From 5c7c1affdda1df2f65104d7daec5b590cf7aeafb Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Thu, 1 Oct 2026 14:57:20 -0700 Subject: [PATCH 1/2] fix(image_uris): reject unsupported container_version instead of bad URI retrieve() concatenated a caller-supplied container_version into the image tag with no validation, so e.g. retrieve("pytorch", version="1.12", image_scope="inference", instance_type="ml.p3.8xlarge", container_version="1.1") returned pytorch-inference:1.12-gpu-py38-1.1 -- an image that does not exist (the "1.1" is a Deep Learning Containers release version, not a tag build component). The call silently succeeded with an unpullable URI (#3702). container_version is only meaningful for frameworks whose tag carries a trailing build component: declared in the config container_version map (which already overrides the caller value), or appended by the graviton, neuron, HuggingFace and pytorch-smp paths. For any other framework a caller-supplied container_version can only produce a non-existent tag, so raise a clear ValueError instead of returning a broken URI. Add a unit test; the existing graviton and HuggingFace container_version tests continue to pass (full image_uris suite: 205 passed). Relates-to #3702 --- .../src/sagemaker/core/image_uris.py | 30 +++++++++++++++++++ .../tests/unit/image_uris/test_retrieve.py | 30 +++++++++++++++++++ 2 files changed, 60 insertions(+) diff --git a/sagemaker-core/src/sagemaker/core/image_uris.py b/sagemaker-core/src/sagemaker/core/image_uris.py index f3c063c34f..0b5d96c1ed 100644 --- a/sagemaker-core/src/sagemaker/core/image_uris.py +++ b/sagemaker-core/src/sagemaker/core/image_uris.py @@ -227,6 +227,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] @@ -291,6 +295,32 @@ def retrieve( _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 != HUGGING_FACE_FRAMEWORK + and framework != "pytorch-smp" + 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, diff --git a/sagemaker-core/tests/unit/image_uris/test_retrieve.py b/sagemaker-core/tests/unit/image_uris/test_retrieve.py index 80cc31bd04..f9a039aa81 100644 --- a/sagemaker-core/tests/unit/image_uris/test_retrieve.py +++ b/sagemaker-core/tests/unit/image_uris/test_retrieve.py @@ -83,6 +83,36 @@ def test_retrieve_unsupported_image_scope(config_for_framework): assert "Supported image scope(s): training, inference, eia." in str(e.value) +@patch("sagemaker.core.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" + + @patch("sagemaker.core.image_uris.config_for_framework", return_value=BASE_CONFIG) def test_retrieve_default_image_scope(config_for_framework, caplog): uri = image_uris.retrieve( From 0f7f5f406bd4dfe3c6a23cff4e8131cee73f758b Mon Sep 17 00:00:00 2001 From: Roja Reddy Sareddy Date: Fri, 2 Oct 2026 14:46:54 -0700 Subject: [PATCH 2/2] fix(image_uris): allow container_version for Spark framework The #3702 guard rejected a caller-supplied container_version for any framework whose config has no container_version map and that is not HuggingFace/graviton/neuron/pytorch-smp. That wrongly broke Spark: PySparkProcessor/SparkJarProcessor call image_uris.retrieve("spark", ..., container_version="v1"), and spark.json has no container_version map, so the guard raised ValueError and broke Spark image retrieval. Replace the ad hoc framework checks with an explicit FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX allowlist (huggingface, pytorch-smp, spark) so the set of frameworks that legitimately append a container_version is self-documenting, and add a Spark regression test. (sagemaker-core has no spark processing unit test, so this guards the path the V2 spark suite caught.) Relates-to #3702 --- sagemaker-core/src/sagemaker/core/image_uris.py | 12 ++++++++++-- .../tests/unit/image_uris/test_retrieve.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/sagemaker-core/src/sagemaker/core/image_uris.py b/sagemaker-core/src/sagemaker/core/image_uris.py index 0b5d96c1ed..a6ddaaf680 100644 --- a/sagemaker-core/src/sagemaker/core/image_uris.py +++ b/sagemaker-core/src/sagemaker/core/image_uris.py @@ -38,6 +38,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" @@ -306,8 +315,7 @@ def retrieve( if ( requested_container_version is not None and not version_config.get("container_version") - and framework != HUGGING_FACE_FRAMEWORK - and framework != "pytorch-smp" + and framework not in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX and repo != f"{framework}-inference-graviton" and "neuron" not in repo ): diff --git a/sagemaker-core/tests/unit/image_uris/test_retrieve.py b/sagemaker-core/tests/unit/image_uris/test_retrieve.py index f9a039aa81..71e9d2996c 100644 --- a/sagemaker-core/tests/unit/image_uris/test_retrieve.py +++ b/sagemaker-core/tests/unit/image_uris/test_retrieve.py @@ -113,6 +113,21 @@ def test_retrieve_rejects_container_version_for_unsupported_framework(config_for 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.core.image_uris.config_for_framework", return_value=BASE_CONFIG) def test_retrieve_default_image_scope(config_for_framework, caplog): uri = image_uris.retrieve(