Generate serving.properties and upload code to s3

Author: siddvenk

src/sagemaker/djl_inference/__init__.py

@@ -0,0 +1,16 @@
+# 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
+
+from sagemaker.djl_inference.model import DJLLargeModel, DJLLargeModelPredictor, DeepSpeedModel, HuggingfaceAccelerateModel

src/sagemaker/djl_inference/defaults.py

@@ -17,7 +17,7 @@
     "bloom",
     "opt",
     "gpt_neox",
-    "gptj",
+    #"gptj",
     "gpt_neo",
     "gpt2",
     "xlm-roberta",

src/sagemaker/djl_inference/model.py

@@ -14,10 +14,13 @@
 from __future__ import absolute_import
 
 import json
+import os.path
+import tempfile
 from enum import Enum
 from typing import Optional, Union, Dict
 
-from sagemaker import s3, Predictor
+import sagemaker
+from sagemaker import s3, Predictor, image_uris, fw_utils
 from sagemaker.deserializers import JSONDeserializer
 from sagemaker.model import FrameworkModel
 from sagemaker.serializers import JSONSerializer
@@ -71,24 +74,49 @@ def __new__(
     def __init__(
         self,
         uncompressed_model_data: str,
+        djl_version: str = None,
         task: str = None,
         data_type: str = None,
         tensor_parallel_degree: int = None,
+        min_workers: int = None,
+        max_workers: int = None,
+        job_queue_size: int = None,
+        parallel_loading: bool = False,
+        model_loading_timeout: int = None,
+        prediction_timeout: 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.djl_version = djl_version
         self.task = task
         self.data_type = data_type
-        self.tensor_parallel_degree = tensor_parallel_degree,
+        self.tensor_parallel_degree = tensor_parallel_degree
+        self.min_workers = min_workers
+        self.max_workers = max_workers
+        self.job_queue_size = job_queue_size
+        self.parallel_loading = parallel_loading
+        self.model_loading_timeout = model_loading_timeout
+        self.prediction_timeout = prediction_timeout
         super(DJLLargeModel, self).__init__(
             None, image_uri, role, entry_point, predictor_cls=predictor_cls,**kwargs
         )
         self.sagemaker_session = self.sagemaker_session or Session()
 
+    def serving_image_uri(self, region_name):
+        if not self.djl_version:
+            self.djl_version = "0.20.0"
+
+        return image_uris.retrieve(
+            self._framework(),
+            region_name,
+            version=self.djl_version,
+        )
+
+
 def _determine_engine_for_model_type(model_type: str):
     if model_type in defaults.DEEPSPEED_RECOMMENDED_ARCHITECTURES:
         return DeepSpeedModel
@@ -102,6 +130,8 @@ def _validate_engine_for_model_type(model_type: str, engine: DJLEngine):
 
 class DeepSpeedModel(DJLLargeModel):
 
+    _framework_name = "djl-deepspeed"
+
     def __init__(
         self,
         uncompressed_model_data: str,
@@ -138,8 +168,92 @@ def __init__(
             **kwargs,
         )
 
+    def prepare_container_def(
+        self,
+        instance_type=None,
+        accelerator_type=None,
+        serverless_inference_config=None,
+    ):
+        if serverless_inference_config is not None:
+            raise ValueError("DJLLargeModel does not support serverless deployment")
+        if accelerator_type is not None:
+            raise ValueError("DJLLargeModel does not support Elastic Inference accelerator")
+
+        deploy_image = self.image_uri
+        if not deploy_image:
+            region_name = self.sagemaker_session.boto_session.region_name
+            deploy_image = self.serving_image_uri(region_name)
+
+        print(f"Deploy image is{deploy_image}")
+        tmp_dir = self._validate_and_write_serving_properties()
+        deploy_key_prefix = fw_utils.model_code_key_prefix(self.key_prefix, self.name, deploy_image)
+        bucket = self.bucket or self.sagemaker_session.default_bucket()
+        print(f"bucket to upload code to is {bucket}")
+        # self.uploaded_code = fw_utils.tar_and_upload_dir(
+        #     session=self.sagemaker_session.boto_session,
+        #     bucket=bucket,
+        #     s3_key_prefix=deploy_key_prefix,
+        #     directory=tmp_dir,
+        #     script=None,
+        # )
+
+    def _validate_and_write_serving_properties(self):
+        serving_properties = {
+            "engine": "DeepSpeed",
+            "option.entryPoint": "djl_python.deepspeed",
+            "option.s3url": self.uncompressed_model_data,
+        }
+        if self.max_tokens:
+            serving_properties["option.max_tokens"] = self.max_tokens
+        if self.low_cpu_mem_usage:
+            serving_properties["option.low_cpu_mem_usage"] = self.low_cpu_mem_usage
+        if self.enable_cuda_graph:
+            if self.tensor_parallel_degree > 1:
+                raise ValueError("enable_cuda_graph is not supported when tensor_parallel_degree > 1")
+            serving_properties["option.enable_cuda_graph"] = self.enable_cuda_graph
+        if self.triangular_masking:
+            serving_properties["option.triangular_masking"] = self.triangular_masking
+        if self.return_tuple:
+            serving_properties["option.return_tuple"] = self.return_tuple
+        if self.deepspeed_checkpoint_file:
+            serving_properties["option.checkpoint"] = self.deepspeed_checkpoint_file
+        if self.tensor_parallel_degree:
+            serving_properties["option.tensor_parallel_degree"] = self.tensor_parallel_degree
+        if self.entry_point:
+            serving_properties["entryPoint"] = self.entry_point
+        if self.task:
+            serving_properties["option.task"] = self.task
+        if self.data_type:
+            serving_properties["option.dtype"] = self.data_type
+        if self.min_workers:
+            serving_properties["minWorkers"] = self.min_workers
+        if self.max_workers:
+            serving_properties["maxWorkers"] = self.max_workers
+        if self.job_queue_size:
+            serving_properties["job_queue_size"] = self.job_queue_size
+        if self.parallel_loading:
+            serving_properties["option.parallel_loading"] = self.parallel_loading
+        if self.model_loading_timeout:
+            serving_properties["option.model_loading_timeout"] = self.model_loading_timeout
+        if self.prediction_timeout:
+            serving_properties["option.prediction_timeout"] = self.prediction_timeout
+
+        local_dir = None if self.sagemaker_session.settings else self.sagemaker_session.settings.local_download_dir
+        tmp_dir = tempfile.mkdtemp(dir=local_dir)
+
+        with open(os.path.join(tmp_dir, "serving.properties"), 'w+') as f:
+            for key, value in serving_properties.items():
+                f.write(f"{key}={value}\n")
+
+        print(f"wrote serving.properties to {tmp_dir}")
+
+        return tmp_dir
+
+
 class HuggingfaceAccelerateModel(DJLLargeModel):
 
+    _framework_name = "djl-deepspeed"
+
     def __init__(
         self,
         uncompressed_model_data: str,
@@ -172,3 +286,21 @@ def __init__(
             **kwargs
         )
 
+if __name__ == "__main__":
+    session = Session()
+    role = "arn:aws:iam::125045733377:role/AmazonSageMaker-ExecutionRole-djl"
+    opt_model = DJLLargeModel(
+        "s3://dlc-deepspeed-test-temp/opt-2.7b/",
+        tensor_parallel_degree=2,
+        data_type="fp32",
+        task="text-generation",
+        max_tokens=2048,
+        parallel_loading=True,
+        role=role,
+        sagemaker_session=session,
+    )
+    opt_model.prepare_container_def()
+    # opt_model.deploy(
+    #     initial_instance_count=1,
+    #     instance_type="ml.g5.12xl"
+    # )