Introduce HF DLC to Model Builder

Author: samruds

src/sagemaker/serve/builder/hf_dlc_builder.py

@@ -0,0 +1,223 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Holds mixin logic to support deployment of Model ID"""
+from __future__ import absolute_import
+import logging
+from typing import Type
+from abc import ABC, abstractmethod
+
+from sagemaker.model import Model
+from sagemaker.serve.utils.local_hardware import (
+    _get_nb_instance,
+    _get_ram_usage_mb,
+    _get_gpu_info,
+    _get_gpu_info_fallback,
+)
+
+from sagemaker.djl_inference.model import _get_model_config_properties_from_hf
+from sagemaker.serve.model_server.djl_serving.utils import (
+    _get_admissible_tensor_parallel_degrees,
+    _get_default_tensor_parallel_degree,
+)
+from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri
+from sagemaker.serve.model_server.hf_dlc.prepare import _create_dir_structure
+from sagemaker.serve.utils.predictors import HfDLCLocalModePredictor
+from sagemaker.serve.utils.types import ModelServer
+from sagemaker.serve.mode.function_pointers import Mode
+from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
+from sagemaker.base_predictor import PredictorBase
+from sagemaker.huggingface.llm_utils import get_huggingface_model_metadata
+
+logger = logging.getLogger(__name__)
+
+
+class HuggingFaceDLC(ABC):
+    """HuggingFace DLC build logic for ModelBuilder()"""
+
+    def __init__(self):
+        self.model = None
+        self.serve_settings = None
+        self.sagemaker_session = None
+        self.model_path = None
+        self.dependencies = None
+        self.modes = None
+        self.mode = None
+        self.model_server = None
+        self.image_uri = None
+        self._original_deploy = None
+        self.hf_model_config = None
+        self._default_data_type = None
+        self._default_max_tokens = None
+        self.pysdk_model = None
+        self.env_vars = None
+        self.nb_instance_type = None
+        self.ram_usage_model_load = None
+        self.secret_key = None
+        self.role_arn = None
+        self.transformers_version = None
+        self.py_version = None
+        self.pytorch_version = None
+        self.tensorflow_version = None
+
+    @abstractmethod
+    def _prepare_for_mode(self):
+        """Placeholder docstring"""
+
+    def _create_hf_dlc_model(self) -> Type[Model]:
+        """Placeholder docstring"""
+
+        hf_model_md = get_huggingface_model_metadata(self.model,
+                                                     self.env_vars.get("HUGGING_FACE_HUB_TOKEN"))
+        if 'pytorch' in hf_model_md.get("tags"):
+            self.pytorch_version = "1.8.1"
+            self.py_version = "py36"
+        elif 'keras' in hf_model_md.get("tags") or 'tensorflow' in hf_model_md.get("tags"):
+            self.py_version = "py37"
+            self.tensorflow_version = "2.4.1"
+
+        self.transformers_version = "4.6.1"
+
+        pysdk_model = HuggingFaceModel(
+            env=self.env_vars,
+            role=self.role_arn,
+            sagemaker_session=self.sagemaker_session,
+            py_version=self.py_version,
+            transformers_version=self.transformers_version,
+            pytorch_version=self.pytorch_version,
+        )
+        self.image_uri = pysdk_model.serving_image_uri(self.sagemaker_session.boto_region_name, "local")
+
+        self._original_deploy = pysdk_model.deploy
+        pysdk_model.deploy = self._hf_dlc_model_builder_deploy_wrapper
+        return pysdk_model
+
+    @_capture_telemetry("hf_dlc.deploy")
+    def _hf_dlc_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBase]:
+        """Placeholder docstring"""
+        timeout = kwargs.get("model_data_download_timeout")
+        if timeout:
+            self.pysdk_model.env.update({"MODEL_LOADING_TIMEOUT": str(timeout)})
+
+        if "mode" in kwargs and kwargs.get("mode") != self.mode:
+            overwrite_mode = kwargs.get("mode")
+            # mode overwritten by customer during model.deploy()
+            logger.warning(
+                "Deploying in %s Mode, overriding existing configurations set for %s mode",
+                overwrite_mode,
+                self.mode,
+            )
+
+            if overwrite_mode == Mode.SAGEMAKER_ENDPOINT:
+                self.mode = self.pysdk_model.mode = Mode.SAGEMAKER_ENDPOINT
+            elif overwrite_mode == Mode.LOCAL_CONTAINER:
+                self._prepare_for_mode()
+                self.mode = self.pysdk_model.mode = Mode.LOCAL_CONTAINER
+            else:
+                raise ValueError("Mode %s is not supported!" % overwrite_mode)
+
+        serializer = self.schema_builder.input_serializer
+        deserializer = self.schema_builder._output_deserializer
+        if self.mode == Mode.LOCAL_CONTAINER:
+            timeout = kwargs.get("model_data_download_timeout")
+
+            predictor = HfDLCLocalModePredictor(
+                self.modes[str(Mode.LOCAL_CONTAINER)], serializer, deserializer
+            )
+
+            ram_usage_before = _get_ram_usage_mb()
+            self.modes[str(Mode.LOCAL_CONTAINER)].create_server(
+                self.image_uri,
+                timeout if timeout else 1800,
+                None,
+                predictor,
+                self.pysdk_model.env,
+                jumpstart=False,
+            )
+            ram_usage_after = _get_ram_usage_mb()
+            self.ram_usage_model_load = max(ram_usage_after - ram_usage_before, 0)
+
+            return predictor
+
+        if "mode" in kwargs:
+            del kwargs["mode"]
+        if "role" in kwargs:
+            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)
+
+        if "endpoint_logging" not in kwargs:
+            kwargs["endpoint_logging"] = True
+
+        if self.nb_instance_type and "instance_type" not in kwargs:
+            kwargs.update({"instance_type": self.nb_instance_type})
+        elif not self.nb_instance_type and "instance_type" not in kwargs:
+            raise ValueError(
+                "Instance type must be provided when deploying " "to SageMaker Endpoint mode."
+            )
+        else:
+            try:
+                tot_gpus = _get_gpu_info(kwargs.get("instance_type"), self.sagemaker_session)
+            except Exception:  # pylint: disable=W0703
+                tot_gpus = _get_gpu_info_fallback(kwargs.get("instance_type"))
+            default_num_shard = _get_default_tensor_parallel_degree(self.hf_model_config, tot_gpus)
+            self.pysdk_model.env.update(
+                {
+                    "NUM_SHARD": str(default_num_shard),
+                    "SHARDED": "true" if default_num_shard > 1 else "false",
+                }
+            )
+
+        if "initial_instance_count" not in kwargs:
+            kwargs.update({"initial_instance_count": 1})
+
+        if "endpoint_logging" not in kwargs:
+            kwargs["endpoint_logging"] = True
+
+        predictor = self._original_deploy(*args, **kwargs)
+
+        self.pysdk_model.env.update({"TRANSFORMERS_OFFLINE": "0"})
+
+        predictor.serializer = serializer
+        predictor.deserializer = deserializer
+        return predictor
+
+    def _build_for_hugging_face_dlc(self):
+        """Placeholder docstring"""
+        self.nb_instance_type = _get_nb_instance()
+
+        _create_dir_structure(self.model_path)
+        if not hasattr(self, "pysdk_model"):
+            self.env_vars.update({"HF_MODEL_ID": self.model})
+            self.hf_model_config = _get_model_config_properties_from_hf(
+                self.model, self.env_vars.get("HUGGING_FACE_HUB_TOKEN")
+            )
+
+        self.pysdk_model = self._create_hf_dlc_model()
+
+        if self.mode == Mode.LOCAL_CONTAINER:
+            self._prepare_for_mode()
+
+        return self.pysdk_model
+
+    def _build_for_hf_dlc(self):
+        """Placeholder docstring"""
+        self.secret_key = None
+
+        self.model_server = ModelServer.HuggingFaceDLC
+
+        self.pysdk_model = self._build_for_hugging_face_dlc()
+        return self.pysdk_model

src/sagemaker/serve/builder/model_builder.py

@@ -34,6 +34,7 @@
 from sagemaker.serve.builder.djl_builder import DJL
 from sagemaker.serve.builder.tgi_builder import TGI
 from sagemaker.serve.builder.jumpstart_builder import JumpStart
+from sagemaker.serve.builder.hf_dlc_builder import HuggingFaceDLC
 from sagemaker.predictor import Predictor
 from sagemaker.serve.save_retrive.version_1_0_0.metadata.metadata import Metadata
 from sagemaker.serve.spec.inference_spec import InferenceSpec
@@ -53,19 +54,21 @@
 from sagemaker.serve.validations.check_image_and_hardware_type import (
     validate_image_uri_and_hardware,
 )
+from sagemaker.huggingface.llm_utils import get_huggingface_model_metadata
 
 logger = logging.getLogger(__name__)
 
 supported_model_server = {
     ModelServer.TORCHSERVE,
     ModelServer.TRITON,
     ModelServer.DJL_SERVING,
+    ModelServer.HuggingFaceDLC,
 }
 
 
 # pylint: disable=attribute-defined-outside-init
 @dataclass
-class ModelBuilder(Triton, DJL, JumpStart, TGI):
+class ModelBuilder(Triton, DJL, JumpStart, TGI, HuggingFaceDLC):
     """Class that builds a deployable model.
 
     Args:
