[serving] deprecate s3Url and replace it with model_id

deepjavalibrary · Mar 23, 2023 · be64fc3 · be64fc3
1 parent b50432e
commit be64fc3
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 36 deletions.
diff --git a/engines/python/setup/djl_python/deepspeed.py b/engines/python/setup/djl_python/deepspeed.py
@@ -109,8 +109,8 @@ def initialize(self, properties: dict):
         self.initialized = True
 
     def _parse_properties(self, properties):
-        # If option.s3url is used, the directory is stored in model_id
-        # If option.s3url is not used but model_id is present, we download from hub
+        # model_id can point to huggingface model_id or local directory.
+        # If option.model_id points to a s3 bucket, we download it and set model_id to the download directory.
         # Otherwise we assume model artifacts are in the model_dir
         self.model_id_or_path = properties.get("model_id") or properties.get(
             "model_dir")

diff --git a/engines/python/setup/djl_python/huggingface.py b/engines/python/setup/djl_python/huggingface.py
@@ -63,8 +63,8 @@ def __init__(self):
         self.initialized = False
 
     def initialize(self, properties: dict):
-        # If option.s3url is used, the directory is stored in model_id
-        # If option.s3url is not used but model_id is present, we download from hub
+        # model_id can point to huggingface model_id or local directory.
+        # If option.model_id points to a s3 bucket, we download it and set model_id to the download directory.
         # Otherwise we assume model artifacts are in the model_dir
         model_id_or_path = properties.get("model_id") or properties.get(
             "model_dir")

diff --git a/engines/python/setup/djl_python/stable-diffusion.py b/engines/python/setup/djl_python/stable-diffusion.py
@@ -46,8 +46,8 @@ def __init__(self):
         self.save_image_dir = None
 
     def initialize(self, properties: dict):
-        # If option.s3url is used, the directory is stored in model_id
-        # If option.s3url is not used but model_id is present, we download from hub
+        # model_id can point to huggingface model_id or local directory.
+        # If option.model_id points to a s3 bucket, we download it and set model_id to the download directory.
         # Otherwise we assume model artifacts are in the model_dir
         self.model_id_or_path = properties.get("model_id") or properties.get(
             "model_dir")

diff --git a/engines/python/src/main/java/ai/djl/python/engine/PyModel.java b/engines/python/src/main/java/ai/djl/python/engine/PyModel.java
@@ -168,15 +168,20 @@ public void load(Path modelPath, String prefix, Map<String, ?> options) throws I
         }
         pyEnv.setEntryPoint(entryPoint);
 
-        String s3Url = pyEnv.getInitParameters().get("s3url");
+        Map<String, String> initParameters = pyEnv.getInitParameters();
+        String s3Url = initParameters.get("s3url");
+        String modelId = pyEnv.getInitParameters().get("model_id");
         if (s3Url != null) {
-            if (pyEnv.getInitParameters().containsKey("model_id")) {
+            // s3url is deprecated, use model_id instead
+            if (modelId != null) {
                 throw new IllegalArgumentException("model_id and s3url could not both set!");
             }
-
-            logger.info("S3 url found, start downloading from {}", s3Url);
+            modelId = s3Url;
+        }
+        if (modelId != null && modelId.startsWith("s3://")) {
+            logger.info("S3 url found, start downloading from {}", modelId);
             String downloadDir = getDownloadDir().toString();
-            downloadS3(s3Url, downloadDir);
+            downloadS3(modelId, downloadDir);
             // point model_id to download directory
             pyEnv.addParameter("model_id", downloadDir);
         }

diff --git a/serving/docker/partition/partition.py b/serving/docker/partition/partition.py
@@ -40,13 +40,14 @@ def __init__(self, props_manager):
         self.download_model_from_s3()
 
     def download_model_from_s3(self):
-        if "s3url" not in self.properties:
+        model_id = self.properties.get("model_id")
+        if not model_id or model_id.startswith("s3://"):
             return
 
         download_dir = os.environ.get("SERVING_DOWNLOAD_DIR",
                                       '/tmp/download/model/')
 
-        s3url = self.properties["s3url"]
+        s3url = model_id
         if Path("/opt/djl/bin/s5cmd").is_file():
             if not s3url.endswith("*"):
                 if s3url.endswith("/"):

diff --git a/serving/docker/partition/properties_manager.py b/serving/docker/partition/properties_manager.py
@@ -71,16 +71,18 @@ def set_and_validate_model_dir(self):
         elif 'model_id' in self.properties:
             self.properties['model_dir'] = self.properties_dir
         elif 's3url' in self.properties:
+            # backward compatible only, should be replaced with model_id
             self.properties['model_dir'] = self.properties_dir
+            self.properties['model_id'] = self.properties['s3url']
+            self.properties.pop('s3url')
         else:
             model_files = glob.glob(os.path.join(self.properties_dir, '*.bin'))
             if model_files:
                 self.properties['model_dir'] = self.properties_dir
             else:
                 raise KeyError(
-                    'Please specify the option.model_dir or option.model_id or option.s3_url or '
-                    'include model '
-                    'files in the model-dir argument.')
+                    'Please specify the option.model_dir or option.model_id or include model files in the model-dir.'
+                )
 
     def generate_properties_file(self):
         checkpoint_path = self.properties.get('save_mp_checkpoint_path')

