feat: jumpstart EULA models (initial commit)

Author: evakravi

src/sagemaker/jumpstart/artifacts/__init__.py

@@ -52,3 +52,4 @@
     _retrieve_supported_accept_types,
     _retrieve_supported_content_types,
 )
+from sagemaker.jumpstart.artifacts.model_packages import _retrieve_model_package_arn

src/sagemaker/jumpstart/artifacts/model_packages.py

@@ -0,0 +1,67 @@
+# 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 obtaining JumpStart model packages."""
+from __future__ import absolute_import
+from copy import deepcopy
+from typing import Dict, List, Optional
+from sagemaker.jumpstart.constants import (
+    JUMPSTART_DEFAULT_REGION_NAME,
+)
+from sagemaker.jumpstart.enums import (
+    JumpStartScriptScope,
+)
+from sagemaker.jumpstart.utils import (
+    verify_model_region_and_return_specs,
+)
+
+
+def _retrieve_model_package_arn(
+    model_id: str,
+    model_version: str,
+    region: Optional[str],
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+) -> Optional[str]:
+    """Retrieves associated model pacakge arn for the model.
+
+    Args:
+        model_id (str): JumpStart model ID of the JumpStart model for which to
+            retrieve the model package arn.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            model package arn.
+        region (Optional[str]): Region for which to retrieve the model package arn.
+        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 versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+
+    Returns:
+        list: the model package arn to use for the model or None.
+    """
+
+    if region is None:
+        region = JUMPSTART_DEFAULT_REGION_NAME
+
+    model_specs = verify_model_region_and_return_specs(
+        model_id=model_id,
+        version=model_version,
+        scope=JumpStartScriptScope.TRAINING,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+    )
+
+    return model_specs.model_package_arn

src/sagemaker/jumpstart/factory/model.py

@@ -25,6 +25,7 @@
     _model_supports_prepacked_inference,
     _retrieve_model_init_kwargs,
     _retrieve_model_deploy_kwargs,
