diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py index ffc543110b..10d2f3e080 100644 --- a/src/sagemaker/image_uris.py +++ b/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, diff --git a/tests/unit/sagemaker/image_uris/test_retrieve.py b/tests/unit/sagemaker/image_uris/test_retrieve.py index 360587677f..5ab6e621a2 100644 --- a/tests/unit/sagemaker/image_uris/test_retrieve.py +++ b/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(