Skip to content

change: add image_uri_config for TensorFlow inference 2.20 and training 2.21 #6230

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mujtaba1747 merged 3 commits into from
Sep 1, 2026
Merged

change: add image_uri_config for TensorFlow inference 2.20 and training 2.21 #6230

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a3d9a47
change: add image_uri_config for TensorFlow inference 2.20 and traini…
Sep 1, 2026
5959ef2
change: strengthen TensorFlow image_uri_config tests
Sep 1, 2026
93cdfe6
change: clarify TensorFlow image_uri_config test comments
Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 96 additions & 2 deletions sagemaker-core/src/sagemaker/core/image_uri_config/tensorflow.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,7 +334,8 @@
"2.14": "2.14.1",
"2.16": "2.16.1",
"2.18": "2.18.0",
"2.19": "2.19.0"
"2.19": "2.19.0",
"2.20": "2.20.0"
},
"versions": {
"1.4.1": {
Expand Down Expand Up @@ -2515,6 +2516,52 @@
"us-west-2": "763104351884"
},
"repository": "tensorflow-inference"
},
"2.20.0": {
"py_versions": [
"py312"
],
"registries": {
"af-south-1": "626614931356",
"ap-east-1": "871362719292",
"ap-east-2": "975050140332",
"ap-northeast-1": "763104351884",
"ap-northeast-2": "763104351884",
"ap-northeast-3": "364406365360",
"ap-south-1": "763104351884",
"ap-south-2": "772153158452",
"ap-southeast-1": "763104351884",
"ap-southeast-2": "763104351884",
"ap-southeast-3": "907027046896",
"ap-southeast-4": "457447274322",
"ap-southeast-5": "550225433462",
"ap-southeast-6": "633930458069",
"ap-southeast-7": "590183813437",
"ca-central-1": "763104351884",
"ca-west-1": "204538143572",
"cn-north-1": "727897471807",
"cn-northwest-1": "727897471807",
"eu-central-1": "763104351884",
"eu-central-2": "380420809688",
"eu-north-1": "763104351884",
"eu-south-1": "692866216735",
"eu-south-2": "503227376785",
"eu-west-1": "763104351884",
"eu-west-2": "763104351884",
"eu-west-3": "763104351884",
"il-central-1": "780543022126",
"me-central-1": "914824155844",
"me-south-1": "217643126080",
"mx-central-1": "637423239942",
"sa-east-1": "763104351884",
"us-east-1": "763104351884",
"us-east-2": "763104351884",
"us-gov-east-1": "446045086412",
"us-gov-west-1": "442386744353",
"us-west-1": "763104351884",
"us-west-2": "763104351884"
},
"repository": "tensorflow-inference"
}
}
},
Expand Down Expand Up @@ -2820,7 +2867,8 @@
"2.14": "2.14.1",
"2.16": "2.16.2",
"2.18": "2.18.0",
"2.19": "2.19.0"
"2.19": "2.19.0",
"2.21": "2.21.0"
},
"versions": {
"1.4.1": {
Expand Down Expand Up @@ -5080,6 +5128,52 @@
"us-west-2": "763104351884"
},
"repository": "tensorflow-training"
},
"2.21.0": {
"py_versions": [
"py312"
],
"registries": {
"af-south-1": "626614931356",
"ap-east-1": "871362719292",
"ap-east-2": "975050140332",
"ap-northeast-1": "763104351884",
"ap-northeast-2": "763104351884",
"ap-northeast-3": "364406365360",
"ap-south-1": "763104351884",
"ap-south-2": "772153158452",
"ap-southeast-1": "763104351884",
"ap-southeast-2": "763104351884",
"ap-southeast-3": "907027046896",
"ap-southeast-4": "457447274322",
"ap-southeast-5": "550225433462",
"ap-southeast-6": "633930458069",
"ap-southeast-7": "590183813437",
"ca-central-1": "763104351884",
"ca-west-1": "204538143572",
"cn-north-1": "727897471807",
"cn-northwest-1": "727897471807",
"eu-central-1": "763104351884",
"eu-central-2": "380420809688",
"eu-north-1": "763104351884",
"eu-south-1": "692866216735",
"eu-south-2": "503227376785",
"eu-west-1": "763104351884",
"eu-west-2": "763104351884",
"eu-west-3": "763104351884",
"il-central-1": "780543022126",
"me-central-1": "914824155844",
"me-south-1": "217643126080",
"mx-central-1": "637423239942",
"sa-east-1": "763104351884",
"us-east-1": "763104351884",
"us-east-2": "763104351884",
"us-gov-east-1": "446045086412",
"us-gov-west-1": "442386744353",
"us-west-1": "763104351884",
"us-west-2": "763104351884"
},
"repository": "tensorflow-training"
}
}
}
Expand Down
174 changes: 174 additions & 0 deletions sagemaker-core/tests/unit/image_uris/test_tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,174 @@
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
from __future__ import absolute_import

