Follow-ups fixes

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -668,8 +668,6 @@ def _build_for_jumpstart(self):
     def _optimize_for_jumpstart(
         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,
@@ -685,9 +683,6 @@ def _optimize_for_jumpstart(
 
         Args:
             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_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
@@ -715,13 +710,13 @@ 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)
+        is_compilation = (not quantization_config) and (
+            (compilation_config is not None) or _is_inferentia_or_trainium(self.instance_type)
         )
 
         pysdk_model_env_vars = dict()
         if is_compilation:
-            pysdk_model_env_vars = self._get_neuron_model_env_vars(instance_type)
+            pysdk_model_env_vars = self._get_neuron_model_env_vars(self.instance_type)
 
         optimization_config, override_env = _extract_optimization_config_and_env(
             quantization_config, compilation_config
@@ -757,8 +752,9 @@ def _optimize_for_jumpstart(
             if self.pysdk_model.deployment_config
             else None
         )
-        self.instance_type = instance_type or deployment_config_instance_type or _get_nb_instance()
-        self.role_arn = role_arn or self.role_arn
+        self.instance_type = (
+            self.instance_type or deployment_config_instance_type or _get_nb_instance()
+        )
 
         create_optimization_job_args = {
             "OptimizationJobName": job_name,
@@ -788,9 +784,10 @@ def _optimize_for_jumpstart(
                 }
 
         if quantization_config or is_compilation:
-            self.pysdk_model.env = _update_environment_variables(
+            optimization_env_vars = _update_environment_variables(
                 optimization_env_vars, override_env
             )
+            self.pysdk_model.env.update(optimization_env_vars)
             return create_optimization_job_args
         return None
 

src/sagemaker/serve/builder/model_builder.py

@@ -73,7 +73,6 @@
     _generate_model_source,
     _extract_optimization_config_and_env,
     _is_s3_uri,
-    _normalize_local_model_path,
     _custom_speculative_decoding,
     _extract_speculative_draft_model_provider,
 )
@@ -833,6 +832,8 @@ def build(  # pylint: disable=R0911
         # until we deprecate HUGGING_FACE_HUB_TOKEN.
         if self.env_vars.get("HUGGING_FACE_HUB_TOKEN") and not self.env_vars.get("HF_TOKEN"):
             self.env_vars["HF_TOKEN"] = self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+        elif self.env_vars.get("HF_TOKEN") and not self.env_vars.get("HUGGING_FACE_HUB_TOKEN"):
+            self.env_vars["HUGGING_FACE_HUB_TOKEN"] = self.env_vars.get("HF_TOKEN")
 
         self.sagemaker_session.settings._local_download_dir = self.model_path
 
@@ -851,7 +852,9 @@ def build(  # pylint: disable=R0911
 
         self._build_validations()
 
-        if not self._is_jumpstart_model_id() and self.model_server:
+        if (
+            not (isinstance(self.model, str) and self._is_jumpstart_model_id())
+        ) and self.model_server:
             return self._build_for_model_server()
 
         if isinstance(self.model, str):
@@ -1216,18 +1219,14 @@ def _model_builder_optimize_wrapper(
             raise ValueError("Quantization config and compilation config are mutually exclusive.")
 
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
-
         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}"
-
         if self._is_jumpstart_model_id():
+            self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
             input_args = self._optimize_for_jumpstart(
                 output_path=output_path,
-                instance_type=instance_type,
-                role_arn=self.role_arn,
                 tags=tags,
                 job_name=job_name,
                 accept_eula=accept_eula,
@@ -1240,10 +1239,13 @@ def _model_builder_optimize_wrapper(
                 max_runtime_in_sec=max_runtime_in_sec,
             )
         else:
+            if self.model_server != ModelServer.DJL_SERVING:
+                logger.info("Overriding model server to DJL_SERVING.")
+                self.model_server = ModelServer.DJL_SERVING
+
+            self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
             input_args = self._optimize_for_hf(
                 output_path=output_path,
-                instance_type=instance_type,
-                role_arn=self.role_arn,
                 tags=tags,
                 job_name=job_name,
                 quantization_config=quantization_config,
@@ -1256,8 +1258,10 @@ def _model_builder_optimize_wrapper(
             )
 
         if input_args:
+            print(input_args)
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
+            print(job_status)
             return _generate_optimized_model(self.pysdk_model, job_status)
 
         self.pysdk_model.remove_tag_with_key(Tag.OPTIMIZATION_JOB_NAME)
@@ -1269,8 +1273,6 @@ def _model_builder_optimize_wrapper(
     def _optimize_for_hf(
         self,
         output_path: str,
-        instance_type: Optional[str] = None,
-        role_arn: Optional[str] = None,
         tags: Optional[Tags] = None,
         job_name: Optional[str] = None,
         quantization_config: Optional[Dict] = None,
@@ -1285,9 +1287,6 @@ def _optimize_for_hf(
 
         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_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``.
@@ -1305,13 +1304,6 @@ def _optimize_for_hf(
         Returns:
             Optional[Dict[str, Any]]: Model optimization job input arguments.
         """
-        if self.model_server != ModelServer.DJL_SERVING:
-            logger.info("Overwriting model server to DJL.")
-            self.model_server = ModelServer.DJL_SERVING
-
-        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
         )
@@ -1371,13 +1363,12 @@ def _optimize_prepare_for_hf(self):
             )
         else:
             if not custom_model_path:
-                custom_model_path = f"/tmp/sagemaker/model-builder/{self.model}/code"
+                custom_model_path = f"/tmp/sagemaker/model-builder/{self.model}"
                 download_huggingface_model_metadata(
                     self.model,
-                    custom_model_path,
+                    os.path.join(custom_model_path, "code"),
                     self.env_vars.get("HUGGING_FACE_HUB_TOKEN"),
                 )
-            custom_model_path = _normalize_local_model_path(custom_model_path)
 
         self.pysdk_model.model_data, env = self._prepare_for_mode(
             model_path=custom_model_path,

src/sagemaker/serve/utils/optimize_utils.py

@@ -282,26 +282,6 @@ def _extract_optimization_config_and_env(
     return None, None
 
 
-def _normalize_local_model_path(local_model_path: Optional[str]) -> Optional[str]:
-    """Normalizes the local model path.
-
-    Args:
-        local_model_path (Optional[str]): The local model path.
-
-    Returns:
-        Optional[str]: The normalized model path.
-    """
-    if local_model_path is None:
-        return local_model_path
-
-    # Removes /code or /code/ path at the end of local_model_path,
-    # as it is appended during artifacts upload.
-    pattern = r"/code/?$"
-    if re.search(pattern, local_model_path):
-        return re.sub(pattern, "", local_model_path)
-    return local_model_path
-
-
 def _custom_speculative_decoding(
     model: Model,
     speculative_decoding_config: Optional[Dict],

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

@@ -28,7 +28,6 @@
     _generate_additional_model_data_sources,
     _generate_channel_name,
     _extract_optimization_config_and_env,
-    _normalize_local_model_path,
     _is_optimized,
     _custom_speculative_decoding,
     _is_inferentia_or_trainium,
@@ -312,19 +311,6 @@ def test_extract_optimization_config_and_env(
     )
 
 
-@pytest.mark.parametrize(
-    "my_path, expected_path",
-    [
-        ("local/path/llama/code", "local/path/llama"),
-        ("local/path/llama/code/", "local/path/llama"),
-        ("local/path/llama/", "local/path/llama/"),
-        ("local/path/llama", "local/path/llama"),
-    ],
-)
-def test_normalize_local_model_path(my_path, expected_path):
-    assert _normalize_local_model_path(my_path) == expected_path
-
-
 class TestCustomSpeculativeDecodingConfig(unittest.TestCase):
 
     @patch("sagemaker.model.Model")