Update mlflow.py

Author: KAVYANSHTYAGI

src/lightning/pytorch/loggers/mlflow.py

@@ -109,6 +109,12 @@ def any_lightning_module_function_or_hook(self):
         ModuleNotFoundError:
             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)
+        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.
+
     """
 
     LOGGER_JOIN_CHAR = "-"
@@ -126,6 +132,7 @@ def __init__(
         run_id: Optional[str] = None,
         synchronous: Optional[bool] = None,
     ):
+
         if not _MLFLOW_AVAILABLE:
             raise ModuleNotFoundError(str(_MLFLOW_AVAILABLE))
         if synchronous is not None and not _MLFLOW_SYNCHRONOUS_AVAILABLE:
@@ -151,6 +158,7 @@ def __init__(
         from mlflow.tracking import MlflowClient
 
         self._mlflow_client = MlflowClient(tracking_uri)
+        self._logged_metrics = set()  # Track (key, step)
 
     @property
     @rank_zero_experiment
@@ -201,6 +209,7 @@ def experiment(self) -> "MlflowClient":
             resolve_tags = _get_resolve_tags()
             run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=resolve_tags(self.tags))
             self._run_id = run.info.run_id
+            self._logged_metrics.clear() 
         self._initialized = True
         return self._mlflow_client
 
@@ -257,7 +266,7 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
             if isinstance(v, str):
                 log.warning(f"Discarding metric with string value {k}={v}.")
                 continue
-
+        
             new_k = re.sub("[^a-zA-Z0-9_/. -]+", "", k)
             if k != new_k:
                 rank_zero_warn(
@@ -266,8 +275,15 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
                     category=RuntimeWarning,
                 )
                 k = new_k
+        
+            metric_id = (k, step or 0)
+            if metric_id in self._logged_metrics:
+                continue 
+            self._logged_metrics.add(metric_id)
+        
             metrics_list.append(Metric(key=k, value=v, timestamp=timestamp_ms, step=step or 0))
 
+
         self.experiment.log_batch(run_id=self.run_id, metrics=metrics_list, **self._log_batch_kwargs)
 
     @override