def test_export_file_is_not_none(self, patches):
patches.path_exists.return_value = False
ModelExportUtils.check_mar_already_exists('some-model',
'/Users/Piyush/', False)
patches.path_exists.assert_called_once_with(
'/Users/Piyush/some-model.mar')
def test_onnx_file_is_none(self, patches):
patches.utils.find_unique.return_value = None
ModelExportUtils.check_custom_model_types(
model_path=self.model_path, model_name=None)
patches.utils.find_unique.assert_called()
patches.utils.convert_onnx_model.assert_not_called()
def test_export_file_already_exists_with_override(self, patches):
patches.path_exists.return_value = True
ModelExportUtils.check_mar_already_exists('some-model', None, True)
patches.path_exists.assert_called_once_with(
'/Users/dummyUser/some-model.mar')
def test_regex_pass(self):
model_names = [
'my-awesome-model', 'Aa.model', 'a', 'aA.model', 'a1234.model',
'a-A-A.model', '123-abc'
]
for m in model_names:
ModelExportUtils.check_model_name_regex_or_exit(m)
def test_regex_fail(self):
model_names = [
'abc%', '123.abc', '12-model-a.model', '##.model', '-.model'
]
for m in model_names:
with pytest.raises(ModelArchiverError):
ModelExportUtils.check_model_name_regex_or_exit(m)
def test_export_file_already_exists_with_override_false(self, patches):
patches.path_exists.return_value = True
with pytest.raises(ModelArchiverError):
ModelExportUtils.check_mar_already_exists('some-model', None, False)
patches.path_exists.assert_called_once_with('/Users/dummyUser/some-model.mar')
def test_export_file_is_none_tar(self, patches):
patches.path_exists.return_value = False
ret_val = ModelExportUtils.check_mar_already_exists(
'some-model', None, False, archive_format='tgz')
patches.path_exists.assert_called_once_with(
"/Users/dummyUser/some-model.tar.gz")
assert ret_val == "/Users/dummyUser"
def test_export_file_is_none(self, patches):
patches.path_exists.return_value = False
ret_val = ModelExportUtils.check_mar_already_exists(
'some-model', None, False)
patches.path_exists.assert_called_once_with(
"/Users/dummyUser/some-model.mar")
assert ret_val == "/Users/dummyUser"
def test_manifest_json(self):
manifest = ModelExportUtils.generate_manifest_json(self.args)
manifest_json = json.loads(manifest)
assert manifest_json['runtime'] == RuntimeType.PYTHON.value
assert 'engine' in manifest_json
assert 'model' in manifest_json
assert 'publisher' in manifest_json
assert 'license' not in manifest_json
def test_export_model_method(self, patches):
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))
patches.export_utils.validate_inputs.assert_called()
patches.export_utils.archive.assert_called()
patches.export_utils.clean_temp_files.assert_called()
def mmcls2torchserve(
config_file: str,
checkpoint_file: str,
output_folder: str,
model_name: str,
model_version: str = '1.0',
force: bool = False,
):
"""Converts mmclassification model (config + checkpoint) to TorchServe
`.mar`.
Args:
config_file:
In MMClassification config format.
The contents vary for each task repository.
checkpoint_file:
In MMClassification checkpoint format.
The contents vary for each task repository.
output_folder:
Folder where `{model_name}.mar` will be created.
The file created will be in TorchServe archive format.
model_name:
If not None, used for naming the `{model_name}.mar` file
that will be created under `output_folder`.
If None, `{Path(checkpoint_file).stem}` will be used.
model_version:
Model's version.
force:
If True, if there is an existing `{model_name}.mar`
file under `output_folder` it will be overwritten.
"""
mmcv.mkdir_or_exist(output_folder)
config = mmcv.Config.fromfile(config_file)
with TemporaryDirectory() as tmpdir:
config.dump(f'{tmpdir}/config.py')
args = Namespace(
**{
'model_file': f'{tmpdir}/config.py',
'serialized_file': checkpoint_file,
'handler': f'{Path(__file__).parent}/mmcls_handler.py',
'model_name': model_name or Path(checkpoint_file).stem,
'version': model_version,
'export_path': output_folder,
'force': force,
'requirements_file': None,
'extra_files': None,
'runtime': 'python',
'archive_format': 'default'
})
manifest = ModelExportUtils.generate_manifest_json(args)
package_model(args, manifest)
def mmgen2torchserver(config_file: str,
checkpoint_file: str,
output_folder: str,
model_name: str,
model_version: str = '1.0',
model_type: str = 'unconditional',
force: bool = False):
"""Converts MMGeneration model (config + checkpoint) to TorchServe `.mar`.
Args:
config_file (str): Path of config file. The config should in
MMGeneration format.
checkpoint_file (str): Path of checkpoint. The checkpoint should in
MMGeneration checkpoint format.
output_folder (str): Folder where `{model_name}.mar` will be created.
The file created will be in TorchServe archive format.
model_name (str): Name of the generated ``'mar'`` file. If not None,
used for naming the `{model_name}.mar` file that will be created
under `output_folder`. If None, `{Path(checkpoint_file).stem}`
will be used.
model_version (str, optional): Model's version. Defaults to '1.0'.
model_type (str, optional): Type of the model to be convert. Handler
named ``{model_type}_handler`` would be used to generate ``mar``
file. Defaults to 'unconditional'.
force (bool, optional): If True, existing `{model_name}.mar` will be
overwritten. Default to False.
"""
mmcv.mkdir_or_exist(output_folder)
config = mmcv.Config.fromfile(config_file)
with TemporaryDirectory() as tmpdir:
config.dump(f'{tmpdir}/config.py')
args = Namespace(
**{
'model_file': f'{tmpdir}/config.py',
'serialized_file': checkpoint_file,
'handler':
f'{Path(__file__).parent}/mmgen_{model_type}_handler.py',
'model_name': model_name or Path(checkpoint_file).stem,
'version': model_version,
'export_path': output_folder,
'force': force,
'requirements_file': None,
'extra_files': None,
'runtime': 'python',
'archive_format': 'default'
})
manifest = ModelExportUtils.generate_manifest_json(args)
package_model(args, manifest)
def test_onnx_file_is_not_none(self, patches):
onnx_file = 'some-file.onnx'
patches.utils.find_unique.return_value = onnx_file
patches.utils.convert_onnx_model.return_value = ('sym', 'param')
temp, exclude = ModelExportUtils.check_custom_model_types(
self.model_path)
patches.utils.convert_onnx_model.assert_called_once_with(
self.model_path, onnx_file, None)
assert len(temp) == 2
assert len(exclude) == 1
assert temp[0] == os.path.join(self.model_path, 'sym')
assert temp[1] == os.path.join(self.model_path, 'param')
assert exclude[0] == onnx_file
def pack(self,
model,
dst_path,
model_name,
model_version,
additional_files=None):
"""
Package a model into the MAR archive so that it can be used for serving using TorchServing
:param model: an object representing a model
:type model: torch.nn.Module
:param dst_path: the destination base path (do not include the filename) of the MAR
:type dst_path: str
:param model_name: the name of the model (will be also used to define the prediction endpoint)
:type model_name: str
:param model_version: a string encoding the version of the model
:type model_version: str
:param additional_files: an optional list of files that should be included in the MAR
:type additional_files: iterable
:return: None
"""
if additional_files is None:
additional_files = []
# save model
torch.save(model, os.path.join(self.tmp_dir, 'model.pt'))
args = ArgClass(
os.path.join(self.tmp_dir, 'model.pt'), self.handler, ','.join([
os.path.join(self.tmp_dir, 'pre_process_tform.pkl'),
os.path.join(self.tmp_dir, 'post_process_tform.pkl'),
os.path.join(self.tmp_dir, 'metadata.json')
] + additional_files), dst_path, model_name, model_version)
manifest = ModelExportUtils.generate_manifest_json(args)
package_model(args, manifest=manifest)
def test_with_exit(self):
files = ['a.onnx', 'b.onnx', 'c.txt']
suffix = '.onnx'
with pytest.raises(ModelArchiverError):
ModelExportUtils.find_unique(files, suffix)
def test_with_starts_with_dot(self):
assert ModelExportUtils.directory_filter('.gitignore', self.unwanted_dirs) is False
def test_with_return_true(self):
assert ModelExportUtils.directory_filter('my-model', self.unwanted_dirs) is True
def test_with_return_true(self):
assert ModelExportUtils.file_filter('abc.mxnet', self.files_to_exclude) is True
def test_with_unwanted_dirs(self):
assert ModelExportUtils.directory_filter('__MACOSX', self.unwanted_dirs) is False
def test_with_pyc(self):
assert ModelExportUtils.file_filter('abc.pyc', self.files_to_exclude) is False
def test_with_ds_store(self):
assert ModelExportUtils.file_filter('.DS_Store', self.files_to_exclude) is False
def test_export_file_already_exists_with_override_false(self, patches):
patches.path_exists.return_value = True
with pytest.raises(ModelArchiverError):
ModelExportUtils.check_mar_already_exists(
'some-model', None, False)
_validate_mar(patches)
def test_with_return_false(self):
assert ModelExportUtils.file_filter('abc.onnx', self.files_to_exclude) is False
def test_model(self):
mod = ModelExportUtils.generate_model(self.args)
assert mod.model_name == self.model_name
assert mod.handler == self.handler
def test_engine(self):
eng = ModelExportUtils.generate_engine(self.args)
assert eng.engine_name == EngineType(self.engine)
def test_publisher(self):
pub = ModelExportUtils.generate_publisher(self.args)
assert pub.email == self.email
assert pub.author == self.author
def test_export_file_already_exists_with_override(self, patches):
patches.path_exists.return_value = True
ModelExportUtils.check_mar_already_exists('some-model', None, True)
_validate_mar(patches)
def test_with_count_zero(self):
files = ['a.txt', 'b.txt', 'c.txt']
suffix = '.mxnet'
val = ModelExportUtils.find_unique(files, suffix)
assert val is None
def mmpose2torchserve(config_file: str,
checkpoint_file: str,
output_folder: str,
model_name: str,
model_version: str = '1.0',
force: bool = False):
"""Converts MMPose model (config + checkpoint) to TorchServe `.mar`.
Args:
config_file:
In MMPose config format.
The contents vary for each task repository.
checkpoint_file:
In MMPose checkpoint format.
The contents vary for each task repository.
output_folder:
Folder where `{model_name}.mar` will be created.
The file created will be in TorchServe archive format.
model_name:
If not None, used for naming the `{model_name}.mar` file
that will be created under `output_folder`.
If None, `{Path(checkpoint_file).stem}` will be used.
model_version:
Model's version.
force:
If True, if there is an existing `{model_name}.mar`
file under `output_folder` it will be overwritten.
"""
mmcv.mkdir_or_exist(output_folder)
config = mmcv.Config.fromfile(config_file)
with TemporaryDirectory() as tmpdir:
model_file = osp.join(tmpdir, 'config.py')
config.dump(model_file)
handler_path = osp.join(osp.dirname(__file__), 'mmpose_handler.py')
model_name = model_name or osp.splitext(
osp.basename(checkpoint_file))[0]
# use mmcv CheckpointLoader if checkpoint is not from a local file
if not osp.isfile(checkpoint_file):
ckpt = CheckpointLoader.load_checkpoint(checkpoint_file)
checkpoint_file = osp.join(tmpdir, 'checkpoint.pth')
with open(checkpoint_file, 'wb') as f:
torch.save(ckpt, f)
args = Namespace(
**{
'model_file': model_file,
'serialized_file': checkpoint_file,
'handler': handler_path,
'model_name': model_name,
'version': model_version,
'export_path': output_folder,
'force': force,
'requirements_file': None,
'extra_files': None,
'runtime': 'python',
'archive_format': 'default'
})
manifest = ModelExportUtils.generate_manifest_json(args)
package_model(args, manifest)
def test_clean_call(self, patches):
temp_files = ['a', 'b', 'c']
ModelExportUtils.clean_temp_files(temp_files)
patches.remove.assert_called()
assert patches.remove.call_count == len(temp_files)