feat: jumpstart retrieve functions (wip)

Author: evakravi

tests/unit/sagemaker/image_uris/jumpstart/conftest.py

@@ -0,0 +1,17 @@
+from mock.mock import Mock
+import pytest
+
+REGION_NAME = "us-west-2"
+BUCKET_NAME = "some-bucket-name"
+
+
+@pytest.fixture(scope="module")
+def session():
+    boto_mock = Mock(region_name=REGION_NAME)
+    sms = Mock(
+        boto_session=boto_mock,
+        boto_region_name=REGION_NAME,
+        config=None,
+    )
+    sms.default_bucket = Mock(return_value=BUCKET_NAME)
+    return sms

tests/unit/sagemaker/image_uris/jumpstart/test_catboost.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_catboost_image_uri(patched_get_model_specs):
+def test_jumpstart_catboost_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -49,6 +49,7 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -71,6 +72,7 @@ def test_jumpstart_catboost_image_uri(patched_get_model_specs):
         py_version=model_specs.training_ecr_specs.py_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_huggingface.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_huggingface_image_uri(patched_get_model_specs):
+def test_jumpstart_huggingface_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -48,6 +48,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs):
         transformers_version=model_specs.hosting_ecr_specs.huggingface_transformers_version,
         pytorch_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -76,6 +77,7 @@ def test_jumpstart_huggingface_image_uri(patched_get_model_specs):
         pytorch_version=model_specs.training_ecr_specs.framework_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert (

tests/unit/sagemaker/image_uris/jumpstart/test_lightgbm.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_lightgbm_image_uri(patched_get_model_specs):
+def test_jumpstart_lightgbm_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -49,6 +49,7 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -71,6 +72,7 @@ def test_jumpstart_lightgbm_image_uri(patched_get_model_specs):
         py_version=model_specs.training_ecr_specs.py_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_mxnet.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_mxnet_image_uri(patched_get_model_specs):
+def test_jumpstart_mxnet_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -49,6 +49,7 @@ def test_jumpstart_mxnet_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -71,6 +72,7 @@ def test_jumpstart_mxnet_image_uri(patched_get_model_specs):
         py_version=model_specs.training_ecr_specs.py_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_pytorch.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_pytorch_image_uri(patched_get_model_specs):
+def test_jumpstart_pytorch_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -49,6 +49,7 @@ def test_jumpstart_pytorch_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -71,6 +72,7 @@ def test_jumpstart_pytorch_image_uri(patched_get_model_specs):
         py_version=model_specs.training_ecr_specs.py_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_sklearn.py

@@ -24,7 +24,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_sklearn_image_uri(patched_get_model_specs):
+def test_jumpstart_sklearn_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -50,6 +50,7 @@ def test_jumpstart_sklearn_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -75,6 +76,7 @@ def test_jumpstart_sklearn_image_uri(patched_get_model_specs):
         instance_type=instance_type,
         instance_count=1,
         image_uri_region=region,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_tensorflow.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_tensorflow_image_uri(patched_get_model_specs):
+def test_jumpstart_tensorflow_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -48,6 +48,7 @@ def test_jumpstart_tensorflow_image_uri(patched_get_model_specs):
         model_data="mock_data",
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -70,6 +71,7 @@ def test_jumpstart_tensorflow_image_uri(patched_get_model_specs):
         py_version=model_specs.training_ecr_specs.py_version,
         instance_type=instance_type,
         instance_count=1,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/jumpstart/test_xgboost.py

@@ -23,7 +23,7 @@
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-def test_jumpstart_xgboost_image_uri(patched_get_model_specs):
+def test_jumpstart_xgboost_image_uri(patched_get_model_specs, session):
 
     patched_get_model_specs.side_effect = get_prototype_model_spec
 
@@ -49,6 +49,7 @@ def test_jumpstart_xgboost_image_uri(patched_get_model_specs):
         entry_point="mock_entry_point",
         framework_version=model_specs.hosting_ecr_specs.framework_version,
         py_version=model_specs.hosting_ecr_specs.py_version,
+        sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
     assert uri == framework_class_uri
@@ -72,6 +73,7 @@ def test_jumpstart_xgboost_image_uri(patched_get_model_specs):
         instance_type=instance_type,
         instance_count=1,
         image_uri_region=region,
+        sagemaker_session=session,
     ).training_image_uri(region=region)
 
     assert uri == framework_class_uri

tests/unit/sagemaker/image_uris/test_jumpstart.py

@@ -0,0 +1,111 @@
+# 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.
+from __future__ import absolute_import
+
+from mock.mock import patch
+
+from sagemaker import image_uris
+import pytest
+
+from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
+from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart import constants as sagemaker_constants
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_script_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_spec_from_base_spec
+    uri = image_uris.retrieve(
+        framework=None,
+        region="us-west-2",
+        image_scope="inference",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="*",
+        instance_type="ml.p2.xlarge",
+    )
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.5.0-gpu-py3"
+    patched_get_model_specs.assert_called_once_with("us-west-2", "pytorch-ic-mobilenet-v2", "*")
+
+    patched_get_model_specs.reset_mock()
+
+    uri = image_uris.retrieve(
+        framework=None,
+        region="us-west-2",
+        image_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="*",
+        instance_type="ml.p2.xlarge",
+    )
+    assert uri == "763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.5.0-gpu-py3"
+    patched_get_model_specs.assert_called_once_with("us-west-2", "pytorch-ic-mobilenet-v2", "*")
+    patched_get_model_specs.reset_mock()
+
+    image_uris.retrieve(
+        framework=None,
+        region="us-west-2",
+        image_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="*",
+        instance_type="ml.p2.xlarge",
+    )
+    patched_get_model_specs.assert_called_once_with(
+        sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
+    )
+
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            image_scope="BAD_SCOPE",
+            model_id="pytorch-ic-mobilenet-v2",
+            model_version="*",
+            instance_type="ml.p2.xlarge",
+        )
+
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework=None,
+            region="mars-south-1",
+            image_scope="training",
+            model_id="pytorch-ic-mobilenet-v2",
+            model_version="*",
+            instance_type="ml.p2.xlarge",
+        )
+
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            model_id="pytorch-ic-mobilenet-v2",
+            model_version="*",
+            instance_type="ml.p2.xlarge",
+        )
+
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            framework=None,
+            region="us-west-2",
+            image_scope="training",
+            model_version="*",
+            instance_type="ml.p2.xlarge",
+        )
+
+    with pytest.raises(ValueError) as e:
+        image_uris.retrieve(
+            region="us-west-2",
+            framework=None,
+            image_scope="training",
+            model_id="pytorch-ic-mobilenet-v2",
+            instance_type="ml.p2.xlarge",
+        )