get_latest_container_image Support for multiple formats of config jsons

Author: nargokul

src/sagemaker/image_utils.py

@@ -1,10 +1,12 @@
 from typing import Optional, Tuple
 
 from sagemaker.image_uris import config_for_framework, retrieve
-
+from packaging.version import Version
 
 def get_latest_container_image(framework: str,
                                image_scope: str,
+                               instance_type: Optional[str] = None,
+                               py_version: Optional[str] = None,
                                region: str = "us-west-2",
                                version: Optional[str] = None) -> Tuple[str, str]:
     try:
@@ -19,9 +21,23 @@ def get_latest_container_image(framework: str,
         version = _fetch_latest_version_from_config(framework_config, image_scope)
     image_uri = retrieve(framework=framework,
                          region=region,
-                         version=version)
+                         version=version,
+                         instance_type=instance_type,
+                         py_version=py_version
+                         )
     return image_uri, version
 
 
 def _fetch_latest_version_from_config(framework_config: dict, image_scope: str) -> str:
-    return framework_config.get(image_scope).get("version_aliases").get("latest")
+    if image_scope in framework_config:
+        if image_scope_config := framework_config[image_scope]:
+            if version_aliases := image_scope_config["version_aliases"]:
+                if latest_version := version_aliases["latest"]:
+                    return latest_version
+    versions = list(framework_config["versions"].keys())
+    top_version = versions[0]
+    bottom_version = versions[-1]
+
+    if Version(top_version) >= Version(bottom_version):
+        return top_version
+    return bottom_version

tests/unit/test_image_utils.py

@@ -5,11 +5,48 @@
 
 
 class TestImageUtils(unittest.TestCase):
+
     @patch('sagemaker.image_utils.config_for_framework')
     @patch('sagemaker.image_utils.retrieve')
     def test_get_latest_container_image(self,
                                         mock_image_retrieve,
                                         mock_config_for_framework):
+        mock_config_for_framework.return_value = {
+            "versions": {
+                "24.03": {
+                    "registries": {
+                        "af-south-1": "626614931356",
+                    },
+                    "repository": "sagemaker-tritonserver",
+                    "tag_prefix": "24.03-py3"
+                },
+                "24.01": {
+                    "registries": {
+                        "af-south-1": "626614931356"
+                    },
+                    "repository": "sagemaker-tritonserver",
+                    "tag_prefix": "24.01-py3"
+                },
+                "23.12": {
+                    "registries": {
+                        "af-south-1": "626614931356"
+                    },
+                    "repository": "sagemaker-tritonserver",
+                    "tag_prefix": "23.12-py3"
+                }
+            }
+        }
+        mock_image_retrieve.return_value = "latest-image"
+
+        image, version = get_latest_container_image("xgboost", "inference")
+        assert image == "latest-image"
+        assert version == "24.03"
+
+    @patch('sagemaker.image_utils.config_for_framework')
+    @patch('sagemaker.image_utils.retrieve')
+    def test_get_latest_container_image_with_alias(self,
+                                                   mock_image_retrieve,
+                                                   mock_config_for_framework):
         mock_config_for_framework.return_value = {
             "inference": {
                 "version_aliases": {
@@ -24,9 +61,7 @@ def test_get_latest_container_image(self,
         assert version == "1"
 
     @patch('sagemaker.image_utils.config_for_framework')
-    @patch('sagemaker.image_utils.retrieve')
     def test_get_latest_container_image_invalid_framework(self,
-                                                          mock_image_retrieve,
                                                           mock_config_for_framework):
         mock_config_for_framework.side_effect = FileNotFoundError
 
@@ -35,12 +70,10 @@ def test_get_latest_container_image_invalid_framework(self,
             assert "No framework config for framework" in str(e.exception)
 
     @patch('sagemaker.image_utils.config_for_framework')
-    @patch('sagemaker.image_utils.retrieve')
     def test_get_latest_container_image_no_framework(self,
-                                                          mock_image_retrieve,
-                                                          mock_config_for_framework):
+                                                     mock_config_for_framework):
         mock_config_for_framework.return_value = {}
 
         with self.assertRaises(ValueError) as e:
             get_latest_container_image("xgboost", "inference")
-            assert "No framework config for framework" in str(e.exception)
+            assert "No framework config for framework" in str(e.exception)