feat: sanitise dict keys for mlflow

psycop/common/model_training_v2/loggers/mlflow_logger.py

@@ -13,6 +13,10 @@
 from psycop.common.model_training_v2.trainer.task.base_metric import CalculatedMetric
 
 
+def sanitise_dict_keys(d: dict[str, Any]) -> dict[str, Any]:
+    return replace_symbols_in_dict_keys(d=d, symbol2replacement={"@": "", "*": "_"})
+
+
 @BaselineRegistry.loggers.register("mlflow_logger")
 class MLFlowLogger(BaselineLogger):
     def __init__(
@@ -64,8 +68,4 @@ def log_metric(self, metric: CalculatedMetric) -> None:
 
     def log_config(self, config: dict[str, Any]):
         config = flatten_nested_dict(config)
-        clean_config = replace_symbols_in_dict_keys(
-            d=config,
-            symbol2replacement={"@": "", "*": "_"},
-        )
-        mlflow.log_params(clean_config)
+        mlflow.log_params(sanitise_dict_keys(config))

psycop/common/sequence_models/train.py

@@ -9,8 +9,11 @@
 import torch
 from torch.utils.data import DataLoader
 
-from psycop.common.global_utils.config_utils import flatten_nested_dict
+from psycop.common.global_utils.config_utils import (
+    flatten_nested_dict,
+)
 
+from ..model_training_v2.loggers.mlflow_logger import sanitise_dict_keys
 from .config_utils import load_config, parse_config
 
 log = logging.getLogger(__name__)
@@ -52,7 +55,7 @@ def train(config_path: Path | None = None) -> None:
             # update config
             log.info("Logging config")
             flat_config = flatten_nested_dict(config_dict)
-            logger.log_hyperparams(flat_config)
+            logger.log_hyperparams(sanitise_dict_keys(flat_config))
 
     # Load and filter dataset
     filter_fn = cfg.model_and_dataset.model.filter_and_reformat