MLflow ONNX asset is incomplete

[CGuichard]

Context

For registered models of type ONNX the default behavior of MLflow on log_model is to set save_as_external_data=True.

This means if you try to load the ONNX file in the MLflow artifacts you will get the error "onnx.onnx_cpp2py_export.checker.ValidationError: Data of TensorProto ( tensor name: XXX) should be stored in XXX, but it doesn't exist or is not accessible.".

IMPORTANT: when loading with mlflow.onnx.load_model MLflow does some processing by reconstructing an out-of-the-box working ONNX from the ONNX proto file and checkpoint file in the artifacts, hence why the problem does not occur.

Docs: https://www.mlflow.org/docs/latest/model/#onnx-onnx

The default behavior for saving ONNX files is to use the ONNX save option save_as_external_data=True in order to support model files that are in excess of 2GB. For edge deployments of small model files, this may create issues. If you need to save a small model as a single file for such deployment considerations, you can set the parameter save_as_external_data=False in either mlflow.onnx.save_model() or mlflow.onnx.log_model() to force the serialization of the model as a small file. Note that if the model is in excess of 2GB, saving as a single file will not work.

Task

In order to load correctly you must download the "checkpoint file" alongside the ONNX "proto" file before loading the latter. Currently SharingHub server only exposes the ONNX "proto" file in the assets, we should also add this checkpoint file in the assets if it exists.

If save_as_external_data=False the ONNX file will not be a proto and will work.