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

Author: evakravi

src/sagemaker/image_uris.py

@@ -19,8 +19,10 @@
 import re
 
 from sagemaker import utils
+from sagemaker.jumpstart.utils import is_jumpstart_model_input
 from sagemaker.spark import defaults
-from sagemaker.jumpstart import accessors as jumpstart_accessors
+from sagemaker.jumpstart import artifacts
+
 
 logger = logging.getLogger(__name__)
 
@@ -81,45 +83,23 @@ def retrieve(
     Raises:
         ValueError: If the combination of arguments specified is not supported.
     """
-    if model_id is not None or model_version is not None:
-        if model_id is None or model_version is None:
-            raise ValueError(
-                "Must specify `model_id` and `model_version` when getting image uri for "
-                "JumpStart models. "
-            )
-        model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
-            region, model_id, model_version
-        )
-        if image_scope is None:
-            raise ValueError(
-                "Must specify `image_scope` argument to retrieve image uri for JumpStart models."
-            )
-        if image_scope == "inference":
-            ecr_specs = model_specs.hosting_ecr_specs
-        elif image_scope == "training":
-            if not model_specs.training_supported:
-                raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
-            ecr_specs = model_specs.training_ecr_specs
-        else:
-            raise ValueError("JumpStart models only support inference and training.")
-
-        if framework is not None and framework != ecr_specs.framework:
-            raise ValueError(
-                f"Bad value for container framework for JumpStart model: '{framework}'."
-            )
-
-        return retrieve(
-            framework=ecr_specs.framework,
-            region=region,
-            version=ecr_specs.framework_version,
-            py_version=ecr_specs.py_version,
-            instance_type=instance_type,
-            accelerator_type=accelerator_type,
-            image_scope=image_scope,
-            container_version=container_version,
-            distribution=distribution,
-            base_framework_version=base_framework_version,
-            training_compiler_config=training_compiler_config,
+    if is_jumpstart_model_input(model_id, model_version):
+        assert model_id is not None
+        assert model_version is not None
+        return artifacts._retrieve_image_uri(
+            model_id,
+            model_version,
+            framework,
+            region,
+            version,
+            py_version,
+            instance_type,
+            accelerator_type,
+            image_scope,
+            container_version,
+            distribution,
+            base_framework_version,
+            training_compiler_config,
         )
 
     if training_compiler_config is None:

src/sagemaker/jumpstart/accessors.py

@@ -103,14 +103,14 @@ def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelS
 
     @staticmethod
     def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:
-        """Sets cache kwargs. Clears the cache.
+        """Sets cache kwargs, clear the cache.
 
         Raises:
             ValueError: If region in `cache_kwargs` is inconsistent with `region` argument.
 
         Args:
             cache_kwargs (str): cache kwargs to validate.
-            region (str): The region to validate along with the kwargs.
+            region (str): Optional. The region to validate along with the kwargs.
         """
         cache_kwargs = JumpStartModelsCache._validate_region_cache_kwargs(cache_kwargs, region)
         JumpStartModelsCache._cache_kwargs = cache_kwargs

src/sagemaker/jumpstart/artifacts.py

@@ -0,0 +1,237 @@
+# 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.
+"""This module contains functions for obtainining JumpStart artifacts."""
+from __future__ import absolute_import
+from typing import Optional
+from sagemaker import image_uris
+from sagemaker.jumpstart.constants import (
+    JUMPSTART_DEFAULT_REGION_NAME,
+    INFERENCE,
+    TRAINING,
+    SUPPORTED_JUMPSTART_SCOPES,
+)
+from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
+from sagemaker.jumpstart import accessors as jumpstart_accessors
+
+
+def _retrieve_image_uri(
+    model_id: str,
+    model_version: 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],
+    training_compiler_config: Optional[str],
+):
+    """Retrieves the container image URI for JumpStart models.
+
+    Only `model_id` and `model_version` are required to be non-None;
+    the rest of the fields are auto-populated.
+
+
+    Args:
+        model_id (str): JumpStart model id for which to retrieve image URI.
+        model_version (str): JumpStart model version for which to retrieve image URI.
+        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
+            more than one supported version for the given framework or algorithm.
+        py_version (str): The Python version. This is required if there is
+            more than one supported Python version for the given framework version.
+        instance_type (str): The SageMaker instance type. For supported types, see
+            https://aws.amazon.com/sagemaker/pricing/instance-types. This is required if
+            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:
+            https://github.com/aws/deep-learning-containers/blob/master/available_images.md
+            (default: None).
+        distribution (dict): A dictionary with information on how to run distributed training
+        training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
+            A configuration class for the SageMaker Training Compiler
+            (default: None).
+
+    Returns:
+        str: the ECR URI for the corresponding SageMaker Docker image.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    assert region is not None
+
+    if image_scope is None:
+        raise ValueError(
+            "Must specify `image_scope` argument to retrieve image uri for JumpStart models."
+        )
+    if image_scope not in SUPPORTED_JUMPSTART_SCOPES:
+        raise ValueError("JumpStart models only support inference and training.")
+
+    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+        region, model_id, model_version
+    )
+
+    if image_scope == INFERENCE:
+        ecr_specs = model_specs.hosting_ecr_specs
+    elif image_scope == TRAINING:
+        if not model_specs.training_supported:
+            raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
+        assert model_specs.training_ecr_specs is not None
+        ecr_specs = model_specs.training_ecr_specs
+
+    if framework is not None and framework != ecr_specs.framework:
+        raise ValueError(f"Bad value for container framework for JumpStart model: '{framework}'.")
+
+    if version is not None and version != ecr_specs.framework_version:
+        raise ValueError(
+            f"Bad value for container framework version for JumpStart model: '{version}'."
+        )
+
+    if py_version is not None and py_version != ecr_specs.py_version:
+        raise ValueError(
+            f"Bad value for container python version for JumpStart model: '{py_version}'."
+        )
+
+    if framework == "huggingface":
+        base_framework_version = ecr_specs.framework_version
+
+    return image_uris.retrieve(
+        framework=ecr_specs.framework,
+        region=region,
+        version=ecr_specs.framework_version,
+        py_version=ecr_specs.py_version,
+        instance_type=instance_type,
+        accelerator_type=accelerator_type,
+        image_scope=image_scope,
+        container_version=container_version,
+        distribution=distribution,
+        base_framework_version=base_framework_version,
+        training_compiler_config=training_compiler_config,
+    )
+
+
+def _retrieve_model_uri(
+    model_id: str,
+    model_version: str,
+    model_scope: Optional[str],
+    region: Optional[str],
+):
+    """Retrieves the model artifact S3 URI for the model matching the given arguments.
+
+    Args:
+        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" and "inference".
+        region (str): Region for which to retrieve model S3 URI.
+    Returns:
+        str: the model artifact S3 URI for the corresponding model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    assert region is not None
+
+    if model_scope is None:
+        raise ValueError(
+            "Must specify `model_scope` argument to retrieve model "
+            "artifact uri for JumpStart models."
+        )
+
+    if model_scope not in SUPPORTED_JUMPSTART_SCOPES:
+        raise ValueError("JumpStart models only support inference and training.")
+
+    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+        region, model_id, model_version
+    )
+    if model_scope == INFERENCE:
+        model_artifact_key = model_specs.hosting_artifact_key
+    elif model_scope == TRAINING:
+        if not model_specs.training_supported:
+            raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
+        assert model_specs.training_artifact_key is not None
+        model_artifact_key = model_specs.training_artifact_key
+
+    bucket = get_jumpstart_content_bucket(region)
+
+    model_s3_uri = f"s3://{bucket}/{model_artifact_key}"
+
+    return model_s3_uri
+
+
+def _retrieve_script_uri(
+    model_id: str,
+    model_version: str,
+    script_scope: Optional[str],
+    region: Optional[str],
+):
+    """Retrieves the model script s3 URI for the model matching the given arguments.
+
+    Args:
+        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.
+        script_scope (str): The script type, i.e. what it is used for.
+            Valid values: "training" and "inference".
+        region (str): Region for which to retrieve model script S3 URI.
+    Returns:
+        str: the model script URI for the corresponding model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    assert region is not None
+
+    if script_scope is None:
+        raise ValueError(
+            "Must specify `script_scope` argument to retrieve model script uri for "
+            "JumpStart models."
+        )
+
+    if script_scope not in SUPPORTED_JUMPSTART_SCOPES:
+        raise ValueError("JumpStart models only support inference and training.")
+
+    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+        region, model_id, model_version
+    )
+    if script_scope == INFERENCE:
+        model_script_key = model_specs.hosting_script_key
+    elif script_scope == TRAINING:
+        if not model_specs.training_supported:
+            raise ValueError(f"JumpStart model id '{model_id}' does not support training.")
+        assert model_specs.training_script_key is not None
+        model_script_key = model_specs.training_script_key
+
+    bucket = get_jumpstart_content_bucket(region)
+
+    script_s3_uri = f"s3://{bucket}/{model_script_key}"
+
+    return script_s3_uri

src/sagemaker/jumpstart/cache.py

@@ -282,7 +282,7 @@ def _select_version(
         semantic_version_str: str,
         available_versions: List[Version],
     ) -> Optional[str]:
-        """Utility to select appropriate version from available versions.
+        """Perform semantic version search on available versions.
 
         Args:
             semantic_version_str (str): the semantic version for which to filter

src/sagemaker/jumpstart/constants.py

@@ -103,10 +103,6 @@
             region_name="cn-north-1",
             content_bucket="jumpstart-cache-prod-cn-north-1",
         ),
-        JumpStartLaunchedRegionInfo(
-            region_name="cn-northwest-1",
-            content_bucket="jumpstart-cache-prod-cn-northwest-1",
-        ),
     ]
 )
 
@@ -118,3 +114,7 @@
 JUMPSTART_DEFAULT_REGION_NAME = boto3.session.Session().region_name
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
+
+INFERENCE = "inference"
+TRAINING = "training"
+SUPPORTED_JUMPSTART_SCOPES = set([INFERENCE, TRAINING])

src/sagemaker/jumpstart/utils.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains utilities related to SageMaker JumpStart."""
 from __future__ import absolute_import
-from typing import Dict, List
+from typing import Dict, List, Optional
 from packaging.version import Version
 import sagemaker
 from sagemaker.jumpstart import constants
@@ -113,3 +113,26 @@ def parse_sagemaker_version() -> str:
     Version(parsed_version)
 
     return parsed_version
+
+
+def is_jumpstart_model_input(model_id: Optional[str], version: Optional[str]) -> bool:
+    """Determines if `model_id` and `version` input are for JumpStart.
+
+    This method returns True if both arguments are not None, false if both arguments
+    are None, and raises an exception if one argument is None but the other isn't.
+
+    Args:
+        model_id (str): Optional. Model id of JumpStart model.
+        version (str): Optional. Version for JumpStart model.
+
+    Raises:
+        ValueError: If only one of the two arguments is None.
+    """
+    if model_id is not None or version is not None:
+        if model_id is None or version is None:
+            raise ValueError(
+                "Must specify `model_id` and `model_version` when getting specs for "
+                "JumpStart models."
+            )
+        return True
+    return False