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

Start an MLFlow server locally

mlflow ui


Run the minimal code example below as is, (with MLFlow logger set to use the default file uri.)
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

  
    
      pytorch-lightning/pytorch_lightning/loggers/mlflow.py
    
    
        Lines 125 to 129
      in
      d438ad8
    
  
  
    

        
          
           @property 
        

        
          
           def run_id(self): 
        

        
          
               # create the experiment if it does not exist to get the run id 
        

        
          
               _ = self.experiment 
        

        
          
               return self._run_id 
        
    
  


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.

  
    
      pytorch-lightning/pytorch_lightning/loggers/mlflow.py
    
    
        Lines 100 to 112
      in
      d438ad8
    
  
  
    

        
          
               @property 
        

        
          
               @rank_zero_experiment 
        

        
          
               def experiment(self) -> MlflowClient: 
        

        
          
                   r""" 
        

        
          
                   Actual MLflow object. To use MLflow features in your 
        

        
          
                   :class:`~pytorch_lightning.core.lightning.LightningModule` do the following. 
        

        
          
            
        

        
          
                   Example:: 
        

        
          
            
        

        
          
                       self.logger.experiment.some_mlflow_function() 
        

        
          
            
        

        
          
                   """ 
        

        
          
                   expt = self._mlflow_client.get_experiment_by_name(self._experiment_name) 
        
    
  


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.