fix(logger): Bugs in AzureMLFlowLogger from #646 (#685)

Finally got around to trying our work in #646, and encountered these two
issues

- a stale reference to `AnemoiMLFlowLogger`
- `getattr` -> `dict.get` 

I'll try to add a test or two that catches these.

Co-authored-by: Ana Prieto Nemesio <[email protected]>

Author: phinate

training/src/anemoi/training/diagnostics/callbacks/checkpoint.py

@@ -122,9 +122,9 @@ def tracker_metadata(self, trainer: pl.Trainer) -> dict:
         if self.config.diagnostics.log.mlflow.enabled:
             self._tracker_name = "mlflow"
 
-            from anemoi.training.diagnostics.mlflow.logger import AnemoiMLflowLogger
+            from anemoi.training.diagnostics.mlflow.logger import BaseAnemoiMLflowLogger
 
-            mlflow_logger = next(logger for logger in trainer.loggers if isinstance(logger, AnemoiMLflowLogger))
+            mlflow_logger = next(logger for logger in trainer.loggers if isinstance(logger, BaseAnemoiMLflowLogger))
             run_id = mlflow_logger.run_id
             run = mlflow_logger._mlflow_client.get_run(run_id)
 

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

@@ -30,12 +30,11 @@ def get_mlflow_logger(config: BaseSchema) -> None:
     del logger_config["enabled"]
 
     # backward compatibility to not break configs
-    logger_config["_target_"] = getattr(
-        logger_config,
-        "_target",
+    logger_config["_target_"] = logger_config.get(
+        "_target_",
         "anemoi.training.diagnostics.mlflow.logger.AnemoiMLflowLogger",
     )
-    logger_config["save_dir"] = getattr(logger_config, "save_dir", str(config.hardware.paths.logs.mlflow))
+    logger_config["save_dir"] = logger_config.get("save_dir", str(config.hardware.paths.logs.mlflow))
 
     logger = instantiate(
         logger_config,