Update mlflow logger to latest aifs-mono version

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

@@ -5,7 +5,6 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-
 import json
 import logging
 import os

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

@@ -1,13 +1,11 @@
-# (C) Copyright 2024 ECMWF.
-#
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
 # This software is licensed under the terms of the Apache Licence Version 2.0
 # which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 # In applying this licence, ECMWF does not waive the privileges and immunities
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
 import io
-import logging
 import os
 import re
 import sys
@@ -26,7 +24,9 @@
 from pytorch_lightning.loggers.mlflow import _flatten_dict
 from pytorch_lightning.utilities.rank_zero import rank_zero_only
 
-LOGGER = logging.getLogger(__name__)
+from anemoi.training.utils.logger import get_code_logger
+
+LOGGER = get_code_logger(__name__)
 
 
 def get_mlflow_run_params(config, tracking_uri):
@@ -213,14 +213,18 @@ def _handle_csi(line):
             # removes the cursor up and down symbols from the line
             # skip tqdm status bar updates ending with "curser up" but the last one in buffer to save space
             def _remove_csi(line):
+                # replacing the leftmost non-overlapping occurrences of pattern ansi_csi_re in string line by the replacement ""
                 return re.sub(ansi_csi_re, b"", line)
 
             for match in ansi_csi_re.finditer(line):
                 arg, command = match.groups()
                 arg = int(arg.decode()) if arg else 1
-                if command == b"A" and (b"0%" not in line and not self._shutdown):  # cursor up
+                if command == b"A" and (
+                    b"0%" not in line and not self._shutdown and b"[INFO]" not in line and b"[DEBUG]" not in line
+                ):  # cursor up
                     # only keep x*10% status updates from tqmd status bars that end with a cursor up
                     # always keep shutdown commands
+                    # always keep logger info and debug prints
                     line = b""
             return _remove_csi(line)
 

src/anemoi/training/utils/logger.py

@@ -0,0 +1,48 @@
+# (C) Copyright 2024 European Centre for Medium-Range Weather Forecasts.
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import logging
+import sys
+import time
+
+
+def get_code_logger(name: str, debug: bool = True) -> logging.Logger:
+    """Returns a logger with a custom level and format.
+
+    We use ISO8601 timestamps and UTC times.
+
+    Parameters
+    ----------
+    name : str
+        Name of logger object
+    debug : bool, optional
+        set logging level to logging.DEBUG; else set to logging.INFO, by default True
+
+    Returns
+    -------
+    logging.Logger
+        Logger object
+    """
+    # create logger object
+    logger = logging.getLogger(name=name)
+    if not logger.hasHandlers():
+        # logging level
+        level = logging.DEBUG if debug else logging.INFO
+        # logging format
+        datefmt = "%Y-%m-%dT%H:%M:%SZ"
+        msgfmt = "[%(asctime)s] [%(filename)s:%(lineno)s - %(funcName).30s] [%(levelname)s] %(message)s"
+        # handler object
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(level)
+        formatter = logging.Formatter(msgfmt, datefmt=datefmt)
+        # record UTC time
+        formatter.converter = time.gmtime
+        handler.setFormatter(formatter)
+        logger.setLevel(level)
+        logger.addHandler(handler)
+
+    return logger