[rsareddy0329]

Issue #, if available: Fixes #3702

Description of changes:

image_uris.retrieve() concatenated a caller-supplied container_version into the image
tag without any validation, so

image_uris.retrieve("pytorch", "us-west-2", version="1.12", image_scope="inference",
                    instance_type="ml.p3.8xlarge", container_version="1.1")
# -> 763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.12-gpu-py38-1.1

returned a URI for an image that does not exist — the call silently succeeded with an
unpullable image (#3702). The "1.1" the reporter used is a Deep Learning Containers
release version, not a tag build component.

container_version is only meaningful for frameworks whose tag carries a trailing build
component — either declared in the config container_version map (which already overrides
the caller value), or appended by the graviton, neuron, HuggingFace and
pytorch-smp code paths. For any other framework a caller-supplied container_version
can only produce a non-existent tag, so this raises a clear ValueError instead of
returning a broken URI.

Why a guard rather than silently dropping it or validating against ECR: the SDK has no
registry of real tags, and valid appends (e.g. cu124 for pytorch-smp, cu110-ubuntu18.04
for HuggingFace) are structurally identical to the invalid one — they differ only by
framework. The allowlist above matches exactly the paths that legitimately append a
container_version, so no working lookup changes behavior.

Testing:

• New test_retrieve_rejects_container_version_for_unsupported_framework: a plain framework
  with container_version raises, and omitting it still returns the valid tag.
• Full tests/unit/image_uris/ suite passes (205 passed, 2 skipped), including the existing
  graviton (ubuntu20.04-sagemaker) and HuggingFace (cu110-ubuntu18.04) container_version
  tests. flake8 and black clean on changed files.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

[github-actions]

Claude finished @rsareddy0329's task in 2m 54s —— View job

---

Code review — fix(image_uris): reject unsupported container_version

• Read the PR diff and the modified image_uris.py / tests
• Trace every code path that appends container_version to a tag
• Check callers across the repo that pass container_version to retrieve()
• Verify the guard's allowlist against those callers
• Post review

---

The intent is good — turning a silently-broken URI into a hard error is the right call for #3702, the comment is clear, and the test is reasonable. But the allowlist is incomplete, and the gap breaks a real, widely-used framework.

🔴 Blocking: the guard breaks all Spark image retrieval

spark is a framework whose tag legitimately carries a trailing container_version (e.g. sagemaker-spark-processing:3.3-cpu-py39-v1), but it is not declared in the config's container_version map and not in the allowlist. Three in-repo callers pass a non-None container_version for framework="spark":

• sagemaker-core/src/sagemaker/core/spark/processing.py:392 — PySparkProcessor / SparkJarProcessor, with container_version = f"v{container_version}"
• sagemaker-core/src/sagemaker/core/remote_function/job.py:874 — container_version=DEFAULT_SPARK_CONTAINER_VERSION ("v1", defined at line 127)
• sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/_image_resolver.py:65 — container_version="v1"

For spark, every condition in the new guard is satisfied:

requested_container_version is not None          # "v1"
and not version_config.get("container_version")  # spark.json has no container_version map
and framework != HUGGING_FACE_FRAMEWORK           # "spark"
and framework != "pytorch-smp"
and repo != f"{framework}-inference-graviton"     # repo == "sagemaker-spark-processing"
and "neuron" not in repo

→ image_uris.retrieve() now raises ValueError for Spark. This breaks PySparkProcessor, SparkJarProcessor, Spark remote functions, and feature-processor image resolution — all of which worked before. (spark.json has no container_version key under its version configs; I confirmed processing → versions → 3.3 carries only repository: sagemaker-spark-processing.)

Why the PR's test run didn't catch it: the tests/unit/image_uris/ suite has no Spark-processing container_version case (test_sparkml.py covers the separate sparkml-serving repo). The Spark processors are tested elsewhere.

Fix: add Spark to the allowlist, e.g. and framework != defaults.SPARK_NAME ("spark"), and add a regression test retrieving a Spark processing image with container_version so this path is locked in.

This also suggests the enumeration should be derived more defensively — e.g. rather than hand-listing frameworks, consider keying off whether any append path actually ran (the container_version value changed from what the caller passed, or the framework/repo matches a known append). The current hard-coded allowlist is easy to drift out of sync with new frameworks.

🟡 Minor: p4d auto-select images with a caller-supplied container_version

_get_image_tag auto-appends a container_version for p4d / smdistributed configs (_should_auto_select_container_version, line 623). A caller who manually passes the correct suffix (e.g. cu110-ubuntu18.04-v3 for a legacy pytorch/tf GPU image) on a p4d instance would previously get a valid URI but will now be rejected, since those frameworks aren't allowlisted. This is narrow (the SDK normally auto-selects it, so callers shouldn't pass it), but worth a note — not a blocker on its own.

