get_latest_container_image utils

Author: nargokul

src/sagemaker/image_utils.py

@@ -0,0 +1,27 @@
+from typing import Optional, Tuple
+
+from sagemaker.image_uris import config_for_framework, retrieve
+
+
+def get_latest_container_image(framework: str,
+                               image_scope: str,
+                               region: str = "us-west-2",
+                               version: Optional[str] = None) -> Tuple[str, str]:
+    try:
+        framework_config = config_for_framework(framework)
+    except FileNotFoundError:
+        raise ValueError("Invalid framework {}".format(framework))
+
+    if not framework_config:
+        raise ValueError("Invalid framework {}".format(framework))
+
+    if not version:
+        version = _fetch_latest_version_from_config(framework_config, image_scope)
+    image_uri = retrieve(framework=framework,
+                         region=region,
+                         version=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")

tests/unit/test_image_utils.py

@@ -0,0 +1,46 @@
+import unittest
+from unittest.mock import patch
+
+from sagemaker.image_utils import get_latest_container_image
+
+
+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 = {
+            "inference": {
+                "version_aliases": {
+                    "latest": "1"
+                }
+            }
+        }
+        mock_image_retrieve.return_value = "latest-image"
+
+        image, version = get_latest_container_image("xgboost", "inference")
+        assert image == "latest-image"
+        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
+
+        with self.assertRaises(ValueError) as e:
+            get_latest_container_image("xgboost", "inference")
+            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.return_value = {}
+
+        with self.assertRaises(ValueError) as e:
+            get_latest_container_image("xgboost", "inference")
+            assert "No framework config for framework" in str(e.exception)