Repository: Lightning-AI/pytorch-lightning
Pull request: Add support for logging the model checkpoints to MLFlowLogger
## What does this PR do?
This PR adds a log_model parameter to MLFlowLogger as the one in WandbLogger (#9138).

The log_model parameter can be set to:
- "all": checkpoints are logged during training.
- True: checkpoints are logged at the end of training, except when 
`:paramref: '~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint.save_top_k' == "-1"` which also logs every checkpoint during training.
- False: no checkpoint is logged.

In particular, the function behavior is not exactly the same as in WandbLogger, since MLFlowClient().log_artifact takes different parameters as input:
- In the WandbLogger, the artifact is logged together with a metadata dictionary and some aliases.
- In the MLFlowLogger, it is not possible to log metadata and aliases together with the artifact, therefore (to keep consistency) a temporary folder is created, metadata.yml and aliases.txt are created inside of it, then the temporary folder is logged on mlflow together with the checkpoint artifact in the model/checkpoints/CHECKPOINT_NAME/ folder.

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

<!-- FILL IN or None -->

### Does your PR introduce any breaking changes? If yes, please list them.
None.

## Before submitting
- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review
Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

## Did you have fun?
Coding this small PR was funny 😃


Problem statement:
Add the WandbLogger log_model parameter to MLFLowLogger
## 🚀 Feature

At the moment I am using both MLFlowLogger and WandbLogger, and I realized that MLFlowLogger is missing the log_model parameter (I also quickly checked the other loggers and it seems that WandbLogger is the only one having this parameter). Will it be possible to add the log_model parameter to MLFlowLogger (and perhaps to all the other loggers where this is applicable)?

### Motivation

Consistency among the different loggers: it would be nice to have this functionality available in all the loggers.

Hints from issue discussion:
haven't tried, but mlflow does support logging artifacts/models, so this seems like a good addition.
https://www.mlflow.org/docs/latest/python_api/mlflow.pytorch.html
https://www.mlflow.org/docs/latest/python_api/mlflow.html#mlflow.log_artifact (with lightning example)
Dear @AlessioQuercia,

Would you be interested to contribute such feature ?

Best,
T.C
Yes, I will try to implement it as close as possible to the one implemented for the [WandbLogger](https://pytorch-lightning.readthedocs.io/en/stable/extensions/generated/pytorch_lightning.loggers.WandbLogger.html#pytorch_lightning.loggers.WandbLogger) during this weekend.

I will keep you posted,

Best regards,
Alessio
Awesome @AlessioQuercia,

I am assigning you to this ticket ! Keep us updated on your progress.

Best,
T.C
I opened a PR (#9187) for this feature @tchaton, @rohitgr7. It seems that there are some failed tests that I did not have locally, I will try to solve them.

The feature is slightly different from the one in WandbLogger, since MLFlowLogger `log_artifact()` cannot log metadata and aliases. Therefore the logic is that metadata.yml and aliases.txt temp files are created and logged together with the checkpoint (3 artifacts per checkpoint) in the model/checkpoint/CHECKPOINT_NAME/ directory.

If you prefer to only log the checkpoint and discard the metadata and aliases info, I can simply delete the temporary folder part of the code and it should still work. Please, let me know if you have any other suggestions on how to manage this problem.

Best regards,
Alessio
I think metadata is fine. Would it be a good idea to merge aliases into metadata as a simple key-value pair? in wandb as well maybe?
It makes totally sense for MLFlow (in fact I wanted to put it inside metadata initially), but I don't know if it makes sense for wandb (there should be a specific space for aliases in wandb if I am not wrong).

Commit message: Updated CHANGELOG

Patch to review:
diff --git a/src/pytorch_lightning/CHANGELOG.md b/src/pytorch_lightning/CHANGELOG.md
index 5f687de11a8da..a6896c11b18bb 100644
--- a/src/pytorch_lightning/CHANGELOG.md
+++ b/src/pytorch_lightning/CHANGELOG.md
@@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added `log_model` parameter to `MLFlowLogger` ([#9187](https://github.com/PyTorchLightning/pytorch-lightning/pull/9187))
 - Added support for requeueing slurm array jobs ([#15040](https://github.com/Lightning-AI/lightning/pull/15040))
 - Added native AMP support for `ddp_fork` (and associated alias strategies) with CUDA GPUs ([#14983](https://github.com/Lightning-AI/lightning/pull/14983))
 - Added `BatchSizeFinder` callback ([#11089](https://github.com/Lightning-AI/lightning/pull/11089))
diff --git a/src/pytorch_lightning/loggers/mlflow.py b/src/pytorch_lightning/loggers/mlflow.py
index 37be8dac9f631..424882565654b 100644
--- a/src/pytorch_lightning/loggers/mlflow.py
+++ b/src/pytorch_lightning/loggers/mlflow.py
@@ -18,14 +18,19 @@
 import logging
 import os
 import re
+import tempfile
 from argparse import Namespace
+from pathlib import Path
 from time import time
 from typing import Any, Dict, Mapping, Optional, Union
 
+import torch
+import yaml
 from lightning_utilities.core.imports import module_available
 
+from pytorch_lightning.callbacks import Checkpoint
 from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment
-from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict
+from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict, _scan_checkpoints
 from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
 
 log = logging.getLogger(__name__)
@@ -108,6 +113,13 @@ def any_lightning_module_function_or_hook(self):
         save_dir: A path to a local directory where the MLflow runs get saved.
             Defaults to `./mlflow` if `tracking_uri` is not provided.
             Has no effect if `tracking_uri` is provided.
+        log_model: Log checkpoints created by :class:`~pytorch_lightning.callbacks.Checkpoint`
+            as MLFlow artifacts.
+            * if ``log_model == 'all'``, checkpoints are logged during training.
+            * if ``log_model == True``, checkpoints are logged at the end of training, except when
+              :paramref:`~pytorch_lightning.callbacks.Checkpoint.save_top_k` ``== -1``
+              which also logs every checkpoint during training.
+            * if ``log_model == False`` (default), no checkpoint is logged.
         prefix: A string to put at the beginning of metric keys.
         artifact_location: The location to store run artifacts. If not provided, the server picks an appropriate
             default.
@@ -127,6 +139,7 @@ def __init__(
         tracking_uri: Optional[str] = os.getenv("MLFLOW_TRACKING_URI"),
         tags: Optional[Dict[str, Any]] = None,
         save_dir: Optional[str] = "./mlruns",
+        log_model: Optional[bool] = False,
         prefix: str = "",
         artifact_location: Optional[str] = None,
         run_id: Optional[str] = None,
@@ -145,6 +158,9 @@ def __init__(
         self._run_name = run_name
         self._run_id = run_id
         self.tags = tags
+        self._log_model = log_model
+        self._logged_model_time = {}
+        self._checkpoint_callback = None
         self._prefix = prefix
         self._artifact_location = artifact_location
 
@@ -261,6 +277,11 @@ def finalize(self, status: str = "success") -> None:
             status = "FINISHED"
         elif status == "failed":
             status = "FAILED"
+
+        # log checkpoints as artifacts
+        if self._checkpoint_callback:
+            self._scan_and_log_checkpoints(self._checkpoint_callback)
+
         if self.experiment.get_run(self.run_id):
             self.experiment.set_terminated(self.run_id, status)
 
@@ -292,3 +313,59 @@ def version(self) -> Optional[str]:
             The run id.
         """
         return self.run_id
+
+    def after_save_checkpoint(self, checkpoint_callback: Checkpoint) -> None:
+        # log checkpoints as artifacts
+        if self._log_model == "all" or self._log_model is True and checkpoint_callback.save_top_k == -1:
+            self._scan_and_log_checkpoints(checkpoint_callback)
+        elif self._log_model is True:
+            self._checkpoint_callback = checkpoint_callback
+
+    def _scan_and_log_checkpoints(self, checkpoint_callback: Checkpoint) -> None:
+        # get checkpoints to be saved with associated score
+        checkpoints = _scan_checkpoints(checkpoint_callback, self._logged_model_time)
+
+        # log iteratively all new checkpoints
+        for t, p, s, tag in checkpoints:
+            metadata = {
+                # Ensure .item() is called to store Tensor contents
+                "score": s.item() if isinstance(s, torch.Tensor) else s,
+                "original_filename": Path(p).name,
+                "Checkpoint": {
+                    k: getattr(checkpoint_callback, k)
+                    for k in [
+                        "monitor",
+                        "mode",
+                        "save_last",
+                        "save_top_k",
+                        "save_weights_only",
+                        "_every_n_train_steps",
+                        "_every_n_val_epochs",
+                    ]
+                    # ensure it does not break if `Checkpoint` args change
+                    if hasattr(checkpoint_callback, k)
+                },
+            }
+            aliases = ["latest", "best"] if p == checkpoint_callback.best_model_path else ["latest"]
+
+            # Artifact path on mlflow
+            artifact_path = f"model/checkpoints/{Path(p).stem}"
+
+            # Log the checkpoint
+            self.experiment.log_artifact(self._run_id, p, artifact_path)
+
+            # Create a temporary directory to log on mlflow
+            with tempfile.TemporaryDirectory(prefix="test", suffix="test", dir=os.getcwd()) as tmp_dir:
+                # Log the metadata
+                with open(f"{tmp_dir}/metadata.yaml", "w") as tmp_file_metadata:
+                    yaml.dump(metadata, tmp_file_metadata, default_flow_style=False)
+
+                # Log the aliases
+                with open(f"{tmp_dir}/aliases.txt", "w") as tmp_file_aliases:
+                    tmp_file_aliases.write(str(aliases))
+
+                # Log the metadata and aliases
+                self.experiment.log_artifacts(self._run_id, tmp_dir, artifact_path)
+
+            # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name)
+            self._logged_model_time[p] = t
diff --git a/tests/tests_pytorch/loggers/test_mlflow.py b/tests/tests_pytorch/loggers/test_mlflow.py
index ad81a5e6cb47f..54f57d77e5359 100644
--- a/tests/tests_pytorch/loggers/test_mlflow.py
+++ b/tests/tests_pytorch/loggers/test_mlflow.py
@@ -18,6 +18,7 @@
 import pytest
 
 from pytorch_lightning import Trainer
+from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.demos.boring_classes import BoringModel
 from pytorch_lightning.loggers import _MLFLOW_AVAILABLE, MLFlowLogger
 from pytorch_lightning.loggers.mlflow import MLFLOW_RUN_NAME, resolve_tags
@@ -277,3 +278,45 @@ def test_mlflow_logger_finalize_when_exception(*_):
     assert logger._initialized
     logger.finalize("failed")
     logger.experiment.set_terminated.assert_called_once_with(logger.run_id, "FAILED")
+
+
+@mock.patch("pytorch_lightning.loggers.mlflow.mlflow")
+@mock.patch("pytorch_lightning.loggers.mlflow.MlflowClient")
+@pytest.mark.parametrize("log_model", ["all", True, False])
+def test_mlflow_log_model(client, mlflow, tmpdir, log_model):
+    """Test that the logger creates the folders and files in the right place."""
+
+    # Reset mocks
+    client.return_value.log_artifact.reset_mock()
+    client.return_value.log_artifacts.reset_mock()
+    # Get model, logger, trainer and train
+    model = BoringModel()
+    logger = MLFlowLogger("test", save_dir=tmpdir, log_model=log_model)
+    checkpoint_callback = ModelCheckpoint(dirpath=tmpdir / "model_checkpoints")
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        callbacks=[checkpoint_callback],
+        logger=logger,
+        max_epochs=2,
+        limit_train_batches=3,
+        limit_val_batches=3,
+    )
+    trainer.fit(model)
+
+    if log_model == "all":
+        # Checkpoint log
+        assert client.return_value.log_artifact.call_count == 2
+        # Metadata and aliases log
+        assert client.return_value.log_artifacts.call_count == 2
+
+    elif log_model is True:
+        # Checkpoint log
+        client.return_value.log_artifact.assert_called_once()
+        # Metadata and aliases log
+        client.return_value.log_artifacts.assert_called_once()
+
+    elif log_model is False:
+        # Checkpoint log
+        assert not client.return_value.log_artifact.called
+        # Metadata and aliases log
+        assert not client.return_value.log_artifacts.called


Code review comments:
diff_hunk: @@ -108,6 +113,13 @@ def any_lightning_module_function_or_hook(self):
         save_dir: A path to a local directory where the MLflow runs get saved.
             Defaults to `./mlflow` if `tracking_uri` is not provided.
             Has no effect if `tracking_uri` is provided.
+        log_model: Log checkpoints created by :class:`~pytorch_lightning.callbacks.Checkpoint`
path: src/pytorch_lightning/loggers/mlflow.py
text: @user1:
```suggestion
        log_model: Log checkpoints created by :class:`~pytorch_lightning.callbacks.model_checkpoint.ModelCheckpoint`
```

diff_hunk: @@ -127,6 +139,7 @@ def __init__(
         tracking_uri: Optional[str] = os.getenv("MLFLOW_TRACKING_URI"),
         tags: Optional[Dict[str, Any]] = None,
         save_dir: Optional[str] = "./mlruns",
+        log_model: Optional[bool] = False,
path: src/pytorch_lightning/loggers/mlflow.py
text: @user1:
The type here needs to reflect the different types of choices listed in the docs above. So I suggest to go for 

```suggestion
        log_model: Literal[True, False, "all"] = False,
```

diff_hunk: @@ -277,3 +278,45 @@ def test_mlflow_logger_finalize_when_exception(*_):
     assert logger._initialized
     logger.finalize("failed")
     logger.experiment.set_terminated.assert_called_once_with(logger.run_id, "FAILED")
+
+
+@mock.patch("pytorch_lightning.loggers.mlflow.mlflow")
+@mock.patch("pytorch_lightning.loggers.mlflow.MlflowClient")
+@pytest.mark.parametrize("log_model", ["all", True, False])
+def test_mlflow_log_model(client, mlflow, tmpdir, log_model):
+    """Test that the logger creates the folders and files in the right place."""
+
+    # Reset mocks
+    client.return_value.log_artifact.reset_mock()
+    client.return_value.log_artifacts.reset_mock()
+    # Get model, logger, trainer and train
+    model = BoringModel()
+    logger = MLFlowLogger("test", save_dir=tmpdir, log_model=log_model)
+    checkpoint_callback = ModelCheckpoint(dirpath=tmpdir / "model_checkpoints")
path: tests/tests_pytorch/loggers/test_mlflow.py
text: @user1:
Is there a specific reason why you chose to create a checkpoint callback instance with this particular path? `Trainer(enable_checkpointing=True)` would also work (default) right?