Optimize support for hf models (aws#1499)

* HF support

* refactoring

* Refactoring

* Refactoing

* HF Refactoring

* Refactoring

* UT

* Fix UT

* Resolving PR comments

* HF Token

* Resolving PR comments

* Fix UT

* Fix JS ModelServer deploy wrapper override

* Fix tests

* fix UT

* Resolve PR comments

* fix doc

---------

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

Author: makungaj1

requirements/extras/test_requirements.txt

@@ -38,3 +38,4 @@ accelerate>=0.24.1,<=0.27.0
 schema==0.7.5
 tensorflow>=2.1,<=2.16
 mlflow>=2.12.2,<2.13
+huggingface_hub>=0.23.4

src/sagemaker/huggingface/llm_utils.py

@@ -13,7 +13,9 @@
 """Functions for generating ECR image URIs for pre-built SageMaker Docker images."""
 from __future__ import absolute_import
 
+import os
 from typing import Optional
+import importlib.util
 
 import urllib.request
 from urllib.error import HTTPError, URLError
@@ -123,3 +125,26 @@ def get_huggingface_model_metadata(model_id: str, hf_hub_token: Optional[str] =
             "Did not find model metadata for the following HuggingFace Model ID %s" % model_id
         )
     return hf_model_metadata_json
+
+
+def download_huggingface_model_metadata(
+    model_id: str, model_local_path: str, hf_hub_token: Optional[str] = None
+) -> None:
+    """Downloads the HuggingFace Model snapshot via HuggingFace API.
+
+    Args:
+        model_id (str): The HuggingFace Model ID
+        model_local_path (str): The local path to save the HuggingFace Model snapshot.
+        hf_hub_token (str): The HuggingFace Hub Token
+
+    Raises:
+        ImportError: If huggingface_hub is not installed.
+    """
+    if not importlib.util.find_spec("huggingface_hub"):
+        raise ImportError("Unable to import huggingface_hub, check if huggingface_hub is installed")
+
+    from huggingface_hub import snapshot_download
+
+    os.makedirs(model_local_path, exist_ok=True)
+    logger.info("Downloading model %s from Hugging Face Hub to %s", model_id, model_local_path)
+    snapshot_download(repo_id=model_id, local_dir=model_local_path, token=hf_hub_token)

src/sagemaker/serve/builder/djl_builder.py

@@ -24,6 +24,7 @@
     LocalModelOutOfMemoryException,
     LocalModelInvocationException,
 )
+from sagemaker.serve.utils.optimize_utils import _is_optimized
 from sagemaker.serve.utils.tuning import (
     _serial_benchmark,
     _concurrent_benchmark,
@@ -214,9 +215,10 @@ def _djl_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
             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)
+        if not _is_optimized(self.pysdk_model):
+            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)
 
         # if the weights have been cached via local container mode -> set to offline
         if str(Mode.LOCAL_CONTAINER) in self.modes:

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -42,10 +42,11 @@
     _update_environment_variables,
     _extract_speculative_draft_model_provider,
     _is_image_compatible_with_optimization_job,
-    _extracts_and_validates_speculative_model_source,
     _generate_channel_name,
-    _generate_additional_model_data_sources,
-    _is_s3_uri,
+    _extract_optimization_config_and_env,
+    _is_optimized,
+    _custom_speculative_decoding,
+    SPECULATIVE_DRAFT_MODEL,
 )
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
@@ -121,7 +122,7 @@ def __init__(self):
         self.speculative_decoding_draft_model_source = None
 
     @abstractmethod
-    def _prepare_for_mode(self):
+    def _prepare_for_mode(self, **kwargs):
         """Placeholder docstring"""
 
     @abstractmethod
@@ -130,6 +131,9 @@ def _get_client_translators(self):
 
     def _is_jumpstart_model_id(self) -> bool:
         """Placeholder docstring"""
+        if self.model is None:
+            return False
+
         try:
             model_uris.retrieve(model_id=self.model, model_version="*", model_scope=_JS_SCOPE)
         except KeyError:
@@ -141,8 +145,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.sagemaker_session = self.sagemaker_session
+        pysdk_model = JumpStartModel(
+            self.model, vpc_config=self.vpc_config, sagemaker_session=self.sagemaker_session
+        )
 
         self._original_deploy = pysdk_model.deploy
         pysdk_model.deploy = self._js_builder_deploy_wrapper