+    _retrieve_model_package_arn,
 )
 from sagemaker.jumpstart.artifacts.resource_names import _retrieve_resource_name_base
 from sagemaker.jumpstart.constants import (
@@ -289,6 +290,21 @@ def _add_env_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelInitKw
     return kwargs
 
 
+def _add_model_package_arn_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelInitKwargs:
+    """Sets model package arn based on default or override, returns full kwargs."""
+
+    model_package_arn = kwargs.model_package_arn or _retrieve_model_package_arn(
+        model_id=kwargs.model_id,
+        model_version=kwargs.model_version,
+        region=kwargs.region,
+        tolerate_deprecated_model=kwargs.tolerate_deprecated_model,
+        tolerate_vulnerable_model=kwargs.tolerate_vulnerable_model,
+    )
+
+    kwargs.model_package_arn = model_package_arn
+    return kwargs
+
+
 def _add_extra_model_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelInitKwargs:
     """Sets extra kwargs based on default or override, returns full kwargs."""
 
@@ -471,6 +487,7 @@ def get_init_kwargs(
     container_log_level: Optional[Union[int, PipelineVariable]] = None,
     dependencies: Optional[List[str]] = None,
     git_config: Optional[Dict[str, str]] = None,
+    model_package_arn: Optional[str] = None,
 ) -> JumpStartModelInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Model` object."""
 
@@ -498,6 +515,7 @@ def get_init_kwargs(
         git_config=git_config,
         tolerate_deprecated_model=tolerate_deprecated_model,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
+        model_package_arn=model_package_arn,
     )
 
     model_init_kwargs = _add_model_version_to_kwargs(kwargs=model_init_kwargs)
@@ -526,4 +544,6 @@ def get_init_kwargs(
     model_init_kwargs = _add_extra_model_kwargs(kwargs=model_init_kwargs)
     model_init_kwargs = _add_role_to_kwargs(kwargs=model_init_kwargs)
 
+    model_init_kwargs = _add_model_package_arn_to_kwargs(kwargs=model_init_kwargs)
+
     return model_init_kwargs

src/sagemaker/jumpstart/model.py

@@ -14,6 +14,7 @@
 
 from __future__ import absolute_import
 import logging
+import re
 
 from typing import Dict, List, Optional, Union
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
@@ -30,7 +31,7 @@
 )
 from sagemaker.jumpstart.utils import is_valid_model_id
 from sagemaker.utils import stringify_object
-from sagemaker.model import Model
+from sagemaker.model import MODEL_PACKAGE_ARN_PATTERN, Model, ModelPackage
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig
 from sagemaker.predictor import PredictorBase
 from sagemaker.serverless.serverless_inference_config import ServerlessInferenceConfig
@@ -71,6 +72,7 @@ def __init__(
         container_log_level: Optional[Union[int, PipelineVariable]] = None,
         dependencies: Optional[List[str]] = None,
         git_config: Optional[Dict[str, str]] = None,
+        model_package_arn: Optional[str] = None,
     ):
         """Initializes a ``JumpStartModel``.
 
@@ -249,6 +251,9 @@ def __init__(
                     >>>               'branch': 'test-branch-git-config',
                     >>>               'commit': '329bfcf884482002c05ff7f44f62599ebc9f445a'}
 
+            model_package_arn (Optional[str]): An existing SageMaker Model Package arn,
+                can be just the name if your account owns the Model Package.
+                ``model_data`` is not required. (Default: None).
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
         """
@@ -291,6 +296,7 @@ def _is_valid_model_id_hook():
             container_log_level=container_log_level,
             dependencies=dependencies,
             git_config=git_config,
+            model_package_arn=model_package_arn,
         )
 
         self.orig_predictor_cls = predictor_cls
@@ -301,9 +307,49 @@ def _is_valid_model_id_hook():
         self.tolerate_vulnerable_model = model_init_kwargs.tolerate_vulnerable_model
         self.tolerate_deprecated_model = model_init_kwargs.tolerate_deprecated_model
         self.region = model_init_kwargs.region
+        self.model_package_arn = model_init_kwargs.model_package_arn
 
         super(JumpStartModel, self).__init__(**model_init_kwargs.to_kwargs_dict())
 
+    def _create_sagemaker_model(self, *args, **kwargs):  # pylint: disable=unused-argument
+        """Create a SageMaker Model Entity
+
+        Args:
+            args: Positional arguments coming from the caller. This class does not require
+                any so they are ignored.
+
+            kwargs: Keyword arguments coming from the caller. This class does not require
+                any so they are ignored.
+        """
+        if self.model_package_arn:
+            # When a ModelPackageArn is provided we just create the Model
+            match = re.match(MODEL_PACKAGE_ARN_PATTERN, self.model_package_arn)
+            if match:
+                model_package_name = match.group(3)
+            else:
+                # model_package_arn can be just the name if your account owns the Model Package
+                model_package_name = self.model_package_arn
+            container_def = {"ModelPackageName": self.model_package_arn}
+
+            if self.env != {}:
+                container_def["Environment"] = self.env
+
+            if self.name is None:
+                self._base_name = model_package_name
+
+            self._set_model_name_if_needed()
+
+            self.sagemaker_session.create_model(
+                self.name,
+                self.role,
+                container_def,
+                vpc_config=self.vpc_config,
+                enable_network_isolation=self.enable_network_isolation(),
+                tags=kwargs.get("tags"),
+            )
+        else:
+            super(JumpStartModel, self)._create_sagemaker_model(*args, **kwargs)
+
     def deploy(
         self,
         initial_instance_count: Optional[int] = None,

src/sagemaker/jumpstart/types.py

@@ -351,6 +351,8 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "inference_enable_network_isolation",
         "training_enable_network_isolation",
         "resource_name_base",
+        "eula_model",
+        "model_package_arn",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -419,6 +421,10 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         )
         self.resource_name_base: bool = json_obj.get("resource_name_base")
 
+        self.eula_model: bool = json_obj.get("eula_model", False)
+
+        self.model_package_arn: Optional[str] = json_obj.get("model_package_arn")
+
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(
                 json_obj["training_ecr_specs"]
@@ -574,6 +580,7 @@ class JumpStartModelInitKwargs(JumpStartKwargs):
         "container_log_level",
         "dependencies",
         "git_config",
+        "model_package_arn",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -583,6 +590,7 @@ class JumpStartModelInitKwargs(JumpStartKwargs):
         "tolerate_vulnerable_model",
         "tolerate_deprecated_model",
         "region",
+        "model_package_arn",
     }
 
     def __init__(
@@ -610,6 +618,7 @@ def __init__(
         git_config: Optional[Dict[str, str]] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
         tolerate_deprecated_model: Optional[bool] = None,
+        model_package_arn: Optional[str] = None,
     ) -> None:
         """Instantiates JumpStartModelInitKwargs object."""
 
@@ -636,6 +645,7 @@ def __init__(
         self.git_config = git_config
         self.tolerate_deprecated_model = tolerate_deprecated_model
         self.tolerate_vulnerable_model = tolerate_vulnerable_model
+        self.model_package_arn = model_package_arn
 
 
 class JumpStartModelDeployKwargs(JumpStartKwargs):

src/sagemaker/jumpstart/utils.py

@@ -402,6 +402,12 @@ def verify_model_region_and_return_specs(
             f"JumpStart model ID '{model_id}' and version '{version}' " "does not support training."
         )
 
+    if model_specs.eula_model:
+        LOGGER.info(
+            "Using model with end-user license agreement (EULA). "
+            "Deploying this model requires accepting EULA terms."
+        )
+
     if model_specs.deprecated:
         if not tolerate_deprecated_model:
             raise DeprecatedJumpStartModelError(model_id=model_id, version=version)

tests/unit/sagemaker/jumpstart/constants.py

@@ -2299,6 +2299,8 @@
     "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
     "training_prepacked_script_key": None,
     "hosting_prepacked_artifact_key": None,
+    "model_package_arn": None,
+    "eula_model": False,
     "hyperparameters": [
         {
             "name": "epochs",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -345,6 +345,7 @@ def test_jumpstart_model_kwargs_match_parent_class(self):
             "tolerate_vulnerable_model",
             "tolerate_deprecated_model",
             "instance_type",
+            "model_package_arn",
         }
         assert parent_class_init_args - js_class_init_args == init_args_to_skip