Merge pull request #182 from metno/add-logging-timing

Add logging and timing

Author: ways

bris/__main__.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import time
 from datetime import datetime, timedelta
 
 from anemoi.utils.dates import frequency_to_seconds
@@ -11,20 +12,22 @@
 from .checkpoint import Checkpoint
 from .inference import Inference
 from .utils import (
+    LOGGER,
     create_config,
     get_all_leadtimes,
     parse_args,
     set_base_seed,
     set_encoder_decoder_num_chunks,
+    setup_logging,
 )
 from .writer import CustomWriter
 
-LOGGER = logging.getLogger(__name__)
-
 
 def main(arg_list: list[str] | None = None):
+    t0 = time.perf_counter()
     args = parse_args(arg_list)
     config = create_config(args["config"], args)
+    setup_logging(config)
 
     models = list(config.checkpoints.keys())
     checkpoints = {
@@ -88,8 +91,8 @@ def main(arg_list: list[str] | None = None):
             ),
             "%Y-%m-%dT%H:%M:%S",
         )
-        LOGGER.info(
-            "No start_date given, setting %s based on start_date and timestep.",
+        LOGGER.warning(
+            "No start_date given, setting %s based on end_date and timestep.",
             config.start_date,
         )
     else:
@@ -143,13 +146,6 @@ def main(arg_list: list[str] | None = None):
     )
     writer = CustomWriter(decoder_outputs, write_interval="batch")
 
-    # Set hydra defaults
-    config.defaults = [
-        {"override hydra/job_logging": "none"},  # disable config parsing logs
-        {"override hydra/hydra_logging": "none"},  # disable config parsing logs
-        "_self_",
-    ]
-
     # Forecaster must know about what leadtimes to output
     model = instantiate(
         config.model,
@@ -180,9 +176,13 @@ def main(arg_list: list[str] | None = None):
     if is_main_thread:
         for decoder_output in decoder_outputs:
             for output in decoder_output["outputs"]:
+                t1 = time.perf_counter()
                 output.finalize()
+                LOGGER.debug(
+                    f"finalizing decoder {decoder_output} output {output.filename_pattern} in {time.perf_counter() - t1:.1f}s"
+                )
 
-        print("Model run completed. 🤖")
+        LOGGER.info(f"Bris completed in {time.perf_counter() - t0:.1f}s. 🤖")
 
 
 if __name__ == "__main__":

bris/conventions/metno.py

@@ -6,6 +6,8 @@
 Additionally, the names of some dimension-variables do not use CF-names
 """
 
+from bris.utils import LOGGER
+
 
 class Metno:
     cf_to_metno = {
@@ -46,12 +48,14 @@ def get_ncname(self, cfname: str, leveltype: str, level: int):
             # This is likely a forcing variable
             return cfname
         else:
-            print(cfname, leveltype, level)
+            LOGGER.error(
+                f"get_ncname not implemented cfname {cfname}, leveltype {leveltype}, level {level}"
+            )
             raise NotImplementedError()
 
         return ncname
 
-    def is_single_level(self, cfname: str, leveltype: str) -> str:
+    def is_single_level(self, cfname: str, leveltype: str) -> bool:
         """Returns true if there should only be a single level in the level dimension for this
         variable.
 
@@ -69,13 +73,13 @@ def is_single_level(self, cfname: str, leveltype: str) -> str:
             "wind_speed",
         ] and leveltype in ["height"]
 
-    def get_name(self, cfname: str):
+    def get_name(self, cfname: str) -> str:
         """Get MetNorway's dimension name from cf standard name"""
         if cfname in self.cf_to_metno:
             return self.cf_to_metno[cfname]
         return cfname
 
-    def get_cfname(self, ncname):
+    def get_cfname(self, ncname) -> str:
         """Get the CF-standard name from a given MetNo name"""
         for k, v in self.cf_to_metno.items():
             if v == ncname:

bris/data/dataset/nativegrid.py

@@ -10,9 +10,7 @@
 from torch.utils.data import IterableDataset
 
 from bris.data.grid_indices import BaseGridIndices
-from bris.utils import get_base_seed, get_usable_indices
-
-LOGGER = logging.getLogger(__name__)
+from bris.utils import LOGGER, get_base_seed, get_usable_indices
 
 
 class NativeGridDataset(IterableDataset):
@@ -200,8 +198,8 @@ def per_worker_init(self, n_workers: int, worker_id: int) -> None:
             f"num_data_parallel = num_nodes * num_gpus_per_node / num_gpus_per_model"
         )
         if len(self.valid_date_indices) % self.ens_comm_num_groups != 0:
