From cbd3d07270502f8767e5981ed21bcabecaa0b251 Mon Sep 17 00:00:00 2001
From: Gert Mertes <13658335+gmertes@users.noreply.github.com>
Date: Wed, 14 Aug 2024 16:39:48 +0100
Subject: [PATCH] Fix null tracker metadata and disable mlflow by default (#18)

* fix: return empty tracker metadata when no tracker

* chore: mlflow disabled by default

* chore: update changelog

* Merge develop into fix/null-tracker
---
 CHANGELOG.md                                           |  1 +
 .../training/config/diagnostics/eval_rollout.yaml      |  2 +-
 src/anemoi/training/diagnostics/callbacks/__init__.py  | 10 ++++------
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a494f12f..3bb2d7d1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -59,6 +59,7 @@ Keep it human-readable, your future self will thank you!
  - Adjusted PyPI release infrastructure to common ECMWF workflow
  - Bumped versions in Pre-commit hooks
  - Fix crash when logging hyperparameters with missing values in the config
+ - Fixed "null" tracker metadata when tracking is disabled, now returns an empty dict
 
 ### Removed
  - Dependency on mlflow-export-import
diff --git a/src/anemoi/training/config/diagnostics/eval_rollout.yaml b/src/anemoi/training/config/diagnostics/eval_rollout.yaml
index 15ae7009..f8bdc0a6 100644
--- a/src/anemoi/training/config/diagnostics/eval_rollout.yaml
+++ b/src/anemoi/training/config/diagnostics/eval_rollout.yaml
@@ -78,7 +78,7 @@ log:
   tensorboard:
     enabled: False
   mlflow:
-    enabled: True
+    enabled: False
     offline: False
     authentication: False
     log_model: False
diff --git a/src/anemoi/training/diagnostics/callbacks/__init__.py b/src/anemoi/training/diagnostics/callbacks/__init__.py
index 2153b7da..8deb75f5 100644
--- a/src/anemoi/training/diagnostics/callbacks/__init__.py
+++ b/src/anemoi/training/diagnostics/callbacks/__init__.py
@@ -789,10 +789,9 @@ def tracker_metadata(self, trainer: pl.Trainer) -> dict:
                     "url": run.url,
                     "project": run.project,
                 }
-            else:
-                self._tracker_metadata = {}
+            return {self._tracker_name: self._tracker_metadata}
 
-        elif self.config.diagnostics.log.mlflow.enabled:
+        if self.config.diagnostics.log.mlflow.enabled:
             self._tracker_name = "mlflow"
 
             from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
@@ -808,10 +807,9 @@ def tracker_metadata(self, trainer: pl.Trainer) -> dict:
                     "url": run.info.artifact_uri,
                     "project": run.info.experiment_id,
                 }
-            else:
-                self._tracker_metadata = {}
+            return {self._tracker_name: self._tracker_metadata}
 
-        return {self._tracker_name: self._tracker_metadata}
+        return {}
 
     def _save_checkpoint(self, trainer: pl.Trainer, lightning_checkpoint_filepath: str) -> None:
         if trainer.is_global_zero: