change: retrieve script uri argument name, comments

Author: evakravi

src/sagemaker/model_uris.py

@@ -36,7 +36,7 @@ def retrieve(
         model_id (str): JumpStart model id for which to retrieve model S3 URI.
         model_version (str): JumpStart model version for which to retrieve model S3 URI.
         model_scope (str): The model type, i.e. what it is used for.
-            Valid values: "training", "inference", "eia".
+            Valid values: "training" and "inference".
     Returns:
         str: the model artifact URI for the corresponding model.
 

src/sagemaker/script_uris.py

@@ -26,16 +26,16 @@ def retrieve(
     region=jumpstart_constants.JUMPSTART_DEFAULT_REGION_NAME,
     model_id=None,
     model_version=None,
-    model_scope=None,
+    script_scope=None,
 ):
     """Retrieves the model script s3 URI for the model matching the given arguments.
 
     Args:
         region (str): Region for which to retrieve model script S3 URI.
         model_id (str): JumpStart model id for which to retrieve model script S3 URI.
         model_version (str): JumpStart model version for which to retrieve model script S3 URI.
-        model_scope (str): The model type, i.e. what it is used for.
-            Valid values: "training", "inference", "eia".
+        script_scope (str): The script type, i.e. what it is used for.
+            Valid values: "training" and "inference".
     Returns:
         str: the model script URI for the corresponding model.
 
@@ -50,13 +50,14 @@ def retrieve(
     model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
         region, model_id, model_version
     )
-    if model_scope is None:
+    if script_scope is None:
         raise ValueError(
-            "Must specify `model_scope` argument to retrieve model script uri for JumpStart models."
+            "Must specify `script_scope` argument to retrieve model script uri for "
+            "JumpStart models."
         )
-    if model_scope == "inference":
+    if script_scope == "inference":
         model_script_key = model_specs.hosting_script_key
-    elif model_scope == "training":
+    elif script_scope == "training":
         if not model_specs.training_supported:
             raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
         model_script_key = model_specs.training_script_key

tests/unit/sagemaker/test_script_uris.py

@@ -28,7 +28,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
     patched_get_model_specs.side_effect = get_spec_from_base_spec
     uri = script_uris.retrieve(
         region="us-west-2",
-        model_scope="inference",
+        script_scope="inference",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="*",
     )
@@ -42,7 +42,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
 
     uri = script_uris.retrieve(
         region="us-west-2",
-        model_scope="training",
+        script_scope="training",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="*",
     )
@@ -54,7 +54,7 @@ def test_jumpstart_script_uri(patched_get_model_specs):
     patched_get_model_specs.reset_mock()
 
     script_uris.retrieve(
-        model_scope="training",
+        script_scope="training",
         model_id="pytorch-ic-mobilenet-v2",
         model_version="*",
     )
@@ -65,15 +65,15 @@ def test_jumpstart_script_uri(patched_get_model_specs):
     with pytest.raises(ValueError):
         script_uris.retrieve(
             region="us-west-2",
-            model_scope="BAD_SCOPE",
+            script_scope="BAD_SCOPE",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
 
     with pytest.raises(ValueError):
         script_uris.retrieve(
             region="mars-south-1",
-            model_scope="training",
+            script_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
             model_version="*",
         )
@@ -86,12 +86,12 @@ def test_jumpstart_script_uri(patched_get_model_specs):
 
     with pytest.raises(ValueError):
         script_uris.retrieve(
-            model_scope="training",
+            script_scope="training",
             model_version="*",
         )
 
     with pytest.raises(ValueError):
         script_uris.retrieve(
-            model_scope="training",
+            script_scope="training",
             model_id="pytorch-ic-mobilenet-v2",
         )