Closed
Description
🚀 The feature
A python API for model archiver
Motivation, pitch
It'd be nice to be able to archive a model directly from the python application, without having to run a bash command with a subshell or others.
Alternatives
I've used a script like this which is almost identical to what the entry point of torch-model-archiver is. The only addition is a pydantic model as config instead of arg parse.
The application then imports ModelArchiverConfig and generate_model_archive to build a new mar file.
One of the limitations here is not being able to pass a custom handler in a clean way when it's part of the library/application.
from typing import Optional, Literal, Union
import os
import logging
import shutil
from pathlib import Path
from pydantic import BaseModel
from model_archiver.model_archiver_error import ModelArchiverError
from model_archiver.model_packaging_utils import ModelExportUtils
from model_archiver.manifest_components.manifest import RuntimeType
class ModelArchiverConfig(BaseModel):
model_name: str
serialized_file: Optional[str] = None
model_file: Optional[str] = None
handler: str
extra_files: Optional[str] = None
runtime: str = RuntimeType.PYTHON.value
export_path: Union[str, Path] = os.getcwd()
archive_format: Literal["default", "tgz", "no-archive"] = "default"
force: bool = False
version: str
requirements_file: Optional[str] = None
def _package_model(config: ModelArchiverConfig, manifest: str):
"""
Internal helper for the exporting model command line interface.
"""
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
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, config.force, config.archive_format
)
# Step 2 : Copy all artifacts to temp directory
artifact_files = {
"model_file": model_file,
"serialized_file": serialized_file,
"handler": handler,
"extra_files": extra_files,
"requirements-file": requirements_file,
}
model_path = ModelExportUtils.copy_artifacts(model_name, **artifact_files)
# Step 2 : Zip 'em all up
ModelExportUtils.archive(
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)
raise e
def generate_model_archive(config: ModelArchiverConfig):
"""
Generate a model archive file
:return:
"""
logging.basicConfig(format="%(levelname)s - %(message)s")
manifest = ModelExportUtils.generate_manifest_json(config)
_package_model(config, manifest=manifest)Additional context
No response