fix: bug for mlflow offline logging (#675)

## Description
After merging the support for azure, we missed a small bug that breaks
the offline logging. This PR fixes the bug and add some tests to capture
this better in the future.

Main issue before was that for forking and resuming the logger was not
setting the tracking_uri to the save_dir.

## What problem does this change solve?
<!-- Describe if it's a bugfix, new feature, doc update, or breaking
change -->

## What issue or task does this change relate to?
<!-- link to Issue Number -->

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)

Author: anaprietonem

training/src/anemoi/training/diagnostics/mlflow/logger.py

@@ -346,16 +346,16 @@ def __init__(
         LOGGER.info("Maximum number of params allowed to be logged is: %s", max_params_length)
 
         self.tracking_uri = tracking_uri
-        if (self._resumed or self._forked) and self.offline:
-            self.tracking_uri = save_dir
-
         # Before creating the run we need to overwrite the tracking_uri and save_dir if offline
         if self.offline:
-            # OFFLINE - When we run offline we can pass a save_dir pointing to a local path
-            self.tracking_uri = None
-            if save_dir is None:
-                # otherwise, by default we create a dir called "None"... not ideal
-                save_dir = "./mlruns"
+            if self._resumed or self._forked:
+                self.tracking_uri = save_dir
+            else:
+                # OFFLINE - When we run offline we can pass a save_dir pointing to a local path
+                self.tracking_uri = None
+                if save_dir is None:
+                    # otherwise, by default we create a dir called "None"... not ideal
+                    save_dir = "./mlruns"
 
         else:
             # ONLINE - When we pass a tracking_uri to mlflow then it will ignore the

training/tests/unit/diagnostics/mlflow/test_mlflow_logger.py

@@ -1,6 +1,9 @@
 from pathlib import Path
 
+import omegaconf
 import pytest
+import yaml
+from hydra.utils import instantiate
 
 from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
 from anemoi.training.schemas.diagnostics import MlflowSchema
@@ -19,6 +22,35 @@ def tmp_uri(monkeypatch: pytest.MonkeyPatch, tmp_path: str) -> Path:
     return uri
 
 
+@pytest.fixture
+def default_offline_config(tmp_path: str) -> omegaconf.DictConfig:
+    base = """
+    diagnostics:
+      log:
+        mlflow:
+            _target_: anemoi.training.diagnostics.mlflow.logger.AnemoiMLflowLogger
+            offline: Trure
+            authentication: False
+            tracking_uri: 'https:test.int'
+            experiment_name: 'anemoi-debug'
+            project_name: 'Anemoi'
+            system: False
+            terminal: True
+            run_name: null
+            on_resume_create_child: True
+            expand_hyperparams:
+                - config
+            http_max_retries: 35
+            max_params_length: 2000
+            save_dir: '/scratch/example'
+    """
+
+    cfg = omegaconf.OmegaConf.create(yaml.safe_load(base))
+    cfg.diagnostics.log.mlflow.save_dir = tmp_path
+
+    return cfg
+
+
 @pytest.fixture
 def default_logger(tmp_path: str, tmp_uri: str) -> AnemoiMLflowLogger:
     return AnemoiMLflowLogger(
@@ -31,6 +63,52 @@ def default_logger(tmp_path: str, tmp_uri: str) -> AnemoiMLflowLogger:
     )
 
 
+def create_run(save_dir: str, experiment_name: str) -> str:
+    import contextlib
+
+    import mlflow
+
+    mlflow.set_tracking_uri(f"file://{save_dir}")
+    with contextlib.suppress(mlflow.exceptions.MlflowException):
+        mlflow.create_experiment(experiment_name)
+    mlflow.set_experiment(experiment_name)
+    with mlflow.start_run():
+        mlflow.log_param("lr", 0.001)
+        return mlflow.active_run().info.run_id
+
+
+def test_offline_logger(default_offline_config: omegaconf.DictConfig) -> None:
+    mlflow_logger = instantiate(default_offline_config.diagnostics.log.mlflow)
+    assert not mlflow_logger.tracking_uri
+
+
+def test_offline_resumed_logger(default_offline_config: omegaconf.DictConfig) -> None:
+
+    run_id = create_run(
+        save_dir=default_offline_config.diagnostics.log.mlflow.save_dir,
+        experiment_name=default_offline_config.diagnostics.log.mlflow.experiment_name,
+    )
+    logger_resumed = instantiate(
+        default_offline_config.diagnostics.log.mlflow,
+        run_id=run_id,
+        fork_run_id=None,
+    )
+    assert logger_resumed.tracking_uri == default_offline_config.diagnostics.log.mlflow.save_dir
+
+
+def test_offline_forked_logger(default_offline_config: omegaconf.DictConfig) -> None:
+    fork_run_id = create_run(
+        save_dir=default_offline_config.diagnostics.log.mlflow.save_dir,
+        experiment_name=default_offline_config.diagnostics.log.mlflow.experiment_name,
+    )
+    logger_forked = instantiate(
+        default_offline_config.diagnostics.log.mlflow,
+        run_id=None,
+        fork_run_id=fork_run_id,
+    )
+    assert logger_forked.tracking_uri == default_offline_config.diagnostics.log.mlflow.save_dir
+
+
 def test_mlflowlogger_params_limit(default_logger: AnemoiMLflowLogger) -> None:
 
     default_logger._max_params_length = 3