import pytest

from sagemaker.core import image_uris
from sagemaker.core.common_utils import ALTERNATE_DOMAINS

GPU_INSTANCE = "ml.g5.2xlarge"
CPU_INSTANCE = "ml.m5.xlarge"
DEFAULT_DOMAIN = "amazonaws.com"

# Representative regions for the exact-URI assertions: three commercial, one China
# (the only partition here with a non-default ECR domain) and one GovCloud. Every other
# region is checked on account/region/repo/tag instead, which keeps the bulk of the
# sweep independent of the host-suffix data in whichever botocore version is installed.
FULL_URI_REGIONS = ["us-east-1", "us-west-2", "eu-west-1", "cn-north-1", "us-gov-west-1"]

# The newest TensorFlow version per scope, with the Python version baked into its tag.
# Asserted below to be the *maximum* registered version, so adding a newer one to
# tensorflow.json fails here until this mapping is updated deliberately.
# tensorflow-inference:2.20.0-{cpu,gpu}-py312
# tensorflow-training:2.21.0-{cpu,gpu}-py312
LATEST = {
"inference": ("2.20.0", "py312"),
"training": ("2.21.0", "py312"),
}

# The already-released version whose registry map a new version is expected to match.
# This is the only account assertion in this file that is not derived from the entry it
# checks, so it catches a one-sided mistake -- a typo in the new version's account, or a
# dropped region -- which asserting against the new entry's own registries cannot. It
# does not catch a change applied to both versions, and it says nothing about whether
# the images are actually published in those regions.
#
# Keep this pointing at an older version. Bumping it to the newly added version makes
# every account assertion in this file tautological again, silently.
REGISTRY_REFERENCE_VERSION = "2.19.0"


def _expected_repo(scope):
return "tensorflow-inference" if scope == "inference" else "tensorflow-training"


def _version_key(version):
return tuple(int(part) for part in version.split("."))


@pytest.mark.parametrize("scope", ["inference", "training"])
@pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
def test_tensorflow_latest_version_is_registered(load_config, scope):
"""The newest version in tensorflow.json is the one this file covers."""
version, py_version = LATEST[scope]
versions = load_config[scope]["versions"]
assert version in versions, f"{version} missing from tensorflow.json {scope}"
newest = max(versions, key=_version_key)
assert newest == version, (
f"tensorflow.json {scope} now registers {newest}, which this file does not cover. "
f"Update LATEST in tests/unit/image_uris/test_tensorflow.py."
)
assert versions[version]["repository"] == _expected_repo(scope)
assert versions[version]["py_versions"] == [py_version]
assert load_config[scope]["processors"] == ["cpu", "gpu"]


