From 30a4a6ae398653c7894edbc1bf550cf43beae386 Mon Sep 17 00:00:00 2001 From: Mohamed Zeidan Date: Mon, 28 Sep 2026 14:26:59 -0700 Subject: [PATCH] feat: add image_uris list_frameworks/list_versions/list_py_versions helpers (#4318) --- .../src/sagemaker/core/image_uris.py | 188 +++++++++++++++++- .../unit/image_uris/test_list_helpers.py | 159 +++++++++++++++ 2 files changed, 345 insertions(+), 2 deletions(-) create mode 100644 sagemaker-core/tests/unit/image_uris/test_list_helpers.py diff --git a/sagemaker-core/src/sagemaker/core/image_uris.py b/sagemaker-core/src/sagemaker/core/image_uris.py index f3c063c34f..e95ff51e40 100644 --- a/sagemaker-core/src/sagemaker/core/image_uris.py +++ b/sagemaker-core/src/sagemaker/core/image_uris.py @@ -18,8 +18,8 @@ import logging import os import re -from typing import Optional -from packaging.version import Version +from typing import List, Optional +from packaging.version import InvalidVersion, Version from sagemaker.core import common_utils as utils from sagemaker.core.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION, JUMPSTART_LOGGER @@ -454,6 +454,190 @@ def config_for_framework(framework): return json.load(f) +def _version_sort_key(version): + """Sort key that orders PEP 440 versions semantically, others lexically. + + Some config version keys are not valid PEP 440 versions (e.g. ``"latest"``, + ``"1.x"``). Parseable versions are grouped first and ordered by + ``packaging.version.Version``; the rest fall back to lexical ordering. + """ + try: + return (0, Version(version)) + except InvalidVersion: + return (1, version) + + +def _has_version_data(config): + """Return ``True`` if a loaded config exposes framework version data. + + Handles both top-level ``versions`` configs and scope-nested configs + (e.g. ``training``/``inference`` -> ``versions``). Non-framework config + files (such as ``instance_gpu_info``) return ``False``. + """ + if not isinstance(config, dict): + return False + if "versions" in config: + return True + return any(isinstance(scope, dict) and "versions" in scope for scope in config.values()) + + +def _iter_config_subtrees(framework, image_scope=None): + """Yield the config subtree(s) that hold ``versions``/``version_aliases``. + + Mirrors :func:`_config_for_framework_and_scope` in supporting both config + shapes: a top-level ``versions`` mapping, or scope-nested mappings keyed by + image scope (``training``, ``inference``, ...). Each yielded subtree is the + dict that owns the ``versions`` (and optional ``version_aliases``) keys. + + When ``image_scope`` is provided for a scope-nested config, only that + scope's subtree is yielded; otherwise every scope is yielded (callers union + across them). For single-map configs ``image_scope`` is validated against + the config's ``scope`` list when present, then the top-level subtree is + yielded (the same images serve every scope). + + Raises: + ValueError: If the framework is unknown, or ``image_scope`` is not a + valid scope for the framework. + """ + try: + config = config_for_framework(framework) + except OSError: + raise ValueError( + "Unsupported framework: {}. Supported framework(s): {}.".format( + framework, ", ".join(list_frameworks()) + ) + ) + + if "versions" in config: + available_scopes = config.get("scope") + if image_scope is not None and available_scopes and image_scope not in available_scopes: + raise ValueError( + "Unsupported image scope: {} for framework: {}. Supported scope(s): {}.".format( + image_scope, framework, ", ".join(sorted(available_scopes)) + ) + ) + yield config + return + + scopes = { + name: scope + for name, scope in config.items() + if isinstance(scope, dict) and "versions" in scope + } + + if image_scope is not None: + if image_scope not in scopes: + raise ValueError( + "Unsupported image scope: {} for framework: {}. Supported scope(s): {}.".format( + image_scope, framework, ", ".join(sorted(scopes)) + ) + ) + yield scopes[image_scope] + return + + for scope in scopes.values(): + yield scope + + +def list_frameworks() -> List[str]: + """Lists the frameworks that have a bundled image URI config. + + Returns: + list[str]: Sorted framework names accepted by :func:`retrieve` and + :func:`list_versions`. Non-framework config files are excluded. + """ + config_dir = os.path.join(os.path.dirname(__file__), "image_uri_config") + frameworks = [] + for fname in os.listdir(config_dir): + if not fname.endswith(".json"): + continue + framework = fname[: -len(".json")] + try: + config = config_for_framework(framework) + except (OSError, ValueError): + continue + if _has_version_data(config): + frameworks.append(framework) + return sorted(frameworks) + + +def list_versions(framework: str, image_scope: Optional[str] = None) -> List[str]: + """Lists the available versions for a framework. + + Both concrete versions and version aliases (e.g. ``"2.0"`` -> ``"2.0.1"``) + are returned, since :func:`retrieve` accepts either as its ``version`` + argument. + + Args: + framework (str): The framework name (see :func:`list_frameworks`). + image_scope (str): Optional image scope (e.g. ``"training"`` or + ``"inference"``). For scope-nested configs, versions are unioned + across all scopes when omitted; for single-map configs it is + validated (when the config declares a ``scope`` list) but does not + change the returned versions. (default: None). + + Returns: + list[str]: Sorted, de-duplicated available versions and aliases. + + Raises: + ValueError: If the framework is unknown, or ``image_scope`` is invalid. + """ + versions = set() + for subtree in _iter_config_subtrees(framework, image_scope): + versions.update(subtree.get("versions", {}).keys()) + versions.update(subtree.get("version_aliases", {}).keys()) + return sorted(versions, key=_version_sort_key) + + +def list_py_versions(framework: str, version: str, image_scope: Optional[str] = None) -> List[str]: + """Lists the available Python versions for a framework version. + + A version alias (e.g. ``"2.0"``) is resolved to its concrete target before + lookup, matching :func:`retrieve`. + + Args: + framework (str): The framework name (see :func:`list_frameworks`). + version (str): The framework version or alias (see + :func:`list_versions`). + image_scope (str): Optional image scope (e.g. ``"training"`` or + ``"inference"``). Only applies to scope-nested configs; when + omitted, Python versions are unioned across all scopes that define + the given ``version``. (default: None). + + Returns: + list[str]: Sorted, de-duplicated Python versions (e.g. ``["py310"]``). + Empty if the version defines no Python versions -- either because the + image has none (e.g. algorithm images) or because the framework nests + Python versions under a base-framework version (e.g. ``huggingface``), + which this helper does not descend into. + + Raises: + ValueError: If the framework, ``image_scope``, or ``version`` is unknown. + """ + py_versions = set() + available_versions = set() + found = False + for subtree in _iter_config_subtrees(framework, image_scope): + versions_dict = subtree.get("versions", {}) + aliases = subtree.get("version_aliases", {}) + available_versions.update(versions_dict.keys()) + available_versions.update(aliases.keys()) + resolved = aliases.get(version, version) + if resolved in versions_dict: + found = True + py_versions.update(versions_dict[resolved].get("py_versions", [])) + + if not found: + raise ValueError( + "Unsupported version: {} for framework: {}. Supported version(s): {}.".format( + version, + framework, + ", ".join(sorted(available_versions, key=_version_sort_key)), + ) + ) + return sorted(py_versions, key=_version_sort_key) + + def _get_final_image_scope(framework, instance_type, image_scope): """Return final image scope based on provided framework and instance type.""" if ( diff --git a/sagemaker-core/tests/unit/image_uris/test_list_helpers.py b/sagemaker-core/tests/unit/image_uris/test_list_helpers.py new file mode 100644 index 0000000000..cf09b4a9a2 --- /dev/null +++ b/sagemaker-core/tests/unit/image_uris/test_list_helpers.py @@ -0,0 +1,159 @@ +# 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. +"""Unit tests for the read-only listing helpers in ``image_uris``. + +These tests run against the REAL bundled image_uri_config JSON files, so they +assert on stable, well-known frameworks/versions rather than mocking config. +""" + +from __future__ import absolute_import + +import pytest + +from sagemaker.core import image_uris + + +def test_list_frameworks_includes_known_frameworks_sorted(): + frameworks = image_uris.list_frameworks() + + assert isinstance(frameworks, list) + assert frameworks == sorted(frameworks) + # No duplicates. + assert len(frameworks) == len(set(frameworks)) + # Well-known frameworks must be present. + for expected in ("pytorch", "tensorflow", "xgboost", "sklearn"): + assert expected in frameworks + # Non-framework config files must be excluded. + assert "instance_gpu_info" not in frameworks + + +def test_list_versions_scoped_framework_is_nonempty_and_sorted(): + # pytorch is a scope-nested config (training/inference/...). + versions = image_uris.list_versions("pytorch") + + assert isinstance(versions, list) + assert len(versions) > 0 + assert len(versions) == len(set(versions)) + # A well-known pytorch training version is present. + assert "2.0.0" in versions + # Versions are returned in the module's sorted order. + assert versions == image_uris.list_versions("pytorch") + + +def test_list_versions_respects_image_scope_argument(): + training = set(image_uris.list_versions("pytorch", image_scope="training")) + inference = set(image_uris.list_versions("pytorch", image_scope="inference")) + union = set(image_uris.list_versions("pytorch")) + + assert training # non-empty + assert inference # non-empty + # The unscoped call unions across scopes, so each scope is a subset. + assert training <= union + assert inference <= union + assert (training | inference) <= union + + +def test_list_versions_top_level_versions_framework(): + # linear-learner ships a top-level ``versions`` config (no scope nesting). + versions = image_uris.list_versions("linear-learner") + + assert isinstance(versions, list) + assert len(versions) > 0 + assert len(versions) == len(set(versions)) + + +def test_list_py_versions_returns_expected_values(): + # pytorch training 2.0.0 ships py310. + py_versions = image_uris.list_py_versions("pytorch", "2.0.0", image_scope="training") + + assert py_versions == ["py310"] + + +def test_list_py_versions_without_scope_unions_across_scopes(): + py_versions = image_uris.list_py_versions("pytorch", "2.0.0") + + assert "py310" in py_versions + + +def test_list_py_versions_version_without_py_versions_returns_empty(): + # Algorithm images (e.g. linear-learner) do not define ``py_versions``. + versions = image_uris.list_versions("linear-learner") + assert versions # sanity + assert image_uris.list_py_versions("linear-learner", versions[0]) == [] + + +def test_list_versions_unknown_framework_raises_clear_error(): + with pytest.raises(ValueError, match="Unsupported framework"): + image_uris.list_versions("not-a-real-framework") + + +def test_list_py_versions_unknown_framework_raises_clear_error(): + with pytest.raises(ValueError, match="Unsupported framework"): + image_uris.list_py_versions("not-a-real-framework", "1.0.0") + + +def test_list_py_versions_unknown_version_raises_clear_error(): + with pytest.raises(ValueError, match="Unsupported version"): + image_uris.list_py_versions("pytorch", "0.0.0-does-not-exist") + + +def test_list_versions_includes_version_aliases(): + # ``retrieve`` accepts aliases like "2.0" (-> "2.0.1"); they must be listed. + versions = image_uris.list_versions("pytorch", image_scope="training") + + assert "2.0" in versions # alias + assert "2.0.1" in versions # concrete target + + +def test_list_py_versions_resolves_version_alias(): + # "2.0" is an alias for the concrete "2.0.1" pytorch training version. + from_alias = image_uris.list_py_versions("pytorch", "2.0", image_scope="training") + from_concrete = image_uris.list_py_versions("pytorch", "2.0.1", image_scope="training") + + assert from_alias == from_concrete + assert from_alias # non-empty + + +def test_list_versions_invalid_scope_raises_clear_error(): + with pytest.raises(ValueError, match="Unsupported image scope"): + image_uris.list_versions("pytorch", image_scope="not-a-scope") + + +def test_list_versions_sorted_semantically_not_lexically(): + versions = image_uris.list_versions("pytorch", image_scope="training") + + # Semantic (not lexical) ordering: 2.9.0 must precede 2.10.0. + numeric = [v for v in versions if v in ("2.9.0", "2.10.0")] + assert numeric == ["2.9.0", "2.10.0"] + + +def test_list_versions_non_pep440_keys_sort_last(): + # data-wrangler ships non-PEP440 version keys ("1.x", "2.x", "3.x"). + versions = image_uris.list_versions("data-wrangler") + + assert any(v.endswith(".x") for v in versions) + # Unparseable keys are grouped after all PEP 440 versions. + from packaging.version import InvalidVersion, Version + + def is_pep440(v): + try: + Version(v) + return True + except InvalidVersion: + return False + + pep440_flags = [is_pep440(v) for v in versions] + # Once we hit the first non-PEP440 key, none after it may be PEP 440. + if False in pep440_flags: + first_non = pep440_flags.index(False) + assert not any(pep440_flags[first_non:])