torch_model_archiver_programmatic_mar_packaging_for_torchserve.py

1import os
2from model_archiver.model_packaging import package_model
3from model_archiver.model_packaging_utils import ModelExportUtils
4from model_archiver.manifest_components.manifest import Manifest
5from model_archiver.arg_parser import ArgParser
6
7# Define the configuration for the model archive
8class ModelArchiverConfig:
9    def __init__(self):
10        self.model_name = "resnet-18"
11        self.handler = "resnet-18" # Can be a path to a python file or a default handler name
12        self.model_file = "model.py"
13        self.serialized_file = "resnet-18.pth"
14        self.version = "1.0"
15        self.export_path = "model_store"
16        self.extra_files = ""
17        self.runtime = "python"
18        self.force = True
19        self.archive_format = "default"
20        self.requirements_file = None
21
22def quickstart_archiver():
23    # 1. Setup configuration
24    args = ModelArchiverConfig()
25    
26    # 2. Ensure export directory exists
27    if not os.path.exists(args.export_path):
28        os.makedirs(args.export_path)
29
30    # 3. Create manifest object
31    manifest = Manifest(args)
32
33    # 4. Package the model
34    # Note: This assumes 'model.py' and 'resnet-18.pth' exist in the current working directory
35    try:
36        package_model(args, manifest)
37        print(f"Successfully archived model to {args.export_path}/{args.model_name}.mar")
38    except Exception as e:
39        print(f"Failed to create model archive: {e}")
40
41if __name__ == "__main__":
42    quickstart_archiver()