feature: jumpstart hyperparameters and environment variables (#2850)

Author: evakravi

src/sagemaker/environment_variables.py

@@ -0,0 +1,53 @@
+# 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.
+"""Accessors to retrieve environment variables for hosting containers."""
+
+from __future__ import absolute_import
+
+import logging
+from typing import Dict
+
+from sagemaker.jumpstart import utils as jumpstart_utils
+from sagemaker.jumpstart import artifacts
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_default(
+    region=None,
+    model_id=None,
+    model_version=None,
+) -> Dict[str, str]:
+    """Retrieves the default container environment variables for the model matching the arguments.
+
+    Args:
+        region (str): Optional. Region for which to retrieve default environment variables.
+            (Default: None).
+        model_id (str): Optional. Model ID of the model for which to
+            retrieve the default environment variables. (Default: None).
+        model_version (str): Optional. Version of the model for which to retrieve the
+            default environment variables. (Default: None).
+    Returns:
+        dict: the variables to use for the model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError("Must specify `model_id` and `model_version` when retrieving script URIs.")
+
+    # mypy type checking require these assertions
+    assert model_id is not None
+    assert model_version is not None
+
+    return artifacts._retrieve_default_environment_variables(model_id, model_version, region)

src/sagemaker/hyperparameters.py

@@ -0,0 +1,58 @@
+# 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.
+"""Accessors to retrieve hyperparameters for training jobs."""
+
+from __future__ import absolute_import
+
+import logging
+from typing import Dict
+
+from sagemaker.jumpstart import utils as jumpstart_utils
+from sagemaker.jumpstart import artifacts
+
+logger = logging.getLogger(__name__)
+
+
+def retrieve_default(
+    region=None,
+    model_id=None,
+    model_version=None,
+    include_container_hyperparameters=False,
+) -> Dict[str, str]:
+    """Retrieves the default training hyperparameters for the model matching the given arguments.
+
+    Args:
+        region (str): Region for which to retrieve default hyperparameters. (Default: None).
+        model_id (str): Model ID of the model for which to
+            retrieve the default hyperparameters. (Default: None).
+        model_version (str): Version of the model for which to retrieve the
+            default hyperparameters. (Default: None).
+        include_container_hyperparameters (bool): True if container hyperparameters
+            should be returned as well. Container hyperparameters are not used to tune
+            the specific algorithm, but rather by SageMaker Training to setup
+            the training container environment. For example, there is a container hyperparameter
+            that indicates the entrypoint script to use. These hyperparameters may be required
+            when creating a training job with boto3, however the ``Estimator`` classes
+            should take care of adding container hyperparameters to the job. (Default: False).
+    Returns:
+        dict: the hyperparameters to use for the model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError("Must specify `model_id` and `model_version` when retrieving script URIs.")
+
+    return artifacts._retrieve_default_hyperparameters(
+        model_id, model_version, region, include_container_hyperparameters
+    )

src/sagemaker/jumpstart/accessors.py

@@ -93,7 +93,9 @@ def get_model_header(region: str, model_id: str, version: str) -> JumpStartModel
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         assert JumpStartModelsAccessor._cache is not None
-        return JumpStartModelsAccessor._cache.get_header(model_id, version)
+        return JumpStartModelsAccessor._cache.get_header(
+            model_id=model_id, semantic_version_str=version
+        )
 
     @staticmethod
     def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelSpecs:
@@ -109,7 +111,9 @@ def get_model_specs(region: str, model_id: str, version: str) -> JumpStartModelS
         )
         JumpStartModelsAccessor._set_cache_and_region(region, cache_kwargs)
         assert JumpStartModelsAccessor._cache is not None
-        return JumpStartModelsAccessor._cache.get_specs(model_id, version)
+        return JumpStartModelsAccessor._cache.get_specs(
+            model_id=model_id, semantic_version_str=version
+        )
 
     @staticmethod
     def set_cache_kwargs(cache_kwargs: Dict[str, Any], region: str = None) -> None:

src/sagemaker/jumpstart/artifacts.py

@@ -12,14 +12,15 @@
 # language governing permissions and limitations under the License.
 """This module contains functions for obtaining JumpStart ECR and S3 URIs."""
 from __future__ import absolute_import
-from typing import Optional
+from typing import Dict, Optional
 from sagemaker import image_uris
 from sagemaker.jumpstart.constants import (
     JUMPSTART_DEFAULT_REGION_NAME,
     INFERENCE,
     TRAINING,
     SUPPORTED_JUMPSTART_SCOPES,
     ModelFramework,
+    VariableScope,
 )
 from sagemaker.jumpstart.utils import get_jumpstart_content_bucket
 from sagemaker.jumpstart import accessors as jumpstart_accessors
