Bug bash fixes (aws#1492)

* HF Optimized

* Revert "HF Optimized"

* MB HF Optimize support

* Refactoring

* HF only s3 upload if optimize

* reuse role if provided in MB

* Refactoring

* New requirements

* Draft

* Refactoring

* Refactoring

* Bug Bash fixes

* UT

* UT

* Fix for parsing optimization output

* Tag fix

* UT

* UT

---------

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

Author: makungaj1

src/sagemaker/model.py

@@ -67,6 +67,7 @@
     format_tags,
     Tags,
     _resolve_routing_config,
+    _validate_new_tags,
 )
 from sagemaker.async_inference import AsyncInferenceConfig
 from sagemaker.predictor_async import AsyncPredictor
@@ -412,10 +413,7 @@ def add_tags(self, tags: Tags) -> None:
         Args:
             tags (Tags): Tags to add.
         """
-        if self._tags and tags:
-            self._tags.update(tags)
-        else:
-            self._tags = tags
+        self._tags = _validate_new_tags(tags, self._tags)
 
     @runnable_by_pipeline
     def register(

src/sagemaker/serve/builder/djl_builder.py

@@ -100,6 +100,7 @@ def __init__(self):
         self.env_vars = None
         self.nb_instance_type = None
         self.ram_usage_model_load = None
+        self.role_arn = None
 
     @abstractmethod
     def _prepare_for_mode(self):
@@ -499,4 +500,8 @@ def _build_for_djl(self):
 
         self.pysdk_model = self._build_for_hf_djl()
         self.pysdk_model.tune = self._tune_for_hf_djl
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -38,11 +38,13 @@
     SkipTuningComboException,
 )
 from sagemaker.serve.utils.optimize_utils import (
-    _extract_model_source,
+    _generate_model_source,
     _update_environment_variables,
     _extract_speculative_draft_model_provider,
     _is_image_compatible_with_optimization_job,
-    _validate_optimization_inputs,
+    _extracts_and_validates_speculative_model_source,
+    _generate_channel_name,
+    _generate_additional_model_data_sources,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -110,6 +112,7 @@ def __init__(self):
         self.ram_usage_model_load = None
         self.model_hub = None
         self.model_metadata = None
+        self.role_arn = None
         self.is_fine_tuned = None
         self.is_gated = None
 
@@ -544,7 +547,7 @@ def _update_model_data_for_fine_tuned_model(self, pysdk_model: Type[Model]) -> T
                 )
             pysdk_model.model_data["S3DataSource"]["S3Uri"] = fine_tuning_model_path
             pysdk_model.add_tags(
-                {"key": Tag.FINE_TUNING_MODEL_PATH, "value": fine_tuning_model_path}
+                {"Key": Tag.FINE_TUNING_MODEL_PATH, "Value": fine_tuning_model_path}
             )
             logger.info(
                 "FINE_TUNING_MODEL_PATH detected. Using fine-tuned model found in %s.",
@@ -633,6 +636,10 @@ def _build_for_jumpstart(self):
                 "with djl-inference, tgi-inference, or mms-inference container."
             )
 
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model
 
     def _optimize_for_jumpstart(
@@ -650,7 +657,7 @@ def _optimize_for_jumpstart(
         vpc_config: Optional[Dict] = None,
         kms_key: Optional[str] = None,
         max_runtime_in_sec: Optional[int] = None,
-    ) -> Dict[str, Any]:
+    ) -> Optional[Dict[str, Any]]:
         """Runs a model optimization job.
 
         Args:
@@ -685,13 +692,9 @@ def _optimize_for_jumpstart(
                 f"Model '{self.model}' requires accepting end-user license agreement (EULA)."
             )
 
-        _validate_optimization_inputs(
-            output_path, instance_type, quantization_config, compilation_config
-        )
-
         optimization_env_vars = None
         pysdk_model_env_vars = None
-        model_source = _extract_model_source(self.pysdk_model.model_data, accept_eula)
+        model_source = _generate_model_source(self.pysdk_model.model_data, accept_eula)
 
         if speculative_decoding_config:
             self._set_additional_model_source(speculative_decoding_config)
