Add ModelBuilder support for JumpStart-provided draft models.

Author: Joseph Zhang

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -49,6 +49,7 @@
     SPECULATIVE_DRAFT_MODEL,
     _is_inferentia_or_trainium,
     _validate_and_set_eula_for_draft_model_sources,
+    _jumpstart_speculative_decoding,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -503,7 +504,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, accept_draft_model_eula: Optional[bool] = False
     ) -> None:
         """Sets the deployment config to apply to the model.
 
@@ -735,6 +736,10 @@ def _optimize_for_jumpstart(
         optimization_config, quantization_override_env, compilation_override_env = (
             _extract_optimization_config_and_env(quantization_config, compilation_config)
         )
+
+        if not optimization_config:
+            optimization_config = {}
+
         if (
             not optimization_config or not optimization_config.get("ModelCompilationConfig")
         ) and is_compilation:
@@ -844,6 +849,7 @@ def _set_additional_model_source(
         """
         if speculative_decoding_config:
             model_provider = _extract_speculative_draft_model_provider(speculative_decoding_config)
+
             channel_name = _generate_channel_name(self.pysdk_model.additional_model_data_sources)
 
             if model_provider == "sagemaker":
@@ -868,17 +874,23 @@ 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"),
-                )
+                    _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}"}
+                    {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"}
                 )
                 self.pysdk_model.add_tags(
                     {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
                 )
+            elif model_provider == "jumpstart":
+                _jumpstart_speculative_decoding(
+                    model=self.pysdk_model,
+                    speculative_decoding_config=speculative_decoding_config,
+                    sagemaker_session=self.sagemaker_session,
+                )
             else:
                 self.pysdk_model = _custom_speculative_decoding(
                     self.pysdk_model, speculative_decoding_config, accept_eula

src/sagemaker/serve/builder/model_builder.py

@@ -76,6 +76,7 @@
     _is_s3_uri,
     _custom_speculative_decoding,
     _extract_speculative_draft_model_provider,
+    _jumpstart_speculative_decoding,
 )
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import (
@@ -99,7 +100,6 @@
     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,
@@ -590,21 +590,6 @@ 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(
@@ -1358,9 +1343,17 @@ def _optimize_for_hf(
         Returns:
             Optional[Dict[str, Any]]: Model optimization job input arguments.
         """
-        self.pysdk_model = _custom_speculative_decoding(
-            self.pysdk_model, speculative_decoding_config, False
-        )
+        if speculative_decoding_config:
+            if speculative_decoding_config.get("ModelProvider", "") == "JumpStart":
+                _jumpstart_speculative_decoding(
+                    model=self.pysdk_model,
+                    speculative_decoding_config=speculative_decoding_config,
+                    sagemaker_session=self.sagemaker_session,
+                )
+            else:
+                self.pysdk_model = _custom_speculative_decoding(
+                    self.pysdk_model, speculative_decoding_config, False
+                )
 
         if quantization_config or compilation_config:
             create_optimization_job_args = {

src/sagemaker/serve/utils/optimize_utils.py

@@ -17,8 +17,10 @@
 import logging
 from typing import Dict, Any, Optional, Union, List, Tuple
 
-from sagemaker import Model
+from sagemaker import Model, Session
 from sagemaker.enums import Tag
+from sagemaker.jumpstart.utils import accessors, get_eula_message
+
 
 logger = logging.getLogger(__name__)
 
@@ -164,6 +166,9 @@ def _extract_speculative_draft_model_provider(
     if speculative_decoding_config is None:
         return None
 
+    if speculative_decoding_config.get("ModelProvider") == "JumpStart":
+        return "jumpstart"
+
     if speculative_decoding_config.get(
         "ModelProvider"
     ) == "Custom" or speculative_decoding_config.get("ModelSource"):
@@ -292,7 +297,7 @@ def _generate_additional_model_data_sources(
         },
     }
     if accept_eula:
-        additional_model_data_source["S3DataSource"]["ModelAccessConfig"] = {"ACCEPT_EULA": True}
+        additional_model_data_source["S3DataSource"]["ModelAccessConfig"] = {"AcceptEula": True}
 
     return [additional_model_data_source]
 
@@ -327,10 +332,10 @@ def _extract_optimization_config_and_env(
     """
     optimization_config = {}
     quantization_override_env = (
-        quantization_config.get("OverrideEnvironment", {}) if quantization_config else None
+        quantization_config.get("OverrideEnvironment") if quantization_config else None
     )
     compilation_override_env = (
-        compilation_config.get("OverrideEnvironment", {}) if compilation_config else None
+        compilation_config.get("OverrideEnvironment") if compilation_config else None
     )
 
     if quantization_config is not None:
@@ -343,7 +348,7 @@ def _extract_optimization_config_and_env(
     if optimization_config:
         return optimization_config, quantization_override_env, compilation_override_env
 
-    return {}, None, None
+    return None, None, None
 
 
 def _custom_speculative_decoding(
@@ -364,7 +369,7 @@ def _custom_speculative_decoding(
             speculative_decoding_config
         )
 
-        accept_eula = speculative_decoding_config.get("AcceptEula", False)
+        accept_eula = speculative_decoding_config.get("AcceptEula", accept_eula)
 
         if _is_s3_uri(additional_model_source):
             channel_name = _generate_channel_name(model.additional_model_data_sources)
@@ -384,6 +389,65 @@ def _custom_speculative_decoding(
     return model
 
 
+def _jumpstart_speculative_decoding(
+    model=Model,
+    speculative_decoding_config: Optional[Dict[str, Any]] = None,
+    sagemaker_session: Optional[Session] = None,
+):
+    """Modifies the given model for speculative decoding config with JumpStart provider.
+
+    Args:
+        model (Model): The model.
+        speculative_decoding_config (Optional[Dict]): The speculative decoding config.
+        sagemaker_session (Optional[Session]): Sagemaker session for execution.
+    """
+    if speculative_decoding_config:
+        js_id = speculative_decoding_config.get("ModelID")
+        if not js_id:
+            raise ValueError(
+                "`ModelID` is a required field in `speculative_decoding_config` when "
+                "using JumpStart as draft model provider."
+            )
+        model_version = speculative_decoding_config.get("ModelVersion", "*")
+        accept_eula = speculative_decoding_config.get("AcceptEula", False)
+        channel_name = _generate_channel_name(model.additional_model_data_sources)
+
+        model_specs = accessors.JumpStartModelsAccessor.get_model_specs(
+            model_id=js_id,
+            version=model_version,
+            region=sagemaker_session.boto_region_name,
+            sagemaker_session=sagemaker_session,
+        )
+        model_spec_json = model_specs.to_json()
+
+        js_bucket = accessors.JumpStartModelsAccessor.get_jumpstart_content_bucket()
+
+        if model_spec_json.get("gated_bucket", False):
+            if not accept_eula:
+                eula_message = get_eula_message(
+                    model_specs=model_specs, region=sagemaker_session.boto_region_name
+                )
+                raise ValueError(
+                    f"{eula_message} Please set `AcceptEula` to True in "
+                    f"speculative_decoding_config once acknowledged."
+                )
+            js_bucket = accessors.JumpStartModelsAccessor.get_jumpstart_gated_content_bucket()
+
+        key_prefix = model_spec_json.get("hosting_prepacked_artifact_key")
+        model.additional_model_data_sources = _generate_additional_model_data_sources(
+            f"s3://{js_bucket}/{key_prefix}",
+            channel_name,
+            accept_eula,
+        )
+
+        model.env.update(
+            {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"}
+        )
+        model.add_tags(
+            {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "jumpstart"},
+        )
+
+
 def _validate_and_set_eula_for_draft_model_sources(
     pysdk_model: Model,
     accept_eula: bool = False,