@@ -93,7 +94,7 @@ def _retrieve_image_uri(
         )
 
     model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
-        region, model_id, model_version
+        region=region, model_id=model_id, version=model_version
     )
 
     if image_scope == INFERENCE:
@@ -110,19 +111,19 @@ 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}' "
-            f"and version {model_version}'."
+            f"and version '{model_version}'."
         )
 
     if version is not None and version != ecr_specs.framework_version:
         raise ValueError(
             f"Incorrect container framework version '{version}' for JumpStart model ID "
-            f"'{model_id}' and version {model_version}'."
+            f"'{model_id}' and version '{model_version}'."
         )
 
     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}' "
-            f"and version {model_version}'."
+            f"and version '{model_version}'."
         )
 
     base_framework_version_override: Optional[str] = None
@@ -201,7 +202,7 @@ def _retrieve_model_uri(
         )
 
     model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
-        region, model_id, model_version
+        region=region, model_id=model_id, version=model_version
     )
     if model_scope == INFERENCE:
         model_artifact_key = model_specs.hosting_artifact_key
@@ -260,7 +261,7 @@ def _retrieve_script_uri(
         )
 
     model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
-        region, model_id, model_version
+        region=region, model_id=model_id, version=model_version
     )
     if script_scope == INFERENCE:
         model_script_key = model_specs.hosting_script_key
@@ -278,3 +279,77 @@ def _retrieve_script_uri(
     script_s3_uri = f"s3://{bucket}/{model_script_key}"
 
     return script_s3_uri
+
+
+def _retrieve_default_hyperparameters(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    include_container_hyperparameters: bool = False,
+):
+    """Retrieves the training hyperparameters for the model matching the given arguments.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the default hyperparameters.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            default hyperparameters.
+        region (str): Region for which to retrieve default hyperparameters.
+        include_container_hyperparameters (bool): True if container hyperparameters
+            should be returned as well. Container hyperparameters are not used to tune
+            the specific algorithm, but rather by SageMaker Training to setup
+            the training container environment. For example, there is a container hyperparameter
+            that indicates the entrypoint script to use. These hyperparameters may be required
+            when creating a training job with boto3, however the ``Estimator`` classes
+            should take care of adding container hyperparameters to the job. (Default: False).
+    Returns:
+        dict: the hyperparameters to use for the model.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    assert region is not None
+
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
+        region=region, model_id=model_id, version=model_version
+    )
+
+    default_hyperparameters: Dict[str, str] = {}
+    for hyperparameter in model_specs.hyperparameters:
+        if (
+            include_container_hyperparameters and hyperparameter.scope == VariableScope.CONTAINER
+        ) or hyperparameter.scope == VariableScope.ALGORITHM:
+            default_hyperparameters[hyperparameter.name] = str(hyperparameter.default)
+    return default_hyperparameters
+
+
+def _retrieve_default_environment_variables(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+):
+    """Retrieves the inference environment variables for the model matching the given arguments.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the default environment variables.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            default environment variables.
+        region (Optional[str]): Region for which to retrieve default environment variables.
+
+    Returns:
+        dict: the inference environment variables to use for the model.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = jumpstart_accessors.JumpStartModelsAccessor.get_model_specs(
+        region=region, model_id=model_id, version=model_version
+    )
+
+    default_environment_variables: Dict[str, str] = {}
+    for environment_variable in model_specs.inference_environment_variables:
+        default_environment_variables[environment_variable.name] = str(environment_variable.default)
+    return default_environment_variables

src/sagemaker/jumpstart/constants.py

@@ -120,6 +120,9 @@
 TRAINING = "training"
 SUPPORTED_JUMPSTART_SCOPES = set([INFERENCE, TRAINING])
 
+INFERENCE_ENTRYPOINT_SCRIPT_NAME = "inference.py"
+TRAINING_ENTRYPOINT_SCRIPT_NAME = "transfer_learning.py"
+
 
 class ModelFramework(str, Enum):
     """Enum class for JumpStart model framework.
@@ -136,3 +139,13 @@ class ModelFramework(str, Enum):
     CATBOOST = "catboost"
     XGBOOST = "xgboost"
     SKLEARN = "sklearn"
+
+
+class VariableScope(str, Enum):
+    """Possible value of the ``scope`` attribute for a hyperparameter or environment variable.
+
+    Used for hosting environment variables and training hyperparameters.
+    """
+
+    CONTAINER = "container"
+    ALGORITHM = "algorithm"