Generate container def for huggingface models

Author: siddvenk

src/sagemaker/djl_inference/model.py

@@ -30,9 +30,9 @@
 
 
 class DJLEngine(Enum):
-    DEEPSPEED = "DeepSpeed"
-    FASTER_TRANSFORMERS = "FasterTransformers"
-    HUGGINGFACE_ACCELERATE = "Python"
+    DEEPSPEED = ("DeepSpeed", "djl_python.deepspeed")
+    FASTER_TRANSFORMERS = ("FasterTransformers", "djl_python.faster_transformers")
+    HUGGINGFACE_ACCELERATE = ("Python", "djl_python.huggingface")
 
 
 class DJLLargeModelPredictor(Predictor):
@@ -93,6 +93,10 @@ def __init__(
         predictor_cls: callable = DJLLargeModelPredictor,
         **kwargs,
     ):
+        super(DJLLargeModel, self).__init__(
+            None, image_uri, role, entry_point, predictor_cls=predictor_cls, **kwargs
+        )
+        self.engine = None
         self.uncompressed_model_data = uncompressed_model_data
         self.djl_version = djl_version
         self.task = task
@@ -104,11 +108,120 @@ def __init__(
         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 package_for_edge(
+        self,
+        output_path,
+        model_name,
+        model_version,
+        role=None,
+        job_name=None,
+        resource_key=None,
+        s3_kms_key=None,
+        tags=None,
+    ):
+        raise NotImplementedError("DJLLargeModels do not support Sagemaker Edge")
+    def compile(
+        self,
+        target_instance_family,
+        input_shape,
+        output_path,
+        role,
+        tags=None,
+        job_name=None,
+        compile_max_run=15 * 60,
+        framework=None,
+        framework_version=None,
+        target_platform_os=None,
+        target_platform_arch=None,
+        target_platform_accelerator=None,
+        compiler_options=None,
+    ):
+        raise NotImplementedError("DJLLargeModels do not currently support compilation with SageMaker Neo")
+
+    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")
+
+        if not self.image_uri:
+            region_name = self.sagemaker_session.boto_session.region_name
+            self.image_uri = self.serving_image_uri(region_name)
+
+        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:
+            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)
+            kwargs_serving_properties = self.generate_serving_properties()
+            existing_serving_properties.update(kwargs_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, self.image_uri
+            )
+            bucket = self.bucket or self.sagemaker_session.default_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(
+                self.image_uri, model_data_url=uploaded_code.s3_prefix, env=self.env
+            )
+
+    def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
+        serving_properties["engine"] = self.engine[0]
+        serving_properties["option.entryPoint"] = self.engine[1]
+        serving_properties["option.s3url"] = self.uncompressed_model_data
+        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
+        return serving_properties
+
     def serving_image_uri(self, region_name):
         if not self.djl_version:
             self.djl_version = "0.20.0"
@@ -167,12 +280,6 @@ def __init__(
         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,
@@ -184,75 +291,16 @@ def __init__(
             predictor_cls=predictor_cls,
             **kwargs,
         )
+        self.engine = DJLEngine.DEEPSPEED
+        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
 
-    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}")
-        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",
-            "option.s3url": self.uncompressed_model_data,
-        }
+    def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
+        serving_properties = super(DeepSpeedModel, self).generate_serving_properties()
         if self.max_tokens:
             serving_properties["option.max_tokens"] = self.max_tokens
         if self.low_cpu_mem_usage:
@@ -269,26 +317,6 @@ def _generate_serving_properties(self):
             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
 
         return serving_properties
 
@@ -312,10 +340,6 @@ def __init__(
         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,
@@ -327,3 +351,22 @@ def __init__(
             predictor_cls=predictor_cls,
             **kwargs,
         )
+        self.engine = DJLEngine.HUGGINGFACE_ACCELERATE
+        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,)
+
+    def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
+        serving_properties = super(HuggingfaceAccelerateModel, self).generate_serving_properties()
+        if self.device_id:
+            serving_properties["option.device_id"] = self.device_id
+        if self.device_map:
+            serving_properties["option.device_map"] = self.device_map
+        if self.load_in_8bit:
+            if not self.data_type == "int8":
+                raise ValueError("Set data_type='int8' to use load_in_8bit")
+            serving_properties["option.load_in_8bit"] = self.load_in_8bit
+        if self.low_cpu_mem_usage:
+            serving_properties["option.low_cpu_mem_usage"] = self.low_cpu_mem_usage
+        return serving_properties