diff --git a/tests/integration/llm/prepare.py b/tests/integration/llm/prepare.py
@@ -8,13 +8,13 @@
 
 ds_aot_list = {
     "opt-6.7b": {
-        "option.s3url": "s3://djl-llm/opt-6b7/",
+        "option.model_id": "s3://djl-llm/opt-6b7/",
         "option.tensor_parallel_degree": 4,
         "option.task": "text-generation",
         "option.dtype": "float16"
     },
     "gpt-neox-20b": {
-        "option.s3url": "s3://djl-llm/gpt-neox-20b/",
+        "option.model_id": "s3://djl-llm/gpt-neox-20b/",
         "option.tensor_parallel_degree": 4,
         "option.task": "text-generation",
         "option.dtype": "float16"
@@ -23,16 +23,16 @@
 
 ds_model_list = {
     "gpt-j-6b": {
-        "option.s3url": "s3://djl-llm/gpt-j-6b/",
+        "option.model_id": "s3://djl-llm/gpt-j-6b/",
         "option.tensor_parallel_degree": 4
     },
     "bloom-7b1": {
-        "option.s3url": "s3://djl-llm/bloom-7b1/",
+        "option.model_id": "s3://djl-llm/bloom-7b1/",
         "option.tensor_parallel_degree": 4,
         "option.dtype": "float16"
     },
     "opt-30b": {
-        "option.s3url": "s3://djl-llm/opt-30b/",
+        "option.model_id": "s3://djl-llm/opt-30b/",
         "option.tensor_parallel_degree": 4
     }
 }
@@ -44,14 +44,14 @@
         "option.tensor_parallel_degree": 2
     },
     "gpt-j-6b": {
-        "option.s3url": "s3://djl-llm/gpt-j-6b/",
+        "option.model_id": "s3://djl-llm/gpt-j-6b/",
         "option.task": "text-generation",
         "option.tensor_parallel_degree": 2,
         "option.device_map": "auto",
         "option.dtype": "fp16"
     },
     "bloom-7b1": {
-        "option.s3url": "s3://djl-llm/bloom-7b1/",
+        "option.model_id": "s3://djl-llm/bloom-7b1/",
         "option.tensor_parallel_degree": 4,
         "option.task": "text-generation",
         "option.load_in_8bit": "TRUE",
@@ -61,19 +61,19 @@
 
 ds_handler_list = {
     "gpt-j-6b": {
-        "option.s3url": "s3://djl-llm/gpt-j-6b/",
+        "option.model_id": "s3://djl-llm/gpt-j-6b/",
         "option.task": "text-generation",
         "option.tensor_parallel_degree": 2,
         "option.dtype": "bf16"
     },
     "bloom-7b1": {
-        "option.s3url": "s3://djl-llm/bloom-7b1/",
+        "option.model_id": "s3://djl-llm/bloom-7b1/",
         "option.tensor_parallel_degree": 4,
         "option.task": "text-generation",
         "option.dtype": "fp16"
     },
     "opt-13b": {
-        "option.s3url": "s3://djl-llm/opt-13b/",
+        "option.model_id": "s3://djl-llm/opt-13b/",
         "option.tensor_parallel_degree": 2,
         "option.task": "text-generation",
         "option.dtype": "fp16"
@@ -82,17 +82,17 @@
 
 sd_handler_list = {
     "stable-diffusion-v1-5": {
-        "option.s3url": "s3://djl-llm/stable-diffusion-v1-5/",
+        "option.model_id": "s3://djl-llm/stable-diffusion-v1-5/",
         "option.tensor_parallel_degree": 4,
         "option.dtype": "fp16"
     },
     "stable-diffusion-2-1-base": {
-        "option.s3url": "s3://djl-llm/stable-diffusion-2-1-base/",
+        "option.model_id": "s3://djl-llm/stable-diffusion-2-1-base/",
         "option.tensor_parallel_degree": 2,
         "option.dtype": "fp16"
     },
     "stable-diffusion-2-depth": {
-        "option.s3url": "s3://djl-llm/stable-diffusion-2-depth/",
+        "option.model_id": "s3://djl-llm/stable-diffusion-2-depth/",
         "option.tensor_parallel_degree": 1,
         "option.dtype": "fp16",
         "gpu.maxWorkers": 1
@@ -101,13 +101,13 @@
 
 ft_handler_list = {
     "bigscience/bloom-3b": {
-        "option.s3url": "s3://djl-llm/bloom-3b/",
+        "option.model_id": "s3://djl-llm/bloom-3b/",
         "option.tensor_parallel_degree": 2,
         "option.dtype": "fp16",
         "gpu.maxWorkers": 1,
     },
     "flan-t5-xxl": {
-        "option.s3url": "s3://djl-llm/flan-t5-xxl/",
+        "option.model_id": "s3://djl-llm/flan-t5-xxl/",
         "option.tensor_parallel_degree": 4,
         "option.dtype": "fp32"
     }
@@ -123,18 +123,18 @@
         "option.tensor_parallel_degree": 1,
     },
     "facebook/opt-6.7b": {
-        "option.s3url": "s3://djl-llm/opt-6b7/",
+        "option.model_id": "s3://djl-llm/opt-6b7/",
         "option.tensor_parallel_degree": 4,
         "option.dtype": "fp16",
     },
     "bigscience/bloom-3b": {
-        "option.s3url": "s3://djl-llm/bloom-3b/",
+        "option.model_id": "s3://djl-llm/bloom-3b/",
         "option.tensor_parallel_degree": 2,
         "option.dtype": "fp16",
         "gpu.maxWorkers": 1,
     },
     "flan-t5-xxl": {
-        "option.s3url": "s3://djl-llm/flan-t5-xxl/",
+        "option.model_id": "s3://djl-llm/flan-t5-xxl/",
         "option.tensor_parallel_degree": 4,
         "option.dtype": "fp32"
     }
@@ -152,15 +152,15 @@
 
 transformers_neuronx_handler_list = {
     "opt-1.3b": {
-        "option.s3url": "s3://djl-llm/opt-1.3b/",
+        "option.model_id": "s3://djl-llm/opt-1.3b/",
         "option.batch_size": 4,
         "option.tensor_parallel_degree": 4,
         "option.n_positions": 256,
         "option.dtype": "fp16",
         "option.model_loading_timeout": 600
     },
     "gpt-j-6b": {
-        "option.s3url": "s3://djl-llm/gpt-j-6b/",
+        "option.model_id": "s3://djl-llm/gpt-j-6b/",
         "option.batch_size": 4,
         "option.tensor_parallel_degree": 8,
         "option.n_positions": 512,