fix(instance-preferences): run integ tests in CI, reject managed spot client-side, stabilize IntPipeVar codegen

sagemaker-core/src/sagemaker/core/modules/utils.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Utils module."""
+
 from __future__ import absolute_import
 
 import os
@@ -200,7 +201,7 @@ def validate_instance_preferences(compute) -> None:
 
     - instance_preferences is mutually exclusive with the classic
       single-cluster fields instance_type / instance_groups /
-      instance_placement_config.
+      instance_placement_config, and with managed spot training.
     - Instance types must not repeat across preferences.
     - Count mode: exactly one of the top-level instance_count (applies to
       whichever preference wins) with no per-preference instance_count, or an
@@ -231,6 +232,12 @@ def _value(obj, field):
                 f"instance_preferences is mutually exclusive with {field}; "
                 "specify either a single fixed cluster or instance_preferences, not both."
             )
+    # Spot is a bool: only an explicit True conflicts (False/None is the default).
+    if _value(compute, "enable_managed_spot_training") is True:
+        raise ValueError(
+            "instance_preferences is mutually exclusive with managed spot training "
+            "(enable_managed_spot_training=True)."
+        )
 
     instance_types = [_value(p, "instance_type") for p in preferences]
     duplicates = sorted(

sagemaker-core/src/sagemaker/core/tools/constants.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Constants used in the code_generator modules."""
+
 import os
 
 CLASS_METHODS = set(["create", "add", "register", "import", "list", "get"])
@@ -134,5 +135,16 @@
         "InstanceCount": "IntPipeVar",
         "VolumeSizeInGB": "IntPipeVar",
         "KeepAlivePeriodInSeconds": "IntPipeVar",
+        "SelectedInstanceCount": "IntPipeVar",
+    },
+    "InstancePreference": {
+        "InstanceCount": "IntPipeVar",
+    },
+    "ProcessingClusterConfig": {
+        "InstanceCount": "IntPipeVar",
+        "SelectedInstanceCount": "IntPipeVar",
+    },
+    "ProcessingInstancePreference": {
+        "InstanceCount": "IntPipeVar",
     },
 }

sagemaker-core/src/sagemaker/core/training/utils.py

@@ -264,7 +264,7 @@ def validate_instance_preferences(compute) -> None:
 
     - instance_preferences is mutually exclusive with the classic
       single-cluster fields instance_type / instance_groups /
-      instance_placement_config.
+      instance_placement_config, and with managed spot training.
     - Instance types must not repeat across preferences.
     - Count mode: exactly one of the top-level instance_count (applies to
       whichever preference wins) with no per-preference instance_count, or an
@@ -295,6 +295,12 @@ def _value(obj, field):
                 f"instance_preferences is mutually exclusive with {field}; "
                 "specify either a single fixed cluster or instance_preferences, not both."
             )
