fix(image_uris): reject unsupported container_version instead of returning a bad URI

src/sagemaker/image_uris.py

@@ -37,6 +37,15 @@
 
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
 HUGGING_FACE_FRAMEWORK = "huggingface"
+# Frameworks whose image tag legitimately carries a caller-supplied container_version
+# component even though their config has no container_version map. For these the SDK's
+# own higher-level classes pass container_version (e.g. SparkProcessor -> "v1",
+# the HuggingFace GPU/trcomp path -> "cu110-ubuntu18.04", pytorch-smp -> "cu124").
+# See issue #3702: for any framework NOT in this set, a caller-supplied container_version
+# can only produce a tag that does not correspond to a published image.
+FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX = frozenset(
+    {HUGGING_FACE_FRAMEWORK, "pytorch-smp", "spark"}
+)
 HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
 HUGGING_FACE_TEI_GPU_FRAMEWORK = "huggingface-tei"
 HUGGING_FACE_TEI_CPU_FRAMEWORK = "huggingface-tei-cpu"
@@ -226,6 +235,10 @@ def retrieve(
         serverless_inference_config,
     )
 
+    # Preserve the caller-supplied container_version so we can detect the case where it
+    # is a framework that does not use a container_version tag component (see guard below).
+    requested_container_version = container_version
+
     # if container version is available in .json file, utilize that
     if version_config.get("container_version"):
         container_version = version_config["container_version"][processor]
@@ -281,6 +294,31 @@ def retrieve(
         container_version = f"{container_version}-sagemaker"
     _validate_instance_deprecation(framework, instance_type, version)
 
+    # A caller-supplied container_version is only meaningful for frameworks whose image
+    # tag actually carries a trailing version component -- either declared in the config
+    # (``container_version`` map, which overrode the value above), or appended by the
+    # graviton/neuron/HuggingFace/SMP code paths. For any other framework the value would
+    # be concatenated into a tag that does not correspond to a published image (e.g.
+    # ``pytorch-inference:1.12-gpu-py38-1.1``), so fail fast instead of returning a URI
+    # that cannot be pulled (see #3702). ``container_version`` is a tag build component
+    # (such as a CUDA/OS suffix), not a Deep Learning Containers release version.
+    if (
+        requested_container_version is not None
+        and not version_config.get("container_version")
+        and framework not in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX
+        and repo != f"{framework}-inference-graviton"
+        and "neuron" not in repo
+    ):
+        raise ValueError(
+            "container_version '{}' is not supported for framework '{}' version '{}'. "
+            "The resulting image tag would not correspond to a published image. "
+            "container_version is an image tag build component (for example a CUDA/OS "
+            "suffix), not a Deep Learning Containers release version; omit it and pass "
+            "the full framework version instead.".format(
+                requested_container_version, framework, version
+            )
+        )
+
     tag = _get_image_tag(
         container_version,
         distribution,

tests/unit/sagemaker/image_uris/test_retrieve.py

@@ -83,6 +83,51 @@ def test_retrieve_unsupported_image_scope(config_for_framework):
     assert "Supported image scope(s): training, inference, eia." in str(e.value)
 
 
+@patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
+def test_retrieve_rejects_container_version_for_unsupported_framework(config_for_framework):
+    """container_version on a framework with no container_version tag component raises (#3702).
+
+    Previously the value was concatenated into the tag, producing a URI for an image
+    that does not exist (e.g. ``dummy:1.0.0-cpu-py3-1.1``).
+    """
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework="useless-string",
+            version="1.0.0",
+            py_version="py3",
+            instance_type="ml.c4.xlarge",
+            region="us-west-2",
+            image_scope="training",
+            container_version="1.1",
+        )
+    assert "container_version '1.1' is not supported" in str(e.value)
+    # Omitting container_version still returns the valid tag (no regression).
+    uri = image_uris.retrieve(
+        framework="useless-string",
+        version="1.0.0",
+        py_version="py3",
+        instance_type="ml.c4.xlarge",
+        region="us-west-2",
+        image_scope="training",
+    )
+    assert uri == "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy:1.0.0-cpu-py3"
+
+
+def test_retrieve_allows_container_version_for_spark():
+    """Spark passes a container_version suffix (e.g. "v1") with no config map (#3702).
+
+    Spark is in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX, so the guard must not reject it.
+    """
+    uri = image_uris.retrieve(
+        framework="spark",
+        version="3.1",
+        instance_type="ml.c5.xlarge",
+        region="us-west-2",
+        container_version="v1",
+    )
+    assert uri.endswith("-v1")
+
+
 @patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
 def test_retrieve_default_image_scope(config_for_framework, caplog):
     uri = image_uris.retrieve(