@@ -151,6 +156,7 @@ def _create_pre_trained_js_model(self) -> Type[Model]:
     @_capture_telemetry("jumpstart.deploy")
     def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
         """Placeholder docstring"""
+        env = {}
         if "mode" in kwargs and kwargs.get("mode") != self.mode:
             overwrite_mode = kwargs.get("mode")
             # mode overwritten by customer during model.deploy()
@@ -167,7 +173,8 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                     or not hasattr(self, "prepared_for_tgi")
                     or not hasattr(self, "prepared_for_mms")
                 ):
-                    self.pysdk_model.model_data, env = self._prepare_for_mode()
+                    if not _is_optimized(self.pysdk_model):
+                        self.pysdk_model.model_data, env = self._prepare_for_mode()
             elif overwrite_mode == Mode.LOCAL_CONTAINER:
                 self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
 
@@ -198,7 +205,6 @@ def _js_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
                     )
 
                 self._prepare_for_mode()
-                env = {}
             else:
                 raise ValueError("Mode %s is not supported!" % overwrite_mode)
 
@@ -726,25 +732,17 @@ def _optimize_for_jumpstart(
                 )
 
         model_source = _generate_model_source(self.pysdk_model.model_data, accept_eula)
-
-        optimization_config = {}
-        if quantization_config:
-            optimization_config["ModelQuantizationConfig"] = quantization_config
-            pysdk_model_env_vars = _update_environment_variables(
-                pysdk_model_env_vars, quantization_config["OverrideEnvironment"]
-            )
-        if compilation_config:
-            optimization_config["ModelCompilationConfig"] = compilation_config
-            pysdk_model_env_vars = _update_environment_variables(
-                pysdk_model_env_vars, compilation_config["OverrideEnvironment"]
-            )
+        optimization_config, env = _extract_optimization_config_and_env(
+            quantization_config, compilation_config
+        )
+        pysdk_model_env_vars = _update_environment_variables(pysdk_model_env_vars, env)
 
         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"
+            instance_type = self.pysdk_model.deployment_config.get("DeploymentArgs", {}).get(
+                "InstanceType", _get_nb_instance()
             )
 
         create_optimization_job_args = {
@@ -771,6 +769,10 @@ def _optimize_for_jumpstart(
             self.pysdk_model.env.update(pysdk_model_env_vars)
         if accept_eula:
             self.pysdk_model.accept_eula = accept_eula
+            if isinstance(self.pysdk_model.model_data, dict):
+                self.pysdk_model.model_data["S3DataSource"]["ModelAccessConfig"] = {
+                    "AcceptEula": True
+                }
 
         if quantization_config or compilation_config:
             return create_optimization_job_args
@@ -806,7 +808,6 @@ 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)
-            speculative_draft_model = f"/opt/ml/additional-model-data-sources/{channel_name}"
 
             if model_provider == "sagemaker":
                 additional_model_data_sources = self.pysdk_model.deployment_config.get(
@@ -825,32 +826,18 @@ def _set_additional_model_source(
                         raise ValueError(
                             "Cannot find deployment config compatible for optimization job."
                         )
+
+                self.pysdk_model.env.update(
+                    {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}"}
+                )
+                self.pysdk_model.add_tags(
+                    {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "sagemaker"},
+                )
             else:
-                model_source = _extracts_and_validates_speculative_model_source(
-                    speculative_decoding_config
+                self.pysdk_model = _custom_speculative_decoding(
+                    self.pysdk_model, speculative_decoding_config, accept_eula
                 )
 
-                if _is_s3_uri(model_source):
-                    self.pysdk_model.additional_model_data_sources = (
-                        _generate_additional_model_data_sources(
-                            model_source, channel_name, accept_eula
-                        )
-                    )
-                else:
-                    speculative_draft_model = model_source
-
-            self.pysdk_model.env = _update_environment_variables(
-                self.pysdk_model.env,
-                {"OPTION_SPECULATIVE_DRAFT_MODEL": speculative_draft_model},
-            )
-            self.pysdk_model.add_tags(
-                {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": model_provider},
-            )
-            if accept_eula and isinstance(self.pysdk_model.model_data, dict):
-                self.pysdk_model.model_data["S3DataSource"]["ModelAccessConfig"] = {
-                    "AcceptEula": True
-                }
-
     def _find_compatible_deployment_config(
         self, speculative_decoding_config: Optional[Dict] = None
     ) -> Optional[Dict[str, Any]]: