Feat: Add optimize to ModelBuilder JS (aws#1474)

* QS JS vanilla model

* Use Alt config for Optimization

* JS Optimize

* Resolve config

* inject additional tags

* Inject tags

* Refactoring

* Refactoring

* Filter Deployment config

* Refactoring

* Refactoring

* Refactoring

* Refactoring

---------

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

Author: makungaj1

src/sagemaker/enums.py

@@ -40,3 +40,9 @@ class RoutingStrategy(Enum):
     """The endpoint routes requests to the specific instances that have
     more capacity to process them.
     """
+
+
+class Tag(str, Enum):
+    """Enum class for tag keys to apply to models."""
+
+    OPTIMIZATION_JOB_NAME = "sagemaker-sdk:optimization-job-name"

src/sagemaker/jumpstart/utils.py

@@ -1336,3 +1336,20 @@ 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/model.py

@@ -404,6 +404,18 @@ def __init__(
         self.content_types = None
         self.response_types = None
         self.accept_eula = None
+        self._tags: Optional[Tags] = None
+
+    def add_tags(self, tags: Tags) -> None:
+        """Add tags to this ``Model``
+
+        Args:
+            tags (Tags): Tags to add.
+        """
+        if self._tags and tags:
+            self._tags.update(tags)
+        else:
+            self._tags = tags
 
     @runnable_by_pipeline
     def register(
@@ -1457,7 +1469,8 @@ def deploy(
             sagemaker_session=self.sagemaker_session,
         )
 
-        tags = format_tags(tags)
+        self.add_tags(tags)
+        tags = format_tags(self._tags)
 
         if (
             getattr(self.sagemaker_session, "settings", None) is not None

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -19,6 +19,8 @@
 from typing import Type, Any, List, Dict, Optional
 import logging
 
+from sagemaker.jumpstart import enums
+from sagemaker.jumpstart.utils import verify_model_region_and_return_specs, get_eula_message
 from sagemaker.model import Model
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
@@ -33,6 +35,11 @@
     LocalModelLoadException,
     SkipTuningComboException,
 )
+from sagemaker.serve.utils.optimize_utils import (
+    _extract_supported_deployment_config,
+    _is_speculation_enabled,
+    _is_compatible_with_optimization_job,
+)
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
     TgiLocalModePredictor,
@@ -53,6 +60,7 @@
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.jumpstart.model import JumpStartModel
+from sagemaker.utils import Tags
 
 _DJL_MODEL_BUILDER_ENTRY_POINT = "inference.py"
 _NO_JS_MODEL_EX = "HuggingFace JumpStart Model ID not detected. Building for HuggingFace Model ID."
@@ -564,6 +572,148 @@ def _build_for_jumpstart(self):
 
         return self.pysdk_model
 
+    def _optimize_for_jumpstart(
+        self,
+        output_path: str,
+        instance_type: Optional[str] = None,
+        role: Optional[str] = None,
+        tags: Optional[Tags] = None,
+        job_name: Optional[str] = None,
+        accept_eula: Optional[bool] = None,
+        quantization_config: Optional[Dict] = None,
+        compilation_config: Optional[Dict] = None,
+        speculative_decoding_config: Optional[Dict] = None,
+        env_vars: Optional[Dict] = None,
+        vpc_config: Optional[Dict] = None,
+        kms_key: Optional[str] = None,
+        max_runtime_in_sec: Optional[int] = None,
+    ) -> None:
+        """Runs a model optimization job.
+
+        Args:
+            output_path (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``.
+            tags (Optional[Tags]): Tags for labeling a model optimization job. Defaults to ``None``.
+            job_name (Optional[str]): The name of the model optimization job. Defaults to ``None``.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False 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
+                models require. (Default: None).
+            quantization_config (Optional[Dict]): Quantization configuration. Defaults to ``None``.
+            compilation_config (Optional[Dict]): Compilation configuration. Defaults to ``None``.
+            speculative_decoding_config (Optional[Dict]): Speculative decoding configuration.
+                Defaults to ``None``
+            env_vars (Optional[Dict]): Additional environment variables to run the optimization
+                container. Defaults to ``None``.
+            vpc_config (Optional[Dict]): The VpcConfig set on the model. Defaults to ``None``.
+            kms_key (Optional[str]): KMS key ARN used to encrypt the model artifacts when uploading
+                to S3. Defaults to ``None``.
+            max_runtime_in_sec (Optional[int]): Maximum job execution time in seconds. Defaults to
+                ``None``.
+        """
+        model_specs = verify_model_region_and_return_specs(
+            region=self.sagemaker_session.boto_region_name,
+            model_id=self.pysdk_model.model_id,
+            version=self.pysdk_model.model_version,
+            sagemaker_session=self.sagemaker_session,
+            scope=enums.JumpStartScriptScope.INFERENCE,
+            model_type=self.pysdk_model.model_type,
+        )
+
+        if model_specs.is_gated_model() and accept_eula is not True:
+            raise ValueError(get_eula_message(model_specs, self.sagemaker_session.boto_region_name))
+
+        if not (self.pysdk_model.model_data and self.pysdk_model.model_data.get("S3DataSource")):
+            raise ValueError("Model Optimization Job only supports model backed by S3.")
+
+        has_alternative_config = self.pysdk_model.deployment_config is not None
+        merged_env_vars = None
+        # TODO: Match Optimization Input Schema
+        model_source = {
+            "S3": {"S3Uri": self.pysdk_model.model_data.get("S3DataSource").get("S3Uri")},
+            "SageMakerModel": {"ModelName": self.model},
+        }
+
+        if has_alternative_config:
+            image_uri = self.pysdk_model.deployment_config.get("DeploymentArgs").get("ImageUri")
+            instance_type = self.pysdk_model.deployment_config.get("InstanceType")
+        else:
+            image_uri = self.pysdk_model.image_uri
+
+        if not _is_compatible_with_optimization_job(instance_type, image_uri) or (
+            speculative_decoding_config
+            and not _is_speculation_enabled(self.pysdk_model.deployment_config)
+        ):
+            deployment_config = _extract_supported_deployment_config(
+                self.pysdk_model.list_deployment_configs(), speculative_decoding_config is None
+            )
+
+            if deployment_config:
+                self.pysdk_model.set_deployment_config(
+                    config_name=deployment_config.get("DeploymentConfigName"),
+                    instance_type=deployment_config.get("InstanceType"),
+                )
+                merged_env_vars = self.pysdk_model.deployment_config.get("Environment")
+
+                if speculative_decoding_config:
+                    # TODO: Match Optimization Input Schema
+                    s3 = {
+                        "S3Uri": self.pysdk_model.additional_model_data_sources[
+                            "SpeculativeDecoding"
+                        ][0]["S3DataSource"]["S3Uri"]
+                    }
+                    model_source["S3"].update(s3)
+            elif speculative_decoding_config:
+                raise ValueError("Can't find deployment config for model optimization job.")
+
+        optimization_config = {}
+        if env_vars:
+            if merged_env_vars:
+                merged_env_vars.update(env_vars)
+            else:
+                merged_env_vars = env_vars
+        if quantization_config:
+            optimization_config["ModelQuantizationConfig"] = quantization_config
+        if compilation_config:
+            optimization_config["ModelCompilationConfig"] = compilation_config
+
+        if accept_eula:
+            self.pysdk_model.accept_eula = accept_eula
+            self.pysdk_model.model_data["S3DataSource"].update(
+                {"ModelAccessConfig": {"AcceptEula": accept_eula}}
+            )
+            model_source["S3"].update({"ModelAccessConfig": {"AcceptEula": accept_eula}})
+
+        output_config = {"S3OutputLocation": output_path}
+        if kms_key:
+            output_config["KmsKeyId"] = kms_key
+
+        create_optimization_job_args = {
+            "OptimizationJobName": job_name,
+            "ModelSource": model_source,
+            "DeploymentInstanceType": instance_type,
+            "Environment": merged_env_vars,
+            "OptimizationConfigs": [optimization_config],
+            "OutputConfig": output_config,
+            "RoleArn": role,
+        }
+
+        if max_runtime_in_sec:
+            create_optimization_job_args["StoppingCondition"] = {
+                "MaxRuntimeInSeconds": max_runtime_in_sec
+            }
+        if tags:
+            create_optimization_job_args["Tags"] = tags
+        if vpc_config:
+            create_optimization_job_args["VpcConfig"] = vpc_config
+
+        self.sagemaker_session.sagemaker_client.create_optimization_job(
+            **create_optimization_job_args
+        )
+
     def _is_gated_model(self, model) -> bool:
         """Determine if ``this`` Model is Gated
 

src/sagemaker/serve/builder/model_builder.py

@@ -62,10 +62,7 @@
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.lineage_utils import _maintain_lineage_tracking_for_mlflow_model
-from sagemaker.serve.utils.optimize_utils import (
-    _is_compatible_with_compilation,
-    _poll_optimization_job,
-)
+from sagemaker.serve.utils.optimize_utils import _poll_optimization_job, _generate_optimized_model
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import (
     _get_gpu_info,
@@ -961,13 +958,15 @@ def optimize(self, *args, **kwargs) -> Type[Model]:
     @_capture_telemetry("optimize")
     def _model_builder_optimize_wrapper(
         self,
-        instance_type: str,
         output_path: str,
+        instance_type: Optional[str] = None,
         role: Optional[str] = None,
         tags: Optional[Tags] = None,
         job_name: Optional[str] = None,
+        accept_eula: Optional[bool] = None,
         quantization_config: Optional[Dict] = None,
         compilation_config: Optional[Dict] = None,
+        speculative_decoding_config: Optional[Dict] = None,
         env_vars: Optional[Dict] = None,
         vpc_config: Optional[Dict] = None,
         kms_key: Optional[str] = None,
@@ -977,13 +976,20 @@ def _model_builder_optimize_wrapper(
         """Runs a model optimization job.
 
         Args:
-            instance_type (str): Target deployment instance type that the model is optimized for.
             output_path (str): Specifies where to store the compiled/quantized model.
+            instance_type (str): Target deployment instance type that the model is optimized for.
             role (Optional[str]): Execution role. Defaults to ``None``.
             tags (Optional[Tags]): Tags for labeling a model optimization job. Defaults to ``None``.
             job_name (Optional[str]): The name of the model optimization job. Defaults to ``None``.
+            accept_eula (bool): For models that require a Model Access Config, specify True or
+                False 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
+                models require. (Default: None).
             quantization_config (Optional[Dict]): Quantization configuration. Defaults to ``None``.
             compilation_config (Optional[Dict]): Compilation configuration. Defaults to ``None``.
+            speculative_decoding_config (Optional[Dict]): Speculative decoding configuration.
+                Defaults to ``None``
             env_vars (Optional[Dict]): Additional environment variables to run the optimization
                 container. Defaults to ``None``.
             vpc_config (Optional[Dict]): The VpcConfig set on the model. Defaults to ``None``.
@@ -999,57 +1005,39 @@ def _model_builder_optimize_wrapper(
             Type[Model]: A deployable ``Model`` object.
         """
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
+        self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
+        job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
 
-        # TODO: inject actual model source location based on different scenarios
-        model_source = {"S3": {"S3Uri": self.model_path, "ModelAccessConfig": {"AcceptEula": True}}}
-
-        optimization_configs = []
-        if quantization_config:
-            optimization_configs.append({"ModelQuantizationConfig": quantization_config})
-        if compilation_config:
-            if _is_compatible_with_compilation(instance_type):
-                optimization_configs.append({"ModelCompilationConfig": compilation_config})
-            else:
-                logger.warning(
-                    "Model compilation is currently only supported for Inferentia and Trainium"
-                    "instances, ignoring `compilation_config'."
-                )
+        if self._is_jumpstart_model_id():
+            self._optimize_for_jumpstart(
+                output_path=output_path,
+                instance_type=instance_type,
+                role=role if role else self.role_arn,
+                tags=tags,
+                job_name=job_name,
+                accept_eula=accept_eula,
+                quantization_config=quantization_config,
+                compilation_config=compilation_config,
+                speculative_decoding_config=speculative_decoding_config,
+                env_vars=env_vars,
+                vpc_config=vpc_config,
+                kms_key=kms_key,
+                max_runtime_in_sec=max_runtime_in_sec,
+            )
 
-        output_config = {"S3OutputLocation": output_path}
-        if kms_key:
-            output_config["KmsKeyId"] = kms_key
+        # TODO: use the wait for job pattern similar to
+        #  https://quip-amazon.com/TKaPAhJck5sD/PySDK-Model-Optimization#temp:C:YcX3f2b103dabb4431090568bca2
+        if not _poll_optimization_job(job_name, self.sagemaker_session):
+            raise Exception("Optimization job timed out.")
 
-        job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
-        create_optimization_job_args = {
-            "OptimizationJobName": job_name,
-            "ModelSource": model_source,
-            "DeploymentInstanceType": instance_type,
-            "OptimizationConfigs": optimization_configs,
-            "OutputConfig": output_config,
-            "RoleArn": role or self.role_arn,
-        }
-
-        if env_vars:
-            create_optimization_job_args["OptimizationEnvironment"] = env_vars
-
-        if max_runtime_in_sec:
-            create_optimization_job_args["StoppingCondition"] = {
-                "MaxRuntimeInSeconds": max_runtime_in_sec
-            }
-
-        # TODO: tag injection if it is a JumpStart model
-        if tags:
-            create_optimization_job_args["Tags"] = tags
-
-        if vpc_config:
-            create_optimization_job_args["VpcConfig"] = vpc_config
-
-        response = self.sagemaker_session.sagemaker_client.create_optimization_job(
-            **create_optimization_job_args
+        describe_optimization_job_res = (
+            self.sagemaker_session.sagemaker_client.describe_optimization_job(
+                OptimizationJobName=job_name
+            )
         )
 
-        if not _poll_optimization_job(job_name, self.sagemaker_session):
-            raise Exception("Optimization job timed out.")
+        self.pysdk_model = _generate_optimized_model(
+            self.pysdk_model, describe_optimization_job_res
+        )
 
-        # TODO: return model created by optimization job
-        return response
+        return self.pysdk_model