@@ -125,7 +128,7 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI):
             in order for model builder to build the artifacts correctly (according
             to the model server). Possible values for this argument are
             ``TORCHSERVE``, ``MMS``, ``TENSORFLOW_SERVING``, ``DJL_SERVING``,
-            ``TRITON``, and ``TGI``.
+            ``TRITON``, ``TGI``, and ``HuggingFaceDLC``.
 
     """
 
@@ -577,12 +580,19 @@ def build(
         )
 
         self.serve_settings = self._get_serve_setting()
+
+        hf_model_md = get_huggingface_model_metadata(self.model,
+                                                     self.env_vars.get("HUGGING_FACE_HUB_TOKEN"))
+        
         if isinstance(self.model, str):
             if self._is_jumpstart_model_id():
                 return self._build_for_jumpstart()
             if self._is_djl():
                 return self._build_for_djl()
-            return self._build_for_tgi()
+            if hf_model_md.get("pipeline_tag") == "text-generation":
+                return self._build_for_tgi()
+            else:
+                return self._build_for_hf_dlc()
 
         self._build_validations()
 

src/sagemaker/serve/mode/local_container_mode.py

@@ -19,6 +19,7 @@
 from sagemaker.serve.model_server.djl_serving.server import LocalDJLServing
 from sagemaker.serve.model_server.triton.server import LocalTritonServer
 from sagemaker.serve.model_server.tgi.server import LocalTgiServing
+from sagemaker.serve.model_server.hf_dlc.server import LocalHFDLCServing
 from sagemaker.session import Session
 
 logger = logging.getLogger(__name__)
@@ -31,7 +32,7 @@
 )
 
 
-class LocalContainerMode(LocalTorchServe, LocalDJLServing, LocalTritonServer, LocalTgiServing):
+class LocalContainerMode(LocalTorchServe, LocalDJLServing, LocalTritonServer, LocalTgiServing, LocalHFDLCServing):
     """A class that holds methods to deploy model to a container in local environment"""
 
     def __init__(
@@ -128,6 +129,15 @@ def create_server(
                 jumpstart=jumpstart,
             )
             self._ping_container = self._tgi_deep_ping
+        elif self.model_server == ModelServer.HuggingFaceDLC:
+            self._start_hf_dlc_serving(
+                client=self.client,
+                image=image,
+                model_path=model_path if model_path else self.model_path,
+                secret_key=secret_key,
+                env_vars=env_vars if env_vars else self.env_vars,
+            )
+            self._ping_container = self._hf_dlc_deep_ping            
 
         # allow some time for container to be ready
         time.sleep(10)

src/sagemaker/serve/mode/sagemaker_endpoint_mode.py

@@ -12,12 +12,13 @@
 from sagemaker.serve.model_server.torchserve.server import SageMakerTorchServe
 from sagemaker.serve.model_server.djl_serving.server import SageMakerDjlServing
 from sagemaker.serve.model_server.tgi.server import SageMakerTgiServing
+from sagemaker.serve.model_server.hf_dlc.server import SageMakerHFDLCServing
 
 logger = logging.getLogger(__name__)
 
 
 class SageMakerEndpointMode(
-    SageMakerTorchServe, SageMakerTritonServer, SageMakerDjlServing, SageMakerTgiServing
+    SageMakerTorchServe, SageMakerTritonServer, SageMakerDjlServing, SageMakerTgiServing, SageMakerHFDLCServing
 ):
     """Holds the required method to deploy a model to a SageMaker Endpoint"""
 
@@ -92,5 +93,12 @@ def prepare(
                 image=image,
                 jumpstart=jumpstart,
             )
+        if self.model_server == ModelServer.HuggingFaceDLC:
+            return self._upload_hf_dlc_artifacts(
+                model_path=model_path,
+                sagemaker_session=sagemaker_session,
+                s3_model_data_url=s3_model_data_url,
+                image=image,
+            )        
 
         raise ValueError("%s model server is not supported" % self.model_server)

src/sagemaker/serve/model_server/hf_dlc/__init__.py

@@ -0,0 +1,38 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Prepare HF DLC Model for Deployment"""
+
+from __future__ import absolute_import
+import logging
+from pathlib import Path
+
+from sagemaker.serve.utils.local_hardware import _check_disk_space, _check_docker_disk_usage
+
+logger = logging.getLogger(__name__)
+
+
+def _create_dir_structure(model_path: str) -> tuple:
+    """Create the expected model directory structure for the HF DLC server"""
+    model_path = Path(model_path)
+    if not model_path.exists():
+        model_path.mkdir(parents=True)
+    elif not model_path.is_dir():
+        raise ValueError("model_dir is not a valid directory")
+
+    code_dir = model_path.joinpath("code")
+    code_dir.mkdir(exist_ok=True, parents=True)
+
+    _check_disk_space(model_path)
+    _check_docker_disk_usage()
+
+    return model_path, code_dir