feat: add image_uris list_frameworks/list_versions/list_py_versions helpers (#4318)

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 (

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:])