Repository: mlflow/mlflow

Issue:
[FR] Parametrize the save as external data for ONNX model logging
### Willingness to contribute

Yes. I would be willing to contribute this feature with guidance from the MLflow community.

### Proposal Summary

Since #7808, ONNX models are systematically saved with the weights saved in a separate file. This introduced a behavior change, and it is needed only for big models. The associated code is here:

https://github.com/mlflow/mlflow/blob/5712024236c45e7830506bb3adba3c6e42131625/mlflow/onnx.py#L175

It would be nice to allow to enable or not this behavior so that a small model can still be serialized as a single file.

### Motivation

> #### What is the use case for this feature?

- Keep the behavior consistent this the previous one
- Serialize small models as a single file

> #### Why is this use case valuable to support for MLflow users in general?

Splitting a model into several files makes using it more complex, primarily when mlflow is not used for inference.

> #### Why is this use case valuable to support for your project(s) or organization?

We never use models bigger than 2GB, and we process models are single-file objects.

> #### Why is it currently difficult to achieve this use case?

It is not possible to log ONNX models as a single file with recent releases of mlflow.

### Details

_No response_

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.

Resolving patch:
diff --git a/mlflow/onnx/__init__.py b/mlflow/onnx/__init__.py
--- a/mlflow/onnx/__init__.py
+++ b/mlflow/onnx/__init__.py
@@ -92,6 +92,7 @@ def save_model(
     onnx_execution_providers=None,
     onnx_session_options=None,
     metadata=None,
+    save_as_external_data=True,
 ):
     """
     Save an ONNX model to a path on the local file system.
@@ -139,6 +140,7 @@ def save_model(
                                  See onnxruntime API for further descriptions:
                                  https://onnxruntime.ai/docs/api/python/api_summary.html#sessionoptions
     :param metadata: Custom metadata dictionary passed to the model and stored in the MLmodel file.
+    :param save_as_external_data: Save tensors to external file(s).
 
                      .. Note:: Experimental: This parameter may change or be removed in a future
                                              release without warning.
@@ -167,7 +169,7 @@ def save_model(
 
     # Save onnx-model
     if Version(onnx.__version__) >= Version("1.9.0"):
-        onnx.save_model(onnx_model, model_data_path, save_as_external_data=True)
+        onnx.save_model(onnx_model, model_data_path, save_as_external_data=save_as_external_data)
     else:
         onnx.save_model(onnx_model, model_data_path)
 
@@ -446,6 +448,7 @@ def log_model(
     onnx_execution_providers=None,
     onnx_session_options=None,
     metadata=None,
+    save_as_external_data=True
 ):
     """
     Log an ONNX model as an MLflow artifact for the current run.
@@ -498,6 +501,7 @@ def log_model(
                                  See onnxruntime API for further descriptions:
                                  https://onnxruntime.ai/docs/api/python/api_summary.html#sessionoptions
     :param metadata: Custom metadata dictionary passed to the model and stored in the MLmodel file.
+    :param save_as_external_data: Save tensors to external file(s).
 
                      .. Note:: Experimental: This parameter may change or be removed in a future
                                              release without warning.
@@ -519,4 +523,5 @@ def log_model(
         onnx_execution_providers=onnx_execution_providers,
         onnx_session_options=onnx_session_options,
         metadata=metadata,
+        save_as_external_data=save_as_external_data
     )


Test patch:
diff --git a/tests/onnx/test_onnx_model_export.py b/tests/onnx/test_onnx_model_export.py
--- a/tests/onnx/test_onnx_model_export.py
+++ b/tests/onnx/test_onnx_model_export.py
@@ -253,16 +253,11 @@ def test_model_save_load(onnx_model, model_path):
 
 
 def test_model_save_load_nonexternal_data(onnx_model, model_path):
-    original_save_model = onnx.save_model
     if Version(onnx.__version__) >= Version("1.9.0"):
+        onnx.convert_model_to_external_data = mock.Mock()
 
-        def onnx_save_nonexternal(
-            model, path, save_as_external_data
-        ):  # pylint: disable=unused-argument
-            original_save_model(model, path, save_as_external_data=False)
-
-        with mock.patch("onnx.save_model", wraps=onnx_save_nonexternal):
-            mlflow.onnx.save_model(onnx_model, model_path)
+        mlflow.onnx.save_model(onnx_model, model_path, save_as_external_data=False)
+        onnx.convert_model_to_external_data.assert_not_called()
 
         # Loading ONNX model
         onnx.checker.check_model = mock.Mock()
@@ -346,10 +341,17 @@ def test_model_save_load_multiple_inputs(onnx_model_multiple_inputs_float64, mod
     assert onnx.checker.check_model.called
 
 
+@pytest.mark.parametrize("save_as_external_data", [True, False])
 def test_model_save_load_evaluate_pyfunc_format_multiple_inputs(
-    onnx_model_multiple_inputs_float64, data_multiple_inputs, predicted_multiple_inputs, model_path
+    onnx_model_multiple_inputs_float64,
+    data_multiple_inputs,
+    predicted_multiple_inputs,
+    model_path,
+    save_as_external_data,
 ):
-    mlflow.onnx.save_model(onnx_model_multiple_inputs_float64, model_path)
+    mlflow.onnx.save_model(
+        onnx_model_multiple_inputs_float64, model_path, save_as_external_data=save_as_external_data
+    )
 
     # Loading pyfunc model
     pyfunc_loaded = mlflow.pyfunc.load_model(model_path)