JS Optimize api ref

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -109,7 +109,7 @@ def __init__(self):
         self.prepared_for_djl = None
         self.prepared_for_mms = None
         self.schema_builder = None
-        self.nb_instance_type = None
+        self.instance_type = None
         self.ram_usage_model_load = None
         self.model_hub = None
         self.model_metadata = None
@@ -138,7 +138,9 @@ 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)
+        pysdk_model = JumpStartModel(
+            self.model, vpc_config=self.vpc_config, instance_type=self.instance_type
+        )
         pysdk_model.sagemaker_session = self.sagemaker_session
 
         self._original_deploy = pysdk_model.deploy
@@ -234,8 +236,8 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
-        if hasattr(self, "nb_instance_type"):
-            kwargs.update({"instance_type": self.nb_instance_type})
+        if self.instance_type:
+            kwargs.update({"instance_type": self.instance_type})
 
         if "mode" in kwargs:
             del kwargs["mode"]
@@ -268,7 +270,7 @@ def _build_for_djl_jumpstart(self):
                 )
             self._prepare_for_mode()
         elif self.mode == Mode.SAGEMAKER_ENDPOINT and hasattr(self, "prepared_for_djl"):
-            self.nb_instance_type = _get_nb_instance()
+            self.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)
@@ -647,7 +649,7 @@ def _optimize_for_jumpstart(
         self,
         output_path: Optional[str] = None,
         instance_type: Optional[str] = None,
-        role: Optional[str] = None,
+        role_arn: Optional[str] = None,
         tags: Optional[Tags] = None,
         job_name: Optional[str] = None,
         accept_eula: Optional[bool] = None,
@@ -665,7 +667,7 @@ def _optimize_for_jumpstart(
             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``.
+            role_arn (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
@@ -735,7 +737,7 @@ def _optimize_for_jumpstart(
             "DeploymentInstanceType": instance_type,
             "OptimizationConfigs": [optimization_config],
             "OutputConfig": output_config,
-            "RoleArn": role,
+            "RoleArn": role_arn,
         }
 
         if optimization_env_vars:

src/sagemaker/serve/builder/model_builder.py

@@ -949,7 +949,7 @@ def optimize(self, *args, **kwargs) -> Model:
             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 (Optional[str]): Execution role. Defaults to ``None``.
+            role_arn (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``.
             quantization_config (Optional[Dict]): Quantization configuration. Defaults to ``None``.
@@ -978,7 +978,7 @@ def _model_builder_optimize_wrapper(
         self,
         output_path: Optional[str] = None,
         instance_type: Optional[str] = None,
-        role: Optional[str] = None,
+        role_arn: Optional[str] = None,
         tags: Optional[Tags] = None,
         job_name: Optional[str] = None,
         accept_eula: Optional[bool] = None,
@@ -996,7 +996,7 @@ def _model_builder_optimize_wrapper(
         Args:
             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``.
+            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
@@ -1030,8 +1030,8 @@ def _model_builder_optimize_wrapper(
 
         if instance_type:
             self.instance_type = instance_type
-        if role:
-            self.role = role
+        if role_arn:
+            self.role_arn = role_arn
 
         self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
         job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
@@ -1041,7 +1041,7 @@ def _model_builder_optimize_wrapper(
             input_args = self._optimize_for_jumpstart(
                 output_path=output_path,
                 instance_type=instance_type,
-                role=role if role else self.role_arn,
+                role_arn=self.role_arn,
                 tags=tags,
                 job_name=job_name,
                 accept_eula=accept_eula,