fix(core): add PyTorch 2.6 graviton (arm64) inference image

[rsareddy0329]

Issue #, if available: Closes #5532

Description of changes:

The inference_graviton scope in pytorch.json stopped at 2.4.0, so
image_uris.retrieve(..., image_scope="inference_graviton") could not resolve a
2.5+/2.6 graviton image even though the regular inference scope already had
2.5.1/2.6.0 (#5532).

AWS publishes the newer graviton PyTorch inference DLC under the renamed ECR
repository pytorch-inference-arm64 (verified against
aws/deep-learning-containers pytorch/inference/buildspec-arm64-2-6-sm.yml and
docs/src/data/pytorch-inference-arm64/2.6-cpu-sagemaker.yml). The published
SageMaker tag is 2.6.0-cpu-py312-ubuntu22.04-sagemaker. Only 2.6 has a SageMaker
arm64 image (there is no 2.5 sagemaker arm64 build), so only 2.6.0 is added.

• Add the 2.6.0 version entry and 2.6 alias to inference_graviton, using the
  pytorch-inference-arm64 repository, py312, ubuntu22.04, and the standard
  DLC registries.
• Recognize the <framework>-inference-arm64 repository when appending the
  -sagemaker tag suffix, alongside the legacy <framework>-inference-graviton
  repository (the suffix was previously hardcoded to the old repo name, which
  produced a tag without -sagemaker for the arm64 repo).
• Update the graviton image_uris test to read the per-version repository from
  the config so it covers both repository names (TensorFlow and older PyTorch
  versions keep the legacy graviton repo default).

Testing:

• image_uris.retrieve(framework="pytorch", version="2.6.0", image_scope="inference_graviton", instance_type="ml.c7g.xlarge") now returns
  pytorch-inference-arm64:2.6.0-cpu-py312-ubuntu22.04-sagemaker, matching the
  published DLC tag; 2.4.0 still resolves to
  pytorch-inference-graviton:2.4.0-cpu-py311-ubuntu22.04-sagemaker.
• tests/unit/image_uris/test_graviton.py (11) pass, now covering 2.6.0 across
  all regions/instance types; the rest of tests/unit/image_uris pass
  (194 passed, 2 skipped).

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 1m 10s —— View job

---

Code Review — PyTorch 2.6 graviton (arm64) inference image

• Read the PR diff (/tmp/pr.diff)
• Verify the image_uris.py suffix change against callers/config loading
• Verify the pytorch.json config entry structure
• Check for cross-framework side effects of the repo-name change
• Review test coverage

---

Overall this is a clean, well-scoped fix. The change correctly resolves image_scope="inference_graviton" for PyTorch 2.6 by pointing the new version entry at the renamed pytorch-inference-arm64 ECR repo and teaching the -sagemaker suffix logic about that repo name. I verified the mechanics end to end and have no blocking concerns.

Correctness — verified

• image_uris.py:222 reads repo = version_config["repository"] per-version, so routing the 2.6.0 entry to pytorch-inference-arm64 works without further plumbing.
• image_uris.py:281 now appends -sagemaker for both -inference-graviton and -inference-arm64, so the resolved tag becomes 2.6.0-cpu-py312-ubuntu22.04-sagemaker, matching the published DLC. Without this the arm64 repo would have produced a tag missing the -sagemaker suffix. Good catch.
• The version_aliases addition ("2.6": "2.6.0") and the versions entry are placed correctly inside the inference_graviton scope.
• No cross-framework side effects: I grepped all configs under image_uri_config/ for inference-arm64 — only pytorch.json uses it, so widening the suffix check does not change behavior for any other framework. TensorFlow and older PyTorch versions keep the legacy -inference-graviton default.

Tests — adequate

• test_graviton.py now reads repository per-version from the config (defaulting to {framework}-inference-graviton), so 2.6.0 is exercised across all regions/instance types with the correct repo. The parameterization over the config's versions means the new entry is picked up automatically.

Minor / non-blocking

• Consider whether a matching JumpStart / regular-inference deprecation note is needed anywhere, but that's out of scope here.
• The absence of a 2.5 graviton entry is intentional and documented in the PR (no 2.5 sagemaker arm64 DLC build) — worth keeping that note in the commit history in case a reviewer wonders about the gap between 2.4 and 2.6.

No injected instructions were found in the PR content.
· fix/pytorch-graviton-image-uris