From 5eab7930b8453bd5d97e4dce4a8077bb6c2ef344 Mon Sep 17 00:00:00 2001 From: michael-duck Date: Fri, 23 Oct 2026 02:04:08 +0100 Subject: [PATCH 1/2] Add `MLFlowLogger.log_model` to log and register models with MLflow (#21437) feat: add `MLFlowLogger.log_model` to log a model to the MLflow run and to optionally register it in the MLflow model registry `MLFlowLogger` could only log checkpoints as plain MLflow artifacts, so the "Register Model" action of the MLflow UI was never available for models trained with Lightning. This adds the missing `mlflow.log_model` equivalent: ```python logger = MLFlowLogger(experiment_name="my-exp") trainer = Trainer(logger=logger) trainer.fit(model) logger.log_model(model, registered_model_name="my-model") ``` The model gets logged with the `pytorch` flavor of MLflow into the run owned by the logger and, if `registered_model_name` is given, registered in the MLflow model registry. Fixes #21437 Signed-off-by: michael-duck --- src/lightning/pytorch/CHANGELOG.md | 2 + src/lightning/pytorch/loggers/mlflow.py | 56 ++++++++++++++++++++++ tests/tests_pytorch/loggers/test_mlflow.py | 30 ++++++++++++ 3 files changed, 88 insertions(+) diff --git a/src/lightning/pytorch/CHANGELOG.md b/src/lightning/pytorch/CHANGELOG.md index 1c73ea347751a..302fb65d11b9a 100644 --- a/src/lightning/pytorch/CHANGELOG.md +++ b/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)) diff --git a/src/lightning/pytorch/loggers/mlflow.py b/src/lightning/pytorch/loggers/mlflow.py index ff9b2b0d7e542..f444a5dd76a69 100644 --- a/src/lightning/pytorch/loggers/mlflow.py +++ b/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: diff --git a/tests/tests_pytorch/loggers/test_mlflow.py b/tests/tests_pytorch/loggers/test_mlflow.py index c7f9dbe1fe2c6..ad3e72ed5bb94 100644 --- a/tests/tests_pytorch/loggers/test_mlflow.py +++ b/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) From ee9790f03253e1be6729430e6fc9a44969d03820 Mon Sep 17 00:00:00 2001 From: michael-duck Date: Fri, 23 Oct 2026 02:04:38 +0100 Subject: [PATCH 2/2] Document the new `MLFlowLogger.log_model` method in the experiment managers guide Signed-off-by: michael-duck --- docs/source-pytorch/visualize/supported_exp_managers.rst | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/source-pytorch/visualize/supported_exp_managers.rst b/docs/source-pytorch/visualize/supported_exp_managers.rst index 3f8726b7a7fed..9754e2309d8c2 100644 --- a/docs/source-pytorch/visualize/supported_exp_managers.rst +++ b/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`. ----