[Michael-Duck]

Summary

Adds MLFlowLogger.log_model(), which logs a model to the MLflow run owned by the logger and optionally registers 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:

from lightning.pytorch import Trainer
from lightning.pytorch.loggers import MLFlowLogger

logger = MLFlowLogger(experiment_name="my-exp")
trainer = Trainer(logger=logger)
trainer.fit(model)

# log the model to the run and register it in the model registry
logger.log_model(model, registered_model_name="my-model")

The model is logged with MLflow's pytorch flavor, so it shows up as a model of the run in the MLflow UI and can be loaded again with mlflow.pyfunc.load_model("models:/my-model/1").

Notes

• The model is logged into the run that the logger manages, via mlflow.models.Model.log(..., run_id=...). mlflow.pytorch.log_model() would log into the globally active MLflow run and terminate that run when it exits, which would interfere with the run owned by the logger.
• serialization_format="pickle" is used by default when the installed MLflow supports it, because the newer pt2 format needs an input example to trace the model, which the logger cannot infer. Pass input_example=... (or serialization_format=...) to override this.
• The log_model argument of the constructor is untouched and keeps controlling whether the checkpoints are logged as artifacts. The new method logs a model object instead of a checkpoint.

Tests

pytest tests/tests_pytorch/loggers/test_mlflow.py

• new test test_mlflow_logger_log_model: without this change it fails with AttributeError: 'MLFlowLogger' object has no attribute 'log_model', with this change it passes (it logs the model into the run, registers it in the model registry and loads it back for inference).
• whole file with this change: 21 passed, 1 failed, 2 errors. The failure (test_mlflow_run_name_setting) and the two errors (the environment variable leak check tripping over MLFLOW_TRACKING_URI, which MLflow 3.x exports) are pre-existing on this machine with MLflow 3.17 and fail identically without this change (20 passed, 1 failed, 2 errors).

pytest tests/tests_pytorch/loggers/ (wider run): 107 passed, 32 failed, 2 skipped, 6 xfailed, 2 errors. All failures are environmental and unrelated to this change: tensorboard/tensorboardX and jsonargparse are not installed, and the ddp_spawn based test_logger_initialization cases cannot spawn processes in this sandbox (PermissionError: [Errno 1] Operation not permitted).

ruff check and ruff format --check on the changed files pass, sphinx-lint passes on the changed documentation.

Fixes #21437