@@ -745,8 +748,12 @@ def _optimize_for_jumpstart(
         if vpc_config:
             create_optimization_job_args["VpcConfig"] = vpc_config
 
-        self.pysdk_model.env.update(pysdk_model_env_vars)
-        return create_optimization_job_args
+        if pysdk_model_env_vars:
+            self.pysdk_model.env.update(pysdk_model_env_vars)
+
+        if quantization_config or compilation_config:
+            return create_optimization_job_args
+        return None
 
     def _is_gated_model(self, model=None) -> bool:
         """Determine if ``this`` Model is Gated
@@ -779,14 +786,13 @@ def _set_additional_model_source(
         """
         if speculative_decoding_config:
             model_provider = _extract_speculative_draft_model_provider(speculative_decoding_config)
+            channel_name = _generate_channel_name(self.pysdk_model.additional_model_data_sources)
 
             if model_provider.lower() == "sagemaker":
-                if (
-                    self.pysdk_model.deployment_config.get("DeploymentArgs", {}).get(
-                        "AdditionalDataSources"
-                    )
-                    is None
-                ):
+                additional_model_data_sources = self.pysdk_model.deployment_config.get(
+                    "DeploymentArgs", {}
+                ).get("AdditionalDataSources")
+                if additional_model_data_sources is None:
                     deployment_config = self._find_compatible_deployment_config(
                         speculative_decoding_config
                     )
@@ -801,28 +807,26 @@ def _set_additional_model_source(
                         )
 
                 self.pysdk_model.add_tags(
-                    {"key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "value": "sagemaker"},
+                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
                 )
             else:
-                s3_uri = speculative_decoding_config.get("ModelSource")
-                if not s3_uri:
-                    raise ValueError("Custom S3 Uri cannot be none.")
-
-                # TODO: Set correct channel name.
-                additional_model_data_source = {
-                    "ChannelName": "DraftModelName",
-                    "S3DataSource": {"S3Uri": s3_uri},
-                }
-                if accept_eula:
-                    additional_model_data_source["S3DataSource"]["ModelAccessConfig"] = {
-                        "ACCEPT_EULA": True
-                    }
-
-                self.pysdk_model.additional_model_data_sources = [additional_model_data_source]
+                s3_uri = _extracts_and_validates_speculative_model_source(
+                    speculative_decoding_config
+                )
+
+                self.pysdk_model.additional_model_data_sources = (
+                    _generate_additional_model_data_sources(s3_uri, channel_name, accept_eula)
+                )
                 self.pysdk_model.add_tags(
-                    {"key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "value": "customer"},
+                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "customer"},
                 )
 
+            speculative_draft_model = f"/opt/ml/additional-model-data-sources/{channel_name}"
+            self.pysdk_model.env = _update_environment_variables(
+                self.pysdk_model.env,
+                {"OPTION_SPECULATIVE_DRAFT_MODEL": speculative_draft_model},
+            )
+
     def _find_compatible_deployment_config(
         self, speculative_decoding_config: Optional[Dict] = None
     ) -> Optional[Dict[str, Any]]:

src/sagemaker/serve/builder/model_builder.py

@@ -63,7 +63,10 @@
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.lineage_utils import _maintain_lineage_tracking_for_mlflow_model
-from sagemaker.serve.utils.optimize_utils import _generate_optimized_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 (
     _get_gpu_info,
@@ -87,7 +90,9 @@
 )
 from sagemaker.utils import Tags
 from sagemaker.workflow.entities import PipelineVariable
-from sagemaker.huggingface.llm_utils import get_huggingface_model_metadata
+from sagemaker.huggingface.llm_utils import (
+    get_huggingface_model_metadata,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -383,7 +388,7 @@ def _get_serve_setting(self):
             sagemaker_session=self.sagemaker_session,
         )
 
-    def _prepare_for_mode(self):
+    def _prepare_for_mode(self, should_upload_artifacts: bool = False):
         """Placeholder docstring"""
         # TODO: move mode specific prepare steps under _model_builder_deploy_wrapper
         self.s3_upload_path = None
@@ -401,6 +406,7 @@ def _prepare_for_mode(self):
                 self.sagemaker_session,
                 self.image_uri,
                 getattr(self, "model_hub", None) == ModelHub.JUMPSTART,
+                should_upload=should_upload_artifacts,
             )
             self.env_vars.update(env_vars_sagemaker)
             return self.s3_upload_path, env_vars_sagemaker
@@ -479,6 +485,10 @@ def _create_model(self):
         self.pysdk_model.mode = self.mode
         self.pysdk_model.modes = self.modes
         self.pysdk_model.serve_settings = self.serve_settings
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
 
         # dynamically generate a method to direct model.deploy() logic based on mode
         # unique method to models created via ModelBuilder()
@@ -935,8 +945,9 @@ def optimize(self, *args, **kwargs) -> Model:
         """Runs a model optimization job.
 
         Args:
-            instance_type (str): Target deployment instance type that the model is optimized for.
-            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.
+            output_path (Optional[str]): Specifies where to store the compiled/quantized model.
             role (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``.
@@ -964,7 +975,7 @@ def optimize(self, *args, **kwargs) -> Model:
     @_capture_telemetry("optimize")
     def _model_builder_optimize_wrapper(
         self,
-        output_path: str,
+        output_path: Optional[str] = None,
         instance_type: Optional[str] = None,
         role: Optional[str] = None,
         tags: Optional[Tags] = None,
@@ -1010,11 +1021,15 @@ def _model_builder_optimize_wrapper(
         Returns:
             Model: A deployable ``Model`` object.
         """
+        _validate_optimization_inputs(
+            output_path, instance_type, quantization_config, compilation_config
+        )
+
         self.sagemaker_session = sagemaker_session or self.sagemaker_session or Session()
         self.build(mode=self.mode, sagemaker_session=self.sagemaker_session)
         job_name = job_name or f"modelbuilderjob-{uuid.uuid4().hex}"
 
-        input_args = {}
+        input_args = None
         if self._is_jumpstart_model_id():
             input_args = self._optimize_for_jumpstart(
                 output_path=output_path,
@@ -1032,8 +1047,9 @@ def _model_builder_optimize_wrapper(
                 max_runtime_in_sec=max_runtime_in_sec,
             )
 
-        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)
+        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 self.pysdk_model

src/sagemaker/serve/builder/tei_builder.py

@@ -162,10 +162,7 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             self.pysdk_model.role = kwargs.get("role")
             del kwargs["role"]
 
-        # set model_data to uncompressed s3 dict
-        self.pysdk_model.model_data, env_vars = self._prepare_for_mode()
-        self.env_vars.update(env_vars)
-        self.pysdk_model.env.update(self.env_vars)
+        self._prepare_for_mode()
 
         # if the weights have been cached via local container mode -> set to offline
         if str(Mode.LOCAL_CONTAINER) in self.modes:
@@ -220,4 +217,8 @@ def _build_for_tei(self):
         self._set_to_tei()
 
         self.pysdk_model = self._build_for_hf_tei()
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model

src/sagemaker/serve/builder/tf_serving_builder.py

@@ -102,6 +102,10 @@ def _create_tensorflow_model(self):
         self.pysdk_model.mode = self.mode
         self.pysdk_model.modes = self.modes
         self.pysdk_model.serve_settings = self.serve_settings
+        if hasattr(self, "role_arn") and self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if hasattr(self, "sagemaker_session") and self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
 
         self._original_deploy = self.pysdk_model.deploy
         self.pysdk_model.deploy = self._model_builder_deploy_wrapper

src/sagemaker/serve/builder/tgi_builder.py

@@ -201,10 +201,7 @@ def _tgi_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             self.pysdk_model.role = kwargs.get("role")
             del kwargs["role"]
 
-        # set model_data to uncompressed s3 dict
-        self.pysdk_model.model_data, env_vars = self._prepare_for_mode()
-        self.env_vars.update(env_vars)
-        self.pysdk_model.env.update(self.env_vars)
+        self._prepare_for_mode()
 
         # if the weights have been cached via local container mode -> set to offline
         if str(Mode.LOCAL_CONTAINER) in self.modes:
@@ -472,4 +469,8 @@ def _build_for_tgi(self):
 
         self.pysdk_model = self._build_for_hf_tgi()
         self.pysdk_model.tune = self._tune_for_hf_tgi
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model

src/sagemaker/serve/builder/transformers_builder.py

@@ -223,10 +223,7 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
             self.pysdk_model.role = kwargs.get("role")
             del kwargs["role"]
 
-        # set model_data to uncompressed s3 dict
-        self.pysdk_model.model_data, env_vars = self._prepare_for_mode()
-        self.env_vars.update(env_vars)
-        self.pysdk_model.env.update(self.env_vars)
+        self._prepare_for_mode()
 
         if "endpoint_logging" not in kwargs:
             kwargs["endpoint_logging"] = True
@@ -303,4 +300,8 @@ def _build_for_transformers(self):
 
         self._build_transformers_env()
 
+        if self.role_arn:
+            self.pysdk_model.role = self.role_arn
+        if self.sagemaker_session:
+            self.pysdk_model.sagemaker_session = self.sagemaker_session
         return self.pysdk_model