Add "Auto" speculative decoding ModelProvider option; add validations to differentiate SageMaker/JumpStart draft models.

Author: Joseph Zhang

src/sagemaker/jumpstart/factory/model.py

@@ -564,10 +564,7 @@ def _add_config_name_to_init_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpSta
 
 
 def _apply_accept_eula_on_model_data_source(
-    model_data_source: Dict[str, Any],
-    model_id: str,
-    region: str,
-    accept_eula: bool
+    model_data_source: Dict[str, Any], model_id: str, region: str, accept_eula: bool
 ):
     """Sets AcceptEula to True for gated speculative decoding models"""
 
@@ -586,7 +583,8 @@ def _apply_accept_eula_on_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_eula=True` once acknowledged."
+                f"/{hosting_eula_key} for terms of use. Please set `accept_draft_model_eula=True` "
+                f"once acknowledged."
             )
         )
 
@@ -608,7 +606,10 @@ def _add_additional_model_data_sources_to_kwargs(
         [
             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,
+                    data_source.to_json(),
+                    kwargs.model_id,
+                    kwargs.region,
+                    kwargs.accept_draft_model_eula,
                 )
             )
             for data_source in speculative_decoding_data_sources

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -49,6 +49,8 @@
     SPECULATIVE_DRAFT_MODEL,
     _is_inferentia_or_trainium,
     _jumpstart_speculative_decoding,
+    _deployment_config_contains_draft_model,
+    _is_draft_model_jumpstart_provided,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -850,7 +852,7 @@ def _set_additional_model_source(
 
             channel_name = _generate_channel_name(self.pysdk_model.additional_model_data_sources)
 
-            if model_provider == "sagemaker":
+            if model_provider in ["sagemaker", "auto"]:
                 additional_model_data_sources = (
                     self.pysdk_model.deployment_config.get("DeploymentArgs", {}).get(
                         "AdditionalDataSources"
@@ -863,6 +865,15 @@ def _set_additional_model_source(
                         speculative_decoding_config
                     )
                     if deployment_config:
+                        if model_provider == "sagemaker" and _is_draft_model_jumpstart_provided(
+                            deployment_config
+                        ):
+                            raise ValueError(
+                                "No `Sagemaker` provided draft model was found for "
+                                f"{self.model}. Try setting `ModelProvider` "
+                                "to `Auto` instead."
+                            )
+
                         try:
                             self.pysdk_model.set_deployment_config(
                                 config_name=deployment_config.get("DeploymentConfigName"),
@@ -878,12 +889,21 @@ def _set_additional_model_source(
                         raise ValueError(
                             "Cannot find deployment config compatible for optimization job."
                         )
+                else:
+                    if model_provider == "sagemaker" and _is_draft_model_jumpstart_provided(
+                        self.pysdk_model.deployment_config
+                    ):
+                        raise ValueError(
+                            "No `Sagemaker` provided draft model was found for "
+                            f"{self.model}. Try setting `ModelProvider` "
+                            "to `Auto` instead."
+                        )
 
                 self.pysdk_model.env.update(
                     {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"}
                 )
                 self.pysdk_model.add_tags(
-                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
+                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": model_provider},
                 )
             elif model_provider == "jumpstart":
                 _jumpstart_speculative_decoding(
@@ -911,15 +931,17 @@ def _find_compatible_deployment_config(
         for deployment_config in self.pysdk_model.list_deployment_configs():
             image_uri = deployment_config.get("deployment_config", {}).get("ImageUri")
 
-            if _is_image_compatible_with_optimization_job(image_uri):
+            if _is_image_compatible_with_optimization_job(
+                image_uri
+            ) and _deployment_config_contains_draft_model(deployment_config):
                 if (
-                    model_provider == "sagemaker"
+                    model_provider in ["sagemaker", "auto"]
                     and deployment_config.get("DeploymentArgs", {}).get("AdditionalDataSources")
                 ) or model_provider == "custom":
                     return deployment_config
 
         # There's no matching config from jumpstart to add sagemaker draft model location
-        if model_provider == "sagemaker":
+        if model_provider in ["sagemaker", "auto"]:
             return None
 
         # fall back to the default jumpstart model deployment config for optimization job

src/sagemaker/serve/utils/optimize_utils.py

@@ -60,6 +60,47 @@ def _is_image_compatible_with_optimization_job(image_uri: Optional[str]) -> bool
     return "djl-inference:" in image_uri and ("-lmi" in image_uri or "-neuronx-" in image_uri)
 
 
+def _deployment_config_contains_draft_model(deployment_config: Optional[Dict]) -> bool:
+    """Checks whether a deployment config contains a speculative decoding draft model.
+
+    Args:
+        deployment_config (Dict): The deployment config to check.
+
+    Returns:
+        bool: Whether the deployment config contains a draft model or not.
+    """
+    if deployment_config is None:
+        return False
+    deployment_args = deployment_config.get("DeploymentArgs", {})
+    additional_data_sources = deployment_args.get("AdditionalDataSources")
+    if not additional_data_sources:
+        return False
+    return additional_data_sources.get("speculative_decoding", False)
+
+
+def _is_draft_model_jumpstart_provided(deployment_config: Optional[Dict]) -> bool:
+    """Checks whether a deployment config's draft model is provided by JumpStart.
+
+    Args:
+        deployment_config (Dict): The deployment config to check.
+
+    Returns:
+        bool: Whether the draft model is provided by JumpStart or not.
+    """
+    if deployment_config is None:
+        return False
+
+    additional_model_data_sources = deployment_config.get("DeploymentArgs", {}).get(
+        "AdditionalDataSources"
+    )
+    for source in additional_model_data_sources.get("speculative_decoding", []):
+        if source["channel_name"] == "draft_model":
+            if source.get("provider", {}).get("name") == "JumpStart":
+                return True
+            continue
+    return False
+
+
 def _generate_optimized_model(pysdk_model: Model, optimization_response: dict) -> Model:
     """Generates a new optimization model.
 
@@ -166,15 +207,18 @@ def _extract_speculative_draft_model_provider(
     if speculative_decoding_config is None:
         return None
 
-    if speculative_decoding_config.get("ModelProvider") == "JumpStart":
+    if speculative_decoding_config.get("ModelProvider").lower() == "jumpstart":
         return "jumpstart"
 
     if speculative_decoding_config.get(
         "ModelProvider"
-    ) == "Custom" or speculative_decoding_config.get("ModelSource"):
+    ).lower() == "custom" or speculative_decoding_config.get("ModelSource"):
         return "custom"
 
-    return "sagemaker"
+    if speculative_decoding_config.get("ModelProvider").lower() == "sagemaker":
+        return "sagemaker"
+
+    return "auto"
 
 
 def _extract_additional_model_data_source_s3_uri(