Refactor api to make parallelism options and corresponding engines more clear

Author: siddvenk

src/sagemaker/djl_inference/model.py

@@ -88,9 +88,12 @@ def __init__(
         )
 
 
-def _determine_engine_for_model_type(model_type: str):
+def _determine_engine_for_model(model_type: str, tensor_parallel_degree: int):
     """Placeholder docstring"""
 
+    if tensor_parallel_degree and tensor_parallel_degree > 1:
+        return DeepSpeedModel
+
     if model_type in defaults.DEEPSPEED_RECOMMENDED_ARCHITECTURES:
         return DeepSpeedModel
     return HuggingFaceAccelerateModel
@@ -126,6 +129,7 @@ class DJLLargeModel(FrameworkModel):
     def __new__(
         cls,
         uncompressed_model_data: str,
+        tensor_parallel_degree: int = None,
         *args,
         **kwargs,
     ):
@@ -143,14 +147,16 @@ def __new__(
 
         model_type = json.loads(s3.S3Downloader.read_file(config_file)).get("model_type")
         cls_to_create = (
-            cls if cls is not DJLLargeModel else _determine_engine_for_model_type(model_type)
+            cls
+            if cls is not DJLLargeModel
+            else _determine_engine_for_model(model_type, tensor_parallel_degree)
         )
         return super(DJLLargeModel, cls).__new__(cls_to_create)
 
     def __init__(
         self,
         uncompressed_model_data: str,
-        role: str,
+        role: str = None,
         djl_version: str = None,
         task: str = None,
         data_type: str = "fp32",
@@ -190,9 +196,10 @@ def __init__(
                 than or equal to the number of gpus available on the instance. Defaults to None.
                 If not provided, no tensor parallel sharding is done. If the provided value is
                 greater than 1, DeepSpeed will be used as the backend.
-            data_parallel_degree (int): The number of copies of the model to instantiate. It should be
+            data_parallel_degree (int): The number of replicas of the model to instantiate. It should be
                 less than or equal to the number of gpus available on the instance. Defaults to None.
-                If not provided, all available gpus will be used.
+                If not provided, all available gpus will be used. If tensor_parallel_degree is set,
+                data_parallel_degree will be computed by DJL Serving based on the number of available GPUs.
             min_workers (int): The minimum number of worker processes. DJL Serving will auto detect
                 the minimum workers if not specified. Defaults to None.
             max_workers (int): The maximum number of worker processes. DJL Serving will auto detect
@@ -536,8 +543,6 @@ def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
         serving_properties["engine"] = self.engine.value[0]
         serving_properties["option.entryPoint"] = self.engine.value[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["option.entryPoint"] = self.entry_point
         if self.task:
@@ -667,6 +672,8 @@ def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
             dict: The model server configuration to use when deploying this model to SageMaker.
         """
         serving_properties = super(DeepSpeedModel, self).generate_serving_properties()
+        if self.tensor_parallel_degree:
+            serving_properties["option.tensor_parallel_degree"] = self.tensor_parallel_degree
         if self.max_tokens:
             serving_properties["option.max_tokens"] = self.max_tokens
         if self.low_cpu_mem_usage:
@@ -681,8 +688,6 @@ def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
             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
 
         return serving_properties
 
@@ -759,7 +764,11 @@ def generate_serving_properties(self, serving_properties={}) -> Dict[str, str]:
             dict: The model server configuration to use when deploying this model to SageMaker.
         """
         serving_properties = super(HuggingFaceAccelerateModel, self).generate_serving_properties()
+        if self.data_parallel_degree:
+            serving_properties["option.tensor_parallel_degree"] = self.data_parallel_degree
         if self.device_id:
+            if self.data_parallel_degree > 1:
+                raise ValueError("device_id cannot be set when data_parallel_degree is > 1")
             serving_properties["option.device_id"] = self.device_id
         if self.device_map:
             serving_properties["option.device_map"] = self.device_map