Repository: Lightning-AI/pytorch-lightning

Issue:
MLFlow Logger slows training steps dramatically, despite only setting metrics to be logged on epoch
## 🐛 Bug

When using the MLFlow logger, with a remote server, logging per step introduces latency which slows the training loop.
I have tried to configure logging of metrics only per epoch, however it seems this still results in much slower performance. I suspect the logger is still communicating with the MLFlow server on each training step.

### To Reproduce
1. Start an MLFlow server locally
```
mlflow ui
```
2. Run the minimal code example below as is, (with MLFlow logger set to use the default file uri.)
3. Uncomment out the `tracking_uri` to use the local MLFlow server and run the code again. You will see a 2-3 times drop in the iterations per second.

#### Code sample
```
import torch
from torch.utils.data import TensorDataset, DataLoader
import pytorch_lightning as pl

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.num_examples = 5000
        self.num_valid = 1000
        self.batch_size = 64
        self.lr = 1e-3
        self.wd = 1e-2
        self.num_features = 2
        self.linear = torch.nn.Linear(self.num_features, 1)
        self.loss_func = torch.nn.MSELoss()
        self.X = torch.rand(self.num_examples, self.num_features)
        self.y = self.X.matmul(torch.rand(self.num_features, 1)) + torch.rand(self.num_examples)
        
    def forward(self, x):
        return self.linear(x)

    def train_dataloader(self): 
        ds = TensorDataset(self.X[:-self.num_valid], self.X[:-self.num_valid])
        dl = DataLoader(ds, batch_size=self.batch_size)
        return dl

    def val_dataloader(self): 
        ds = TensorDataset(self.X[-self.num_valid:], self.X[-self.num_valid:])
        dl = DataLoader(ds, batch_size=self.batch_size)
        return dl

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=self.lr, weight_decay=self.wd)

    def training_step(self, batch, batch_idx):
        x, y = batch
        yhat = self(x)
        loss = self.loss_func(yhat, y)
        result = pl.TrainResult(minimize=loss)
        result.log('train_loss', loss, on_epoch=True, on_step=False)
        return result

    def validation_step(self, batch, batch_idx):
        x, y = batch
        yhat = self(x)
        loss = self.loss_func(yhat, y)
        result = pl.EvalResult(early_stop_on=loss)
        result.log('val_loss', loss, on_epoch=True, on_step=False)
        return result

if __name__ == '__main__':
    from pytorch_lightning.loggers import TensorBoardLogger, MLFlowLogger
    mlf_logger = MLFlowLogger(
        experiment_name=f"MyModel",
        # tracking_uri="http://localhost:5000"
    )
    trainer = pl.Trainer(
        min_epochs=5,
        max_epochs=50,
        early_stop_callback=True,
        logger=mlf_logger
    )
    model = MyModel()
    trainer.fit(model)
```

### Expected behavior

When using the TrainResult and EvalResult, or manually handling metric logging using the `training_epoch_end` and `validation_epoch_end` callbacks. It should be possible to avoid the MLFlow logger from communicating with the server in each training loop. 
This would make it feasible to implement the MLFlow when a remote server is used for experiment tracking.

### Environment
```
* CUDA:
	- GPU:
	- available:         False
	- version:           None
* Packages:
	- numpy:             1.18.2
	- pyTorch_debug:     False
	- pyTorch_version:   1.6.0+cpu
	- pytorch-lightning: 0.9.0
	- tensorboard:       2.2.0
	- tqdm:              4.48.2
* System:
	- OS:                Linux
	- architecture:
		- 64bit
		-
	- processor:         x86_64
	- python:            3.7.9
	- version:           #1 SMP Tue May 26 11:42:35 UTC 2020
```
### Additional context

We host a MLFlow instance in AWS and would like to be able to track experiments without affecting the training speed. 
It appears that in general the MLFlow logger is much less performant than the default Tensorboard Logger, but this would not be much of a problem if we could avoid calls to the logger during the training loop.

### Solution
I've done a bit of debugging in the codebase and have been able to isolate the cause in two places
https://github.com/PyTorchLightning/pytorch-lightning/blob/d438ad8a8db3e76d3ed4e3c6bc9b91d6b3266b8e/pytorch_lightning/loggers/mlflow.py#L125-L129
Here `self.experiment` is called regardless of whether `self._run_id` exists. If we add an `if not self._run_id` here we avoid calling `self._mlflow_client.get_experiment_by_name(self._experiment_name)` on each step.
However we still call it each time we log metrics to MFlow, because of the property `self.experiment`.

