Apply suggestions from code review

Co-authored-by: Nicki Skafte Detlefsen <[email protected]>

Author: Borda

requirements/pytorch/test.txt

@@ -18,4 +18,4 @@ fastapi  # for `ServableModuleValidator`  # not setting version as re-defined in
 uvicorn  # for `ServableModuleValidator`  # not setting version as re-defined in App
 
 tensorboard >=2.9.1, <2.21.0  # for `TensorBoardLogger`
-mlflow
+mlflow >=3.0.0, <4.0   # for `MLFlowLogger 

src/lightning/pytorch/loggers/mlflow.py

@@ -110,10 +110,10 @@ def any_lightning_module_function_or_hook(self):
             If required MLFlow package is not installed on the device.
 
     Note:
-        As of vX.XX, MLFlowLogger will skip logging any metric (same name and step)
+        MLFlowLogger will skip logging any metric (same name and step)
         more than once per run, to prevent database unique constraint violations on
         some MLflow backends (such as PostgreSQL). Only the first value for each (metric, step)
-        pair will be logged per run. This improves robustness for all users.
+        pair will be logged per run.
 
     """
 
@@ -157,7 +157,7 @@ def __init__(
         from mlflow.tracking import MlflowClient
 
         self._mlflow_client = MlflowClient(tracking_uri)
-        self._logged_metrics = set()  # Track (key, step)
+        self._logged_metrics: Set[Tuple[str, float]] = set()  # Track (key, step)
 
     @property
     @rank_zero_experiment
@@ -275,6 +275,7 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
                 )
                 k = new_k
 
+            # Prevent trying to log duplicate metric value
             metric_id = (k, step or 0)
             if metric_id in self._logged_metrics:
                 continue