update: Add optimize to ModelBuilder JS (aws#1485)

* MB JS Optimize

* UT

* Refactore

* UT

* UT

* refactore

* refactore

---------

Co-authored-by: Jonathan Makunga <[email protected]>

Author: makungaj1

src/sagemaker/jumpstart/types.py

@@ -2568,6 +2568,7 @@ class DeploymentArgs(BaseDeploymentConfigDataHolder):
         "compute_resource_requirements",
         "model_data_download_timeout",
         "container_startup_health_check_timeout",
+        "additional_data_sources",
     ]
 
     def __init__(
@@ -2597,6 +2598,7 @@ def __init__(
             self.supported_instance_types = resolved_config.get(
                 "supported_inference_instance_types"
             )
+            self.additional_data_sources = resolved_config.get("hosting_additional_data_sources")
 
 
 class DeploymentConfigMetadata(BaseDeploymentConfigDataHolder):

src/sagemaker/jumpstart/utils.py

@@ -1364,20 +1364,3 @@ def wrapped_f(*args, **kwargs):
     if _func is None:
         return wrapper_cache
     return wrapper_cache(_func)
-
-
-def _extract_image_tag_and_version(image_uri: str) -> Tuple[Optional[str], Optional[str]]:
-    """Extract Image tag and version from image URI.
-
-    Args:
-        image_uri (str): Image URI.
-
-    Returns:
-        Tuple[Optional[str], Optional[str]]: The tag and version of the image.
-    """
-    if image_uri is None:
-        return None, None
-
-    tag = image_uri.split(":")[-1]
-
-    return tag, tag.split("-")[0]

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -38,9 +38,11 @@
     SkipTuningComboException,
 )
 from sagemaker.serve.utils.optimize_utils import (
-    _is_compatible_with_optimization_job,
     _extract_model_source,
     _update_environment_variables,
+    _extract_speculative_draft_model_provider,
+    _is_image_compatible_with_optimization_job,
+    _validate_optimization_inputs,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -628,7 +630,7 @@ def _build_for_jumpstart(self):
 
     def _optimize_for_jumpstart(
         self,
-        output_path: str,
+        output_path: Optional[str] = None,
         instance_type: Optional[str] = None,
         role: Optional[str] = None,
         tags: Optional[Tags] = None,
@@ -645,7 +647,7 @@ def _optimize_for_jumpstart(
         """Runs a model optimization job.
 
         Args:
-            output_path (str): Specifies where to store the compiled/quantized model.
+            output_path (Optional[str]): Specifies where to store the compiled/quantized model.
             instance_type (Optional[str]): Target deployment instance type that
                 the model is optimized for.
             role (Optional[str]): Execution role. Defaults to ``None``.
@@ -673,40 +675,30 @@ def _optimize_for_jumpstart(
         """
         if self._is_gated_model() and accept_eula is not True:
             raise ValueError(
-                f"ValueError: Model '{self.model}' "
-                f"requires accepting end-user license agreement (EULA)."
+                f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
+        _validate_optimization_inputs(
+            output_path, instance_type, quantization_config, compilation_config
+        )
+
         optimization_env_vars = None
         pysdk_model_env_vars = None
         model_source = _extract_model_source(self.pysdk_model.model_data, accept_eula)
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
-            optimization_env_vars = self.pysdk_model.deployment_config.get("DeploymentArgs").get(
-                "Environment"
-            )
+            optimization_env_vars = self.pysdk_model.deployment_config.get(
+                "DeploymentArgs", {}
+            ).get("Environment")
         else:
-            image_uri = None
-            if quantization_config and quantization_config.get("Image"):
-                image_uri = quantization_config.get("Image")
-            elif compilation_config and compilation_config.get("Image"):
-                image_uri = compilation_config.get("Image")
-            instance_type = (
-                instance_type
-                or self.pysdk_model.deployment_config.get("DeploymentArgs").get("InstanceType")
-                or _get_nb_instance()
-            )
-            if not _is_compatible_with_optimization_job(instance_type, image_uri):
-                deployment_config = self._find_compatible_deployment_config(None)
-                if deployment_config:
-                    optimization_env_vars = deployment_config.get("DeploymentArgs").get(
-                        "Environment"
-                    )
-                    self.pysdk_model.set_deployment_config(
-                        config_name=deployment_config.get("DeploymentConfigName"),
-                        instance_type=deployment_config.get("InstanceType"),
-                    )
+            deployment_config = self._find_compatible_deployment_config(None)
+            if deployment_config:
+                optimization_env_vars = deployment_config.get("DeploymentArgs").get("Environment")
+                self.pysdk_model.set_deployment_config(
+                    config_name=deployment_config.get("DeploymentConfigName"),
+                    instance_type=deployment_config.get("InstanceType"),
+                )
 
         optimization_env_vars = _update_environment_variables(optimization_env_vars, env_vars)
 
@@ -736,7 +728,7 @@ def _optimize_for_jumpstart(
         }
 
         if optimization_env_vars:
-            create_optimization_job_args["Environment"] = optimization_env_vars
+            create_optimization_job_args["OptimizationEnvironment"] = optimization_env_vars
         if max_runtime_in_sec:
             create_optimization_job_args["StoppingCondition"] = {
                 "MaxRuntimeInSeconds": max_runtime_in_sec
@@ -766,18 +758,26 @@ def _is_gated_model(self, model=None) -> bool:
         return "private" in s3_uri
 
     def _set_additional_model_source(
-        self, speculative_decoding_config: Optional[Dict[str, Any]] = None
+        self,
+        speculative_decoding_config: Optional[Dict[str, Any]] = None,
+        accept_eula: Optional[bool] = None,
     ) -> None:
         """Set Additional Model Source to ``this`` model.
 
         Args:
             speculative_decoding_config (Optional[Dict[str, Any]]): Speculative decoding config.
+            accept_eula (Optional[bool]): For models that require a Model Access Config.
         """
         if speculative_decoding_config:
-            model_provider: str = speculative_decoding_config["ModelProvider"]
+            model_provider = _extract_speculative_draft_model_provider(speculative_decoding_config)
 
             if model_provider.lower() == "sagemaker":
-                if not self._is_speculation_enabled(self.pysdk_model.deployment_config):
+                if (
+                    self.pysdk_model.deployment_config.get("DeploymentArgs", {}).get(
+                        "AdditionalDataSources"
+                    )
+                    is None
+                ):
                     deployment_config = self._find_compatible_deployment_config(
                         speculative_decoding_config
                     )
@@ -786,21 +786,30 @@ def _set_additional_model_source(
                             config_name=deployment_config.get("DeploymentConfigName"),
                             instance_type=deployment_config.get("InstanceType"),
                         )
-                        self.pysdk_model.add_tags(
-                            {"key": Tag.SPECULATIVE_DRAFT_MODL_PROVIDER, "value": "sagemaker"},
-                        )
                     else:
                         raise ValueError(
                             "Cannot find deployment config compatible for optimization job."
                         )
+
+                self.pysdk_model.add_tags(
+                    {"key": Tag.SPECULATIVE_DRAFT_MODL_PROVIDER, "value": "sagemaker"},
+                )
             else:
                 s3_uri = speculative_decoding_config.get("ModelSource")
                 if not s3_uri:
                     raise ValueError("Custom S3 Uri cannot be none.")
 
-                self.pysdk_model.additional_model_data_sources["speculative_decoding"][0][
-                    "s3_data_source"
-                ]["s3_uri"] = s3_uri
+                # TODO: Set correct channel name.
+                additional_model_data_source = {
+                    "ChannelName": "DraftModelName",
+                    "S3DataSource": {"S3Uri": s3_uri},
+                }
+                if accept_eula:
+                    additional_model_data_source["S3DataSource"]["ModelAccessConfig"] = {
+                        "ACCEPT_EULA": True
+                    }
+
+                self.pysdk_model.additional_model_data_sources = [additional_model_data_source]
                 self.pysdk_model.add_tags(
                     {"key": Tag.SPECULATIVE_DRAFT_MODL_PROVIDER, "value": "customer"},
                 )
@@ -816,36 +825,20 @@ def _find_compatible_deployment_config(
         Returns:
             Optional[Dict[str, Any]]: A compatible model deployment config for optimization job.
         """
+        model_provider = _extract_speculative_draft_model_provider(speculative_decoding_config)
         for deployment_config in self.pysdk_model.list_deployment_configs():
-            instance_type = deployment_config.get("deployment_config").get("InstanceType")
-            image_uri = deployment_config.get("deployment_config").get("ImageUri")
-
-            if _is_compatible_with_optimization_job(instance_type, image_uri):
-                if not speculative_decoding_config:
-                    return deployment_config
+            image_uri = deployment_config.get("deployment_config", {}).get("ImageUri")
 
-                if self._is_speculation_enabled(deployment_config):
+            if _is_image_compatible_with_optimization_job(image_uri):
+                if (
+                    model_provider == "sagemaker"
+                    and deployment_config.get("DeploymentArgs", {}).get("AdditionalDataSources")
+                ) or model_provider == "custom":
                     return deployment_config
 
-        return None
-
-    def _is_speculation_enabled(self, deployment_config: Optional[Dict[str, Any]]) -> bool:
-        """Checks whether speculative is enabled for the given deployment config.
+        # There's no matching config from jumpstart to add sagemaker draft model location
+        if model_provider == "sagemaker":
+            return None
 
-        Args:
-            deployment_config (Dict[str, Any]): A deployment config.
-
-        Returns:
-            bool: Whether speculative is enabled for this deployment config.
-        """
-        if deployment_config is None:
-            return False
-
-        acceleration_configs = deployment_config.get("AccelerationConfigs")
-        if acceleration_configs:
-            for acceleration_config in acceleration_configs:
-                if acceleration_config.get(
-                    "type", "default"
-                ).lower() == "speculative" and acceleration_config.get("enabled"):
-                    return True
-        return False
+        # fall back to the default jumpstart model deployment config for optimization job
+        return self.pysdk_model.deployment_config

src/sagemaker/serve/utils/optimize_utils.py

@@ -19,7 +19,6 @@
 
 from sagemaker import Model
 from sagemaker.enums import Tag
-from sagemaker.fw_utils import _is_gpu_instance
 
 
 logger = logging.getLogger(__name__)
@@ -42,30 +41,19 @@ def _is_inferentia_or_trainium(instance_type: Optional[str]) -> bool:
     return False
 
 
-def _is_compatible_with_optimization_job(
-    instance_type: Optional[str], image_uri: Optional[str]
-) -> bool:
+def _is_image_compatible_with_optimization_job(image_uri: Optional[str]) -> bool:
     """Checks whether an instance is compatible with an optimization job.
 
     Args:
-        instance_type (str): The instance type used for the compilation job.
         image_uri (str): The image URI of the optimization job.
 
     Returns:
         bool: Whether the given instance type is compatible with an optimization job.
     """
-    if not instance_type:
-        return False
-
-    compatible_image = True
-    if image_uri:
-        compatible_image = "djl-inference:" in image_uri and (
-            "-lmi" in image_uri or "-neuronx-" in image_uri
-        )
-
-    return (
-        _is_gpu_instance(instance_type) or _is_inferentia_or_trainium(instance_type)
-    ) and compatible_image
+    # TODO: Use specific container type instead.
+    if image_uri is None:
+        return True
+    return "djl-inference:" in image_uri and ("-lmi" in image_uri or "-neuronx-" in image_uri)
 
 
 def _generate_optimized_model(pysdk_model: Model, optimization_response: dict) -> Model:
@@ -89,28 +77,6 @@ def _generate_optimized_model(pysdk_model: Model, optimization_response: dict) -
     return pysdk_model
 
 
-def _is_speculation_enabled(deployment_config: Optional[Dict[str, Any]]) -> bool:
-    """Checks whether speculation is enabled for this deployment config.
-
-    Args:
-        deployment_config (Dict[str, Any]): A deployment config.
-
-    Returns:
-        bool: Whether the speculation is enabled for this deployment config.
-    """
-    if deployment_config is None:
-        return False
-
-    acceleration_configs = deployment_config.get("AccelerationConfigs")
-    if acceleration_configs:
-        for acceleration_config in acceleration_configs:
-            if acceleration_config.get("type").lower() == "speculation" and acceleration_config.get(
-                "enabled"
-            ):
-                return True
-    return False
-
-
 def _extract_model_source(
     model_data: Optional[Union[Dict[str, Any], str]], accept_eula: Optional[bool]
 ) -> Optional[Dict[str, Any]]:
@@ -129,7 +95,6 @@ def _extract_model_source(
     if isinstance(s3_uri, dict):
         s3_uri = s3_uri.get("S3DataSource").get("S3Uri")
 
-    # Todo: Inject fine-tune data source
     model_source = {"S3": {"S3Uri": s3_uri}}
     if accept_eula:
         model_source["S3"]["ModelAccessConfig"] = {"AcceptEula": True}
@@ -154,3 +119,61 @@ def _update_environment_variables(
         else:
             env = new_env
     return env
+
+
+def _extract_speculative_draft_model_provider(
+    speculative_decoding_config: Optional[Dict] = None,
+) -> Optional[str]:
+    """Extracts speculative draft model provider from speculative decoding config.
+
+    Args:
+        speculative_decoding_config (Optional[Dict]): A speculative decoding config.
+
+    Returns:
+        Optional[str]: The speculative draft model provider.
+    """
+    if speculative_decoding_config is None:
+        return None
+
+    if speculative_decoding_config.get(
+        "ModelProvider"
+    ) == "Custom" or speculative_decoding_config.get("ModelSource"):
+        return "custom"
+
+    return "sagemaker"
+
+
+def _validate_optimization_inputs(
+    output_path: Optional[str] = None,
+    instance_type: Optional[str] = None,
+    quantization_config: Optional[Dict] = None,
+    compilation_config: Optional[Dict] = None,
+) -> None:
+    """Validates optimization inputs.
+
+    Args:
+        output_path (Optional[str]): The output path.
+        instance_type (Optional[str]): The instance type.
+        quantization_config (Optional[Dict]): The quantization config.
+        compilation_config (Optional[Dict]): The compilation config.
+
+    Raises:
+        ValueError: If an optimization input is invalid.
+    """
+    if quantization_config and compilation_config:
+        raise ValueError("Quantization config and compilation config are mutually exclusive.")
+
+    instance_type_msg = "Please provide an instance type for %s optimization job."
+    output_path_msg = "Please provide an output path for %s optimization job."
+
+    if quantization_config:
+        if not instance_type:
+            raise ValueError(instance_type_msg.format("quantization"))
+        if not output_path:
+            raise ValueError(output_path_msg.format("quantization"))
+
+    if compilation_config:
+        if not instance_type:
+            raise ValueError(instance_type_msg.format("compilation"))
+        if not output_path:
+            raise ValueError(output_path_msg.format("compilation"))