change: improve jumpstart retrieve fx impl, cleanup tests, comments, and code

Author: evakravi

tests/unit/sagemaker/image_uris/test_jumpstart.py

@@ -11,11 +11,11 @@
 # 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
+import pytest
 
 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 import constants as sagemaker_constants
@@ -62,6 +62,20 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "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="1.*",
+        instance_type="ml.p2.xlarge",
+    )
+    patched_get_model_specs.assert_called_once_with(
+        sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "1.*"
+    )
+
     with pytest.raises(ValueError):
         image_uris.retrieve(
             framework=None,

tests/unit/sagemaker/test_model_uris.py

@@ -11,12 +11,12 @@
 # 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
+import pytest
 
 from mock.mock import patch
 
 from sagemaker import model_uris
 from sagemaker.jumpstart import constants as sagemaker_constants
-import pytest
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 
@@ -60,6 +60,16 @@ def test_jumpstart_model_uri(patched_get_model_specs):
     patched_get_model_specs.assert_called_once_with(
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
     )
+    patched_get_model_specs.reset_mock()
+
+    model_uris.retrieve(
+        model_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="1.*",
+    )
+    patched_get_model_specs.assert_called_once_with(
+        sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "1.*"
+    )
 
     with pytest.raises(ValueError):
         model_uris.retrieve(

tests/unit/sagemaker/test_script_uris.py

@@ -11,11 +11,11 @@
 # 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
+import pytest
 
 from mock.mock import patch
 
 from sagemaker import script_uris
-import pytest
 
 from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
@@ -62,6 +62,17 @@ def test_jumpstart_script_uri(patched_get_model_specs):
         sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "*"
     )
 
+    patched_get_model_specs.reset_mock()
+
+    script_uris.retrieve(
+        script_scope="training",
+        model_id="pytorch-ic-mobilenet-v2",
+        model_version="1.*",
+    )
+    patched_get_model_specs.assert_called_once_with(
+        sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME, "pytorch-ic-mobilenet-v2", "1.*"
+    )
+
     with pytest.raises(ValueError):
         script_uris.retrieve(
             region="us-west-2",