Checkpoint Export Functionality Should Require Bundle Directory Path Only

[ericspod]

Is your feature request related to a problem? Please describe.
Currently to export a bundle as a Torchscript file one has to specify the files to include individually, eg:

python -m monai.bundle ckpt_export network_def \
    --filepath foo.ts \
    --ckpt_file foo/models/model.pt \
    --meta_file foo/configs/metadata.json \
    --config_file foo/configs/inference.json

Describe the solution you'd like
I think the optimal behaviour is to specify export of a whole bundle directory which will include all the config files in the configs directory and use models/model.pt as the model weight file implicitly. The network definition network_def is in config files in the tutorials example but perhaps it should be moved to the metadata.json file since there will really only be the one primary network definition per bundle. The actual command line to export would be :

python -m monai.bundle ckpt_export --bundle_path /path/to/bundle/dir --out_path foo.ts 

(I would also suggest changing filepath to out_path to be clearer about what it is.)

Describe alternatives you've considered
The current command works just fine so this is a QoL improvement.

[Nic-Ma]

Thanks for your feedback, maybe let's try to collect feedback from first users and adjust APIs of bundle soon.

[ericspod]

Also we should allow any number of json files to be stored, so anything that's in the configs directory gets kept in the exported Torchscript file without their names being changed, currently the config file gets stored as "config". I think we want to support using downloaded models that have multiple configs in them with other software which takes as a parameter a config name to use to do something.

[Nic-Ma]

Hi @ericspod ,

Currently, the ckpt_export API can support multiple config files but will merge them together and save with "config":
https://github.com/Project-MONAI/MONAI/blob/dev/monai/bundle/scripts.py#L374
So you mean we should save the configs separately instead of merging them before exporting?

Thanks.

[ericspod]

I would suggest we should have them as separate names so that the user can choose one or the other depending on the task. The files should also retain their extensions since the parser relies on these to determine if they're json or yaml.

A separate enhance that would be help is being able to load metadata and configuration from the zip file directly, I'm currently having to save files out to a temporary directory and read them from there but we can have a cleaner process I think.

[Nic-Ma]

Sure, let me try to investigate deeper soon.

Thanks for the feedback and ideas.

[Nic-Ma]

Hi @ericspod ,

About your suggestion "load metadata and configuration from the zip file directly", the "zip" here I think you mean the torchscript file with embedded config & metadata? We can support that by enhancing your previous API:
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/torchscript_utils.py#L103
Or you mean just from zip bundle in Github storage? I think the monai.bundle download API(@yiheng-wang-nv is working) will automatically extract the downloaded zip.

Thanks.

[ericspod]

By zip I mean the Torchscript zip file. My current workaround in the deploy context is this:

def get_bundle_config(bundle_path):
    """
    Get the configuration parser from the specified Torchscript bundle file path.
    """
    name, _ = os.path.splitext(os.path.basename(bundle_path))
    parser = ConfigParser()
    
    with tempfile.TemporaryDirectory() as td:
        archive = zipfile.ZipFile(bundle_path, "r")
        archive.extract(name + "/extra/metadata.json", td)
        archive.extract(name + "/extra/config", td)

        os.rename(f"{td}/{name}/extra/config", f"{td}/{name}/extra/config.json")

        parser.read_meta(f=f"{td}/{name}/extra/metadata.json")
        parser.read_config(f=f"{td}/{name}/extra/config.json")

        parser.parse()

    return parser

I would like to be able to read meta or config data from the zip via a data stream I can open using the zipfile API, all this would then go into an API function we'd provide. I'd also like to be able to select different json files since the Torchscript zip file could contain multiples like inference.json and training.json just like the bundle directory rather than putting multiple things together into config. I had to add the extension here as well to get things to parse. We could also support parsing from a zip file that contains the bundle (ie. is not a Torchscript object) as well by specifying paths in a similar way.

[Nic-Ma]

Hi @ericspod ,

Thanks for your explanation.
About export / load configs with TorchScript model, I think maybe we can narrow down the requirements for current version and here is my draft proposal:

1. Enhance the ckpt_export API to save configs separately in the TorchScript model.
2. Enhance the load_net_with_metadata API to get expected config files from TorchScript model by more_extra_files:
   https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/torchscript_utils.py#L103
   Maybe we can also wrapper this API in @yiheng-wang-nv 's new load API of bundle scripts for TorchScript case.

For the config files:

1. One option is to keep the extension in the name, but I didn't find a way to serialize a local file into TorchScript model directly to avoid calling load/dumps, do you know some idea?
2. Another option is to simplify the logic: load JSON or YAML configs, but always serialize to JSON format string in the TorchScript model, and don't keep the extension in the name.

What do you think?

Thanks.

[Nic-Ma]

Hi @ericspod ,

I implemented the above proposal idea with simplified option 2 to save config files in the TorchScript in PR: #4159.
And added unit tests to load TorchScript model and meta, configs.
Could you please help review it?

Thanks in advance.