feat: jumpstart instance types (#3686)

Author: evakravi

src/sagemaker/environment_variables.py

@@ -27,6 +27,8 @@ def retrieve_default(
     region=None,
     model_id=None,
     model_version=None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
 ) -> Dict[str, str]:
     """Retrieves the default container environment variables for the model matching the arguments.
 
@@ -37,6 +39,13 @@ def retrieve_default(
             retrieve the default environment variables. (Default: None).
         model_version (str): Optional. The version of the model for which to retrieve the
             default environment variables. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
     Returns:
         dict: The variables to use for the model.
 
@@ -48,4 +57,6 @@ def retrieve_default(
             "Must specify `model_id` and `model_version` when retrieving environment variables."
         )
 
-    return artifacts._retrieve_default_environment_variables(model_id, model_version, region)
+    return artifacts._retrieve_default_environment_variables(
+        model_id, model_version, region, tolerate_vulnerable_model, tolerate_deprecated_model
+    )

src/sagemaker/hyperparameters.py

@@ -30,6 +30,8 @@ def retrieve_default(
     model_id=None,
     model_version=None,
     include_container_hyperparameters=False,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
 ) -> Dict[str, str]:
     """Retrieves the default training hyperparameters for the model matching the given arguments.
 
@@ -47,6 +49,13 @@ def retrieve_default(
             that indicates the entrypoint script to use. These hyperparameters may be required
             when creating a training job with boto3, however the ``Estimator`` classes
             add required container hyperparameters to the job. (Default: False).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
     Returns:
         dict: The hyperparameters to use for the model.
 
@@ -59,7 +68,12 @@ def retrieve_default(
         )
 
     return artifacts._retrieve_default_hyperparameters(
-        model_id, model_version, region, include_container_hyperparameters
+        model_id,
+        model_version,
+        region,
+        include_container_hyperparameters,
+        tolerate_vulnerable_model,
+        tolerate_deprecated_model,
     )
 
 
@@ -68,7 +82,7 @@ def validate(
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
     hyperparameters: Optional[dict] = None,
-    validation_mode: Optional[HyperparameterValidationMode] = None,
+    validation_mode: HyperparameterValidationMode = HyperparameterValidationMode.VALIDATE_PROVIDED,
 ) -> None:
     """Validates hyperparameters for models.
 
@@ -99,6 +113,9 @@ def validate(
             "Must specify `model_id` and `model_version` when validating hyperparameters."
         )
 
+    if model_id is None or model_version is None:
+        raise RuntimeError("Model id and version must both be non-None")
+
     if hyperparameters is None:
         raise ValueError("Must specify hyperparameters.")
 

src/sagemaker/instance_types.py

@@ -0,0 +1,121 @@
+# 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 instance types."""
+
+from __future__ import absolute_import
+
+import logging
+from typing import List
+
+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,
+    scope=None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> str:
+    """Retrieves the default instance type for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the default instance type.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the default instance type. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            default instance type. (Default: None).
+        scope (str): The model type, i.e. what it is used for.
+            Valid values: "training" and "inference".
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
+    Returns:
+        str: The default instance type 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 instance types."
+        )
+
+    if scope is None:
+        raise ValueError("Must specify scope for instance types.")
+
+    return artifacts._retrieve_default_instance_type(
+        model_id,
+        model_version,
+        scope,
+        region,
+        tolerate_vulnerable_model,
+        tolerate_deprecated_model,
+    )
+
+
+def retrieve(
+    region=None,
+    model_id=None,
+    model_version=None,
+    scope=None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> List[str]:
+    """Retrieves the supported training instance types for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the supported instance types.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the supported instance types. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            supported instance types. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
+    Returns:
+        list: The supported instance types 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 instance types."
+        )
+
+    if scope is None:
+        raise ValueError("Must specify scope for instance types.")
+
+    return artifacts._retrieve_instance_types(
+        model_id,
+        model_version,
+        scope,
+        region,
+        tolerate_vulnerable_model,
+        tolerate_deprecated_model,
+    )