Refactoring

Author: Jonathan Makunga

src/sagemaker/jumpstart/types.py

@@ -2569,6 +2569,8 @@ class DeploymentArgs(BaseDeploymentConfigDataHolder):
         "model_data_download_timeout",
         "container_startup_health_check_timeout",
         "additional_data_sources",
+        "neuron_model_id",
+        "neuron_model_version",
     ]
 
     def __init__(
@@ -2599,6 +2601,8 @@ def __init__(
                 "supported_inference_instance_types"
             )
             self.additional_data_sources = resolved_config.get("hosting_additional_data_sources")
+            self.neuron_model_id = resolved_config.get("hosting_neuron_model_id")
+            self.neuron_model_version = resolved_config.get("hosting_neuron_model_version")
 
 
 class DeploymentConfigMetadata(BaseDeploymentConfigDataHolder):

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -110,6 +110,7 @@ def __init__(self):
         self.prepared_for_mms = None
         self.schema_builder = None
         self.instance_type = None
+        self.nb_instance_type = None
         self.ram_usage_model_load = None
         self.model_hub = None
         self.model_metadata = None
@@ -236,8 +237,8 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
-        if self.instance_type:
-            kwargs.update({"instance_type": self.instance_type})
+        if hasattr(self, "nb_instance_type"):
+            kwargs.update({"instance_type": self.nb_instance_type})
 
         if "mode" in kwargs:
             del kwargs["mode"]
@@ -270,7 +271,7 @@ def _build_for_djl_jumpstart(self):
                 )
             self._prepare_for_mode()
         elif self.mode == Mode.SAGEMAKER_ENDPOINT and hasattr(self, "prepared_for_djl"):
-            self.instance_type = self.instance_type or _get_nb_instance()
+            self.nb_instance_type = self.instance_type or _get_nb_instance()
             self.pysdk_model.model_data, env = self._prepare_for_mode()
 
         self.pysdk_model.env.update(env)
@@ -695,25 +696,29 @@ def _optimize_for_jumpstart(
                 f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
-        optimization_env_vars = None
-        pysdk_model_env_vars = None
-        model_source = _generate_model_source(self.pysdk_model.model_data, accept_eula)
+        if compilation_config:
+            neuro_model_id = self.pysdk_model.deployment_config.get("DeploymentArgs").get(
+                "NeuronModelId"
+            )
+            self.model = neuro_model_id
+            self.pysdk_model = self._create_pre_trained_js_model()
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
-            optimization_env_vars = self.pysdk_model.deployment_config.get(
-                "DeploymentArgs", {}
-            ).get("Environment")
         else:
             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"),
                 )
 
+        model_source = _generate_model_source(self.pysdk_model.model_data, accept_eula)
+        optimization_env_vars = self.pysdk_model.deployment_config.get("DeploymentArgs", {}).get(
+            "Environment"
+        )
         optimization_env_vars = _update_environment_variables(optimization_env_vars, env_vars)
+        pysdk_model_env_vars = env_vars
 
         optimization_config = {}
         if quantization_config:
@@ -730,6 +735,10 @@ def _optimize_for_jumpstart(
         output_config = {"S3OutputLocation": output_path}
         if kms_key:
             output_config["KmsKeyId"] = kms_key
+        if not instance_type:
+            instance_type = self.pysdk_model.deployment_config.get("DeploymentArgs").get(
+                "InstanceType"
+            )
 
         create_optimization_job_args = {
             "OptimizationJobName": job_name,

src/sagemaker/serve/builder/model_builder.py

@@ -65,7 +65,6 @@
 from sagemaker.serve.utils.lineage_utils import _maintain_lineage_tracking_for_mlflow_model
 from sagemaker.serve.utils.optimize_utils import (
     _generate_optimized_model,
-    _validate_optimization_inputs,
 )
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import (
@@ -238,7 +237,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers, TensorflowServing,
         metadata={"help": "Define the s3 location where you want to upload the model package"},
     )
     instance_type: Optional[str] = field(
-        default="ml.c5.xlarge",
+        default=None,
         metadata={"help": "Define the instance_type of the endpoint"},
     )
     schema_builder: Optional[SchemaBuilder] = field(
@@ -1022,9 +1021,8 @@ def _model_builder_optimize_wrapper(
         Returns:
             Model: A deployable ``Model`` object.
         """
-        _validate_optimization_inputs(
-            output_path, instance_type, quantization_config, compilation_config
-        )
+        if quantization_config and compilation_config:
+            raise ValueError("Quantization config and compilation config are mutually exclusive.")
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
 

src/sagemaker/serve/utils/optimize_utils.py

@@ -160,42 +160,6 @@ def _extracts_and_validates_speculative_model_source(
     return s3_uri
 
 
-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"))
-
-
 def _generate_channel_name(additional_model_data_sources: Optional[List[Dict]]) -> str:
     """Generates a channel name.
 

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

@@ -22,7 +22,6 @@
     _update_environment_variables,
     _is_image_compatible_with_optimization_job,
     _extract_speculative_draft_model_provider,
-    _validate_optimization_inputs,
     _extracts_and_validates_speculative_model_source,
     _is_s3_uri,
     _generate_additional_model_data_sources,
@@ -168,31 +167,6 @@ def test_extract_speculative_draft_model_provider(
     )
 
 
-@pytest.mark.parametrize(
-    "output_path, instance, quantization_config, compilation_config",
-    [
-        (
-            None,
-            None,
-            {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}},
-            {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}},
-        ),
-        (None, None, {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}}, None),
-        (None, None, None, {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}}),
-        ("output_path", None, None, {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}}),
-        (None, "instance_type", None, {"OverrideEnvironment": {"TENSOR_PARALLEL_DEGREE": 4}}),
-    ],
-)
-def test_validate_optimization_inputs(
-    output_path, instance, quantization_config, compilation_config
-):
-
-    with pytest.raises(ValueError):
-        _validate_optimization_inputs(
-            output_path, instance, quantization_config, compilation_config
-        )
-
-
 def test_extract_speculative_draft_model_s3_uri():
     res = _extracts_and_validates_speculative_model_source({"ModelSource": "s3://"})
     assert res == "s3://"