Nits (non-blocking)

• The guard reads version for the error message after it's been resolved from any alias; that's fine, just be aware the message may show a resolved version rather than what the caller typed.

Summary

One blocking issue: the allowlist omits spark, which turns previously-valid Spark image lookups into ValueError and breaks PySparkProcessor/SparkJarProcessor, Spark remote functions, and feature-processor image resolution. Please add Spark to the allowlist and add a Spark container_version regression test. No prompt-injection or security concerns found in the diff.
· branch fix/image-uris-container-version

[github-actions]

Claude finished @rsareddy0329's task in 1m 34s —— View job

---

Code review — fix(image_uris): reject unsupported container_version

• Read PR diff and modified image_uris.py / tests
• Trace every code path that appends container_version
• Verify Spark fix addresses prior blocking issue
• Check remaining concerns (p4d auto-select, neuron/graviton gaps)
• Post review

The prior blocking issue is resolved. spark is now in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX and there's a test_retrieve_allows_container_version_for_spark regression test. I re-traced every path that appends a container_version and the guard's exclusions now line up with them:

Append path	Guard exclusion
config container_version map (image_uris.py:244)	not version_config.get("container_version")
neuron sdk_version prefix (:255, :274)	"neuron" not in repo
graviton -sagemaker suffix (:295)	repo != f"{framework}-inference-graviton"
HuggingFace / pytorch-smp / spark	framework not in FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX

I confirmed SPARK_NAME = "spark" (sagemaker-core/src/sagemaker/core/spark/defaults.py:17), so the three in-repo Spark callers (PySparkProcessor/SparkJarProcessor, Spark remote functions, feature-processor image resolution) are no longer broken by the guard. 👍

No security concerns in the diff. No prompt-injection attempts in the PR content.

🟡 Minor (non-blocking) — p4d / smdistributed manual container_version

_get_image_tag auto-appends a container_version for p4d / smdistributed configs (_should_auto_select_container_version, image_uris.py:631), but the frameworks that use it (pytorch, tensorflow) are not allowlisted. A caller who manually passes the correct suffix (e.g. cu110-ubuntu18.04-v3) on a p4d instance would previously get a valid URI and will now be rejected. This is narrow — the SDK auto-selects this value, so callers normally shouldn't pass it, and the auto-select path itself is unaffected (it keys off the computed tag, not the caller value). Worth a note, not a blocker.

Nits (non-blocking)

• FRAMEWORKS_WITH_CONTAINER_VERSION_SUFFIX hardcodes "spark" and "pytorch-smp" as string literals. defaults.SPARK_NAME is already imported and used at image_uris.py:659; using defaults.SPARK_NAME (and a shared constant for "pytorch-smp", which also appears at :756) would avoid drift. Cosmetic.
• The error message reads version after alias resolution, so it may show a resolved version rather than exactly what the caller typed. Harmless.
• test_retrieve_allows_container_version_for_spark uses the real spark.json config (not BASE_CONFIG); good for exercising the true path, just note it couples the test to that config file staying as-is.

Summary

The update correctly fixes the previously-blocking Spark regression and adds a matching test. The guard's allowlist now matches every path that legitimately appends a container_version. Remaining items are minor/cosmetic. Looks good to merge from a correctness standpoint.
· branch fix/image-uris-container-version