Retrieve image uri from s3 (#1694)

* Retrieve image uri from s3

* Add intelligent default test

---------

Co-authored-by: pintaoz <pintaoz@amazon.com>

Author: pintaoz-aws

sagemaker_utils/src/sagemaker/utils/image_retriever/__init__.py

@@ -1 +1,3 @@
 from __future__ import absolute_import
+
+from sagemaker.utils.image_retriever.image_retriever import ImageRetriever  # noqa: F401

sagemaker_utils/src/sagemaker/utils/image_retriever/image_retriever.py

@@ -1,13 +1,27 @@
+# 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.
+"""Functions for generating ECR image URIs for pre-built SageMaker Docker images."""
+from __future__ import absolute_import
+
 import re
 from typing import Optional
 from graphene.utils.str_converters import to_camel_case
 
-# TODO: Update these dependencies after they are moved to corresponding submodule
-from legacy.src.sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
-from legacy.src.sagemaker.training_compiler.config import TrainingCompilerConfig
+from sagemaker.utils.serverless_inference_config import ServerlessInferenceConfig
+from sagemaker.utils.training_compiler_config import TrainingCompilerConfig
 from sagemaker.utils.utils import _botocore_resolver
 from sagemaker.utils.workflow import is_pipeline_variable
-from image_retriever_utils import (
+from sagemaker.utils.image_retriever.image_retriever_utils import (
     _config_for_framework_and_scope,
     _get_final_image_scope,
     _get_image_tag,
@@ -26,7 +40,7 @@
     config_for_framework,
 )
 from sagemaker.utils.workflow.utilities import override_pipeline_parameter_var
-from sagemaker.utils.config.config_schema import IMAGE_RETRIEVER, MODULES, SAGEMAKER, _simple_path
+from sagemaker.utils.config.config_schema import IMAGE_RETRIEVER, MODULES, PYTHON_SDK, SAGEMAKER, _simple_path
 from sagemaker.utils.config.config_manager import SageMakerConfig
 
 ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
@@ -109,10 +123,13 @@ def retrieve_hugging_face_uri(
             str: The ECR URI for the corresponding SageMaker Docker image.
         """
         args = dict(locals())
+        config = SageMakerConfig()
         for name, val in args.items():
             if name in CONFIGURABLE_ATTRIBUTES and not val:
-                default_value = SageMakerConfig.resolve_value_from_config(
-                    config_path=_simple_path(SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name))
+                default_value = config.resolve_value_from_config(
+                    config_path=_simple_path(
+                        SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
+                    )
                 )
                 if default_value is not None:
                     locals()[name] = default_value
@@ -497,10 +514,13 @@ def retrieve(
             DeprecatedJumpStartModelError: If the version of the model is deprecated.
         """
         args = dict(locals())
+        config = SageMakerConfig()
         for name, val in args.items():
             if name in CONFIGURABLE_ATTRIBUTES and not val:
-                default_value = SageMakerConfig.resolve_value_from_config(
-                    config_path=_simple_path(SAGEMAKER, MODULES, IMAGE_RETRIEVER, to_camel_case(name))
+                default_value = config.resolve_value_from_config(
+                    config_path=_simple_path(
+                        SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, to_camel_case(name)
+                    )
                 )
                 if default_value is not None:
                     locals()[name] = default_value

sagemaker_utils/src/sagemaker/utils/image_retriever/image_retriever_utils.py

@@ -13,9 +13,7 @@
 """Functions for generating ECR image URIs for pre-built SageMaker Docker images."""
 from __future__ import absolute_import
 
-import json
 import logging
-import os
 from typing import Optional
 from packaging.version import Version
 import requests
@@ -44,6 +42,8 @@
 DATA_WRANGLER_FRAMEWORK = "data-wrangler"
 STABILITYAI_FRAMEWORK = "stabilityai"
 SAGEMAKER_TRITONSERVER_FRAMEWORK = "sagemaker-tritonserver"
+# TODO: Update to use a bucket in prod account when GA
+S3_URL_TEMPLATE = "https://image_uri_configs_beta.s3.amazonaws.com/{framework}.json"
 
 
 def _get_image_tag(
@@ -185,11 +185,13 @@ def _validate_for_suppported_frameworks_and_instance_type(framework, instance_ty
 
 def config_for_framework(framework):
     """Loads the JSON config for the given framework."""
-    response = requests.get(s3_url)
-    return response.json()
-    fname = os.path.join(os.path.dirname(__file__), "image_uri_config", "{}.json".format(framework))
-    with open(fname) as f:
-        return json.load(f)
+    try:
+        s3_url = S3_URL_TEMPLATE.format(framework=framework)
+        response = requests.get(s3_url)
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        print(f"Error retrieving config from S3: {e}")
+        raise e
 
 
 def _get_final_image_scope(framework, instance_type, image_scope):

sagemaker_utils/tests/integ/image_retriever/__init__.py

@@ -0,0 +1,117 @@
+from unittest.mock import patch
+import pytest
+
+from sagemaker.utils.config.config_schema import (
+    IMAGE_RETRIEVER,
+    MODULES,
+    PYTHON_SDK,
+    SAGEMAKER,
+    _simple_path,
+)
+from sagemaker.utils.image_retriever import ImageRetriever
+from sagemaker.utils.config.config_manager import SageMakerConfig
+
+
+@pytest.mark.integ
+def test_retrieve_image_uri():
+    image_uri = ImageRetriever.retrieve("clarify", "us-west-2")
+    assert (
+        image_uri == "306415355426.dkr.ecr.us-west-2.amazonaws.com/sagemaker-clarify-processing:1.0"
+    )
+
+    image_uri = ImageRetriever.retrieve(
+        framework="sagemaker-distribution",
+        image_scope="inference",
+        instance_type="ml.g5.4xlarge",
+        region="us-west-1",
+    )
+    assert (
+        image_uri
+        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.0.0-gpu"
+    )
+
+    image_uri = ImageRetriever.retrieve(
+        "xgboost",
+        "eu-west-1",
+        version="0.90-1",
+        instance_type="ml.m5.xlarge",
+        image_scope="inference",
+    )
+    assert (
+        image_uri == "141502667606.dkr.ecr.eu-west-1.amazonaws.com/sagemaker-xgboost:0.90-1-cpu-py3"
+    )
+
+    image_uri = ImageRetriever.retrieve(
+        framework="tensorflow",
+        region="us-west-2",
+        version="2.3",
+        py_version="py37",
+        instance_type="ml.p4d.24xlarge",
+        image_scope="training",
+    )
+    assert (
+        image_uri
+        == "763104351884.dkr.ecr.us-west-2.amazonaws.com/tensorflow-training:2.3-gpu-py37-cu110-ubuntu18.04-v3"
+    )
+
+
+@pytest.mark.integ
+def test_retrieve_pytorch_uri():
+    image_uri = ImageRetriever.retrieve_pytorch_uri(
+        region="us-west-2",
+        version="1.6",
+        py_version="py3",
+        instance_type="ml.p4d.24xlarge",
+        image_scope="training",
+    )
+    assert (
+        image_uri
+        == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.6-gpu-py3-cu110-ubuntu18.04-v3"
+    )
+
+
+@pytest.mark.integ
+def test_retrieve_hugging_face_uri():
+    image_uri = ImageRetriever.retrieve_hugging_face_uri(
+        version="4.28.1",
+        py_version="py310",
+        instance_type="ml.p2.xlarge",
+        region="us-east-1",
+        image_scope="training",
+        base_framework_version="pytorch2.0.0",
+        container_version="cu110-ubuntu20.04",
+    )
+    assert image_uri == "763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training"
+    ":2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04"
+
+
+@pytest.mark.integ
+def test_retrieve_base_python_image_uri():
+    image_uri = ImageRetriever.retrieve_base_python_image_uri()
+    assert image_uri == "236514542706.dkr.ecr.us-west-2.amazonaws.com/sagemaker-base-python-310:1.0"
+
+
+@pytest.mark.integ
+@patch.object(SageMakerConfig, "resolve_value_from_config")
+def test_retrieve_image_uri_intelligent_default(mock_load_config):
+    def custom_return(config_path):
+        if config_path == _simple_path(
+            SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, "ImageScope"
+        ):
+            return "inference"
+        if config_path == _simple_path(
+            SAGEMAKER, PYTHON_SDK, MODULES, IMAGE_RETRIEVER, "InstanceType"
+        ):
+            return "ml.g5.4xlarge"
+
+    mock_load_config.side_effect = custom_return
+
+    # Will get image_scope="inference" and instance_type="ml.g5.4xlarge" from intelligent default
+    image_uri = ImageRetriever.retrieve(
+        framework="sagemaker-distribution",
+        region="us-west-1",
+    )
+    assert (
+        image_uri
+        == "053634841547.dkr.ecr.us-west-1.amazonaws.com/sagemaker-distribution-prod:3.0.0-gpu"
+    )