Require EULA acceptance when using a gated 1p draft model via ModelBuilder.

Joseph Zhang · Joseph Zhang · commit cf70f596bf9e · 2024-11-07T18:12:42.000-08:00
diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -48,6 +48,7 @@
     _custom_speculative_decoding,
     SPECULATIVE_DRAFT_MODEL,
     _is_inferentia_or_trainium,
+    _validate_and_set_eula_for_draft_model_sources,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -733,10 +734,6 @@ def _optimize_for_jumpstart(
         if (
             not optimization_config or not optimization_config.get("ModelCompilationConfig")
         ) and is_compilation:
-            # Ensure optimization_config exists
-            if not optimization_config:
-                optimization_config = {}
-
             # Fallback to default if override_env is None or empty
             if not compilation_override_env:
                 compilation_override_env = pysdk_model_env_vars
@@ -867,6 +864,11 @@ def _set_additional_model_source(
                             "Cannot find deployment config compatible for optimization job."
                         )
 
+                _validate_and_set_eula_for_draft_model_sources(
+                    pysdk_model=self.pysdk_model,
+                    accept_eula=speculative_decoding_config.get("AcceptEula"),
+                )
+
                 self.pysdk_model.env.update(
                     {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}"}
                 )
diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
@@ -99,6 +99,7 @@
     validate_image_uri_and_hardware,
 )
 from sagemaker.utils import Tags
