Fix public optimize api signature (aws#1507)

* Fix public optimize api signature

* JS Compilation fix

* Refactoring

* Refactoring

---------

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

Author: makungaj1

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -47,6 +47,7 @@
     _is_optimized,
     _custom_speculative_decoding,
     SPECULATIVE_DRAFT_MODEL,
+    _is_inferentia_or_trainium,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -714,10 +715,25 @@ def _optimize_for_jumpstart(
                 f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
+        is_compilation = (quantization_config is None) and (
+            (compilation_config is not None) or _is_inferentia_or_trainium(instance_type)
+        )
+
         pysdk_model_env_vars = dict()
-        if compilation_config:
+        if is_compilation:
             pysdk_model_env_vars = self._get_neuron_model_env_vars(instance_type)
 
+        optimization_config, override_env = _extract_optimization_config_and_env(
+            quantization_config, compilation_config
+        )
+        if not optimization_config and is_compilation:
+            override_env = override_env or pysdk_model_env_vars
+            optimization_config = {
+                "ModelCompilationConfig": {
+                    "OverrideEnvironment": override_env,
+                }
+            }
+
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
         else:
@@ -732,10 +748,6 @@ def _optimize_for_jumpstart(
         model_source = _generate_model_source(self.pysdk_model.model_data, accept_eula)
         optimization_env_vars = _update_environment_variables(pysdk_model_env_vars, env_vars)
 
-        optimization_config, override_env = _extract_optimization_config_and_env(
-            quantization_config, compilation_config
-        )
-
         output_config = {"S3OutputLocation": output_path}
         if kms_key:
             output_config["KmsKeyId"] = kms_key
@@ -775,7 +787,7 @@ def _optimize_for_jumpstart(
                     "AcceptEula": True
                 }
 
-        if quantization_config or compilation_config:
+        if quantization_config or is_compilation:
             self.pysdk_model.env = _update_environment_variables(
                 optimization_env_vars, override_env
             )

src/sagemaker/serve/builder/model_builder.py

@@ -1083,25 +1083,47 @@ def _try_fetch_gpu_info(self):
                 f"Unable to determine single GPU size for instance: [{self.instance_type}]"
             )
 
-    def optimize(self, *args, **kwargs) -> Model:
-        """Runs a model optimization job.
+    def optimize(
+        self,
+        output_path: Optional[str] = None,
+        instance_type: Optional[str] = None,
+        role_arn: 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] = 36000,
+        sagemaker_session: Optional[Session] = None,
+    ) -> Model:
+        """Create an optimized deployable ``Model`` instance with ``ModelBuilder``.
 
         Args:
-            instance_type (Optional[str]): Target deployment instance type that the
-                model is optimized for.
-            output_path (Optional[str]): Specifies where to store the compiled/quantized model.
-            role_arn (Optional[str]): Execution role. Defaults to ``None``.
+            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_arn (Optional[str]): Execution role arn. 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``.
+                36000 seconds.
             sagemaker_session (Optional[Session]): Session object which manages interactions
                 with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
                 function creates one using the default AWS configuration chain.
@@ -1113,7 +1135,22 @@ def optimize(self, *args, **kwargs) -> Model:
         # need to get telemetry_opt_out info before telemetry decorator is called
         self.serve_settings = self._get_serve_setting()
 
-        return self._model_builder_optimize_wrapper(*args, **kwargs)
+        return self._model_builder_optimize_wrapper(
+            output_path=output_path,
+            instance_type=instance_type,
+            role_arn=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,
+            sagemaker_session=sagemaker_session,
+        )
 
     @_capture_telemetry("optimize")
     def _model_builder_optimize_wrapper(
@@ -1178,10 +1215,8 @@ def _model_builder_optimize_wrapper(
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
 
-        if instance_type:
-            self.instance_type = instance_type
-        if role_arn:
-            self.role_arn = role_arn
+        self.instance_type = instance_type or self.instance_type
+        self.role_arn = role_arn or self.role_arn
 
         self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
         job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
@@ -1266,7 +1301,7 @@ def _optimize_for_hf(
                 ``None``.
 
         Returns:
-            Dict[str, Any]: Model optimization job input arguments.
+            Optional[Dict[str, Any]]: Model optimization job input arguments.
         """
         if self.model_server != ModelServer.DJL_SERVING:
             logger.info("Overwriting model server to DJL.")
@@ -1275,6 +1310,10 @@ def _optimize_for_hf(
         self.role_arn = role_arn or self.role_arn
         self.instance_type = instance_type or self.instance_type
 
+        self.pysdk_model = _custom_speculative_decoding(
+            self.pysdk_model, speculative_decoding_config, False
+        )
+
         if quantization_config or compilation_config:
             create_optimization_job_args = {
                 "OptimizationJobName": job_name,
@@ -1290,10 +1329,6 @@ def _optimize_for_hf(
             model_source = _generate_model_source(self.pysdk_model.model_data, False)
             create_optimization_job_args["ModelSource"] = model_source
 
-            self.pysdk_model = _custom_speculative_decoding(
-                self.pysdk_model, speculative_decoding_config, False
-            )
-
             optimization_config, override_env = _extract_optimization_config_and_env(
                 quantization_config, compilation_config
             )

src/sagemaker/serve/utils/optimize_utils.py

@@ -26,6 +26,23 @@
 SPECULATIVE_DRAFT_MODEL = "/opt/ml/additional-model-data-sources"
 
 
+def _is_inferentia_or_trainium(instance_type: Optional[str]) -> bool:
+    """Checks whether an instance is compatible with Inferentia.
+
+    Args:
+        instance_type (str): The instance type used for the compilation job.
+
+    Returns:
+        bool: Whether the given instance type is Inferentia or Trainium.
+    """
+    if isinstance(instance_type, str):
+        match = re.match(r"^ml[\._]([a-z\d]+)\.?\w*$", instance_type)
+        if match:
+            if match[1].startswith("inf") or match[1].startswith("trn"):
+                return True
+    return False
+
+
 def _is_image_compatible_with_optimization_job(image_uri: Optional[str]) -> bool:
     """Checks whether an instance is compatible with an optimization job.
 
@@ -169,11 +186,11 @@ def _extracts_and_validates_speculative_model_source(
     Raises:
         ValueError: If model source is none.
     """
-    s3_uri: str = speculative_decoding_config.get("ModelSource")
+    model_source: str = speculative_decoding_config.get("ModelSource")
 
-    if not s3_uri:
+    if not model_source:
         raise ValueError("ModelSource must be provided in speculative decoding config.")
-    return s3_uri
+    return model_source
 
 
 def _generate_channel_name(additional_model_data_sources: Optional[List[Dict]]) -> str:

tests/unit/sagemaker/serve/utils/test_optimize_utils.py

@@ -31,6 +31,7 @@
     _normalize_local_model_path,
     _is_optimized,
     _custom_speculative_decoding,
+    _is_inferentia_or_trainium,
 )
 
 mock_optimization_job_output = {
@@ -81,6 +82,18 @@
 }
 
 
+@pytest.mark.parametrize(
+    "instance, expected",
+    [
+        ("ml.trn1.2xlarge", True),
+        ("ml.inf2.xlarge", True),
+        ("ml.c7gd.4xlarge", False),
+    ],
+)
+def test_is_inferentia_or_trainium(instance, expected):
+    assert _is_inferentia_or_trainium(instance) == expected
+
+
 @pytest.mark.parametrize(
     "image_uri, expected",
     [