https://github.com/PyTorchLightning/pytorch-lightning/blob/d438ad8a8db3e76d3ed4e3c6bc9b91d6b3266b8e/pytorch_lightning/loggers/mlflow.py#L100-L112
Here if we store `expt` within the logger and only call `self._mlflow_client.get_experiment_by_name` when it does not exist, we eliminate all overhead, it runs as fast as fast as the tensorboard logger and all the mlflow logging appears to be working as expected.

I'd be happy to raise a PR for this fix.


Discussion:
Hi! thanks for your contribution!, great first issue!
have you tried to just increase the row_log_interval, its a trainer flag that controls how often logs are sent to the logger.
I mean, your network is a single linear layer, you probably run through epochs super fast.
I am not yet convinced it is a bug, but I'll try your example code
hey @awaelchli, Thanks for replying!
The model above is a contrived example, upon further testing I have realised that the performance difference between MFLow logger and the Tensorboard logger is not inherent to the MLFlow client.

I've done some debugging and added a solution section to the issue. It appears to be in in the `experiment` property of the MLFlowLogger. Each time `.experiment` is accessed, `self._mlflow_client.get_experiment_by_name(self._experiment_name)` is called, which communicates with the MLFlow server.

It seems we can store the response of this method thereby needing to call it only once, and this seems to resolve the dramatic difference between the Tensorboard and MLFlow Logger.
oh ok, that makes sense. Would you like to send a PR with your suggestion and see if the tests pass? Happy to review it. 
yeah sure, I'll link it here shortly.
Did you encounter this #3392 problem as well?

Resolving patch:
diff --git a/pytorch_lightning/loggers/mlflow.py b/pytorch_lightning/loggers/mlflow.py
--- a/pytorch_lightning/loggers/mlflow.py
+++ b/pytorch_lightning/loggers/mlflow.py
@@ -109,15 +109,15 @@ def experiment(self) -> MlflowClient:
             self.logger.experiment.some_mlflow_function()
 
         """
-        expt = self._mlflow_client.get_experiment_by_name(self._experiment_name)
-
-        if expt:
-            self._experiment_id = expt.experiment_id
-        else:
-            log.warning(f'Experiment with name {self._experiment_name} not found. Creating it.')
-            self._experiment_id = self._mlflow_client.create_experiment(name=self._experiment_name)
-
-        if not self._run_id:
+        if self._experiment_id is None:
+            expt = self._mlflow_client.get_experiment_by_name(self._experiment_name)
+            if expt is not None:
+                self._experiment_id = expt.experiment_id
+            else:
+                log.warning(f'Experiment with name {self._experiment_name} not found. Creating it.')
+                self._experiment_id = self._mlflow_client.create_experiment(name=self._experiment_name)
+
+        if self._run_id is None:
             run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=self.tags)
             self._run_id = run.info.run_id
         return self._mlflow_client


Test patch:
diff --git a/tests/loggers/test_mlflow.py b/tests/loggers/test_mlflow.py
--- a/tests/loggers/test_mlflow.py
+++ b/tests/loggers/test_mlflow.py
@@ -1,5 +1,8 @@
 import os
 
+from unittest import mock
+from mlflow.tracking import MlflowClient
+
 from pytorch_lightning import Trainer
 from pytorch_lightning.loggers import MLFlowLogger
 from tests.base import EvalModelTemplate
@@ -39,3 +42,13 @@ def test_mlflow_logger_dirs_creation(tmpdir):
     assert set(os.listdir(tmpdir / exp_id / run_id / 'params')) == model.hparams.keys()
     assert trainer.checkpoint_callback.dirpath == (tmpdir / exp_id / run_id / 'checkpoints')
     assert set(os.listdir(trainer.checkpoint_callback.dirpath)) == {'epoch=0.ckpt'}
+
+
+def test_mlflow_experiment_id_retrieved_once(tmpdir):
+    logger = MLFlowLogger('test', save_dir=tmpdir)
+    get_experiment_name = logger._mlflow_client.get_experiment_by_name
+    with mock.patch.object(MlflowClient, 'get_experiment_by_name', wraps=get_experiment_name) as mocked:
+        _ = logger.experiment
+        _ = logger.experiment
+        _ = logger.experiment
+        assert mocked.call_count == 1