+from sagemaker.serve.utils.optimize_utils import _validate_and_set_eula_for_draft_model_sources
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.huggingface.llm_utils import (
     get_huggingface_model_metadata,
@@ -589,6 +590,21 @@ def _model_builder_deploy_wrapper(
                     model_server=self.model_server,
                 )
 
+            if self.deployment_config:
+                accept_draft_model_eula = kwargs.get("accept_draft_model_eula", False)
+                try:
+                    _validate_and_set_eula_for_draft_model_sources(
+                        pysdk_model=self,
+                        accept_eula=accept_draft_model_eula,
+                    )
+                except ValueError as e:
+                    logger.error(
+                        "This deployment tried to use a gated draft model but the EULA was not "
+                        "accepted. Please review the EULA, set accept_draft_model_eula to True, "
+                        "and try again."
+                    )
+                    raise e
+
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
         predictor = self._original_deploy(
diff --git a/src/sagemaker/serve/utils/optimize_utils.py b/src/sagemaker/serve/utils/optimize_utils.py
@@ -172,6 +172,60 @@ def _extract_speculative_draft_model_provider(
     return "sagemaker"
 
 
+def _extract_additional_model_data_source_s3_uri(
+    additional_model_data_source: Optional[Dict] = None,
+) -> Optional[str]:
+    """Extracts model data source s3 uri from a model data source in Pascal case.
+
+    Args:
+        additional_model_data_source (Optional[Dict]): A model data source.
+
+    Returns:
+        str: S3 uri of the model resources.
+    """
+    if (
+        additional_model_data_source is None
+        or additional_model_data_source.get("S3DataSource", None) is None
+    ):
+        return None
+
+    return additional_model_data_source.get("S3DataSource").get("S3Uri", None)
+
+
+def _extract_deployment_config_additional_model_data_source_s3_uri(
+    additional_model_data_source: Optional[Dict] = None,
+) -> Optional[str]:
+    """Extracts model data source s3 uri from a model data source in snake case.
+
+    Args:
+        additional_model_data_source (Optional[Dict]): A model data source.
+
+    Returns:
+        str: S3 uri of the model resources.
+    """
+    if (
+        additional_model_data_source is None
+        or additional_model_data_source.get("s3_data_source", None) is None
+    ):
+        return None
+
+    return additional_model_data_source.get("s3_data_source").get("s3_uri", None)
+
+
+def _is_draft_model_gated(
+    draft_model_config: Optional[Dict] = None,
+) -> bool:
+    """Extracts model gated-ness from draft model data source.
+
+    Args:
+        draft_model_config (Optional[Dict]): A model data source.
+
+    Returns:
+        bool: Whether the draft model is gated or not.
+    """
+    return draft_model_config.get("hosting_eula_key", None)
+
+
 def _extracts_and_validates_speculative_model_source(
     speculative_decoding_config: Dict,
 ) -> str:
@@ -289,7 +343,7 @@ def _extract_optimization_config_and_env(
     if optimization_config:
         return optimization_config, quantization_override_env, compilation_override_env
 
-    return None, None, None
+    return {}, None, None
 
 
 def _custom_speculative_decoding(
@@ -310,6 +364,8 @@ def _custom_speculative_decoding(
             speculative_decoding_config
         )
 
+        accept_eula = speculative_decoding_config.get("AcceptEula", False)
+
         if _is_s3_uri(additional_model_source):
             channel_name = _generate_channel_name(model.additional_model_data_sources)
             speculative_draft_model = f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}"
@@ -326,3 +382,78 @@ def _custom_speculative_decoding(
         )
 
     return model
+
+
+def _validate_and_set_eula_for_draft_model_sources(
+    pysdk_model: Model,
+    accept_eula: bool = False,
+):
+    """Validates whether the EULA has been accepted for gated additional draft model sources.
+
+    If accepted, updates the model data source's model access config.
+
+    Args:
+        pysdk_model (Model): The model whose additional model data sources to check.
+        accept_eula (bool): EULA acceptance for the draft model.
+    """
+    if not pysdk_model:
+        return
+
+    deployment_config_draft_model_sources = (
+        pysdk_model.deployment_config.get("DeploymentArgs", {})
+        .get("AdditionalDataSources", {})
+        .get("speculative_decoding", [])
+        if pysdk_model.deployment_config
+        else None
+    )
+    pysdk_model_additional_model_sources = pysdk_model.additional_model_data_sources
+
+    if not deployment_config_draft_model_sources or not pysdk_model_additional_model_sources:
+        return
+
+    # Gated/ungated classification is only available through deployment_config.
+    # Thus we must check each draft model in the deployment_config and see if it is set
+    # as an additional model data source on the PySDK model itself.
+    model_access_config_updated = False
+    for source in deployment_config_draft_model_sources:
+        if source.get("channel_name") != "draft_model":
+            continue
+
+        if not _is_draft_model_gated(source):
+            continue
+
+        deployment_config_draft_model_source_s3_uri = (
+            _extract_deployment_config_additional_model_data_source_s3_uri(source)
+        )
+
+        # If EULA is accepted, proceed with modifying the draft model data source
+        for additional_source in pysdk_model_additional_model_sources:
+            if additional_source.get("ChannelName") != "draft_model":
+                continue
+
+            # Verify the pysdk model source and deployment config model source match
+            pysdk_model_source_s3_uri = _extract_additional_model_data_source_s3_uri(
+                additional_source
+            )
+            if deployment_config_draft_model_source_s3_uri not in pysdk_model_source_s3_uri:
+                continue
+
+            if not accept_eula:
+                raise ValueError(
+                    "Gated draft model requires accepting end-user license agreement (EULA)."
+                )
+
+            # Set ModelAccessConfig.AcceptEula to True
+            updated_source = additional_source.copy()
+            updated_source["S3DataSource"]["ModelAccessConfig"] = {"AcceptEula": True}
+
+            index = pysdk_model.additional_model_data_sources.index(additional_source)
+            pysdk_model.additional_model_data_sources[index] = updated_source
+
+            model_access_config_updated = True
+            break
+
+        if model_access_config_updated:
+            break
+
+    return
diff --git a/tests/unit/sagemaker/serve/builder/test_js_builder.py b/tests/unit/sagemaker/serve/builder/test_js_builder.py
@@ -25,7 +25,10 @@
     LocalModelOutOfMemoryException,
     LocalModelInvocationException,
 )
-from tests.unit.sagemaker.serve.constants import DEPLOYMENT_CONFIGS
+from tests.unit.sagemaker.serve.constants import (
+    DEPLOYMENT_CONFIGS,
+    OPTIMIZED_DEPLOYMENT_CONFIG_WITH_GATED_DRAFT_MODEL,
+)
 
 mock_model_id = "huggingface-llm-amazon-falconlite"
 mock_t5_model_id = "google/flan-t5-xxl"
@@ -1198,6 +1201,51 @@ def test_optimize_quantize_for_jumpstart(
 
         self.assertIsNotNone(out_put)
 
+    @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
+    @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
+    def test_optimize_gated_draft_model_for_jumpstart_with_accept_eula_false(
+        self,
+        mock_serve_settings,
+        mock_telemetry,
+    ):
+        mock_sagemaker_session = Mock()
+
+        mock_pysdk_model = Mock()
+        mock_pysdk_model.env = {"SAGEMAKER_ENV": "1"}
+        mock_pysdk_model.model_data = mock_model_data
+        mock_pysdk_model.image_uri = mock_tgi_image_uri
+        mock_pysdk_model.list_deployment_configs.return_value = DEPLOYMENT_CONFIGS
+        mock_pysdk_model.deployment_config = OPTIMIZED_DEPLOYMENT_CONFIG_WITH_GATED_DRAFT_MODEL
+
+        sample_input = {
+            "inputs": "The diamondback terrapin or simply terrapin is a species "
+            "of turtle native to the brackish coastal tidal marshes of the",
+            "parameters": {"max_new_tokens": 1024},
+        }
+        sample_output = [
+            {
+                "generated_text": "The diamondback terrapin or simply terrapin is a "
+                "species of turtle native to the brackish coastal "
+                "tidal marshes of the east coast."
+            }
+        ]
+
+        model_builder = ModelBuilder(
+            model="meta-textgeneration-llama-3-70b",
+            schema_builder=SchemaBuilder(sample_input, sample_output),
+            sagemaker_session=mock_sagemaker_session,
+        )
+
+        model_builder.pysdk_model = mock_pysdk_model
+
+        self.assertRaises(
+            ValueError,
+            model_builder._optimize_for_jumpstart(
+                accept_eula=True,
+                speculative_decoding_config={"Provider": "sagemaker", "AcceptEula": False},
+            ),
+        )
+
     @patch("sagemaker.serve.builder.jumpstart_builder._capture_telemetry", side_effect=None)
     @patch.object(ModelBuilder, "_get_serve_setting", autospec=True)
     def test_optimize_quantize_and_compile_for_jumpstart(
@@ -1248,10 +1296,6 @@ def test_optimize_quantize_and_compile_for_jumpstart(
                 "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
             },
             compilation_config={"OverrideEnvironment": {"OPTION_TENSOR_PARALLEL_DEGREE": "2"}},
-            env_vars={
-                "OPTION_TENSOR_PARALLEL_DEGREE": "1",
-                "OPTION_MAX_ROLLING_BATCH_SIZE": "2",
-            },
             output_path="s3://bucket/code/",
         )
 
diff --git a/tests/unit/sagemaker/serve/constants.py b/tests/unit/sagemaker/serve/constants.py
@@ -165,3 +165,105 @@
         },
     },
 ]
+OPTIMIZED_DEPLOYMENT_CONFIG_WITH_GATED_DRAFT_MODEL = {
+    "DeploymentConfigName": "lmi-optimized",
+    "DeploymentArgs": {
+        "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+        "djl-inference:0.29.0-lmi11.0.0-cu124",
+        "ModelData": {
+            "S3DataSource": {
+                "S3Uri": "s3://jumpstart-private-cache-alpha-us-west-2/meta-textgeneration/"
+                "meta-textgeneration-llama-3-1-70b/artifacts/inference-prepack/v2.0.0/",
+                "S3DataType": "S3Prefix",
+                "CompressionType": "None",
+            }
+        },
+        "ModelPackageArn": None,
+        "Environment": {
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "ENDPOINT_SERVER_TIMEOUT": "3600",
+            "MODEL_CACHE_ROOT": "/opt/ml/model",
+            "SAGEMAKER_ENV": "1",
+            "HF_MODEL_ID": "/opt/ml/model",
+            "OPTION_SPECULATIVE_DRAFT_MODEL": "/opt/ml/additional-model-data-sources/draft_model",
+            "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+        },
+        "InstanceType": "ml.g6.2xlarge",
+        "ComputeResourceRequirements": {
+            "MinMemoryRequiredInMb": 131072,
+            "NumberOfAcceleratorDevicesRequired": 1,
+        },
+        "ModelDataDownloadTimeout": 1200,
+        "ContainerStartupHealthCheckTimeout": 1200,
+        "AdditionalDataSources": {
+            "speculative_decoding": [
+                {
+                    "channel_name": "draft_model",
+                    "provider": {"name": "JumpStart", "classification": "gated"},
+                    "artifact_version": "v1",
+                    "hosting_eula_key": "fmhMetadata/eula/llama3_2Eula.txt",
+                    "s3_data_source": {
+                        "s3_uri": "meta-textgeneration/meta-textgeneration-llama-3-2-1b/artifacts/"
+                        "inference-prepack/v1.0.0/",
+                        "compression_type": "None",
+                        "s3_data_type": "S3Prefix",
+                    },
+                }
+            ]
+        },
+    },
+    "AccelerationConfigs": [
+        {
+            "type": "Compilation",
+            "enabled": False,
+            "diy_workflow_overrides": {
+                "gpu-lmi-trt": {
+                    "enabled": False,
+                    "reason": "TRT-LLM 0.11.0 in LMI v11 does not support llama 3.1",
+                }
+            },
+        },
+        {
+            "type": "Speculative-Decoding",
+            "enabled": True,
+            "diy_workflow_overrides": {
+                "gpu-lmi-trt": {
+                    "enabled": False,
+                    "reason": "LMI v11 does not support Speculative Decoding for TRT",
+                }
+            },
+        },
+        {
+            "type": "Quantization",
+            "enabled": False,
+            "diy_workflow_overrides": {
+                "gpu-lmi-trt": {
+                    "enabled": False,
+                    "reason": "TRT-LLM 0.11.0 in LMI v11 does not support llama 3.1",
+                }
+            },
+        },
+    ],
+    "BenchmarkMetrics": {"ml.g6.2xlarge": None},
+}
+GATED_DRAFT_MODEL_CONFIG = {
+    "channel_name": "draft_model",
+    "provider": {"name": "JumpStart", "classification": "gated"},
+    "artifact_version": "v1",
+    "hosting_eula_key": "fmhMetadata/eula/llama3_2Eula.txt",
+    "s3_data_source": {
+        "s3_uri": "meta-textgeneration/meta-textgeneration-llama-3-2-1b/artifacts/"
+        "inference-prepack/v1.0.0/",
+        "compression_type": "None",
+        "s3_data_type": "S3Prefix",
+    },
+}
+NON_GATED_DRAFT_MODEL_CONFIG = {
+    "channel_name": "draft_model",
+    "s3_data_source": {
+        "compression_type": "None",
+        "s3_data_type": "S3Prefix",
+        "s3_uri": "s3://sagemaker-sd-models-beta-us-west-2/"
+        "sagemaker-speculative-decoding-llama3-small-v3/",
+    },
+}
diff --git a/tests/unit/sagemaker/serve/utils/test_optimize_utils.py b/tests/unit/sagemaker/serve/utils/test_optimize_utils.py
