From 4f40b5d0e384311ae72f10b1303ccce6bfa01adb Mon Sep 17 00:00:00 2001 From: Lucas Jia Date: Tue, 6 Oct 2026 11:37:53 -0700 Subject: [PATCH] test(image_uris): derive pytorch-amzn2023 versions from config test_pytorch_amzn2023_training_uris compared the config's version list against a hardcoded ["2.11", "2.12", "2.13"]. The automated image_uri_configs update (#6374) added 2.14, so the exact-equality assertion broke sagemaker-core unit tests on master. Iterate over the versions present in the config, require a minimum set of known versions, and check that the "latest" alias points to a defined version, so future version additions do not break the test. --- .../tests/unit/image_uris/test_pytorch_amzn2023.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/sagemaker-core/tests/unit/image_uris/test_pytorch_amzn2023.py b/sagemaker-core/tests/unit/image_uris/test_pytorch_amzn2023.py index 328cf84880..fd9b9aec9e 100644 --- a/sagemaker-core/tests/unit/image_uris/test_pytorch_amzn2023.py +++ b/sagemaker-core/tests/unit/image_uris/test_pytorch_amzn2023.py @@ -25,16 +25,20 @@ # encodes CUDA directly (e.g. 2.13-cu133-amzn2023-sagemaker) with no "gpu" token, # so these versions set "processor_in_tag": false and bake the accelerator into # container_version. -AMZN2023_VERSIONS = ["2.11", "2.12", "2.13"] +# Versions that must stay resolvable. Newer versions added by the automated +# image_uri_configs updates are picked up from the config automatically, so this +# set only needs to grow when a version must be guaranteed to remain available. +REQUIRED_AMZN2023_VERSIONS = {"2.11", "2.12", "2.13", "2.14"} @pytest.mark.parametrize("load_config", ["pytorch-amzn2023.json"], indirect=True) def test_pytorch_amzn2023_training_uris(load_config): """pytorch-amzn2023 resolves both cpu and gpu; the gpu tag carries cuNNN, not "gpu".""" training = load_config["training"] - assert sorted(training["versions"]) == sorted(AMZN2023_VERSIONS) - for version in AMZN2023_VERSIONS: - version_config = training["versions"][version] + versions = training["versions"] + assert REQUIRED_AMZN2023_VERSIONS <= set(versions) + assert training["version_aliases"]["latest"] in versions + for version, version_config in versions.items(): assert version_config["repository"] == "pytorch" assert version_config["processor_in_tag"] is False container_version = version_config["container_version"]