Add timings, split out functions, add logging with configurable levels

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,21 @@
 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):
     args = parse_args(arg_list)
     config = create_config(args["config"], args)
+    setup_logging(config)
 
     models = list(config.checkpoints.keys())
     checkpoints = {
@@ -77,7 +79,7 @@ def main(arg_list: list[str] | None = None):
     try:
         multistep = checkpoints["forecaster"].config.training.multistep_input
     except KeyError:
-        LOGGER.debug("Multistep not found in checkpoint")
+        LOGGER.error("Multistep not found in checkpoint")
 
     # If no start_date given, calculate as end_date-((multistep-1)*timestep)
     if "start_date" not in config or config.start_date is None:
@@ -88,7 +90,7 @@ def main(arg_list: list[str] | None = None):
             ),
             "%Y-%m-%dT%H:%M:%S",
         )
-        LOGGER.info(
+        LOGGER.error(
             "No start_date given, setting %s based on start_date and timestep.",
             config.start_date,
         )
@@ -143,13 +145,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 +175,14 @@ def main(arg_list: list[str] | None = None):
     if is_main_thread:
         for decoder_output in decoder_outputs:
             for output in decoder_output["outputs"]:
+                t0 = time.perf_counter()
                 output.finalize()
+                LOGGER.debug(
+                    f"finalizing decoder {decoder_output} output {output.filename_pattern} in %d.1s"
+                    % (time.perf_counter() - t0)
+                )
 
-        print("Model run completed. 🤖")
+        LOGGER.info("Model run completed. 🤖")
 
 
 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

@@ -200,8 +200,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

@@ -20,12 +20,15 @@
 
 from bris.utils import get_base_seed
 
-LOGGER = logging.getLogger(__name__)
+from .utils import LOGGER
 
 
 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 +92,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 +236,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
@@ -282,15 +292,15 @@ def process_dataloader(
         )
 
         dataloader.dataset.set_comm_group_info(
-            self.global_rank,
-            model_comm_group_id,
-            model_comm_group_rank,
-            model_comm_num_groups,
-            ens_comm_group_id,
-            ens_comm_group_rank,
-            ens_comm_num_groups,
-            reader_group_rank,
-            self.read_group_size,
+            global_rank=self.global_rank,
+            model_comm_group_id=model_comm_group_id,
+            model_comm_group_rank=model_comm_group_rank,
+            model_comm_num_groups=model_comm_num_groups,
+            ens_comm_group_id=ens_comm_group_id,
+            ens_comm_group_rank=ens_comm_group_rank,
+            ens_comm_num_groups=ens_comm_num_groups,
+            reader_group_rank=reader_group_rank,
+            reader_group_size=self.read_group_size,
         )
 
         return dataloader

bris/inference.py

@@ -1,4 +1,5 @@
 import logging
+import time
 from functools import cached_property
 from typing import Any, Optional
 
@@ -9,8 +10,7 @@
 from bris.ddp_strategy import DDPGroupStrategy
 
 from .data.datamodule import DataModule
-
-LOGGER = logging.getLogger(__name__)
+from .utils import LOGGER
 
 
 class Inference:
@@ -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("bris/Inference.run: %d.1s", (time.perf_counter() - t0))

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,11 +1,14 @@
 import copy
+import time
 from typing import Optional
 
 import numpy as np
 
 from bris import sources
 from bris.predict_metadata import PredictMetadata
 
+from ..utils import LOGGER
+
 
 def instantiate(name: str, predict_metadata: PredictMetadata, workdir: str, init_args):
     """Creates an object of type name with config
@@ -100,13 +103,13 @@ def add_forecast(self, times: list, ensemble_member: int, pred: np.ndarray):
             ensemble_member: Which ensemble member is this?
             pred: 3D numpy array with dimensions (leadtime, location, variable)
         """
-
-        # Append extra variables to prediction
+        # Append extra variables to prediction (quick)
         for name in self.extra_variables:
             if name not in self.pm.variables:
                 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 = []
@@ -120,6 +123,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)
@@ -130,7 +136,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"""