Model archiver api

model-archiver/model_archiver/__init__.py

@@ -1,9 +1,10 @@
-
-
 """
 This module does the following:
 Exports the model folder to generate a Model Archive file out of it in .mar format
 """
 from . import version
 
 __version__ = version.__version__
+
+from .model_archiver import ModelArchiver
+from .model_archiver_config import ModelArchiverConfig

model-archiver/model_archiver/arg_parser.py

@@ -6,6 +6,7 @@
 import os
 
 from .manifest_components.manifest import RuntimeType
+from .model_archiver_config import ModelArchiverConfig
 
 
 # noinspection PyTypeChecker
@@ -154,4 +155,4 @@ def export_model_args_parser():
             help="Path to a yaml file containing model configuration eg. batch_size.",
         )
 
-        return parser_export
+        return ModelArchiverConfig.from_args(parser_export.parse_args())

model-archiver/model_archiver/model_archiver.py

@@ -0,0 +1,17 @@
+"""
+Helper class to generate a model archive file
+"""
+
+from model_archiver.model_archiver_config import ModelArchiverConfig
+from model_archiver.model_packaging import generate_model_archive
+
+
+class ModelArchiver:
+    @staticmethod
+    def generate_model_archive(config: ModelArchiverConfig) -> None:
+        """
+        Generate a model archive file
+        :param config: Model Archiver Config object
+        :return:
+        """
+        generate_model_archive(config)

model-archiver/model_archiver/model_archiver_config.py

@@ -0,0 +1,28 @@
+import os
+from argparse import Namespace
+from dataclasses import dataclass, fields
+from typing import Literal, Optional
+
+from model_archiver.manifest_components.manifest import RuntimeType
+
+
+@dataclass
+class ModelArchiverConfig:
+    model_name: str
+    handler: str
+    version: str
+    serialized_file: Optional[str] = None
+    model_file: Optional[str] = None
+    extra_files: Optional[str] = None
+    runtime: str = RuntimeType.PYTHON.value
+    export_path: str = os.getcwd()
+    archive_format: Literal["default", "tgz", "no-archive"] = "default"
+    force: bool = False
+    requirements_file: Optional[str] = None
+    config_file: Optional[str] = None
+
+    @classmethod
+    def from_args(cls, args: Namespace) -> "ModelArchiverConfig":
+        params = {field.name: getattr(args, field.name) for field in fields(cls)}
+        config = cls(**params)
+        return config

model-archiver/model_archiver/model_packaging.py

@@ -4,31 +4,32 @@
 
 import logging
 import shutil
-import sys
+from typing import Optional
 
 from model_archiver.arg_parser import ArgParser
+from model_archiver.model_archiver_config import ModelArchiverConfig
 from model_archiver.model_archiver_error import ModelArchiverError
 from model_archiver.model_packaging_utils import ModelExportUtils
 
 
-def package_model(args, manifest):
+def package_model(config: ModelArchiverConfig, manifest: str):
     """
     Internal helper for the exporting model command line interface.
     """
-    model_file = args.model_file
-    serialized_file = args.serialized_file
-    model_name = args.model_name
-    handler = args.handler
-    extra_files = args.extra_files
-    export_file_path = args.export_path
-    requirements_file = args.requirements_file
-    config_file = args.config_file
+    model_file = config.model_file
+    serialized_file = config.serialized_file
+    model_name = config.model_name
+    handler = config.handler
+    extra_files = config.extra_files
+    export_file_path = config.export_path
+    requirements_file = config.requirements_file
+    config_file = config.config_file
 
     try:
         ModelExportUtils.validate_inputs(model_name, export_file_path)
         # Step 1 : Check if .mar already exists with the given model name
         export_file_path = ModelExportUtils.check_mar_already_exists(
-            model_name, export_file_path, args.force, args.archive_format
+            model_name, export_file_path, config.force, config.archive_format
         )
 
         # Step 2 : Copy all artifacts to temp directory
@@ -45,27 +46,28 @@ def package_model(args, manifest):
 
         # Step 2 : Zip 'em all up
         ModelExportUtils.archive(
-            export_file_path, model_name, model_path, manifest, args.archive_format
+            export_file_path, model_name, model_path, manifest, config.archive_format
         )
         shutil.rmtree(model_path)
         logging.info(
             "Successfully exported model %s to file %s", model_name, export_file_path
         )
     except ModelArchiverError as e:
         logging.error(e)
