Repository: mlflow/mlflow
Pull request: Add 'save_format' param to Keras autolog()
Signed-off-by: Bruno Alvisio <bruno.alvisio@gmail.com>

<!-- 🚨 We recommend pull requests be filed from a non-master branch on a repository fork (e.g. <username>:fix-xxx). 🚨 -->

## Related Issues/PRs

Resolve https://github.com/mlflow/mlflow/issues/7024

## What changes are proposed in this pull request?

Allow configuring the format in which the model is saved in `keras.autolog()` 

## How is this patch tested?

Added test to load model in 'HDF5' format
<!--
If you're unsure about what to test, where to add tests, or how to run tests, please feel free to ask. We'd be happy to help.
-->
- [x] I have written tests (not required for typo or doc fix) and confirmed the proposed feature/bug-fix/change works.

<!--
Please describe how you confirmed the proposed feature/bug-fix/change works here. For example, if you fixed an MLflow client API, you could attach the code that didn't work prior to the fix but works now, or if you added a new feature on MLflow UI, you could attach a video that demonstrates the feature.
-->

## Does this PR change the documentation?

- [ ] No. You can skip the rest of this section.
- [x] Yes. Make sure the changed pages / sections render correctly in the documentation preview.

## Release Notes

### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [x] Yes. Give a description of this change to be included in the release notes for MLflow users.

(Details in 1-2 sentences. You can just refer to another PR with a description if this PR is part of a larger change.)

### What component(s), interfaces, languages, and integrations does this PR affect?
Components
- [ ] `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
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline 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
- [x] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface
- [ ] `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

Language
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->
<a name="release-note-category"></a>
### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [x] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
[FR] Allowing changing the Keras serialization format when using `mlflow.keras.autolog()`
### Willingness to contribute

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

### Proposal Summary

Currently, it is not possible to change the serialization format of the Keras models (tf) to HDF% when using the `mlflow.keras.autolog()`. It would be great if this could be configured. One way could be for `autolog()` to take a `kwargs` that is passed to inner functions so that serialization configuration can be set when using `autolog()` 

### Motivation

> #### What is the use case for this feature?

> #### Why is this use case valuable to support for MLflow users in general?

> #### Why is this use case valuable to support for your project(s) or organization?

> #### Why is it currently difficult to achieve this use case?


### 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
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/pipelines`: Pipelines, Pipeline APIs, Pipeline configs, Pipeline 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
- [X] `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

Hints from issue discussion:
Hi @balvisio, thank you for your feature request. I think that introducing a `save_format` argument to `mlflow.tensorflow.autolog()` makes sense to me, which may be slightly easier to implement and easier for users to reason about than generic **kwargs for `mlflow.tensorflow.autolog()`. We would be very excited to review a pull request that implements this capability. 
Thanks @dbczumar. Will work in the PR in the coming weeks.

Commit message: Add 'save_format' param to Keras autolog()

Signed-off-by: Bruno Alvisio <bruno.alvisio@gmail.com>

Patch to review:
diff --git a/mlflow/keras.py b/mlflow/keras.py
index 0c6604098f88c..3bfa41d1daaaf 100644
--- a/mlflow/keras.py
+++ b/mlflow/keras.py
@@ -603,6 +603,7 @@ def autolog(
     disable_for_unsupported_versions=False,
     silent=False,
     registered_model_name=None,
+    save_format="tf",
 ):  # pylint: disable=unused-argument
     # pylint: disable=E0611
     """
@@ -668,6 +669,9 @@ def autolog(
     :param registered_model_name: If given, each time a model is trained, it is registered as a
                                   new model version of the registered model with this name.
                                   The registered model is created if it does not already exist.
+    :param save_format: File format to save the model at the end of the training. Allowed
+                        values are: ``tf`` for saving the model in ``SavedModel`` format
+                        or ``h5`` for saving the model in ``HDF5`` format.
     """
     import keras
 
@@ -681,6 +685,12 @@ def autolog(
             stacklevel=2,
         )
 
