feat: jumpstart hyperparameters and env variables

Author: evakravi

src/sagemaker/environment_variables.py

@@ -0,0 +1,52 @@
+# 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 to run pretrained ML models."""
+
+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 environment variables for the model matching the given arguments.
+
+    Args:
+        region (str): Region for which to retrieve default environment variables.
+        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.
+    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,57 @@
+# 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 to run pretrained ML models."""
+
+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.
+        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.
+        include_container_hyperparameters (bool): True if container hyperparameters
+            should be returned as well. (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.")
+
+    # mypy type checking require these assertions
+    assert model_id is not None
+    assert model_version is not None
+
+    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

@@ -93,7 +93,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:
@@ -201,7 +201,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 +260,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 +278,80 @@ 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. (Default: False)
+    Returns:
+        dict: the hyperparameters to use for the 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
+
+    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+        region=region, model_id=model_id, version=model_version
+    )
+
+    default_hyperparameters = {}
+    for hyperparameter in model_specs.hyperparameters:
+        if (
+            include_container_hyperparameters and hyperparameter.scope == "container"
+        ) or hyperparameter.scope == "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 (str): Region for which to retrieve default environment variables.
+
+    Returns:
+        dict: the inference environment variables to use for the 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
+
+    model_specs = jumpstart_accessors.JumpStartModelsCache.get_model_specs(
+        region=region, model_id=model_id, version=model_version
+    )
+
+    default_environment_variables = {}
+    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.

src/sagemaker/jumpstart/types.py

@@ -170,6 +170,93 @@ def to_json(self) -> Dict[str, Any]:
         return json_obj
 
 
+class JumpStartHyperparameter(JumpStartDataHolderType):
+    """Data class for JumpStart hyperparameter."""
+
+    __slots__ = {
+        "name",
+        "type",
+        "options",
+        "default",
+        "scope",
+        "min",
+        "max",
+    }
+
+    def __init__(self, spec: Dict[str, Any]):
+        """Initializes a JumpStartHyperparameter object from its json representation.
+
+        Args:
+            spec (Dict[str, Any]): Dictionary representation of hyperparameter.
+        """
+        self.from_json(spec)
+
+    def from_json(self, json_obj: Dict[str, Any]) -> None:
+        """Sets fields in object based on json.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of hyperparameter.
+        """
+
+        self.name = json_obj["name"]
+        self.type = json_obj["type"]
+        self.default = json_obj["default"]
+        self.scope = json_obj["scope"]
+
+        options = json_obj.get("options")
+        if options is not None:
+            self.options = options
+
+        min_val = json_obj.get("min")
+        if min_val is not None:
+            self.min = min_val
+
+        max_val = json_obj.get("max")
+        if max_val is not None:
+            self.max = max_val
+
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of JumpStartHyperparameter object."""
+        json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
+        return json_obj
+
+
+class JumpStartEnvironmentVariable(JumpStartDataHolderType):
+    """Data class for JumpStart environment variable."""
+
+    __slots__ = {
+        "name",
+        "type",
+        "default",
+        "scope",
+    }
+
+    def __init__(self, spec: Dict[str, Any]):
+        """Initializes a JumpStartEnvironmentVariable object from its json representation.
+
+        Args:
+            spec (Dict[str, Any]): Dictionary representation of environment variable.
+        """
+        self.from_json(spec)
+
+    def from_json(self, json_obj: Dict[str, Any]) -> None:
+        """Sets fields in object based on json.
+
+        Args:
+            json_obj (Dict[str, Any]): Dictionary representation of environment variable.
+        """
+
+        self.name = json_obj["name"]
+        self.type = json_obj["type"]
+        self.default = json_obj["default"]
+        self.scope = json_obj["scope"]
+
+    def to_json(self) -> Dict[str, Any]:
+        """Returns json representation of JumpStartEnvironmentVariable object."""
+        json_obj = {att: getattr(self, att) for att in self.__slots__ if hasattr(self, att)}
+        return json_obj
+
+
 class JumpStartModelSpecs(JumpStartDataHolderType):
     """Data class JumpStart model specs."""
 
@@ -186,6 +273,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_artifact_key",
         "training_script_key",
         "hyperparameters",
+        "inference_environment_variables",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -210,22 +298,37 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.hosting_artifact_key: str = json_obj["hosting_artifact_key"]
         self.hosting_script_key: str = json_obj["hosting_script_key"]
         self.training_supported: bool = bool(json_obj["training_supported"])
+        self.inference_environment_variables = [
+            JumpStartEnvironmentVariable(env_variable)
+            for env_variable in json_obj["inference_environment_variables"]
+        ]
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]
             )
             self.training_artifact_key: str = json_obj["training_artifact_key"]
             self.training_script_key: str = json_obj["training_script_key"]
-            self.hyperparameters: Dict[str, Any] = json_obj.get("hyperparameters", {})
+            hyperparameters = json_obj.get("hyperparameters")
+            if hyperparameters is not None:
+                self.hyperparameters = [
+                    JumpStartHyperparameter(hyperparameter) for hyperparameter in hyperparameters
+                ]
 
     def to_json(self) -> Dict[str, Any]:
         """Returns json representation of JumpStartModelSpecs object."""
         json_obj = {}
         for att in self.__slots__:
             if hasattr(self, att):
                 cur_val = getattr(self, att)
-                if isinstance(cur_val, JumpStartECRSpecs):
+                if issubclass(type(cur_val), JumpStartDataHolderType):
                     json_obj[att] = cur_val.to_json()
+                elif isinstance(cur_val, list):
+                    json_obj[att] = []
+                    for obj in cur_val:
+                        if issubclass(type(obj), JumpStartDataHolderType):
+                            json_obj[att].append(obj.to_json())
+                        else:
+                            json_obj[att].append(obj)
                 else:
                     json_obj[att] = cur_val
         return json_obj