-        sys.exit(1)
+        raise e
 
 
-def generate_model_archive():
+def generate_model_archive(config: Optional[ModelArchiverConfig] = None):
     """
     Generate a model archive file
     :return:
     """
 
     logging.basicConfig(format="%(levelname)s - %(message)s")
-    args = ArgParser.export_model_args_parser().parse_args()
-    manifest = ModelExportUtils.generate_manifest_json(args)
-    package_model(args, manifest=manifest)
+    if config is None:
+        config = ArgParser.export_model_args_parser()
+    manifest = ModelExportUtils.generate_manifest_json(config)
+    package_model(config, manifest=manifest)
 
 
 if __name__ == "__main__":

model-archiver/model_archiver/model_packaging_utils.py

@@ -15,6 +15,7 @@
 
 from .manifest_components.manifest import Manifest
 from .manifest_components.model import Model
+from .model_archiver_config import ModelArchiverConfig
 from .model_archiver_error import ModelArchiverError
 
 archiving_options = {
@@ -107,29 +108,29 @@ def find_unique(files, suffix):
             )
 
     @staticmethod
-    def generate_model(modelargs):
+    def generate_model(modelcfg: ModelArchiverConfig):
         model = Model(
-            model_name=modelargs.model_name,
-            serialized_file=modelargs.serialized_file,
-            model_file=modelargs.model_file,
-            handler=modelargs.handler,
-            model_version=modelargs.version,
-            requirements_file=modelargs.requirements_file,
-            config_file=modelargs.config_file,
+            model_name=modelcfg.model_name,
+            serialized_file=modelcfg.serialized_file,
+            model_file=modelcfg.model_file,
+            handler=modelcfg.handler,
+            model_version=modelcfg.version,
+            requirements_file=modelcfg.requirements_file,
+            config_file=modelcfg.config_file,
         )
         return model
 
     @staticmethod
-    def generate_manifest_json(args):
+    def generate_manifest_json(config: ModelArchiverConfig) -> str:
         """
         Function to generate manifest as a json string from the inputs provided by the user in the command line
         :param args:
         :return:
         """
 
-        model = ModelExportUtils.generate_model(args)
+        model = ModelExportUtils.generate_model(config)
 
-        manifest = Manifest(runtime=args.runtime, model=model)
+        manifest = Manifest(runtime=config.runtime, model=model)
 
         return str(manifest)
 

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

@@ -40,12 +40,11 @@ def delete_file_path(path):
         pass
 
 