+    if save_format != "tf" and save_format != "h5":
+        raise Exception(
+            "Invalid value for `save_format` argument. "
+            "Valid values for `save_format` are `tf` or `h5` only."
+        )
+
     def getKerasCallback(metrics_logger):
         class __MLflowKerasCallback(keras.callbacks.Callback, metaclass=ExceptionSafeClass):
             """
@@ -733,6 +743,7 @@ def on_train_end(self, logs=None):
                         self.model,
                         artifact_path="model",
                         registered_model_name=registered_model_name,
+                        **{"save_format": save_format},
                     )
 
             # As of Keras 2.4.0, Keras Callback implementations must define the following
diff --git a/tests/keras/test_keras_autolog.py b/tests/keras/test_keras_autolog.py
index b0fd27e4f2a80..1806ce4227e00 100644
--- a/tests/keras/test_keras_autolog.py
+++ b/tests/keras/test_keras_autolog.py
@@ -410,3 +410,20 @@ def test_autolog_registering_model(random_train_data, random_one_hot_labels):
 
         registered_model = MlflowClient().get_registered_model(registered_model_name)
         assert registered_model.name == registered_model_name
+
+
+def test_autolog_load_saved_hdf5_model(random_train_data, random_one_hot_labels):
+    mlflow.keras.autolog(save_format="h5")
+
+    data = random_train_data
+    labels = random_one_hot_labels
+
+    model = create_model()
+    with mlflow.start_run() as run:
+        model.fit(data, labels, epochs=10)
+        mlflow.keras.load_model("runs:/" + run.info.run_id + "/model")
+
+
+def test_autolog_invalid_model_save():
+    with pytest.raises(Exception, match="Invalid value for `save_format` argument."):
+        mlflow.keras.autolog(save_format="foo")


Code review comments:
diff_hunk: @@ -681,6 +685,12 @@ def autolog(
             stacklevel=2,
         )
 
+    if save_format != "tf" and save_format != "h5":
+        raise Exception(
+            "Invalid value for `save_format` argument. "
path: mlflow/keras.py
text: @user1:
Can we include the invalid value in the message?

@user1:
```suggestion
            "Invalid value for `save_format` argument: {save_format}. "
```

diff_hunk: @@ -681,6 +685,12 @@ def autolog(
             stacklevel=2,
         )
 
+    if save_format != "tf" and save_format != "h5":
+        raise Exception(
+            "Invalid value for `save_format` argument. "
+            "Valid values for `save_format` are `tf` or `h5` only."
+        )
path: mlflow/keras.py
text: @user1:
```suggestion
        raise MlflowException(
            "Invalid value for `save_format` argument: {save_format}. "
            "Valid values for `save_format` are 'tf' or 'h5' only.",
            INVALID_PARAMETER,
        )
```

diff_hunk: @@ -733,6 +743,7 @@ def on_train_end(self, logs=None):
                         self.model,
                         artifact_path="model",
                         registered_model_name=registered_model_name,
+                        **{"save_format": save_format},
path: mlflow/keras.py
text: @user1:
```suggestion
                        save_format=save_format,
```

diff_hunk: @@ -410,3 +410,20 @@ def test_autolog_registering_model(random_train_data, random_one_hot_labels):
 
         registered_model = MlflowClient().get_registered_model(registered_model_name)
         assert registered_model.name == registered_model_name
+
+
+def test_autolog_load_saved_hdf5_model(random_train_data, random_one_hot_labels):
+    mlflow.keras.autolog(save_format="h5")
+
+    data = random_train_data
+    labels = random_one_hot_labels
+
+    model = create_model()
+    with mlflow.start_run() as run:
+        model.fit(data, labels, epochs=10)
+        mlflow.keras.load_model("runs:/" + run.info.run_id + "/model")
path: tests/keras/test_keras_autolog.py
text: @user1:
Is there a way to check the model is saved in the specified format?

diff_hunk: @@ -681,6 +685,12 @@ def autolog(
             stacklevel=2,
         )
 
+    if save_format != "tf" and save_format != "h5":
path: mlflow/keras.py
text: @user1:
```suggestion
    if save_format not in ["tf", "h5"]:
```