-            print(
-                f"Warning: Dataloader has {len(self.valid_date_indices)} samples, which is not divisible by "
+            LOGGER.warning(
+                f"Dataloader has {len(self.valid_date_indices)} samples, which is not divisible by "
                 f"{self.ens_comm_num_groups} data parallel workers. This will lead to "
                 f"{len(self.valid_date_indices) % self.ens_comm_num_groups} unprocessed samples.",
                 "num_data_parallel = num_nodes * num_gpus_per_node / num_gpus_per_model",

bris/ddp_strategy.py

@@ -18,14 +18,15 @@
 from pytorch_lightning.strategies.ddp import DDPStrategy
 from pytorch_lightning.trainer.states import TrainerFn
 
-from bris.utils import get_base_seed
-
-LOGGER = logging.getLogger(__name__)
+from bris.utils import LOGGER, get_base_seed
 
 
 class DDPGroupStrategy(DDPStrategy):
     """Distributed Data Parallel strategy with group communication."""
 
+    # Define type of model, set in DDPStrategy somewhere
+    model: pl.LightningModule
+
     def __init__(
         self,
         num_gpus_per_model: int,
@@ -89,6 +90,13 @@ def setup(self, trainer: pl.Trainer) -> None:
 
         # set up reader groups by further splitting model_comm_group_ranks with read_group_size:
 
+        LOGGER.debug(
+            "world_size %d, model_comm_group_size %d, read_group_size %d",
+            self.world_size,
+            self.model_comm_group_size,
+            self.read_group_size,
+        )
+
         assert self.model_comm_group_size % self.read_group_size == 0, (
             f"Number of GPUs per model ({self.model_comm_group_size}) must be divisible by read_group_size "
             f"({self.read_group_size})."
@@ -226,7 +234,7 @@ def get_my_model_comm_group(self, num_gpus_per_model: int) -> tuple[int, int, in
 
     def get_my_reader_group(
         self, model_comm_group_rank: int, read_group_size: int
-    ) -> tuple[int, int, int]:
+    ) -> tuple[int, int, int, int]:
         """Determine tasks that work together and from a reader group.
 
         Parameters

bris/inference.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from functools import cached_property
 from typing import Any, Optional
 
@@ -7,11 +8,10 @@
 from anemoi.utils.config import DotDict
 
 from bris.ddp_strategy import DDPGroupStrategy
+from bris.utils import LOGGER
 
 from .data.datamodule import DataModule
 
-LOGGER = logging.getLogger(__name__)
-
 
 class Inference:
     def __init__(
@@ -41,7 +41,7 @@ def device(self) -> str:
                 LOGGER.info("Specified device not set. Found GPU")
                 return "cuda"
 
-            LOGGER.info("Specified device not set. Could not find gpu, using CPU")
+            LOGGER.warning("Specified device not set. Could not find gpu, using CPU")
             return "cpu"
 
         LOGGER.info("Using specified device: %s", self._device)
@@ -74,6 +74,9 @@ def trainer(self) -> pl.Trainer:
         return trainer
 
     def run(self):
+        t0 = time.perf_counter()
+        LOGGER.debug("Bris/Inference/run Predicting")
         self.trainer.predict(
             self.model, datamodule=self.datamodule, return_predictions=False
         )
+        LOGGER.debug(f"bris/Inference.run: {time.perf_counter() - t0:.1f}s")

bris/model/brispredictor.py

@@ -18,14 +18,13 @@
     get_dynamic_forcings,
 )
 from ..utils import (
+    LOGGER,
     check_anemoi_training,
     timedelta64_from_timestep,
 )
 from .basepredictor import BasePredictor
 from .model_utils import get_model_static_forcings, get_variable_indices
 
-LOGGER = logging.getLogger(__name__)
-
 
 class BrisPredictor(BasePredictor):
     """

bris/outputs/__init__.py

@@ -1,10 +1,12 @@
 import copy
+import time
 from typing import Optional
 
 import numpy as np
 
 from bris import sources
 from bris.predict_metadata import PredictMetadata
+from bris.utils import LOGGER
 
 
 def instantiate(name: str, predict_metadata: PredictMetadata, workdir: str, init_args):
@@ -113,6 +115,7 @@ def add_forecast(self, times: list, ensemble_member: int, pred: np.ndarray):
                 self.pm.variables.append(name)
 
         # only do this once. For multiple members, intermediate calls this several times
+        t0 = time.perf_counter()
         if pred.shape[2] != len(self.pm.variables):
             # Append extra variables to prediction
             extra_pred = []
@@ -126,6 +129,9 @@ def add_forecast(self, times: list, ensemble_member: int, pred: np.ndarray):
                     raise ValueError(f"No recipe to compute {name}")
 
             pred = np.concatenate([pred] + extra_pred, axis=2)
+        LOGGER.debug(
+            f"outputs.add_forecast Calculate ws in {time.perf_counter() - t0:.1f}s"
+        )
 
         assert pred.shape[0] == self.pm.num_leadtimes
         assert pred.shape[1] == len(self.pm.lats)
@@ -136,7 +142,11 @@ def add_forecast(self, times: list, ensemble_member: int, pred: np.ndarray):
         assert ensemble_member >= 0
         assert ensemble_member < self.pm.num_members
 
+        t1 = time.perf_counter()
         self._add_forecast(times, ensemble_member, pred)
+        LOGGER.debug(
+            f"outputs.add_forecast called _add_forecast in {time.perf_counter() - t1:.1f}s"
+        )
 
     def _add_forecast(self, times: list, ensemble_member: int, pred: np.ndarray):
         """Subclasses should implement this"""

bris/outputs/intermediate.py

@@ -1,5 +1,6 @@
 import glob
 import os
+import time
 from typing import Optional
 
 import numpy as np
@@ -24,17 +25,21 @@ def __init__(
         super().__init__(predict_metadata, extra_variables)
         self.workdir = workdir
 
-    def _add_forecast(self, times, ensemble_member, pred):
+    def _add_forecast(self, times, ensemble_member, pred) -> None:
+        t0 = time.perf_counter()
         filename = self.get_filename(times[0], ensemble_member)
         utils.create_directory(filename)
 
         np.save(filename, pred)
+        utils.LOGGER.debug(
+            f"Intermediate._add_forecast for {filename} in {time.perf_counter() - t0:.1f}s"
+        )
 
-    def get_filename(self, forecast_reference_time, ensemble_member):
+    def get_filename(self, forecast_reference_time, ensemble_member) -> str:
         frt_ut = utils.datetime_to_unixtime(forecast_reference_time)
         return f"{self.workdir}/{frt_ut:.0f}_{ensemble_member:.0f}.npy"
 
-    def get_forecast_reference_times(self):
+    def get_forecast_reference_times(self) -> list[np.datetime64]:
         """Returns all forecast reference times that have been saved"""
         filenames = self.get_filenames()
         frts = []
@@ -48,7 +53,9 @@ def get_forecast_reference_times(self):
 
         return frts
 
-    def get_forecast(self, forecast_reference_time, ensemble_member=None):
+    def get_forecast(
+        self, forecast_reference_time, ensemble_member=None
+    ) -> np.ndarray | None:
         """Fetches forecasts from stored numpy files
 
         Args:
@@ -61,6 +68,7 @@ def get_forecast(self, forecast_reference_time, ensemble_member=None):
                       4D otherwise (leadtime, points, variables, members)
         """
 
+        t0 = time.perf_counter()
         if ensemble_member is None:
             shape = [
                 self.pm.num_leadtimes,
@@ -73,16 +81,21 @@ def get_forecast(self, forecast_reference_time, ensemble_member=None):
                 filename = self.get_filename(forecast_reference_time, e)
                 if os.path.exists(filename):
                     pred[..., e] = np.load(filename)
+                utils.LOGGER.debug(
+                    f"Intermediate.get_forecast for {filename} in {time.perf_counter() - t0:.1f}s"
+                )
         else:
             assert isinstance(ensemble_member, int)
 
             filename = self.get_filename(forecast_reference_time, ensemble_member)
             pred = np.load(filename) if os.path.exists(filename) else None
-
+            utils.LOGGER.debug(
+                f"Intermediate.get_forecast for {filename} in {time.perf_counter() - t0:.1f}s"
+            )
         return pred
 
     @property
-    def num_members(self):
+    def num_members(self) -> int:
         filenames = self.get_filenames()
 
         max_member = 0
@@ -91,22 +104,23 @@ def num_members(self):
             max_member = max(int(member), max_member)
         return max_member + 1
 
-    def get_filenames(self):
+    def get_filenames(self) -> list[str]:
         return glob.glob(f"{self.workdir}/*_*.npy")
 
-    def cleanup(self):
+    def cleanup(self) -> None:
         """Removes up all intermediate files and removes the workdir. Called in finalize of the main output."""
-
+        t0 = time.perf_counter()
         for _filename in self.get_filenames():
             try:
                 os.remove(_filename)
             except OSError as e:
-                print(f"Error during cleanup of {_filename}: {e}")
+                utils.LOGGER.warning(f"Error during cleanup of {_filename}: {e}")
 
         try:
             os.rmdir(self.workdir)
         except OSError as e:
-            print(f"Error removing workdir {self.workdir}: {e}")
+            utils.LOGGER.warning(f"Error removing workdir {self.workdir}: {e}")
+        utils.LOGGER.debug(f"Intermediate.cleanup in {time.perf_counter() - t0:.1f}s")
 
     def finalize(self):
         pass