feat: optimization technique related validations.

src/sagemaker/jumpstart/model.py

@@ -665,7 +665,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,
+        model_access_configs: Optional[Dict[str, ModelAccessConfig]] = None,
     ) -> PredictorBase:
         """Creates endpoint by calling base ``Model`` class `deploy` method.
 
@@ -766,7 +766,7 @@ def deploy(
                 ModelAccessConfig, provide a `{ "model_id", ModelAccessConfig(accept_eula=True) }`
                 to indicate whether model terms of use have been accepted. The `accept_eula` value
                 must be explicitly defined as `True` in order to accept the end-user license
-                agreement (EULA) that some. (Default: None)
+                agreement (EULA) that some models require. (Default: None)
 
         Raises:
             MarketplaceModelSubscriptionError: If the caller is not subscribed to the model.
@@ -817,12 +817,14 @@ def deploy(
                 f"{EndpointType.INFERENCE_COMPONENT_BASED} is not supported for Proprietary models."
             )
 
+        print(self.additional_model_data_sources)
         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,
         )
+        print(self.additional_model_data_sources)
 
         try:
             predictor = super(JumpStartModel, self).deploy(**deploy_kwargs.to_kwargs_dict())

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 = {
-        "artifact_version"
-    } | AdditionalModelDataSource.SERIALIZATION_EXCLUSION_SET
+    SERIALIZATION_EXCLUSION_SET = AdditionalModelDataSource.SERIALIZATION_EXCLUSION_SET.union(
+        {"artifact_version"}
+    )
 
     __slots__ = list(SERIALIZATION_EXCLUSION_SET) + AdditionalModelDataSource.__slots__
 

src/sagemaker/jumpstart/utils.py

@@ -558,12 +558,12 @@ 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_template(
+    return get_formatted_eula_message_template(
         model_id=model_specs.model_id, region=region, hosting_eula_key=model_specs.hosting_eula_key
     )
 
 
-def format_eula_message_template(model_id: str, region: str, hosting_eula_key: str):
+def get_formatted_eula_message_template(model_id: str, region: str, hosting_eula_key: str) -> str:
     """Returns a formatted EULA message."""
     return (
         f"Model '{model_id}' requires accepting end-user license agreement (EULA). "
@@ -1542,17 +1542,32 @@ def _add_model_access_configs_to_model_data_sources(
     model_access_configs: Dict[str, ModelAccessConfig],
     model_id: str,
     region: str,
-):
-    """Sets AcceptEula to True for gated speculative decoding models"""
+) -> List[Dict[str, any]]:
+    """Iterate over the accept EULA configs to ensure all channels are matched
 
+    Args:
+        model_data_sources (DeploymentConfigMetadata): Model data sources that will be updated
+        model_access_configs (DeploymentConfigMetadata): Config holding accept_eula field
+        model_id (DeploymentConfigMetadata): Jumpstart mode id.
+        region (str): Region where the user is operating in.
+    Returns:
+        List[Dict[str, Any]]: List of model data sources with accept EULA configs applied
+    Raise:
+        ValueError if at least one channel that requires EULA acceptance as not passed.
+    """
     if not model_data_sources:
         return model_data_sources
 
     acked_model_data_sources = []
     for model_data_source in model_data_sources:
         hosting_eula_key = model_data_source.get("HostingEulaKey")
+        mutable_model_data_source = model_data_source.copy()
         if hosting_eula_key:
-            if not model_access_configs or not model_access_configs.get(model_id):
+            if (
+                not model_access_configs
+                or not model_access_configs.get(model_id)
+                or not model_access_configs.get(model_id).accept_eula
+            ):
                 eula_message_template = (
                     "{model_source}{base_eula_message}{model_access_configs_message}"
                 )
@@ -1562,24 +1577,28 @@ def _add_model_access_configs_to_model_data_sources(
                 raise ValueError(
                     eula_message_template.format(
                         model_source="Additional " if model_data_source.get("ChannelName") else "",
-                        base_eula_message=format_eula_message_template(
+                        base_eula_message=get_formatted_eula_message_template(
                             model_id=model_id, region=region, hosting_eula_key=hosting_eula_key
                         ),
                         model_access_configs_message=(
-                            " Please add a ModelAccessConfig entry:"
+                            "Please add a ModelAccessConfig entry:"
                             f" {model_access_config_entry} "
-                            "to model_access_configs to acknowledge the EULA."
+                            "to model_access_configs to accept the EULA."
                         ),
                     )
                 )
-            acked_model_data_source = model_data_source.copy()
-            acked_model_data_source.pop("HostingEulaKey")
-            acked_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
+            mutable_model_data_source.pop(
+                "HostingEulaKey"
+            )  # pop when model access config is applied
+            mutable_model_data_source["S3DataSource"]["ModelAccessConfig"] = (
                 camel_case_to_pascal_case(model_access_configs.get(model_id).model_dump())
             )
-            acked_model_data_sources.append(acked_model_data_source)
+            acked_model_data_sources.append(mutable_model_data_source)
         else:
-            acked_model_data_sources.append(model_data_source)
+            mutable_model_data_source.pop(
+                "HostingEulaKey"
+            )  # pop when model access config is not applicable
+            acked_model_data_sources.append(mutable_model_data_source)
     return acked_model_data_sources
 
 

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,
@@ -1283,7 +1283,7 @@ def _model_builder_optimize_wrapper(
             # TRTLLM is used by Neo if the following are provided:
             #  1) a GPU instance type
             #  2) compilation config
-            gpu_instance_families = ["g4", "g5", "p4d"]
+            gpu_instance_families = ["g5", "g6", "p4d", "p4de", "p5"]
             is_gpu_instance = optimization_instance_type and any(
                 gpu_instance_family in optimization_instance_type
                 for gpu_instance_family in gpu_instance_families
@@ -1296,8 +1296,16 @@ def _model_builder_optimize_wrapper(
                 keyword in self.model.lower() for keyword in llama_3_1_keywords
             )
 
-            if is_gpu_instance and self.model and is_llama_3_1 and self.is_compiled:
-                raise ValueError("Compilation is not supported for Llama-3.1 with a GPU instance.")
+            if is_gpu_instance and self.model and self.is_compiled:
+                if is_llama_3_1:
+                    raise ValueError(
+                        "Compilation is not supported for Llama-3.1 with a GPU instance."
+                    )
+                if speculative_decoding_config:
+                    raise ValueError(
+                        "Compilation is not supported with speculative decoding with "
+                        "a GPU instance."
+                    )
 
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)

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