feat: Keynote1 - Gated Models (aws#1246)

Co-authored-by: evakravi <[email protected]>

Author: knikure

src/sagemaker/jumpstart/factory/model.py

@@ -466,6 +466,7 @@ def get_deploy_kwargs(
     tolerate_vulnerable_model: Optional[bool] = None,
     tolerate_deprecated_model: Optional[bool] = None,
     sagemaker_session: Optional[Session] = None,
+    accept_eula: Optional[bool] = None,
 ) -> JumpStartModelDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Model` object."""
 
@@ -493,6 +494,7 @@ def get_deploy_kwargs(
         tolerate_deprecated_model=tolerate_deprecated_model,
         tolerate_vulnerable_model=tolerate_vulnerable_model,
         sagemaker_session=sagemaker_session,
+        accept_eula=accept_eula,
     )
 
     deploy_kwargs = _add_sagemaker_session_to_kwargs(kwargs=deploy_kwargs)

src/sagemaker/jumpstart/model.py

@@ -448,6 +448,7 @@ def deploy(
         container_startup_health_check_timeout: Optional[int] = None,
         inference_recommendation_id: Optional[str] = None,
         explainer_config: Optional[ExplainerConfig] = None,
+        accept_eula: Optional[bool] = None,
     ) -> PredictorBase:
         """Creates endpoint by calling base ``Model`` class `deploy` method.
 
@@ -526,7 +527,11 @@ def deploy(
                 (Default: None).
             explainer_config (Optional[sagemaker.explainer.ExplainerConfig]): Specifies online
                 explainability configuration for use with Amazon SageMaker Clarify. (Default: None).
-
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False to indicate whether model terms of use have been accepted.
+                The `accept_eula` value must be explicitly defined as `True` in order to
+                accept the end-user license agreement (EULA) that some
+                models require. (Default: None).
         """
 
         deploy_kwargs = get_deploy_kwargs(
@@ -552,6 +557,7 @@ def deploy(
             container_startup_health_check_timeout=container_startup_health_check_timeout,
             inference_recommendation_id=inference_recommendation_id,
             explainer_config=explainer_config,
+            accept_eula=accept_eula,
         )
 
         predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())

src/sagemaker/jumpstart/types.py

@@ -1065,6 +1065,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "tolerate_vulnerable_model",
         "tolerate_deprecated_model",
         "sagemaker_session",
+        "accept_eula",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -1101,6 +1102,7 @@ def __init__(
         tolerate_deprecated_model: Optional[bool] = None,
         tolerate_vulnerable_model: Optional[bool] = None,
         sagemaker_session: Optional[Session] = None,
+        accept_eula: Optional[bool] = None,
     ) -> None:
         """Instantiates JumpStartModelDeployKwargs object."""
 
@@ -1127,6 +1129,7 @@ def __init__(
         self.tolerate_vulnerable_model = tolerate_vulnerable_model
         self.tolerate_deprecated_model = tolerate_deprecated_model
         self.sagemaker_session = sagemaker_session
+        self.accept_eula = accept_eula
 
 
 class JumpStartEstimatorInitKwargs(JumpStartKwargs):

src/sagemaker/model.py

@@ -379,6 +379,7 @@ def __init__(
         self.repacked_model_data = None
         self.content_types = None
         self.response_types = None
+        self.accept_eula = None
 
     @runnable_by_pipeline
     def register(
@@ -634,6 +635,7 @@ def prepare_container_def(
             self.repacked_model_data or self.model_data,
             deploy_env,
             image_config=self.image_config,
+            accept_eula=getattr(self, "accept_eula", None),
         )
 
     def is_repack(self) -> bool:
@@ -1260,6 +1262,7 @@ def deploy(
         container_startup_health_check_timeout=None,
         inference_recommendation_id=None,
         explainer_config=None,
+        accept_eula: Optional[bool] = None,
         **kwargs,
     ):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
@@ -1342,6 +1345,11 @@ def deploy(
                 a list of ``RealtimeInferenceRecommendations`` within ``DeploymentRecommendation``
             explainer_config (sagemaker.explainer.ExplainerConfig): Specifies online explainability
                 configuration for use with Amazon SageMaker Clarify. Default: None.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False to indicate whether model terms of use have been accepted.
+                The `accept_eula` value must be explicitly defined as `True` in order to
+                accept the end-user license agreement (EULA) that some
+                models require. (Default: None).
         Raises:
              ValueError: If arguments combination check failed in these circumstances:
                 - If no role is specified or
@@ -1355,6 +1363,8 @@ def deploy(
                 ``self.predictor_cls`` on the created endpoint name, if ``self.predictor_cls``
                 is not None. Otherwise, return None.
         """
+        self.accept_eula = accept_eula
+
         removed_kwargs("update_endpoint", kwargs)
 
         self._init_sagemaker_session_if_does_not_exist(instance_type)

src/sagemaker/session.py

@@ -6128,7 +6128,14 @@ def update_args(args: Dict[str, Any], **kwargs):
             args.update({key: value})
 
 
-def container_def(image_uri, model_data_url=None, env=None, container_mode=None, image_config=None):
+def container_def(
+    image_uri,
+    model_data_url=None,
+    env=None,
+    container_mode=None,
+    image_config=None,
+    accept_eula=None,
+):
     """Create a definition for executing a container as part of a SageMaker model.
 
     Args:
@@ -6145,6 +6152,11 @@ def container_def(image_uri, model_data_url=None, env=None, container_mode=None,
         image_config (dict[str, str]): Specifies whether the image of model container is pulled
             from ECR, or private registry in your VPC. By default it is set to pull model
             container image from ECR. (default: None).
+        accept_eula (bool): For models that require a Model Access Config, specify True or
+            False to indicate whether model terms of use have been accepted.
+            The `accept_eula` value must be explicitly defined as `True` in order to
+            accept the end-user license agreement (EULA) that some
+            models require. (Default: None).
 
     Returns:
         dict[str, str]: A complete container definition object usable with the CreateModel API if
@@ -6154,9 +6166,28 @@ def container_def(image_uri, model_data_url=None, env=None, container_mode=None,
         env = {}
     c_def = {"Image": image_uri, "Environment": env}
 
-    if isinstance(model_data_url, dict):
-        c_def["ModelDataSource"] = model_data_url
-    elif model_data_url:
+    if isinstance(model_data_url, str) and (
+        not (model_data_url.startswith("s3://") and model_data_url.endswith("tar.gz"))
+        or accept_eula is None
+    ):
+        c_def["ModelDataUrl"] = model_data_url
+
+    elif isinstance(model_data_url, (dict, str)):
+        if isinstance(model_data_url, dict):
+            c_def["ModelDataSource"] = model_data_url
+        else:
+            c_def["ModelDataSource"] = {
+                "S3DataSource": {
+                    "S3Uri": model_data_url,
+                    "S3DataType": "S3Object",
+                    "CompressionType": "Gzip",
+                }
+            }
+        if accept_eula is not None:
+            c_def["ModelDataSource"]["S3DataSource"]["ModelAccessConfig"] = {
+                "AcceptEula": accept_eula
+            }
+    elif model_data_url is not None:
         c_def["ModelDataUrl"] = model_data_url
 
     if container_mode: