Skip to content

Commit 2d0804d

Browse files
rsareddy0329Roja Reddy Sareddy
andauthored
fix(image_uris): reject unsupported container_version instead of returning a bad URI (#6377)
* 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. This is the V2 (master-v2) counterpart of #6376. Add a unit test; existing graviton and HuggingFace container_version tests continue to pass. Relates-to #3702 * 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 (test_processing.py::test_retrieve_image_uri). 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. Relates-to #3702 --------- Co-authored-by: Roja Reddy Sareddy <rsareddy@amazon.com>
1 parent 7c37579 commit 2d0804d

2 files changed

Lines changed: 83 additions & 0 deletions

File tree

‎src/sagemaker/image_uris.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,15 @@
3737

3838
ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
3939
HUGGING_FACE_FRAMEWORK = "huggingface"
40+
# Frameworks whose image tag legitimately carries a caller-supplied container_version
41+
# component even though their config has no container_version map. For these the SDK's
42+
# own higher-level classes pass container_version (e.g. SparkProcessor -> "v1",
43+
# the HuggingFace GPU/trcomp path -> "cu110-ubuntu18.04", pytorch-smp -> "cu124").
44+
# See issue #3702: for any framework NOT in this set, a caller-supplied container_version
45+
# can only produce a tag that does not correspond to a published image.
46+
FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX = frozenset(
47+
{HUGGING_FACE_FRAMEWORK, "pytorch-smp", "spark"}
48+
)
4049
HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
4150
HUGGING_FACE_TEI_GPU_FRAMEWORK = "huggingface-tei"
4251
HUGGING_FACE_TEI_CPU_FRAMEWORK = "huggingface-tei-cpu"
@@ -226,6 +235,10 @@ def retrieve(
226235
serverless_inference_config,
227236
)
228237

238+
# Preserve the caller-supplied container_version so we can detect the case where it
239+
# is a framework that does not use a container_version tag component (see guard below).
240+
requested_container_version = container_version
241+
229242
# if container version is available in .json file, utilize that
230243
if version_config.get("container_version"):
231244
container_version = version_config["container_version"][processor]
@@ -281,6 +294,31 @@ def retrieve(
281294
container_version = f"{container_version}-sagemaker"
282295
_validate_instance_deprecation(framework, instance_type, version)
283296

297+
# A caller-supplied container_version is only meaningful for frameworks whose image
298+
# tag actually carries a trailing version component -- either declared in the config
299+
# (``container_version`` map, which overrode the value above), or appended by the
300+
# graviton/neuron/HuggingFace/SMP code paths. For any other framework the value would
301+
# be concatenated into a tag that does not correspond to a published image (e.g.
302+
# ``pytorch-inference:1.12-gpu-py38-1.1``), so fail fast instead of returning a URI
303+
# that cannot be pulled (see #3702). ``container_version`` is a tag build component
304+
# (such as a CUDA/OS suffix), not a Deep Learning Containers release version.
305+
if (
306+
requested_container_version is not None
307+
and not version_config.get("container_version")
308+
and framework not in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX
309+
and repo != f"{framework}-inference-graviton"
310+
and "neuron" not in repo
311+
):
312+
raise ValueError(
313+
"container_version '{}' is not supported for framework '{}' version '{}'. "
314+
"The resulting image tag would not correspond to a published image. "
315+
"container_version is an image tag build component (for example a CUDA/OS "
316+
"suffix), not a Deep Learning Containers release version; omit it and pass "
317+
"the full framework version instead.".format(
318+
requested_container_version, framework, version
319+
)
320+
)
321+
284322
tag = _get_image_tag(
285323
container_version,
286324
distribution,

‎tests/unit/sagemaker/image_uris/test_retrieve.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,51 @@ def test_retrieve_unsupported_image_scope(config_for_framework):
8383
assert "Supported image scope(s): training, inference, eia." in str(e.value)
8484

8585

86+
@patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
87+
def test_retrieve_rejects_container_version_for_unsupported_framework(config_for_framework):
88+
"""container_version on a framework with no container_version tag component raises (#3702).
89+
90+
Previously the value was concatenated into the tag, producing a URI for an image
91+
that does not exist (e.g. ``dummy:1.0.0-cpu-py3-1.1``).
92+
"""
93+
with pytest.raises(ValueError) as e:
94+
image_uris.retrieve(
95+
framework="useless-string",
96+
version="1.0.0",
97+
py_version="py3",
98+
instance_type="ml.c4.xlarge",
99+
region="us-west-2",
100+
image_scope="training",
101+
container_version="1.1",
102+
)
103+
assert "container_version '1.1' is not supported" in str(e.value)
104+
# Omitting container_version still returns the valid tag (no regression).
105+
uri = image_uris.retrieve(
106+
framework="useless-string",
107+
version="1.0.0",
108+
py_version="py3",
109+
instance_type="ml.c4.xlarge",
110+
region="us-west-2",
111+
image_scope="training",
112+
)
113+
assert uri == "123412341234.dkr.ecr.us-west-2.amazonaws.com/dummy:1.0.0-cpu-py3"
114+
115+
116+
def test_retrieve_allows_container_version_for_spark():
117+
"""Spark passes a container_version suffix (e.g. "v1") with no config map (#3702).
118+
119+
Spark is in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX, so the guard must not reject it.
120+
"""
121+
uri = image_uris.retrieve(
122+
framework="spark",
123+
version="3.1",
124+
instance_type="ml.c5.xlarge",
125+
region="us-west-2",
126+
container_version="v1",
127+
)
128+
assert uri.endswith("-v1")
129+
130+
86131
@patch("sagemaker.image_uris.config_for_framework", return_value=BASE_CONFIG)
87132
def test_retrieve_default_image_scope(config_for_framework, caplog):
88133
uri = image_uris.retrieve(

0 commit comments

Comments
 (0)