move the accept eula configurations into deploy flow

gwang111 · gwang111 · commit 779f6d69ad34 · 2024-11-12T01:22:38.000Z
diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
@@ -16,6 +16,7 @@
 
 
 from typing import Any, Dict, List, Optional, Union
+from sagemaker_core.shapes import ModelAccessConfig
 from sagemaker import environment_variables, image_uris, instance_types, model_uris, script_uris
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
@@ -58,7 +59,6 @@
     update_dict_if_key_not_present,
     resolve_model_sagemaker_config_field,
     verify_model_region_and_return_specs,
-    get_jumpstart_content_bucket,
 )
 
 from sagemaker.jumpstart.factory.utils import (
@@ -563,37 +563,6 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
     return kwargs
 
 
-def _apply_accept_eula_on_model_data_source(
-    model_data_source: Dict[str, Any], model_id: str, region: str, accept_eula: bool
-):
-    """Sets AcceptEula to True for gated speculative decoding models"""
-
-    mutable_model_data_source = model_data_source.copy()
-
-    hosting_eula_key = mutable_model_data_source.get("hosting_eula_key")
-    del mutable_model_data_source["hosting_eula_key"]
-
-    if not hosting_eula_key:
-        return mutable_model_data_source
-
-    if not accept_eula:
-        raise ValueError(
-            (
-                f"The set deployment config comes optimized with an additional model data source "
-                f"'{model_id}' that requires accepting end-user license agreement (EULA). "
-                f"See https://{get_jumpstart_content_bucket(region=region)}.s3.{region}."
-                f"{get_domain_for_region(region)}"
-                f"/{hosting_eula_key} for terms of use. Please set `accept_draft_model_eula=True` "
-                f"once acknowledged."
-            )
-        )
-
-    mutable_model_data_source["s3_data_source"]["model_access_config"] = {
-        "accept_eula": accept_eula
-    }
-    return mutable_model_data_source
-
-
 def _add_additional_model_data_sources_to_kwargs(
     kwargs: JumpStartModelInitKwargs,
 ) -> JumpStartModelInitKwargs:
@@ -606,14 +575,7 @@ def _add_additional_model_data_sources_to_kwargs(
         data_source.s3_data_source.set_bucket(get_neo_content_bucket(region=kwargs.region))
     api_shape_additional_model_data_sources = (
         [
-            camel_case_to_pascal_case(
-                _apply_accept_eula_on_model_data_source(
-                    data_source.to_json(),
-                    kwargs.model_id,
-                    kwargs.region,
-                    kwargs.accept_draft_model_eula,
-                )
-            )
+            camel_case_to_pascal_case(data_source.to_json())
             for data_source in speculative_decoding_data_sources
         ]
         if specs.get_speculative_decoding_s3_data_sources()
@@ -693,6 +655,7 @@ def get_deploy_kwargs(
     training_config_name: Optional[str] = None,
     config_name: Optional[str] = None,
     routing_config: Optional[Dict[str, Any]] = None,
+    model_access_configs: Optional[List[ModelAccessConfig]] = None,
 ) -> JumpStartModelDeployKwargs:
     """Returns kwargs required to call `deploy` on `sagemaker.estimator.Model` object."""
 
@@ -729,6 +692,7 @@ def get_deploy_kwargs(
         resources=resources,
         config_name=config_name,
         routing_config=routing_config,
+        model_access_configs=model_access_configs,
     )
     deploy_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(kwargs=deploy_kwargs)
     deploy_kwargs.specs = verify_model_region_and_return_specs(
@@ -903,7 +867,6 @@ def get_init_kwargs(
     resources: Optional[ResourceRequirements] = None,
     config_name: Optional[str] = None,
     additional_model_data_sources: Optional[Dict[str, Any]] = None,
-    accept_draft_model_eula: Optional[bool] = None,
 ) -> JumpStartModelInitKwargs:
     """Returns kwargs required to instantiate `sagemaker.estimator.Model` object."""
 
@@ -938,7 +901,6 @@ def get_init_kwargs(
         resources=resources,
         config_name=config_name,
         additional_model_data_sources=additional_model_data_sources,
-        accept_draft_model_eula=accept_draft_model_eula,
     )
     model_init_kwargs, orig_session = _set_temp_sagemaker_session_if_not_set(
         kwargs=model_init_kwargs
diff --git a/src/sagemaker/jumpstart/model.py b/src/sagemaker/jumpstart/model.py
@@ -18,6 +18,7 @@
 import pandas as pd
 from botocore.exceptions import ClientError
 
+from sagemaker_core.shapes import ModelAccessConfig
 from sagemaker import payloads
 from sagemaker.async_inference.async_inference_config import AsyncInferenceConfig
 from sagemaker.base_deserializers import BaseDeserializer
@@ -51,6 +52,7 @@
     add_instance_rate_stats_to_benchmark_metrics,
     deployment_config_response_data,
     _deployment_config_lru_cache,
+    _add_model_access_configs_to_model_data_sources,
 )
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION, JUMPSTART_LOGGER
 from sagemaker.jumpstart.enums import JumpStartModelType
@@ -111,7 +113,6 @@ def __init__(
         resources: Optional[ResourceRequirements] = None,
         config_name: Optional[str] = None,
         additional_model_data_sources: Optional[Dict[str, Any]] = None,
-        accept_draft_model_eula: Optional[bool] = None,
     ):
         """Initializes a ``JumpStartModel``.
 
@@ -302,10 +303,6 @@ def __init__(
                 optionally applied to the model.
             additional_model_data_sources (Optional[Dict[str, Any]]): Additional location
                 of SageMaker model data (default: None).
-            accept_draft_model_eula (bool): For draft models that require a Model Access Config, specify True or
-                False to indicate whether model terms of use have been accepted.
-                The `accept_draft_model_eula` value must be explicitly defined as `True` in order to
-                accept the end-user license agreement (EULA) that some
         Raises:
             ValueError: If the model ID is not recognized by JumpStart.
         """
@@ -365,7 +362,6 @@ def _validate_model_id_and_type():
             resources=resources,
             config_name=config_name,
             additional_model_data_sources=additional_model_data_sources,
-            accept_draft_model_eula=accept_draft_model_eula
         )
 
         self.orig_predictor_cls = predictor_cls
@@ -463,7 +459,7 @@ def retrieve_example_payload(self) -> JumpStartSerializablePayload:
         )
 
     def set_deployment_config(
-        self, config_name: str, instance_type: str, accept_draft_model_eula: Optional[bool] = False
+        self, config_name: str, instance_type: str
     ) -> None:
         """Sets the deployment config to apply to the model.
 
@@ -483,8 +479,7 @@ def set_deployment_config(
             instance_type=instance_type,
             config_name=config_name,
             sagemaker_session=self.sagemaker_session,
-            role=self.role,
-            accept_draft_model_eula=accept_draft_model_eula,
+            role=self.role
         )
 
     @property
@@ -674,6 +669,7 @@ def deploy(
         managed_instance_scaling: Optional[str] = None,
         endpoint_type: EndpointType = EndpointType.MODEL_BASED,
         routing_config: Optional[Dict[str, Any]] = None,
+        model_access_configs: Optional[List[ModelAccessConfig]] = None,
     ) -> PredictorBase:
         """Creates endpoint by calling base ``Model`` class `deploy` method.
 
@@ -770,6 +766,11 @@ def deploy(
                 (Default: EndpointType.MODEL_BASED).
             routing_config (Optional[Dict]): Settings the control how the endpoint routes
                 incoming traffic to the instances that the endpoint hosts.
+            model_access_configs (Optional[List[ModelAccessConfig]]): For models that require Model Access Configs,
+                provide one or multiple ModelAccessConfig objects to indicate whether model terms of use have been accepted.
+                The `AcceptEula` value must be explicitly defined as `True` in order to
+                accept the end-user license agreement (EULA) that some.
+                (Default: None)
 
         Raises:
             MarketplaceModelSubscriptionError: If the caller is not subscribed to the model.
@@ -810,6 +811,7 @@ def deploy(
             model_type=self.model_type,
             config_name=self.config_name,
             routing_config=routing_config,
+            model_access_configs=model_access_configs,
         )
         if (
             self.model_type == JumpStartModelType.PROPRIETARY
@@ -819,6 +821,13 @@ def deploy(
                 f"{EndpointType.INFERENCE_COMPONENT_BASED} is not supported for Proprietary models."
             )
 
+        self.additional_model_data_sources = _add_model_access_configs_to_model_data_sources(
+            self.additional_model_data_sources,
+            deploy_kwargs.model_access_configs,
+            deploy_kwargs.model_id,
+            deploy_kwargs.region,
+        )
+
         try:
             predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())
         except ClientError as e:
@@ -1058,7 +1067,6 @@ def _get_deployment_configs(
                 region=self.region,
                 model_version=self.model_version,
                 hub_arn=self.hub_arn,
-                accept_draft_model_eula=True,
             )
             deploy_kwargs = get_deploy_kwargs(
                 model_id=self.model_id,
diff --git a/src/sagemaker/jumpstart/types.py b/src/sagemaker/jumpstart/types.py
@@ -17,6 +17,7 @@
 from copy import deepcopy
 from enum import Enum
 from typing import Any, Dict, List, Optional, Set, Union
+from sagemaker_core.shapes import ModelAccessConfig as CoreModelAccessConfig
 from sagemaker.model_card.model_card import ModelCard, ModelPackageModelCard
 from sagemaker.utils import (
     S3_PREFIX,
@@ -2117,7 +2118,6 @@ class JumpStartModelInitKwargs(JumpStartKwargs):
         "hub_content_type",
         "model_reference_arn",
         "specs",
-        "accept_draft_model_eula",
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2133,7 +2133,6 @@ class JumpStartModelInitKwargs(JumpStartKwargs):
         "training_instance_type",
         "config_name",
         "hub_content_type",
-        "accept_draft_model_eula",
     }
 
     def __init__(
@@ -2168,7 +2167,6 @@ def __init__(
         resources: Optional[ResourceRequirements] = None,
         config_name: Optional[str] = None,
         additional_model_data_sources: Optional[Dict[str, Any]] = None,
-        accept_draft_model_eula: Optional[bool] = False
     ) -> None:
         """Instantiates JumpStartModelInitKwargs object."""
 
@@ -2202,7 +2200,6 @@ def __init__(
         self.resources = resources
         self.config_name = config_name
         self.additional_model_data_sources = additional_model_data_sources
-        self.accept_draft_model_eula = accept_draft_model_eula
 
 
 class JumpStartModelDeployKwargs(JumpStartKwargs):
@@ -2244,6 +2241,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "config_name",
         "routing_config",
         "specs",
+        "model_access_configs"
     ]
 
     SERIALIZATION_EXCLUSION_SET = {
@@ -2257,6 +2255,7 @@ class JumpStartModelDeployKwargs(JumpStartKwargs):
         "sagemaker_session",
         "training_instance_type",
         "config_name",
+        "model_access_configs"
     }
 
     def __init__(
@@ -2295,6 +2294,7 @@ def __init__(
         endpoint_type: Optional[EndpointType] = None,
         config_name: Optional[str] = None,
         routing_config: Optional[Dict[str, Any]] = None,
+        model_access_configs: Optional[List[CoreModelAccessConfig]] = None
     ) -> None:
         """Instantiates JumpStartModelDeployKwargs object."""
 
@@ -2332,6 +2332,7 @@ def __init__(
         self.endpoint_type = endpoint_type
         self.config_name = config_name
         self.routing_config = routing_config
+        self.model_access_configs = model_access_configs
 
 
 class JumpStartEstimatorInitKwargs(JumpStartKwargs):
diff --git a/src/sagemaker/jumpstart/utils.py b/src/sagemaker/jumpstart/utils.py
@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains utilities related to SageMaker JumpStart."""
 from __future__ import absolute_import
+
 from copy import copy
 import logging
 import os
@@ -22,6 +23,7 @@
 from botocore.exceptions import ClientError
 from packaging.version import Version
 import botocore
+from sagemaker_core.shapes import ModelAccessConfig
 import sagemaker
 from sagemaker.config.config_schema import (
     MODEL_ENABLE_NETWORK_ISOLATION_PATH,
@@ -55,6 +57,7 @@
     TagsDict,
     get_instance_rate_per_hour,
     get_domain_for_region,
+    camel_case_to_pascal_case,
 )
 from sagemaker.workflow import is_pipeline_variable
 from sagemaker.user_agent import get_user_agent_extra_suffix
@@ -555,11 +558,17 @@ def get_eula_message(model_specs: JumpStartModelSpecs, region: str) -> str:
     """Returns EULA message to display if one is available, else empty string."""
     if model_specs.hosting_eula_key is None:
         return ""
+    return format_eula_message_from_specs(
+        model_id=model_specs.model_id, region=region, hosting_eula_key=model_specs.hosting_eula_key
+    )
+
+
+def format_eula_message_from_specs(model_id: str, region: str, hosting_eula_key: str):
     return (
-        f"Model '{model_specs.model_id}' requires accepting end-user license agreement (EULA). "
+        f"Model '{model_id}' requires accepting end-user license agreement (EULA). "
         f"See https://{get_jumpstart_content_bucket(region=region)}.s3.{region}."
         f"{get_domain_for_region(region)}"
-        f"/{model_specs.hosting_eula_key} for terms of use."
+        f"/{hosting_eula_key} for terms of use."
     )
 
 
@@ -1525,3 +1534,41 @@ def wrapped_f(*args, **kwargs):
     if _func is None:
         return wrapper_cache
     return wrapper_cache(_func)
+
+
+def _add_model_access_configs_to_model_data_sources(
+    model_data_sources: List[Dict[str, any]],
+    model_access_configs: List[ModelAccessConfig],
+    model_id: str,
+    region: str,
+):
+    """Sets AcceptEula to True for gated speculative decoding models"""
+
+    if not model_data_sources:
+        return model_data_sources
+
+    acked_model_data_sources = []
+    acked_model_access_configs = 0
+    for model_data_source in model_data_sources:
+        hosting_eula_key = model_data_source.pop("HostingEulaKey", None)
+        if hosting_eula_key:
+            if not model_access_configs or acked_model_access_configs == len(model_access_configs):
+                eula_message_template = "{model_source}{base_eula_message}{model_access_configs_message}"
+                raise ValueError(eula_message_template.format(
+                    model_source="Draft " if model_data_source.get("ChannelName") else "",
+                    base_eula_message=format_eula_message_from_specs(
+                        model_id=model_id, region=region, hosting_eula_key=hosting_eula_key
+                    ),
+                    model_access_configs_message=(
+                        " Please add a ModelAccessConfig with AcceptEula=True"
+                        " to model_access_configs to acknowledge the EULA."
+                    )
+                ))
+            acked_model_data_source = model_data_source.copy()
+            acked_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
+                camel_case_to_pascal_case(model_access_configs[acked_model_access_configs].model_dump())
+            )
+            acked_model_data_sources.append(acked_model_data_source)
+        else:
+            acked_model_data_sources.append(model_data_source)
+    return acked_model_data_sources
diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -505,7 +505,7 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
         )
 
     def set_deployment_config(
-        self, config_name: str, instance_type: str, accept_draft_model_eula: Optional[bool] = False
+        self, config_name: str, instance_type: str
     ) -> None:
         """Sets the deployment config to apply to the model.
 
@@ -522,7 +522,7 @@ def set_deployment_config(
         if not hasattr(self, "pysdk_model") or self.pysdk_model is None:
             raise Exception("Cannot set deployment config to an uninitialized model.")
 
-        self.pysdk_model.set_deployment_config(config_name, instance_type, accept_draft_model_eula)
+        self.pysdk_model.set_deployment_config(config_name, instance_type)
         self.deployment_config_name = config_name
 
         self.instance_type = instance_type
