change: minor updates to jumpstart retrieve functions

Author: evakravi

src/sagemaker/image_uris.py

@@ -95,13 +95,13 @@ def retrieve(
         return artifacts._retrieve_image_uri(
             model_id,
             model_version,
+            image_scope,
             framework,
             region,
             version,
             py_version,
             instance_type,
             accelerator_type,
-            image_scope,
             container_version,
             distribution,
             base_framework_version,

src/sagemaker/jumpstart/artifacts.py

@@ -28,13 +28,13 @@
 def _retrieve_image_uri(
     model_id: str,
     model_version: str,
+    image_scope: str,
     framework: Optional[str],
     region: Optional[str],
     version: Optional[str],
     py_version: Optional[str],
     instance_type: Optional[str],
     accelerator_type: Optional[str],
-    image_scope: Optional[str],
     container_version: Optional[str],
     distribution: Optional[str],
     base_framework_version: Optional[str],
@@ -50,6 +50,9 @@ def _retrieve_image_uri(
         model_id (str): JumpStart model ID for which to retrieve image URI.
         model_version (str): Version of the JumpStart model for which to retrieve
             the image URI (default: None).
+        image_scope (str): The image type, i.e. what it is used for.
+            Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
+            ``image_scope`` is ignored.
         framework (str): The name of the framework or algorithm.
         region (str): The AWS region.
         version (str): The framework or algorithm version. This is required if there is
@@ -61,9 +64,6 @@ def _retrieve_image_uri(
             there are different images for different processor types.
         accelerator_type (str): Elastic Inference accelerator type. For more, see
             https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
-        image_scope (str): The image type, i.e. what it is used for.
-            Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
-            ``image_scope`` is ignored.
         container_version (str): the version of docker image.
             Ideally the value of parameter should be created inside the framework.
             For custom use, see the list of supported container versions:
@@ -112,7 +112,7 @@ def _retrieve_image_uri(
     if framework is not None and framework != ecr_specs.framework:
         raise ValueError(
             f"Incorrect container framework '{framework}' for JumpStart model ID '{model_id}' "
-            "and version {model_version}'."
+            f"and version {model_version}'."
         )
 
     if version is not None and version != ecr_specs.framework_version:
@@ -124,7 +124,7 @@ def _retrieve_image_uri(
     if py_version is not None and py_version != ecr_specs.py_version:
         raise ValueError(
             f"Incorrect python version '{py_version}' for JumpStart model ID '{model_id}' "
-            "and version {model_version}'."
+            f"and version {model_version}'."
         )
 
     base_framework_version_override: Optional[str] = None

src/sagemaker/script_uris.py

@@ -10,9 +10,8 @@
 # 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.
-"""Accessors to retrieve the script S3 URI to be run pretrained ML models
-in SageMaker containers.
-"""
+"""Accessors to retrieve the script S3 URI to run pretrained ML models."""
+
 from __future__ import absolute_import
 
 import logging

tests/integ/sagemaker/jumpstart/retrieve_uri/inference.py

@@ -51,15 +51,15 @@ def __init__(
         region=JUMPSTART_DEFAULT_REGION_NAME,
         boto_config=Config(retries={"max_attempts": 10, "mode": "standard"}),
         base_name="jumpstart-inference-job",
-        execution_role=Session().get_caller_identity_arn(),
+        execution_role=None,
     ) -> None:
 
         self.suffix = suffix
         self.test_suite_id = os.environ[ENV_VAR_JUMPSTART_SDK_TEST_SUITE_ID]
         self.region = region
         self.config = boto_config
         self.base_name = base_name
-        self.execution_role = execution_role
+        self.execution_role = execution_role or Session().get_caller_identity_arn()
         self.account_id = boto3.client("sts").get_caller_identity()["Account"]
         self.image_uri = image_uri
         self.script_uri = script_uri

tests/integ/sagemaker/jumpstart/retrieve_uri/training.py

@@ -44,7 +44,7 @@ def __init__(
         region=JUMPSTART_DEFAULT_REGION_NAME,
         boto_config=Config(retries={"max_attempts": 10, "mode": "standard"}),
         base_name="jumpstart-training-job",
-        execution_role=Session().get_caller_identity_arn(),
+        execution_role=None,
     ) -> None:
 
         self.account_id = boto3.client("sts").get_caller_identity()["Account"]
@@ -53,7 +53,7 @@ def __init__(
         self.region = region
         self.config = boto_config
         self.base_name = base_name
-        self.execution_role = execution_role
+        self.execution_role = execution_role or Session().get_caller_identity_arn()
         self.image_uri = image_uri
         self.script_uri = script_uri
         self.model_uri = model_uri

tests/unit/sagemaker/jumpstart/utils.py

@@ -48,6 +48,9 @@ def get_header_from_base_header(
 def get_prototype_model_spec(
     region: str = None, model_id: str = None, version: str = None
 ) -> JumpStartModelSpecs:
+    """This function mocks cache accessor functions. For this mock,
+    we only retrieve model specs based on the model id.
+    """
 
     specs = JumpStartModelSpecs(PROTOTYPICAL_MODEL_SPECS_DICT[model_id])
     return specs