torchserve_model_archiver_package_model_to_mar_file.py

1import os
2from model_archiver.model_packaging import package_model
3from model_archiver.model_archiver_config import ModelArchiverConfig
4
5# 1. Define the configuration for the model archiver
6# Note: Ensure that the 'model_file', 'handler', and 'serialized_file' exist at the specified paths.
7config = ModelArchiverConfig(
8    model_name="my_model",
9    version="1.0",
10    model_file="model.py",            # Path to the file containing the model architecture
11    handler="handler.py",            # Path to the custom handler file
12    serialized_file="model.pt",      # Path to the serialized parameter file (.pt or .pth)
13    export_path="./model_store",     # Directory where the .mar file will be saved
14    runtime="python",                # Runtime of the model (default is python)
15    force=True                       # Overwrite existing .mar file if it exists
16)
17
18# 2. Create the export directory if it doesn't exist
19if not os.path.exists(config.export_path):
20    os.makedirs(config.export_path)
21
22# 3. Package the model into a .mar file
23try:
24    package_model(config)
25    print(f"Successfully archived model to {config.export_path}/{config.model_name}.mar")
26except Exception as e:
27    print(f"Failed to archive model: {e}")