Add generic DJL Large model and engine specific model classes with constructors

Author: siddvenk

src/sagemaker/djl_inference/__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.
+"""Placeholder docstring"""
+from __future__ import absolute_import
+
+DEEPSPEED_RECOMMENDED_ARCHITECTURES = {
+    "bloom",
+    "opt",
+    "gpt_neox",
+    "gptj",
+    "gpt_neo",
+    "gpt2",
+    "xlm-roberta",
+    "roberta",
+    "bert",
+}
+
+DEEPSPEED_SUPPORTED_ARCHITECTURES = {
+    "bloom",
+    "opt",
+    "gpt_neox",
+    "gptj",
+    "gpt_neo",
+    "gpt2",
+    "xlm-roberta",
+    "roberta",
+    "bert",
+}

src/sagemaker/djl_inference/defaults.py

@@ -0,0 +1,174 @@
+# 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.
+
+from __future__ import absolute_import
+
+import json
+from enum import Enum
+from typing import Optional, Union, Dict
+
+from sagemaker import s3, Predictor
+from sagemaker.deserializers import JSONDeserializer
+from sagemaker.model import FrameworkModel
+from sagemaker.serializers import JSONSerializer
+from sagemaker.session import Session
+import defaults
+from sagemaker.workflow.entities import PipelineVariable
+
+
+class DJLEngine(Enum):
+    DEEPSPEED = "DeepSpeed"
+    FASTER_TRANSFORMERS = "FasterTransformers"
+    HUGGINGFACE_ACCELERATE = "Python"
+
+
+class DJLLargeModelPredictor(Predictor):
+
+    def __init__(
+        self,
+        endpoint_name,
+        sagemaker_session=None,
+        serializer=JSONSerializer(),
+        deserializer=JSONDeserializer(),
+    ):
+        super(DJLLargeModelPredictor, self).__init__(
+            endpoint_name,
+            sagemaker_session,
+            serializer=serializer,
+            deserializer=deserializer,
+        )
+
+
+class DJLLargeModel(FrameworkModel):
+
+    def __new__(
+        cls,
+        uncompressed_model_data: str,
+        *args,
+        **kwargs,
+    ):
+        if not uncompressed_model_data.startswith("s3://"):
+            raise ValueError("DJLLargeModel only supports loading model artifacts from s3")
+        if uncompressed_model_data.endswith("/"):
+            config_file = uncompressed_model_data + "config.json"
+        else:
+            config_file = uncompressed_model_data + "/config.json"
+
+        model_type = json.loads(s3.S3Downloader.read_file(config_file)).get("model_type")
+        cls_to_create = _determine_engine_for_model_type(model_type)
+        return super(DJLLargeModel, cls).__new__(cls_to_create)
+
+    def __init__(
+        self,
+        uncompressed_model_data: str,
+        task: str = None,
+        data_type: str = None,
+        tensor_parallel_degree: int = None,
+        role: str = None,
+        entry_point: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        predictor_cls: callable = DJLLargeModelPredictor,
+        **kwargs
+    ):
+        self.uncompressed_model_data = uncompressed_model_data
+        self.task = task
+        self.data_type = data_type
+        self.tensor_parallel_degree = tensor_parallel_degree,
+        super(DJLLargeModel, self).__init__(
+            None, image_uri, role, entry_point, predictor_cls=predictor_cls,**kwargs
+        )
+        self.sagemaker_session = self.sagemaker_session or Session()
+
+def _determine_engine_for_model_type(model_type: str):
+    if model_type in defaults.DEEPSPEED_RECOMMENDED_ARCHITECTURES:
+        return DeepSpeedModel
+    return HuggingfaceAccelerateModel
+
+def _validate_engine_for_model_type(model_type: str, engine: DJLEngine):
+    if engine == DJLEngine.DEEPSPEED:
+        if model_type not in defaults.DEEPSPEED_SUPPORTED_ARCHITECTURES:
+            raise ValueError(f"{model_type} is not supported by DeepSpeed. " \
+                      f"Supported model_types are {defaults.DEEPSPEED_SUPPORTED_ARCHITECTURES}")
+
+class DeepSpeedModel(DJLLargeModel):
+
+    def __init__(
+        self,
+        uncompressed_model_data: str,
+        max_tokens: int = None,
+        low_cpu_mem_usage: bool = True,
+        enable_cuda_graph: bool = False,
+        triangular_masking: bool = True,
+        return_tuple = True,
+        deepspeed_checkpoint_file = None,
+        task: str = None,
+        data_type: str = None,
+        tensor_parallel_degree: int = None,
+        role: str = None,
+        entry_point: Optional[str] = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        predictor_cls: callable = DJLLargeModelPredictor,
+        **kwargs,
+    ):
+        self.max_tokens = max_tokens
+        self.low_cpu_mem_usage = low_cpu_mem_usage
+        self.enable_cuda_graph = enable_cuda_graph
+        self.triangular_masking = triangular_masking
+        self.return_tuple = return_tuple
+        self.deepspeed_checkpoint_file = deepspeed_checkpoint_file
+        super(DeepSpeedModel, self).__init__(
+            uncompressed_model_data,
+            role=role,
+            task=task,
+            data_type=data_type,
+            tensor_parallel_degree=tensor_parallel_degree,
+            entry_point=entry_point,
+            image_uri=image_uri,
+            predictor_cls=predictor_cls,
+            **kwargs,
+        )
+
+class HuggingfaceAccelerateModel(DJLLargeModel):
+
+    def __init__(
+        self,
+        uncompressed_model_data: str,
+        device_id: int = None,
+        device_map: Union[str, Dict[str, str]] = None,
+        load_in_8bit: bool = False,
+        low_cpu_mem_usage: bool = True,
+        task: str = None,
+        data_type: str = None,
+        tensor_parallel_degree: int = None,
+        role: str = None,
+        entry_point: str = None,
+        image_uri: Optional[Union[str, PipelineVariable]] = None,
+        predictor_cls: callable = DJLLargeModelPredictor,
+        **kwargs
+    ):
+        self.device_id = device_id
+        self.device_map = device_map
+        self.load_in_8bit = load_in_8bit,
+        self.low_cpu_mem_usage = low_cpu_mem_usage,
+        super(HuggingfaceAccelerateModel, self).__init__(
+            uncompressed_model_data,
+            role=role,
+            task=task,
+            data_type=data_type,
+            tensor_parallel_degree=tensor_parallel_degree,
+            entry_point=entry_point,
+            image_uri=image_uri,
+            predictor_cls=predictor_cls,
+            **kwargs
+        )
+