change: improve jumpstart retrieve uri unit tests, fix logic for image uris

Author: evakravi

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

@@ -23,7 +23,6 @@
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
 def test_jumpstart_common_image_uri(patched_get_model_specs):
-
     patched_get_model_specs.side_effect = get_spec_from_base_spec
 
     image_uris.retrieve(

tests/unit/sagemaker/image_uris/test_jumpstart.py

@@ -0,0 +1,49 @@
+# 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 model_uris
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+def test_jumpstart_catboost_model_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    # inference
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id="catboost-classification-model",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/catboost-infer/"
+        "infer-catboost-classification-model.tar.gz"
+    )
+
+    # training
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="training",
+        model_id="catboost-classification-model",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/catboost-training/"
+        "train-catboost-classification-model.tar.gz"
+    )

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

@@ -0,0 +1,49 @@
+# 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 model_uris
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+def test_jumpstart_huggingface_model_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    # inference
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id="huggingface-spc-bert-base-cased",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/huggingface-infer/"
+        "infer-huggingface-spc-bert-base-cased.tar.gz"
+    )
+
+    # training
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="training",
+        model_id="huggingface-spc-bert-base-cased",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/huggingface-training/"
+        "train-huggingface-spc-bert-base-cased.tar.gz"
+    )

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

@@ -0,0 +1,49 @@
+# 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 model_uris
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+def test_jumpstart_lightgbm_model_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    # inference
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id="lightgbm-classification-model",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/lightgbm-infer/"
+        "infer-lightgbm-classification-model.tar.gz"
+    )
+
+    # training
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="training",
+        model_id="lightgbm-classification-model",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/lightgbm-training/"
+        "train-lightgbm-classification-model.tar.gz"
+    )

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

@@ -0,0 +1,49 @@
+# 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 model_uris
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+def test_jumpstart_mxnet_model_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    # inference
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id="mxnet-semseg-fcn-resnet50-ade",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/mxnet-infer/"
+        "infer-mxnet-semseg-fcn-resnet50-ade.tar.gz"
+    )
+
+    # training
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="training",
+        model_id="mxnet-semseg-fcn-resnet50-ade",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/mxnet-training/"
+        "train-mxnet-semseg-fcn-resnet50-ade.tar.gz"
+    )

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

@@ -0,0 +1,49 @@
+# 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 model_uris
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsCache.get_model_specs")
+def test_jumpstart_pytorch_model_uri(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    # inference
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="inference",
+        model_id="pytorch-eqa-bert-base-cased",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/pytorch-infer/"
+        "infer-pytorch-eqa-bert-base-cased.tar.gz"
+    )
+
+    # training
+    uri = model_uris.retrieve(
+        region="us-west-2",
+        model_scope="training",
+        model_id="pytorch-eqa-bert-base-cased",
+        model_version="*",
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/pytorch-training/"
+        "train-pytorch-eqa-bert-base-cased.tar.gz"
+    )