Add zip-store option to model-archiver tool

model-archiver/README.md

@@ -56,10 +56,11 @@ torch-model-archiver --model-name densenet161 --version 1.0 --model-file example
 
 ```
 $ torch-model-archiver -h
-usage: torch-model-archiver [-h] --model-name MODEL_NAME  --version MODEL_VERSION_NUMBER
-                      --model-file MODEL_FILE_PATH --serialized-file MODEL_SERIALIZED_PATH
-                      --handler HANDLER [--runtime {python,python3}]
-                      [--export-path EXPORT_PATH] [-f] [--requirements-file]
+usage: torch-model-archiver [-h] --model-name MODEL_NAME --handler HANDLER --version MODEL_VERSION_NUMBER
+                      [--model-file MODEL_FILE_PATH] [--serialized-file MODEL_SERIALIZED_PATH]
+                      [--runtime {python,python3}] [--export-path EXPORT_PATH]
+                      [--archive-format {tgz,no-archive,zip-store,default}]
+                      [-f] [--requirements-file]
 
 Model Archiver Tool
 
@@ -93,7 +94,7 @@ optional arguments:
                         is an optional parameter. If --export-path is not
                         specified, the file will be saved in the current
                         working directory.
-  --archive-format {tgz,default}
+  --archive-format {tgz, no-archive, zip-store, default}
                         The format in which the model artifacts are archived.
                         "tgz": This creates the model-archive in <model-name>.tar.gz format.
                         If platform hosting requires model-artifacts to be in ".tar.gz"
@@ -102,6 +103,9 @@ optional arguments:
                         at "export-path/{model-name}" location. As a result of this choice,
                         MANIFEST file will be created at "export-path/{model-name}" location
                         without archiving these model files
+                        "zip-store": This creates the model-archive in <model-name>.mar format
+                        but will skip deflating the files to speed up creation. Mainly used
+                        for testing purposes
                         "default": This creates the model-archive in <model-name>.mar format.
                         This is the default archiving format. Models archived in this format
                         will be readily hostable on TorchServe.
@@ -182,7 +186,7 @@ This will package all the model artifacts files and output `densenet_161.mar` in
 
 
 ### Model specific custom python requirements
-Custom models/handlers may depend on different python packages which are not installed by-default as a part of `TorchServe` setup. 
+Custom models/handlers may depend on different python packages which are not installed by-default as a part of `TorchServe` setup.
 Supply a [python requirements](https://pip.pypa.io/en/stable/user_guide/#requirements-files) file containing the list of required python packages to be installed by `TorchServe` for seamless model serving using `--requirements-file` parameter while creating the model-archiver.
 
 Example:

model-archiver/model_archiver/arg_parser.py

@@ -100,7 +100,7 @@ def export_model_args_parser():
             required=False,
             type=str,
             default="default",
-            choices=["tgz", "no-archive", "default"],
+            choices=["tgz", "no-archive", "zip-store", "default"],
             help="The format in which the model artifacts are archived.\n"
             '"tgz": This creates the model-archive in <model-name>.tar.gz format.\n'
             'If platform hosting TorchServe requires model-artifacts to be in ".tar.gz"\n'
@@ -109,6 +109,9 @@ def export_model_args_parser():
             'at "export-path/{model-name}" location. As a result of this choice, \n'
             'MANIFEST file will be created at "export-path/{model-name}" location\n'
             "without archiving these model files\n"
+            '"zip-store": This creates the model-archive in <model-name>.mar format\n'
+            "but will skip deflating the files to speed up creation. Mainly used\n"
+            "for testing purposes\n"
             '"default": This creates the model-archive in <model-name>.mar format.\n'
             "This is the default archiving format. Models archived in this format\n"
             "will be readily hostable on native TorchServe.\n",

model-archiver/model_archiver/model_packaging_utils.py

@@ -16,7 +16,12 @@
 from .manifest_components.model import Model
 from .model_archiver_error import ModelArchiverError
 
-archiving_options = {"tgz": ".tar.gz", "no-archive": "", "default": ".mar"}
+archiving_options = {
+    "tgz": ".tar.gz",
+    "no-archive": "",
+    "zip-store": ".mar",
+    "default": ".mar",
+}
 
 
 model_handlers = {
@@ -217,7 +222,12 @@ def archive(
                 with open(os.path.join(manifest_path, MANIFEST_FILE_NAME), "w") as f:
                     f.write(manifest)
             else:
-                with zipfile.ZipFile(mar_path, "w", zipfile.ZIP_DEFLATED) as z:
+                zip_mode = (
+                    zipfile.ZIP_STORED
+                    if archive_format == "zip-store"
+                    else zipfile.ZIP_DEFLATED
+                )
+                with zipfile.ZipFile(mar_path, "w", zip_mode) as z:
                     ModelExportUtils.archive_dir(
                         model_path, z, archive_format, model_name
                     )

model-archiver/model_archiver/tests/conftest.py

@@ -0,0 +1,50 @@
+import json
+from pathlib import Path
+
+import pytest
+
+INTEG_TEST_CONFIG_FILE = "integ_tests/configuration.json"
+DEFAULT_HANDLER_CONFIG_FILE = "integ_tests/default_handler_configuration.json"
+
+TEST_ROOT_DIR = Path(__file__).parent
+MODEL_ARCHIVER_ROOT_DIR = Path(__file__).parents[2]
+
+
+def make_paths_absolute(test, keys):
+    def make_absolute(paths):
+        if "," in paths:
+            return ",".join([make_absolute(p) for p in paths.split(",")])
+        return MODEL_ARCHIVER_ROOT_DIR.joinpath(paths).as_posix()
+
+    for k in keys:
+        test[k] = make_absolute(test[k])
+
+    return test
+
+
+@pytest.fixture(name="integ_tests")
+def load_integ_tests():
+    with open(TEST_ROOT_DIR.joinpath(INTEG_TEST_CONFIG_FILE), "r") as f:
+        tests = json.loads(f.read())
+    keys = (
+        "model-file",
+        "serialized-file",
+        "handler",
+        "extra-files",
+    )
+    return [make_paths_absolute(t, keys) for t in tests]
+
+
+@pytest.fixture(name="default_handler_tests")
+def load_default_handler_tests():
+    with open(TEST_ROOT_DIR.joinpath(DEFAULT_HANDLER_CONFIG_FILE), "r") as f:
+        default_handler_tests = json.loads(f.read())
+    keys = (
+        "model-file",
+        "serialized-file",
+        "extra-files",
+    )
+    default_handler_tests = [
+        make_paths_absolute(t, keys) for t in default_handler_tests
+    ]
+    return default_handler_tests

model-archiver/model_archiver/tests/integ_tests/configuration.json

@@ -60,6 +60,19 @@
     "version": "1.0",
     "force": true
   },
+  {
+    "name": "packaging_zip_store_mar",
+    "model-name": "model",
+    "model-file": "model_archiver/tests/integ_tests/resources/regular_model/test_model.py",
+    "serialized-file": "model_archiver/tests/integ_tests/resources/regular_model/test_serialized_file.pt",
+    "handler": "model_archiver/tests/integ_tests/resources/regular_model/test_handler.py",
+    "extra-files": "model_archiver/tests/integ_tests/resources/regular_model/test_index_to_name.json",
+    "export-path": "model",
+    "archive-format": "zip-store",
+    "iterations": 2,
+    "version": "1.0",
+    "force": true
+  },
   {
     "name": "packaging_mar_with_handler_name",
     "model-name": "model",
@@ -85,4 +98,4 @@
     "iterations": 2,
     "version": "1.0",
     "force": true
-  }]
+  }]

model-archiver/model_archiver/tests/integ_tests/test_integration_model_archiver.py

@@ -180,62 +180,68 @@ def make_absolute(paths):
     return test
 
 
-def test_model_archiver():
-    with open(TEST_ROOT_DIR.joinpath(INTEG_TEST_CONFIG_FILE), "r") as f:
-        tests = json.loads(f.read())
-        keys = (
-            "model-file",
-            "serialized-file",
-            "handler",
-            "extra-files",
-        )
-        tests = [make_paths_absolute(t, keys) for t in tests]
-        for test in tests:
-            # tar.gz format problem on windows hence ignore
-            if platform.system() == "Windows" and test["archive-format"] == "tgz":
-                continue
-            try:
-                test["export-path"] = os.path.join(
-                    tempfile.gettempdir(), test["export-path"]
-                )
-                delete_file_path(test.get("export-path"))
-                create_file_path(test.get("export-path"))
-                test["runtime"] = test.get("runtime", DEFAULT_RUNTIME)
-                test["model-name"] = (
-                    test["model-name"] + "_" + str(int(time.time() * 1000.0))
-                )
-                cmd = build_cmd(test)
-                if test.get("force"):
-                    cmd += " -f"
-
-                if run_test(test, cmd):
-                    validate(test)
-            finally:
-                delete_file_path(test.get("export-path"))
-
-
-def test_default_handlers():
-    with open(TEST_ROOT_DIR.joinpath(DEFAULT_HANDLER_CONFIG_FILE), "r") as f:
-        tests = json.loads(f.read())
-        keys = (
-            "model-file",
-            "serialized-file",
-            "extra-files",
-        )
-        tests = [make_paths_absolute(t, keys) for t in tests]
-        for test in tests:
+def test_model_archiver(integ_tests):
+    for test in integ_tests:
+        # tar.gz format problem on windows hence ignore
+        if platform.system() == "Windows" and test["archive-format"] == "tgz":
+            continue
+        try:
+            test["export-path"] = os.path.join(
+                tempfile.gettempdir(), test["export-path"]
+            )
+            delete_file_path(test.get("export-path"))
+            create_file_path(test.get("export-path"))
+            test["runtime"] = test.get("runtime", DEFAULT_RUNTIME)
+            test["model-name"] = (
+                test["model-name"] + "_" + str(int(time.time() * 1000.0))
+            )
             cmd = build_cmd(test)
-            try:
-                delete_file_path(test.get("export-path"))
-                create_file_path(test.get("export-path"))
+            if test.get("force"):
+                cmd += " -f"
+
+            if run_test(test, cmd):
+                validate(test)
+        finally:
+            delete_file_path(test.get("export-path"))
+
+
+def test_default_handlers(default_handler_tests):
+    for test in default_handler_tests:
+        cmd = build_cmd(test)
+        try:
+            delete_file_path(test.get("export-path"))
+            create_file_path(test.get("export-path"))
+
+            if test.get("force"):
+                cmd += " -f"
+
+            if run_test(test, cmd):
+                validate(test)
+        finally:
+            delete_file_path(test.get("export-path"))
+
+
+def test_zip_store(tmp_path, integ_tests):
+    integ_tests = list(
+        filter(lambda t: t["name"] == "packaging_zip_store_mar", integ_tests)
+    )
+    assert len(integ_tests) == 1
+    test = integ_tests[0]
+
+    test["export-path"] = tmp_path
+    test["iterations"] = 1
+
+    test["model-name"] = "zip-store"
+    run_test(test, build_cmd(test))
+
+    test["model-name"] = "zip"
+    test["archive-format"] = "default"
+    run_test(test, build_cmd(test))
 
-                if test.get("force"):
-                    cmd += " -f"
+    stored_size = Path(tmp_path).joinpath("zip-store.mar").stat().st_size
+    zipped_size = Path(tmp_path).joinpath("zip.mar").stat().st_size
 
-                if run_test(test, cmd):
-                    validate(test)
-            finally:
-                delete_file_path(test.get("export-path"))
+    assert zipped_size < stored_size
 
 
 if __name__ == "__main__":