Refactoring

Author: Jonathan Makunga

src/sagemaker/jumpstart/types.py

@@ -1098,6 +1098,8 @@ class JumpStartMetadataBaseFields(JumpStartDataHolderType):
         "gated_bucket",
         "model_subscription_link",
         "hosting_additional_data_sources",
+        "hosting_neuron_model_id",
+        "hosting_neuron_model_version",
     ]
 
     def __init__(self, fields: Dict[str, Any]):
@@ -1208,6 +1210,10 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
             if json_obj.get("hosting_additional_data_sources")
             else None
         )
+        self.hosting_neuron_model_id: Optional[str] = json_obj.get("hosting_neuron_model_id")
+        self.hosting_neuron_model_version: Optional[str] = json_obj.get(
+            "hosting_neuron_model_version"
+        )
 
         if self.training_supported:
             self.training_ecr_specs: Optional[JumpStartECRSpecs] = (
@@ -2569,8 +2575,6 @@ class DeploymentArgs(BaseDeploymentConfigDataHolder):
         "model_data_download_timeout",
         "container_startup_health_check_timeout",
         "additional_data_sources",
-        "neuron_model_id",
-        "neuron_model_version",
     ]
 
     def __init__(
@@ -2601,8 +2605,6 @@ 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

@@ -139,9 +139,7 @@ def _is_jumpstart_model_id(self) -> bool:
 
     def _create_pre_trained_js_model(self) -> Type[Model]:
         """Placeholder docstring"""
-        pysdk_model = JumpStartModel(
-            self.model, vpc_config=self.vpc_config, instance_type=self.instance_type
-        )
+        pysdk_model = JumpStartModel(self.model, vpc_config=self.vpc_config)
         pysdk_model.sagemaker_session = self.sagemaker_session
 
         self._original_deploy = pysdk_model.deploy
@@ -696,12 +694,12 @@ def _optimize_for_jumpstart(
                 f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
+        optimization_env_vars = env_vars
+        pysdk_model_env_vars = env_vars
+
         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()
+            neuron_env = self._get_neuron_model_env_vars(instance_type)
+            optimization_env_vars = _update_environment_variables(neuron_env, optimization_env_vars)
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
@@ -714,11 +712,6 @@ def _optimize_for_jumpstart(
                 )
 
         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:
@@ -874,3 +867,26 @@ def _find_compatible_deployment_config(
 
         # fall back to the default jumpstart model deployment config for optimization job
         return self.pysdk_model.deployment_config
+
+    def _get_neuron_model_env_vars(
+        self, instance_type: Optional[str] = None
+    ) -> Optional[Dict[str, Any]]:
+        """Gets Neuron model env vars.
+
+        Args:
+            instance_type (Optional[str]): Instance type.
+
+        Returns:
+            Optional[Dict[str, Any]]: Neuron Model environment variables.
+        """
+        metadata_config = self.pysdk_model._metadata_configs.get(self.pysdk_model.config_name)
+        resolve_config = metadata_config.resolved_config
+        if instance_type not in resolve_config.get("supported_inference_instance_types", []):
+            neuro_model_id = resolve_config.get("hosting_neuron_model_id")
+            neuro_model_version = resolve_config.get("hosting_neuron_model_version")
+            if neuro_model_id:
+                job_model = JumpStartModel(
+                    neuro_model_id, model_version=neuro_model_version, vpc_config=self.vpc_config
+                )
+                return job_model.env
+        return None

src/sagemaker/serve/builder/model_builder.py

@@ -237,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=None,
+        default="ml.c5.xlarge",
         metadata={"help": "Define the instance_type of the endpoint"},
     )
     schema_builder: Optional[SchemaBuilder] = field(
@@ -1055,6 +1055,6 @@ def _model_builder_optimize_wrapper(
         if input_args:
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
-            self.pysdk_model = _generate_optimized_model(self.pysdk_model, job_status)
+            return _generate_optimized_model(self.pysdk_model, job_status)
 
         return self.pysdk_model