+    # Spot is a bool: only an explicit True conflicts (False/None is the default).
+    if _value(compute, "enable_managed_spot_training") is True:
+        raise ValueError(
+            "instance_preferences is mutually exclusive with managed spot training "
+            "(enable_managed_spot_training=True)."
+        )
 
     instance_types = [_value(p, "instance_type") for p in preferences]
     duplicates = sorted(

sagemaker-core/tests/integ/processing/test_instance_preferences_processing.py

@@ -24,14 +24,11 @@
   ``SelectedInstanceCount`` report the resolved winner, which must be one of
   the submitted preferences.
 
-Configuration (environment variables; SKIPPED when unset):
-
-- ``PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN``  - execution role ARN
-- ``PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI`` - processing image
-- ``PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` - optional comma-separated
-  candidate types (default: ml.m5.xlarge, ml.m5.large)
-- ``SAGEMAKER_ENDPOINT``                             - optional endpoint
-  override.
+Runs in the standard integration-test account: the execution role is the
+suite's ``SageMakerRole`` and the image is resolved through ``image_uris``, the
+same way the other ``sagemaker-core`` integration tests obtain theirs.
+``PROCESSING_INSTANCE_PREFERENCES_TEST_INSTANCE_TYPES`` (comma-separated)
+overrides the candidate types for accounts where the defaults lack quota.
 """
 
 from __future__ import absolute_import
@@ -40,11 +37,12 @@
 import time
 import uuid
 
-import pytest
+import boto3
 
-ROLE_ARN = os.environ.get("PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN")
-IMAGE_URI = os.environ.get("PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI")
-ENDPOINT = os.environ.get("SAGEMAKER_ENDPOINT")
+from sagemaker.core import image_uris
+
+ROLE = "SageMakerRole"
+REGION = os.environ.get("AWS_REGION", os.environ.get("AWS_DEFAULT_REGION", "us-west-2"))
 
 PREFERENCE_TYPES = [
     t.strip()
@@ -55,36 +53,40 @@
 ]
 WAIT_TIMEOUT_SECONDS = 30 * 60
 POLL_SECONDS = 30
-
-pytestmark = pytest.mark.skipif(
-    not (ROLE_ARN and IMAGE_URI),
-    reason=(
-        "Processing Instance Preferences integ test requires "
-        "PROCESSING_INSTANCE_PREFERENCES_TEST_ROLE_ARN and "
-        "PROCESSING_INSTANCE_PREFERENCES_TEST_IMAGE_URI"
-    ),
-)
-
-
-def _sagemaker_client():
-    """Build a boto3 SageMaker client, honoring SAGEMAKER_ENDPOINT."""
-    import boto3
-
-    region = os.environ.get("AWS_REGION", "us-west-2")
-    return boto3.client("sagemaker", region_name=region, endpoint_url=ENDPOINT)
+TERMINAL = ("Completed", "Failed", "Stopped")
+
+
+def _stop_quietly(client, job_name):
+    """The job has done its part once a winner is selected; do not leave it
+    running to max_runtime on shared quota."""
+    try:
+        if (
+            client.describe_processing_job(ProcessingJobName=job_name)["ProcessingJobStatus"]
+            in TERMINAL
+        ):
+            return
+        client.stop_processing_job(ProcessingJobName=job_name)
+    except Exception:  # pylint: disable=broad-except
+        pass
+
+
+def _processing_image():
+    return image_uris.retrieve(
+        "sklearn", REGION, version="1.2-1", py_version="py3", instance_type=PREFERENCE_TYPES[0]
+    )
 
 
 def test_processor_instance_preferences_e2e():
     """Create a real processing job with InstancePreferences and verify the winner."""
     from sagemaker.core.helper.session_helper import Session
     from sagemaker.core.processing import Processor
 
-    client = _sagemaker_client()
+    client = boto3.client("sagemaker", region_name=REGION)
     session = Session(sagemaker_client=client)
 
     processor = Processor(
-        role=ROLE_ARN,
-        image_uri=IMAGE_URI,
+        role=ROLE,
+        image_uri=_processing_image(),
         instance_count=1,
         instance_preferences=[{"InstanceType": t} for t in PREFERENCE_TYPES],
         volume_size_in_gb=30,
@@ -95,39 +97,43 @@ def test_processor_instance_preferences_e2e():
     job_name = f"instance-prefs-proc-integ-{uuid.uuid4().hex[:8]}"
     processor.run(wait=False, logs=False, job_name=job_name)
 
-    # --- Create accepted; Describe echoes the request contract -------------
-    described = client.describe_processing_job(ProcessingJobName=job_name)
-    cluster_config = described["ProcessingResources"]["ClusterConfig"]
-    assert [p["InstanceType"] for p in cluster_config.get("InstancePreferences", [])] == (
-        PREFERENCE_TYPES
-    ), f"Describe did not echo InstancePreferences: {cluster_config}"
-    # The customer never set the top-level InstanceType; it must not come
-    # back populated on Describe.
-    assert "InstanceType" not in cluster_config, cluster_config
-
-    # --- Wait for a terminal-or-resolved state ------------------------------
-    deadline = time.time() + WAIT_TIMEOUT_SECONDS
-    status = described["ProcessingJobStatus"]
-    while time.time() < deadline:
+    try:
+
+        # --- Create accepted; Describe echoes the request contract -------------
         described = client.describe_processing_job(ProcessingJobName=job_name)
+        cluster_config = described["ProcessingResources"]["ClusterConfig"]
+        assert [p["InstanceType"] for p in cluster_config.get("InstancePreferences", [])] == (
+            PREFERENCE_TYPES
+        ), f"Describe did not echo InstancePreferences: {cluster_config}"
+        # The customer never set the top-level InstanceType; it must not come
+        # back populated on Describe.
+        assert "InstanceType" not in cluster_config, cluster_config
+
+        # --- Wait for a terminal-or-resolved state ------------------------------
+        deadline = time.time() + WAIT_TIMEOUT_SECONDS
         status = described["ProcessingJobStatus"]
+        while time.time() < deadline:
+            described = client.describe_processing_job(ProcessingJobName=job_name)
+            status = described["ProcessingJobStatus"]
+            cluster_config = described["ProcessingResources"]["ClusterConfig"]
+            if status in TERMINAL:
+                break
+            if status == "InProgress" and cluster_config.get("SelectedInstanceType"):
+                break
+            time.sleep(POLL_SECONDS)
+
+        failure_reason = described.get("FailureReason", "")
+
+        # --- Full contract: resolved winner is surfaced and is a submitted pref -
         cluster_config = described["ProcessingResources"]["ClusterConfig"]
-        if status in ("Completed", "Failed", "Stopped"):
-            break
-        if status == "InProgress" and cluster_config.get("SelectedInstanceType"):
-            break
-        time.sleep(POLL_SECONDS)
-
-    failure_reason = described.get("FailureReason", "")
-
-    # --- Full contract: resolved winner is surfaced and is a submitted pref -
-    cluster_config = described["ProcessingResources"]["ClusterConfig"]
-    selected_type = cluster_config.get("SelectedInstanceType")
-    selected_count = cluster_config.get("SelectedInstanceCount")
-    assert selected_type in PREFERENCE_TYPES, (
-        f"SelectedInstanceType={selected_type!r} not among submitted preferences "
-        f"{PREFERENCE_TYPES} (job={job_name}, status={status}, "
-        f"failure={failure_reason!r})"
-    )
-    assert selected_count == 1
-    assert "InstanceType" not in cluster_config, cluster_config
+        selected_type = cluster_config.get("SelectedInstanceType")
+        selected_count = cluster_config.get("SelectedInstanceCount")
+        assert selected_type in PREFERENCE_TYPES, (
+            f"SelectedInstanceType={selected_type!r} not among submitted preferences "
+            f"{PREFERENCE_TYPES} (job={job_name}, status={status}, "
+            f"failure={failure_reason!r})"
+        )
+        assert selected_count == 1
+        assert "InstanceType" not in cluster_config, cluster_config
+    finally:
+        _stop_quietly(client, job_name)

sagemaker-core/tests/unit/test_compute_configs.py

@@ -1,4 +1,5 @@
 """Unit tests for Compute and HyperPodCompute config classes."""
+
 import pytest
 from sagemaker.core.training.configs import Compute, HyperPodCompute
 
@@ -169,6 +170,26 @@ def test_instance_type_rejected_with_preferences(self, compute_cls):
                 instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")],
             )
 
+    def test_managed_spot_rejected_with_preferences(self, compute_cls):
+        from sagemaker.core.shapes.shapes import InstancePreference
+
+        with pytest.raises(ValueError, match="mutually exclusive with managed spot training"):
+            compute_cls(
+                instance_count=1,
+                enable_managed_spot_training=True,
+                instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")],
+            )
+
+    def test_managed_spot_false_allowed_with_preferences(self, compute_cls):
+        from sagemaker.core.shapes.shapes import InstancePreference
+
+        compute = compute_cls(
+            instance_count=1,
+            enable_managed_spot_training=False,
+            instance_preferences=[InstancePreference(instance_type="ml.m5.xlarge")],
+        )
+        assert compute.enable_managed_spot_training is False
+
     def test_uniform_count_rejected_with_per_preference_counts(self, compute_cls):
         from sagemaker.core.shapes.shapes import InstancePreference
 
@@ -239,9 +260,7 @@ def test_whole_job_plan_allowed_without_per_preference_plans(self, compute_cls):
 
         compute_cls(
             instance_count=1,
-            training_plan_arn=(
-                "arn:aws:sagemaker:us-west-2:111122223333:training-plan/whole-job"
-            ),
+            training_plan_arn=("arn:aws:sagemaker:us-west-2:111122223333:training-plan/whole-job"),
             instance_preferences=[
                 InstancePreference(instance_type="ml.p5.48xlarge"),
                 InstancePreference(instance_type="ml.p4d.24xlarge"),
@@ -260,9 +279,7 @@ def test_modules_compute_instance_preferences_round_trip(self):
             InstancePreference(instance_type="ml.p5.48xlarge"),
             InstancePreference(instance_type="ml.p4d.24xlarge"),
         ]
-        rc = ModulesCompute(
-            instance_preferences=prefs, instance_count=2
-        )._to_resource_config()
+        rc = ModulesCompute(instance_preferences=prefs, instance_count=2)._to_resource_config()
         assert [p.instance_type for p in rc.instance_preferences] == [
             "ml.p5.48xlarge",
             "ml.p4d.24xlarge",
@@ -271,9 +288,7 @@ def test_modules_compute_instance_preferences_round_trip(self):
     def test_modules_compute_single_type_still_works(self):
         from sagemaker.core.modules.configs import Compute as ModulesCompute
 
-        rc = ModulesCompute(
-            instance_type="ml.m5.xlarge", instance_count=1
-        )._to_resource_config()
+        rc = ModulesCompute(instance_type="ml.m5.xlarge", instance_count=1)._to_resource_config()
         assert rc.instance_type == "ml.m5.xlarge"
 
 

sagemaker-core/tests/unit/tools/test_shapes_extractor.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Unit tests for sagemaker.core.tools.shapes_extractor module."""
+
 from __future__ import absolute_import
 
 import pytest
@@ -343,3 +344,42 @@ def test_get_required_members_none(self, mock_reformat):
         result = extractor.get_required_members("TestShape")
 
         assert len(result) == 0
+
+
+class TestInstancePreferencesPipeVarOverrides:
+    """The IntPipeVar annotations on instance-preferences count members come from
+    PIPE_VAR_OVERRIDES, not the service model. If an override is dropped, codegen
+    silently narrows the member back to int and pipeline variables stop being
+    accepted -- so assert the generated type directly."""
+
+    @pytest.fixture
+    def extractor(self, tmp_path):
+        # Constructing the extractor regenerates shape_dag.py; point that write
+        # at a temp file so the unit test leaves the checked-in file alone.
+        with (
+            patch("sagemaker.core.tools.shapes_extractor.reformat_file_with_black"),
+            patch(
+                "sagemaker.core.tools.shapes_extractor.SHAPE_DAG_FILE_PATH",
+                str(tmp_path / "shape_dag.py"),
+            ),
+        ):
+            return ShapesExtractor()
+
+    @pytest.mark.parametrize(
+        "shape, member",
+        [
+            ("ResourceConfig", "instance_count"),
+            ("ResourceConfig", "selected_instance_count"),
+            ("InstancePreference", "instance_count"),
+            ("ProcessingClusterConfig", "instance_count"),
+            ("ProcessingClusterConfig", "selected_instance_count"),
+            ("ProcessingInstancePreference", "instance_count"),
+        ],
+    )
+    def test_count_members_generate_as_int_pipe_var(self, extractor, shape, member):
+        members = extractor.generate_shape_members(shape)
+        assert member in members, f"{shape}.{member} missing from generated members"
+        assert "IntPipeVar" in members[member], (
+            f"{shape}.{member} generated as {members[member]!r}; "
+            "expected IntPipeVar via PIPE_VAR_OVERRIDES"
+        )