@pytest.mark.parametrize("scope", ["inference", "training"])
@pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
def test_tensorflow_latest_version_registries_match_previous_release(load_config, scope):
"""The newest version ships in the same regions and accounts as the previous release."""
version, _ = LATEST[scope]
versions = load_config[scope]["versions"]
assert versions[version]["registries"] == versions[REGISTRY_REFERENCE_VERSION]["registries"]


@pytest.mark.parametrize("scope", ["inference", "training"])
@pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
def test_tensorflow_latest_version_uris(load_config, scope):
"""Every (processor, region) for the newest version resolves to the expected tag."""
version, py_version = LATEST[scope]
version_config = load_config[scope]["versions"][version]
repo = _expected_repo(scope)
for processor, instance_type in (("cpu", CPU_INSTANCE), ("gpu", GPU_INSTANCE)):
expected_tag = f"{version}-{processor}-{py_version}"
for region, account in version_config["registries"].items():
uri = image_uris.retrieve(
framework="tensorflow",
region=region,
version=version,
image_scope=scope,
instance_type=instance_type,
)
assert uri.startswith(f"{account}.dkr.ecr.{region}."), uri
assert uri.endswith(f"/{repo}:{expected_tag}"), uri


@pytest.mark.parametrize("scope", ["inference", "training"])
@pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
def test_tensorflow_latest_version_full_uri(load_config, scope):
"""Exact URI (including domain) for representative commercial/China/GovCloud regions."""
version, py_version = LATEST[scope]
version_config = load_config[scope]["versions"][version]
repo = _expected_repo(scope)
for region in FULL_URI_REGIONS:
account = version_config["registries"][region]
domain = ALTERNATE_DOMAINS.get(region, DEFAULT_DOMAIN)
for processor, instance_type in (("cpu", CPU_INSTANCE), ("gpu", GPU_INSTANCE)):
uri = image_uris.retrieve(
framework="tensorflow",
region=region,
version=version,
image_scope=scope,
instance_type=instance_type,
)
expected_tag = f"{version}-{processor}-{py_version}"
assert uri == f"{account}.dkr.ecr.{region}.{domain}/{repo}:{expected_tag}"


@pytest.mark.parametrize("scope", ["inference", "training"])
@pytest.mark.parametrize("load_config", ["tensorflow.json"], indirect=True)
def test_tensorflow_minor_alias_resolves_to_newest_patch(load_config, scope):
"""The minor alias (2.20 / 2.21) points at its newest patch and keeps the py suffix."""
version, py_version = LATEST[scope]
alias = version.rsplit(".", 1)[0]
assert load_config[scope]["version_aliases"][alias] == version
for processor, instance_type in (("cpu", CPU_INSTANCE), ("gpu", GPU_INSTANCE)):
uri = image_uris.retrieve(
framework="tensorflow",
region="us-west-2",
version=alias,
image_scope=scope,
instance_type=instance_type,
)
# The alias is used verbatim as the tag prefix, matching the published
# `<minor>-<processor>-py312` tags.
assert uri.endswith(f"/{_expected_repo(scope)}:{alias}-{processor}-{py_version}"), uri


@pytest.mark.parametrize("scope", ["inference", "training"])
def test_tensorflow_latest_version_rejects_other_python_versions(scope):
"""Only py312 is offered for the newest version, so any other py_version is an error."""
version, _ = LATEST[scope]
with pytest.raises(ValueError) as error:
image_uris.retrieve(
framework="tensorflow",
region="us-west-2",
version=version,
py_version="py310",
image_scope=scope,
instance_type=CPU_INSTANCE,
)
assert "Unsupported Python version: py310." in str(error.value)


def test_tensorflow_inference_2_19_keeps_tag_without_python_version():
"""2.19 and earlier inference images have no py suffix; adding 2.20 must not change that."""
uri = image_uris.retrieve(
framework="tensorflow",
region="us-west-2",
version="2.19.0",
image_scope="inference",
instance_type=CPU_INSTANCE,
)
assert uri.endswith("/tensorflow-inference:2.19.0-cpu"), uri
Loading