Add MLFlowLogger.log_model to log and register models with MLflow

docs/source-pytorch/visualize/supported_exp_managers.rst

@@ -84,6 +84,12 @@ Access the mlflow logger from any function (except the LightningModule *init*) t
             fake_images = torch.Tensor(32, 3, 28, 28)
             mlf_logger.add_image("generated_images", fake_images, 0)
 
+Log the model to the run and optionally register it in the MLflow model registry:
+
+.. code-block:: python
+
+    mlf_logger.log_model(model, registered_model_name="my-model")
+
 Here's the full documentation for the :class:`~lightning.pytorch.loggers.MLFlowLogger`.
 
 ----

src/lightning/pytorch/CHANGELOG.md

@@ -10,6 +10,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `MLFlowLogger.log_model` to log a model to the MLflow run and optionally register it in the MLflow model registry ([#21437](https://github.com/Lightning-AI/pytorch-lightning/issues/21437))
+
 - Added `suggest_integrations` flag to `Trainer` to control whether optional integration suggestions (e.g., litmodels, litlogger) are shown in logs ([#21632](https://github.com/Lightning-AI/pytorch-lightning/pull/21632))
 
 - Fixed ``ModelParallelStrategy`` single-file checkpointing when ``torch.compile`` wraps the model so optimizer states no longer raise ``KeyError`` during save ([#21357](https://github.com/Lightning-AI/pytorch-lightning/issues/21357))

src/lightning/pytorch/loggers/mlflow.py

@@ -16,6 +16,7 @@
 -------------
 """
 
+import inspect
 import logging
 import os
 import re
@@ -29,6 +30,7 @@
 import yaml
 from lightning_utilities.core.imports import RequirementCache
 from torch import Tensor
+from torch.nn import Module
 from typing_extensions import override
 
 from lightning.fabric.utilities.logger import _add_prefix, _convert_params, _flatten_dict
@@ -270,6 +272,60 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
 
         self.experiment.log_batch(run_id=self.run_id, metrics=metrics_list, **self._log_batch_kwargs)
 
+    @rank_zero_only
+    def log_model(
+        self,
+        model: Module,
+        artifact_path: str = "model",
+        registered_model_name: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Log a model to the MLflow run.
+
+        The model gets logged with the ``pytorch`` flavor of MLflow. This makes the model available in the MLflow UI
+        and allows registering it in the MLflow model registry.
+
+        This is unrelated to the ``log_model`` argument of the logger, which controls whether the checkpoints created
+        by the :class:`~lightning.pytorch.callbacks.model_checkpoint.ModelCheckpoint` callback get logged as MLflow
+        artifacts.
+
+        Args:
+            model: The model to log.
+            artifact_path: The run relative path the model gets logged to.
+            registered_model_name: If given, the model gets registered in the MLflow model registry under this name.
+            **kwargs: Additional keyword arguments passed to the MLflow ``pytorch`` flavor, for example
+                ``input_example`` or ``signature``.
+
+        """
+        import mlflow.pytorch
+        from mlflow.models import Model
+
+        if (
+            "input_example" not in kwargs
+            and "serialization_format" not in kwargs
+            and "serialization_format" in inspect.signature(mlflow.pytorch.save_model).parameters
+        ):
+            # The `pt2` format requires an input example to trace the model, which cannot be inferred here.
+            # The `pickle` format supports any model.
+            kwargs["serialization_format"] = "pickle"
+
+        # `artifact_path` got deprecated in MLflow 3.0 in favor of `name`
+        if "name" in inspect.signature(Model.log).parameters:
+            kwargs["artifact_path"] = None
+            kwargs["name"] = artifact_path
+        else:
+            kwargs["artifact_path"] = artifact_path
+
+        # `mlflow.pytorch.log_model` logs the model to the currently active run, but this logger owns the run.
+        # `Model.log` allows logging the model to the run created by this logger.
+        Model.log(
+            flavor=mlflow.pytorch,
+            pytorch_model=model,
+            registered_model_name=registered_model_name,
+            run_id=self.run_id,
+            **kwargs,
+        )
+
     @override
     @rank_zero_only
     def finalize(self, status: str = "success") -> None:

tests/tests_pytorch/loggers/test_mlflow.py

@@ -15,6 +15,7 @@
 from unittest import mock
 from unittest.mock import MagicMock, Mock
 
+import numpy as np
 import pytest
 
 from lightning.pytorch import Trainer
@@ -427,3 +428,32 @@ def test_set_tracking_uri(mlflow_mock):
     mlflow_mock.set_tracking_uri.assert_not_called()
     _ = logger.experiment
     mlflow_mock.set_tracking_uri.assert_called_with("the_tracking_uri")
+
+
+def test_mlflow_logger_log_model(tmp_path):
+    """Test that a model can be logged to the run and registered in the MLflow model registry."""
+    if not _MLFLOW_AVAILABLE:
+        pytest.skip("test for explicit file creation requires mlflow dependency to be installed.")
+
+    import mlflow.pyfunc
+    from mlflow.tracking import MlflowClient
+
+    logger = MLFlowLogger("test", tracking_uri=f"sqlite:///{tmp_path}/mlflow.db")
+    model = BoringModel()
+
+    logger.log_model(model, "my-model", registered_model_name="test-model")
+
+    client = MlflowClient(logger._tracking_uri)
+    model_versions = client.search_model_versions("name='test-model'")
+    assert len(model_versions) == 1
+    assert model_versions[0].run_id == logger.run_id
+    assert model_versions[0].status == "READY"
+    # logging a model should not terminate the run
+    assert client.get_run(logger.run_id).info.status == "RUNNING"
+
+    # the logged model can be loaded and used for inference
+    loaded_model = mlflow.pyfunc.load_model(f"models:/test-model/{model_versions[0].version}")
+    assert loaded_model.predict(np.zeros((1, 32), dtype="float32")).shape == (1, 2)
+
+    # MLflow exports the tracking URI to the environment, which would leak out of the test
+    os.environ.pop("MLFLOW_TRACKING_URI", None)