Generate container definition for deepspeed models

Author: siddvenk

src/sagemaker/djl_inference/__init__.py

@@ -13,4 +13,9 @@
 """Placeholder docstring"""
 from __future__ import absolute_import
 
-from sagemaker.djl_inference.model import DJLLargeModel, DJLLargeModelPredictor, DeepSpeedModel, HuggingfaceAccelerateModel
+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",
@@ -35,4 +35,4 @@
     "xlm-roberta",
     "roberta",
     "bert",
-}
+}

src/sagemaker/djl_inference/model.py

@@ -15,17 +15,17 @@
 
 import json
 import os.path
-import tempfile
 from enum import Enum
 from typing import Optional, Union, Dict
 
+import defaults
 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
 from sagemaker.session import Session
-import defaults
+from sagemaker.utils import _tmpdir, _create_or_update_code_dir
 from sagemaker.workflow.entities import PipelineVariable
 
 
@@ -36,7 +36,6 @@ class DJLEngine(Enum):
 
 
 class DJLLargeModelPredictor(Predictor):
-
     def __init__(
         self,
         endpoint_name,
@@ -53,7 +52,6 @@ def __init__(
 
 
 class DJLLargeModel(FrameworkModel):
-
     def __new__(
         cls,
         uncompressed_model_data: str,
@@ -62,6 +60,11 @@ def __new__(
     ):
         if not uncompressed_model_data.startswith("s3://"):
             raise ValueError("DJLLargeModel only supports loading model artifacts from s3")
+        if uncompressed_model_data.endswith("tar.gz"):
+            raise ValueError(
+                "DJLLargeModel does not support model artifacts in tar.gz format."
+                "Please store the model in uncompressed format and provide the s3 uri of the folder"
+            )
         if uncompressed_model_data.endswith("/"):
             config_file = uncompressed_model_data + "config.json"
         else:
@@ -88,7 +91,7 @@ def __init__(
         entry_point: Optional[str] = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         predictor_cls: callable = DJLLargeModelPredictor,
-        **kwargs
+        **kwargs,
     ):
         self.uncompressed_model_data = uncompressed_model_data
         self.djl_version = djl_version
@@ -102,7 +105,7 @@ def __init__(
         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
+            None, image_uri, role, entry_point, predictor_cls=predictor_cls, **kwargs
         )
         self.sagemaker_session = self.sagemaker_session or Session()
 
@@ -122,14 +125,28 @@ def _determine_engine_for_model_type(model_type: str):
         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}")
+            raise ValueError(
+                f"{model_type} is not supported by DeepSpeed. "
+                f"Supported model_types are {defaults.DEEPSPEED_SUPPORTED_ARCHITECTURES}"
+            )
+
+
+def _read_existing_serving_properties(directory: str):
+    serving_properties_path = os.path.join(directory, "serving.properties")
+    properties = {}
+    if os.path.exists(serving_properties_path):
+        with open(serving_properties_path, "r") as f:
+            for line in f:
+                key, val = line.split("=", 1)
+                properties[key] = val
+    return properties
 
-class DeepSpeedModel(DJLLargeModel):
 
+class DeepSpeedModel(DJLLargeModel):
     _framework_name = "djl-deepspeed"
 
     def __init__(
@@ -139,8 +156,8 @@ def __init__(
         low_cpu_mem_usage: bool = True,
         enable_cuda_graph: bool = False,
         triangular_masking: bool = True,
-        return_tuple = True,
-        deepspeed_checkpoint_file = None,
+        return_tuple=True,
+        deepspeed_checkpoint_file=None,
         task: str = None,
         data_type: str = None,
         tensor_parallel_degree: int = None,
@@ -185,19 +202,52 @@ def prepare_container_def(
             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):
+        local_download_dir = (
+            None
+            if self.sagemaker_session.settings is None
+            or self.sagemaker_session.settings.local_download_dir is None
+            else self.sagemaker_session.settings.local_download_dir
+        )
+        with _tmpdir(directory=local_download_dir) as tmp:
+            # Check to see if we need to bundle user provided code with serving.properties and re upload
+            if self.source_dir or self.entry_point:
+                _create_or_update_code_dir(
+                    tmp,
+                    self.entry_point,
+                    self.source_dir,
+                    self.dependencies,
+                    self.sagemaker_session,
+                    tmp,
+                )
+            existing_serving_properties = _read_existing_serving_properties(tmp)
+            provided_serving_properties = self._generate_serving_properties()
+            # provided kwargs take precedence over existing serving.properties file
+            existing_serving_properties.update(provided_serving_properties)
+            # self._validate_serving_properties(existing_serving_properties)
+
+            with open(os.path.join(tmp, "serving.properties"), "w+") as f:
+                for key, val in existing_serving_properties.items():
+                    f.write(f"{key}={val}\n")
+
+            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}")
+            uploaded_code = fw_utils.tar_and_upload_dir(
+                self.sagemaker_session.boto_session,
+                bucket,
+                deploy_key_prefix,
+                self.entry_point,
+                directory=tmp,
+                dependencies=self.dependencies,
+                kms_key=self.model_kms_key,
+            )
+            return sagemaker.container_def(
+                deploy_image, model_data_url=uploaded_code.s3_prefix, env=self.env
+            )
+
+    def _generate_serving_properties(self):
         serving_properties = {
             "engine": "DeepSpeed",
             "option.entryPoint": "djl_python.deepspeed",
@@ -209,7 +259,9 @@ def _validate_and_write_serving_properties(self):
             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")
+                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
@@ -238,20 +290,10 @@ def _validate_and_write_serving_properties(self):
         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
+        return serving_properties
 
 
 class HuggingfaceAccelerateModel(DJLLargeModel):
-
     _framework_name = "djl-deepspeed"
 
     def __init__(
@@ -268,12 +310,12 @@ def __init__(
         entry_point: str = None,
         image_uri: Optional[Union[str, PipelineVariable]] = None,
         predictor_cls: callable = DJLLargeModelPredictor,
-        **kwargs
+        **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,
+        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,
@@ -283,24 +325,5 @@ def __init__(
             entry_point=entry_point,
             image_uri=image_uri,
             predictor_cls=predictor_cls,
-            **kwargs
+            **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"
-    # )