ModelBuilder speculative decoding UTs and minor fixes.

Author: Joseph Zhang

src/sagemaker/jumpstart/types.py

@@ -1122,9 +1122,9 @@ def to_json(self, exclude_keys=True) -> Dict[str, Any]:
 class JumpStartModelDataSource(AdditionalModelDataSource):
     """Data class JumpStart additional model data source."""
 
-    SERIALIZATION_EXCLUSION_SET = {
+    SERIALIZATION_EXCLUSION_SET = AdditionalModelDataSource.SERIALIZATION_EXCLUSION_SET.union(
         "artifact_version"
-    } | AdditionalModelDataSource.SERIALIZATION_EXCLUSION_SET
+    )
 
     __slots__ = list(SERIALIZATION_EXCLUSION_SET) + AdditionalModelDataSource.__slots__
 

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -737,9 +737,7 @@ def _optimize_for_jumpstart(
         if not optimization_config:
             optimization_config = {}
 
-        if (
-            not optimization_config or not optimization_config.get("ModelCompilationConfig")
-        ) and is_compilation:
+        if not optimization_config.get("ModelCompilationConfig") and is_compilation:
             # Fallback to default if override_env is None or empty
             if not compilation_override_env:
                 compilation_override_env = pysdk_model_env_vars
@@ -907,7 +905,9 @@ def _set_additional_model_source(
                 )
             else:
                 self.pysdk_model = _custom_speculative_decoding(
-                    self.pysdk_model, speculative_decoding_config, speculative_decoding_config.get("AcceptEula", False)
+                    self.pysdk_model,
+                    speculative_decoding_config,
+                    speculative_decoding_config.get("AcceptEula", False),
                 )
 
     def _find_compatible_deployment_config(

src/sagemaker/serve/builder/model_builder.py

@@ -591,7 +591,7 @@ def _model_builder_deploy_wrapper(
                 )
 
         if "endpoint_logging" not in kwargs:
-            kwargs["endpoint_logging"] = True
+            kwargs["endpoint_logging"] = False
         predictor = self._original_deploy(
             *args,
             instance_type=instance_type,

src/sagemaker/serve/utils/optimize_utils.py

@@ -73,9 +73,8 @@ def _deployment_config_contains_draft_model(deployment_config: Optional[Dict]) -
         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)
+
+    return "speculative_decoding" in additional_data_sources if additional_data_sources else False
 
 
 def _is_draft_model_jumpstart_provided(deployment_config: Optional[Dict]) -> bool:
@@ -207,15 +206,15 @@ def _extract_speculative_draft_model_provider(
     if speculative_decoding_config is None:
         return None
 
-    if speculative_decoding_config.get("ModelProvider").lower() == "jumpstart":
+    model_provider = speculative_decoding_config.get("ModelProvider", "").lower()
+
+    if model_provider == "jumpstart":
         return "jumpstart"
 
-    if speculative_decoding_config.get(
-        "ModelProvider"
-    ).lower() == "custom" or speculative_decoding_config.get("ModelSource"):
+    if model_provider == "custom" or speculative_decoding_config.get("ModelSource"):
         return "custom"
 
-    if speculative_decoding_config.get("ModelProvider").lower() == "sagemaker":
+    if model_provider == "sagemaker":
         return "sagemaker"
 
     return "auto"
@@ -238,7 +237,7 @@ def _extract_additional_model_data_source_s3_uri(
     ):
         return None
 
-    return additional_model_data_source.get("S3DataSource").get("S3Uri", None)
+    return additional_model_data_source.get("S3DataSource").get("S3Uri")
 
 
 def _extract_deployment_config_additional_model_data_source_s3_uri(
@@ -272,7 +271,7 @@ def _is_draft_model_gated(
     Returns:
         bool: Whether the draft model is gated or not.
     """
-    return draft_model_config.get("hosting_eula_key", None)
+    return "hosting_eula_key" in draft_model_config if draft_model_config else False
 
 
 def _extracts_and_validates_speculative_model_source(
@@ -371,7 +370,7 @@ def _extract_optimization_config_and_env(
         compilation_config (Optional[Dict]): The compilation config.
 
     Returns:
-        Optional[Tuple[Optional[Dict], Optional[Dict]]]:
+        Optional[Tuple[Optional[Dict], Optional[Dict], Optional[Dict]]]:
             The optimization config and environment variables.
     """
     optimization_config = {}
@@ -388,7 +387,7 @@ def _extract_optimization_config_and_env(
     if compilation_config is not None:
         optimization_config["ModelCompilationConfig"] = compilation_config
 
-    # Return both dicts and environment variable if either is present
+    # Return optimization config dict and environment variables if either is present
     if optimization_config:
         return optimization_config, quantization_override_env, compilation_override_env
 

tests/unit/sagemaker/serve/builder/test_js_builder.py

@@ -28,6 +28,7 @@
 from tests.unit.sagemaker.serve.constants import (
     DEPLOYMENT_CONFIGS,
     OPTIMIZED_DEPLOYMENT_CONFIG_WITH_GATED_DRAFT_MODEL,
+    CAMEL_CASE_ADDTL_DRAFT_MODEL_DATA_SOURCES,
 )
 
 mock_model_id = "huggingface-llm-amazon-falconlite"
@@ -1203,19 +1204,34 @@ def test_optimize_quantize_for_jumpstart(
 
     @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(
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._is_jumpstart_model_id",
+        return_value=True,
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder.JumpStart._create_pre_trained_js_model",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "sagemaker.serve.builder.jumpstart_builder._jumpstart_speculative_decoding",
+        return_value=True,
+    )
+    def test_jumpstart_model_provider_calls_jumpstart_speculative_decoding(
         self,
+        mock_js_speculative_decoding,
+        mock_pretrained_js_model,
+        mock_is_js_model,
         mock_serve_settings,
-        mock_telemetry,
+        mock_capture_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
+        mock_pysdk_model.additional_model_data_sources = CAMEL_CASE_ADDTL_DRAFT_MODEL_DATA_SOURCES
 
         sample_input = {
             "inputs": "The diamondback terrapin or simply terrapin is a species "
@@ -1238,14 +1254,17 @@ def test_optimize_gated_draft_model_for_jumpstart_with_accept_eula_false(
 
         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},
-            ),
+        model_builder._optimize_for_jumpstart(
+            accept_eula=True,
+            speculative_decoding_config={
+                "ModelProvider": "JumpStart",
+                "ModelID": "meta-textgeneration-llama-3-2-1b",
+                "AcceptEula": False,
+            },
         )
 
+        mock_js_speculative_decoding.assert_called_once()
+
     @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(

tests/unit/sagemaker/serve/constants.py

@@ -165,6 +165,43 @@
         },
     },
 ]
+NON_OPTIMIZED_DEPLOYMENT_CONFIG = {
+    "ConfigName": "neuron-inference",
+    "BenchmarkMetrics": [
+        {"name": "Latency", "value": "100", "unit": "Tokens/S"},
+        {"name": "Throughput", "value": "1867", "unit": "Tokens/S"},
+    ],
+    "DeploymentArgs": {
+        "ModelDataDownloadTimeout": 1200,
+        "ContainerStartupHealthCheckTimeout": 1200,
+        "ImageUri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4"
+        ".0-gpu-py310-cu121-ubuntu20.04",
+        "ModelData": {
+            "S3DataSource": {
+                "S3Uri": "s3://jumpstart-private-cache-prod-us-west-2/meta-textgeneration/meta-textgeneration"
+                "-llama-2-7b/artifacts/inference-prepack/v1.0.0/",
+                "S3DataType": "S3Prefix",
+                "CompressionType": "None",
+            }
+        },
+        "InstanceType": "ml.p2.xlarge",
+        "Environment": {
+            "SAGEMAKER_PROGRAM": "inference.py",
+            "ENDPOINT_SERVER_TIMEOUT": "3600",
+            "MODEL_CACHE_ROOT": "/opt/ml/model",
+            "SAGEMAKER_ENV": "1",
+            "HF_MODEL_ID": "/opt/ml/model",
+            "MAX_INPUT_LENGTH": "4095",
+            "MAX_TOTAL_TOKENS": "4096",
+            "SM_NUM_GPUS": "1",
+            "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
+        },
+        "ComputeResourceRequirements": {
+            "MinMemoryRequiredInMb": 16384,
+            "NumberOfAcceleratorDevicesRequired": 1,
+        },
+    },
+}
 OPTIMIZED_DEPLOYMENT_CONFIG_WITH_GATED_DRAFT_MODEL = {
     "DeploymentConfigName": "lmi-optimized",
     "DeploymentArgs": {
@@ -267,3 +304,14 @@
         "sagemaker-speculative-decoding-llama3-small-v3/",
     },
 }
+CAMEL_CASE_ADDTL_DRAFT_MODEL_DATA_SOURCES = [
+    {
+        "ChannelName": "draft_model",
+        "S3DataSource": {
+            "S3Uri": "meta-textgeneration/meta-textgeneration-llama-3-2-1b/artifacts/"
+            "inference-prepack/v1.0.0/",
+            "CompressionType": "None",
+            "S3DataType": "S3Prefix",
+        },
+    }
+]