docs(instance-preferences): resolve example images from a candidate type; make the PipeVar override test tox-safe

docs/ml_ops/index.rst

@@ -321,8 +321,14 @@ Run data preprocessing with ``ScriptProcessor`` (sklearn) or ``FrameworkProcesso
 
 .. code-block:: python
 
+   from sagemaker.core import image_uris
    from sagemaker.core.processing import Processor
 
+   # Resolve the image from one of the candidates; every candidate must be able to run it.
+   processing_image = image_uris.retrieve(
+       framework="sklearn", region=region, version="1.2-1", instance_type="ml.m5.4xlarge"
+   )
+
    processor = Processor(
        role=role, image_uri=processing_image, volume_size_in_gb=100,
        instance_preferences=[
@@ -333,7 +339,7 @@ Run data preprocessing with ``ScriptProcessor`` (sklearn) or ``FrameworkProcesso
 
    processor.run(job_name="instance-prefs-processing")
 
-Up to 5 candidates are allowed, each instance type at most once, and exactly one is selected; the list is mutually exclusive with ``instance_type``. Counts use exactly one of two modes — a top-level ``instance_count`` shared by whichever candidate wins, or an ``InstanceCount`` on every candidate — and mixed, partial, or omitted counts are rejected. Selection is based on capacity, not on workload fit, so list only types the job can genuinely run on. The winner is reported as ``SelectedInstanceType`` / ``SelectedInstanceCount`` on the job's ``ClusterConfig``, and billing is for that type and count. Supported on ``Processor``, ``ScriptProcessor``, ``PySparkProcessor``, and ``SparkJarProcessor``; training plans are training-only and do not apply to processing.
+Up to 5 candidates are allowed, each instance type at most once, and exactly one is selected; the list is mutually exclusive with ``instance_type``. Counts use exactly one of two modes — a top-level ``instance_count`` shared by whichever candidate wins, or an ``InstanceCount`` on every candidate — and mixed, partial, or omitted counts are rejected. Selection is based on capacity, not on workload fit, so list only types that can run the job's ``image_uri`` (the image is fixed at submission time; the instance type is not). The winner is reported as ``SelectedInstanceType`` / ``SelectedInstanceCount`` on the job's ``ClusterConfig``, and billing is for that type and count. Supported on ``Processor``, ``ScriptProcessor``, ``PySparkProcessor``, and ``SparkJarProcessor``; training plans are training-only and do not apply to processing.
 
 :doc:`Instance Preferences example <../v3-examples/ml-ops-examples/v3-processing-instance-preferences>`
 

docs/training/index.rst

@@ -398,10 +398,22 @@ Provide an ordered list of candidate instance types and the platform launches th
 
 .. code-block:: python
 
+   from sagemaker.core import image_uris
    from sagemaker.train.model_trainer import ModelTrainer
    from sagemaker.core.training.configs import Compute, SourceCode
    from sagemaker.core.shapes import InstancePreference
 
+   # The image is fixed at submission time while the instance type is not:
+   # resolve it from one of the candidates and list only types that can run it.
+   gpu_training_image = image_uris.retrieve(
+       framework="pytorch",
+       region=region,
+       version="2.0.0",
+       py_version="py310",
+       instance_type="ml.p5.48xlarge",
+       image_scope="training",
+   )
+
    compute = Compute(
        instance_preferences=[
            InstancePreference(instance_type="ml.p5.48xlarge"),
@@ -412,7 +424,7 @@ Provide an ordered list of candidate instance types and the platform launches th
    )
 
    model_trainer = ModelTrainer(
-       training_image=training_image,
+       training_image=gpu_training_image,
        source_code=SourceCode(source_dir="./source", entry_script="train.py"),
        compute=compute,
        base_job_name="instance-preferences-training",

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

@@ -350,7 +350,54 @@ 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."""
+    accepted -- so assert the generated type directly.
+
+    Uses a minimal in-memory model mirroring the real member -> shape wiring;
+    the packaged service JSON is not available in every test environment."""
+
+    _COUNT = {"type": "integer", "min": 1}
+    _STRING = {"type": "string"}
+    _MODEL = {
+        "TrainingInstanceCount": _COUNT,
+        "ProcessingInstanceCount": _COUNT,
+        "TrainingInstanceType": _STRING,
+        "ProcessingInstanceType": _STRING,
+        "ResourceConfig": {
+            "type": "structure",
+            "members": {
+                "InstanceType": {"shape": "TrainingInstanceType"},
+                "InstanceCount": {"shape": "TrainingInstanceCount"},
+                "SelectedInstanceCount": {"shape": "TrainingInstanceCount"},
+            },
+        },
+        "InstancePreference": {
+            "type": "structure",
+            "members": {
+                "InstanceType": {"shape": "TrainingInstanceType"},
+                "InstanceCount": {"shape": "TrainingInstanceCount"},
+            },
+        },
+        "ProcessingClusterConfig": {
+            "type": "structure",
+            "members": {
+                "InstanceType": {"shape": "ProcessingInstanceType"},
+                "InstanceCount": {"shape": "ProcessingInstanceCount"},
+                "SelectedInstanceCount": {"shape": "ProcessingInstanceCount"},
+            },
+        },
+        "ProcessingInstancePreference": {
+            "type": "structure",
+            "members": {
+                "InstanceType": {"shape": "ProcessingInstanceType"},
+                "InstanceCount": {"shape": "ProcessingInstanceCount"},
+            },
+        },
+        # Control: same integer shape, no override registered -> must stay int.
+        "UnrelatedConfig": {
+            "type": "structure",
+            "members": {"InstanceCount": {"shape": "TrainingInstanceCount"}},
+        },
+    }
 
     @pytest.fixture
     def extractor(self, tmp_path):
@@ -363,7 +410,7 @@ def extractor(self, tmp_path):
                 str(tmp_path / "shape_dag.py"),
             ),
         ):
-            return ShapesExtractor()
+            return ShapesExtractor(combined_shapes=self._MODEL)
 
     @pytest.mark.parametrize(
         "shape, member",
@@ -383,3 +430,8 @@ def test_count_members_generate_as_int_pipe_var(self, extractor, shape, member):
             f"{shape}.{member} generated as {members[member]!r}; "
             "expected IntPipeVar via PIPE_VAR_OVERRIDES"
         )
+
+    def test_override_is_targeted_not_blanket(self, extractor):
+        members = extractor.generate_shape_members("UnrelatedConfig")
+        assert "IntPipeVar" not in members["instance_count"]
+        assert "int" in members["instance_count"]

sagemaker-train/tests/integ/train/shallow/README.md

@@ -75,6 +75,7 @@ of any deep test is easy to find:
 | Shallow file | Deep counterpart |
 |---|---|
 | `test_model_trainer.py` | `test_model_trainer.py` |
+| `test_instance_preferences.py` | `test_instance_preferences.py` |
 | `test_sft_trainer.py` | `test_sft_trainer_integration.py` |
 | `test_dpo_trainer.py` | `test_dpo_trainer_integration.py` |
 | `test_rlvr_trainer.py` | `test_rlvr_trainer_integration.py` |
@@ -108,6 +109,7 @@ below accounts for all of them.
 | Deep test | Shallow equivalent |
 |---|---|
 | `test_model_trainer.py` — 8 tests (tar source, py/sh entry, MPI, torchrun, HP json/yaml, custom driver) | `test_model_trainer.py` — `TestSourceCodePackaging`, `TestPayloadShaping`, `TestComputeConfiguration` |
+| `test_instance_preferences.py::test_instance_preferences_select_a_winner_and_complete` (runs to a selected winner and completion) | `test_instance_preferences.py` — `TestInstancePreferencesAccepted` (submit, Describe echo, stop) + `TestInstancePreferencesRejected` |
 | `test_sft_trainer_integration.py::test_sft_trainer_lora_complete_workflow` | `test_minimal_request_is_accepted` + `test_mlflow_resource_arn` |
 | `::test_sft_trainer_with_validation_dataset` | `test_with_validation_dataset` |
 | `::test_sft_trainer_lora_with_sequence_length` | `test_sft_trainer.py::test_sequence_length_is_accepted` |

v3-examples/ml-ops-examples/v3-processing-instance-preferences.ipynb

@@ -39,7 +39,7 @@
     "    framework=\"sklearn\",\n",
     "    region=region,\n",
     "    version=\"1.2-1\",\n",
-    "    instance_type=\"ml.m5.xlarge\",\n",
+    "    instance_type=\"ml.m5.4xlarge\",  # a Step 2 candidate; every candidate must be able to run the image\n",
     "    image_scope=\"training\",\n",
     ")\n"
    ]

v3-examples/training-examples/instance-preferences-example.ipynb

@@ -17,7 +17,11 @@
    "source": [
     "## Step 1: Setup Session\n",
     "\n",
-    "Initialize the SageMaker session, execution role, and a training image.\n"
+    "Initialize the SageMaker session, execution role, and a training image.\n",
+    "\n",
+    "The image is fixed at submission time while the instance type is not, so resolve it\n",
+    "from one of the candidate types you will list in Step 2 and make sure every candidate\n",
+    "can run it. Here both candidates are GPU types and share the same GPU image.\n"
    ]
   },
   {
@@ -42,7 +46,7 @@
     "    region=region,\n",
     "    version=\"2.0.0\",\n",
     "    py_version=\"py310\",\n",
-    "    instance_type=\"ml.m5.xlarge\",\n",
+    "    instance_type=\"ml.p5.48xlarge\",  # a Step 2 candidate; ml.p4d.24xlarge resolves to the same image\n",
     "    image_scope=\"training\",\n",
     ")\n"
    ]