-def run_test(test, args, mocker):
-    m = mocker.patch(
+def run_test(test, config, mocker):
+    mocker.patch(
         "model_archiver.model_packaging.ArgParser.export_model_args_parser",
+        return_value=config,
     )
-    m.return_value.parse_args.return_value = args
-    mocker.patch("sys.exit", side_effect=Exception())
     from model_archiver.model_packaging import generate_model_archive
 
     it = test.get("iterations", 1)
@@ -179,7 +178,9 @@ def build_namespace(test):
 
     args = Namespace(**{k.replace("-", "_"): test[k] for k in keys})
 
-    return args
+    config = model_archiver.ModelArchiverConfig.from_args(args)
+
+    return config
 
 
 def make_paths_absolute(test, keys):

model-archiver/model_archiver/tests/unit_tests/test_model_archiver.py

@@ -0,0 +1,68 @@
+from argparse import Namespace
+from collections import namedtuple
+
+import pytest
+from model_archiver import ModelArchiver, ModelArchiverConfig
+from model_archiver.manifest_components.manifest import RuntimeType
+
+
+# noinspection PyClassHasNoInit
+class TestModelArchiver:
+    model_name = "my-model"
+    model_file = "my-model/"
+    serialized_file = "my-model/"
+    handler = "a.py::my-awesome-func"
+    export_path = "/Users/dummyUser/"
+    version = "1.0"
+    requirements_file = "requirements.txt"
+    config_file = None
+
+    config = ModelArchiverConfig(
+        model_name=model_name,
+        handler=handler,
+        runtime=RuntimeType.PYTHON.value,
+        model_file=model_file,
+        serialized_file=serialized_file,
+        extra_files=None,
+        export_path=export_path,
+        force=False,
+        archive_format="default",
+        version=version,
+        requirements_file=requirements_file,
+        config_file=None,
+    )
+
+    @pytest.fixture()
+    def patches(self, mocker):
+        Patches = namedtuple("Patches", ["arg_parse", "export_utils", "export_method"])
+        patches = Patches(
+            mocker.patch("model_archiver.arg_parser.ArgParser"),
+            mocker.patch("model_archiver.model_packaging.ModelExportUtils"),
+            mocker.patch("model_archiver.model_packaging.package_model"),
+        )
+        mocker.patch("shutil.rmtree")
+
+        return patches
+
+    def test_gen_model_archive(self, patches):
+        ModelArchiver.generate_model_archive(self.config)
+        patches.export_method.assert_called()
+
+    def test_model_archiver_config_from_args(self):
+        args = Namespace(
+            model_name=self.model_name,
+            handler=self.handler,
+            runtime=RuntimeType.PYTHON.value,
+            model_file=self.model_file,
+            serialized_file=self.serialized_file,
+            extra_files=None,
+            export_path=self.export_path,
+            force=False,
+            archive_format="default",
+            version=self.version,
+            requirements_file=self.requirements_file,
+            config_file=None,
+        )
+        config = ModelArchiverConfig.from_args(args)
+
+        assert config == self.config

model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py

@@ -1,20 +1,14 @@
 from collections import namedtuple
 
 import pytest
+from model_archiver import ModelArchiverConfig
 from model_archiver.manifest_components.manifest import RuntimeType
 from model_archiver.model_packaging import generate_model_archive, package_model
 from model_archiver.model_packaging_utils import ModelExportUtils
 
 
 # noinspection PyClassHasNoInit
 class TestModelPackaging:
-    class Namespace:
-        def __init__(self, **kwargs):
-            self.__dict__.update(kwargs)
-
-        def update(self, **kwargs):
-            self.__dict__.update(kwargs)
-
     model_name = "my-model"
     model_file = "my-model/"
     serialized_file = "my-model/"
@@ -23,9 +17,8 @@ def update(self, **kwargs):
     version = "1.0"
     requirements_file = "requirements.txt"
     config_file = None
-    source_vocab = None
 
-    args = Namespace(
+    config = ModelArchiverConfig(
         model_name=model_name,
         handler=handler,
         runtime=RuntimeType.PYTHON.value,
@@ -35,9 +28,7 @@ def update(self, **kwargs):
         export_path=export_path,
         force=False,
         archive_format="default",
-        convert=False,
         version=version,
-        source_vocab=source_vocab,
         requirements_file=requirements_file,
         config_file=None,
     )
@@ -55,7 +46,7 @@ def patches(self, mocker):
         return patches
 
     def test_gen_model_archive(self, patches):
-        patches.arg_parse.export_model_args_parser.parse_args.return_value = self.args
+        patches.arg_parse.export_model_args_parser.parse_args.return_value = self.config
         generate_model_archive()
         patches.export_method.assert_called()
 
@@ -67,32 +58,31 @@ def test_export_model_method(self, patches):
         )
         patches.export_utils.zip.return_value = None
 
-        package_model(self.args, ModelExportUtils.generate_manifest_json(self.args))
+        package_model(self.config, ModelExportUtils.generate_manifest_json(self.config))
         patches.export_utils.validate_inputs.assert_called()
         patches.export_utils.archive.assert_called()
 
     def test_export_model_method_tar(self, patches):
-        self.args.update(archive_format="tar")
+        self.config.archive_format = "tgz"
         patches.export_utils.check_mar_already_exists.return_value = "/Users/dummyUser/"
         patches.export_utils.check_custom_model_types.return_value = (
             "/Users/dummyUser",
             ["a.txt", "b.txt"],
         )
-        patches.export_utils.zip.return_value = None
 
-        package_model(self.args, ModelExportUtils.generate_manifest_json(self.args))
+        package_model(self.config, ModelExportUtils.generate_manifest_json(self.config))
         patches.export_utils.validate_inputs.assert_called()
         patches.export_utils.archive.assert_called()
 
     def test_export_model_method_noarchive(self, patches):
-        self.args.update(archive_format="no-archive")
+        self.config.archive_format = "no-archive"
         patches.export_utils.check_mar_already_exists.return_value = "/Users/dummyUser/"
         patches.export_utils.check_custom_model_types.return_value = (
             "/Users/dummyUser",
             ["a.txt", "b.txt"],
         )
         patches.export_utils.zip.return_value = None
 
-        package_model(self.args, ModelExportUtils.generate_manifest_json(self.args))
+        package_model(self.config, ModelExportUtils.generate_manifest_json(self.config))
         patches.export_utils.validate_inputs.assert_called()
         patches.export_utils.archive.assert_called()