Merge main into ensembles. There is still a known bug for ensemble runs, related to the ens_comm_group: Inference is running, but does not give outputs

Signed-off-by: evenmn <evenmn@fys.uio.no>

Author: evenmn

.github/workflows/python-package.yml

@@ -52,14 +52,6 @@ jobs:
       run: |
         tox -e ruff_checkformat
 
-    - name: Run tox unit tests
-      run: |
-        tox -e py310
-
-    - name: Run tox security check
-      run: |
-        tox -e bandit
-
     - name: Run ruff check
       run: |
         tox -e ruff_check
@@ -68,6 +60,14 @@ jobs:
       run: |
         tox -e trainingdata
 
+    - name: Run tox unit tests
+      run: |
+        tox -e py310
+
+    - name: Run tox security check
+      run: |
+        tox -e bandit
+
     - name: Run inference
       run: |
         tox -e inference_CI

README.md

@@ -36,7 +36,7 @@ See [Dockerfile](https://gitlab.met.no/yrop/bris-cicd/-/blob/main/Dockerfile?ref
 
 ## How to run tests
 
-    pip install ".[tests]"
+    pip install -e '.[dev]'
     tox
 
 When pushing to github, default tests will be run automatically and must succeed.

bris/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.3"
+__version__ = "0.2.0"
 
 
 def main():

bris/__main__.py

@@ -14,6 +14,7 @@
 from .utils import (
     create_config,
     get_all_leadtimes,
+    parse_args,
     set_base_seed,
     set_encoder_decoder_num_chunks,
 )
@@ -22,12 +23,9 @@
 LOGGER = logging.getLogger(__name__)
 
 
-def main():
-    parser = ArgumentParser()
-    parser.add_argument("--debug", action="store_true")
-    parser.add_argument("--config", type=str, required=True)
-
-    config = create_config(parser)
+def main(arg_list: list[str] | None = None):
+    args = parse_args(arg_list)
+    config = create_config(args["config"], args)
 
     models = list(config.checkpoints.keys())
     checkpoints = {
@@ -38,6 +36,9 @@ def main():
         for model in models
     }
     set_encoder_decoder_num_chunks(getattr(config, "inference_num_chunks", 1))
+    if "release_cache" not in config or not isinstance(config["release_cache"], bool):
+        config["release_cache"] = False
+
     set_base_seed()
 
     # Get timestep from checkpoint. Also store a version in seconds for local use.
@@ -168,7 +169,7 @@ def main():
             for output in decoder_output["outputs"]:
                 output.finalize()
 
-    print("Hello world")
+    print("Model run completed. 🤖")
 
 
 if __name__ == "__main__":

bris/checkpoint.py

@@ -153,7 +153,6 @@ def graph(self) -> HeteroData:
 
         return:
             HeteroData graph object
-
         """
         return (
             self._model_instance.graph_data

bris/conventions/cf.py

@@ -100,7 +100,7 @@ def get_metadata(anemoi_variable: str) -> dict:
 #     raise ValueError(f"Unknown leveltype: {leveltype}")
 
 
-def get_attributes(cfname):
+def get_attributes(cfname: str) -> dict[str, str]:
     ret = {"standard_name": cfname}
 
     # Coordinate variables
@@ -127,6 +127,9 @@ def get_attributes(cfname):
         ret["description"] = "height above ground"
         ret["long_name"] = "height"
         ret["positive"] = "up"
+    elif cfname in ["thunder_event"]:
+        ret["standard_name"] = "thunderstorm_probability"
+        ret["units"] = "1"
 
     # Data variables
     elif cfname in [
@@ -139,7 +142,7 @@ def get_attributes(cfname):
         ret["units"] = "m/s"
     elif cfname in ["air_temperature", "dew_point_temperature"]:
         ret["units"] = "K"
-    elif cfname == "land_sea_mask":
+    elif cfname in ["land_sea_mask", "area_fraction"]:
         ret["units"] = "1"
     elif cfname in ["geopotential", "surface_geopotential"]:
         ret["units"] = "m^2/s^2"
@@ -151,8 +154,6 @@ def get_attributes(cfname):
         ret["units"] = "kg/kg"
     elif cfname in ["cloud_base_altitude", "visibility_in_air"]:
         ret["units"] = "m"
-    elif "area_fraction" in cfname:
-        ret["units"] = "1"
     elif cfname in [
         "integral_of_surface_downwelling_longwave_flux_in_air_wrt_time",
         "integral_of_surface_downwelling_shortwave_flux_in_air_wrt_time",

bris/data/dataset/__init__.py

@@ -0,0 +1,37 @@
+import logging
+
+from torch.utils.data import get_worker_info
+
+# Make NativeGridDataset and ZipDataset importable from bris.data.dataset:
+from .nativegrid import NativeGridDataset  # noqa
+from .zip import ZipDataset  # noqa
+
+LOGGER = logging.getLogger(__name__)
+
+
+def worker_init_func(worker_id: int) -> None:
+    """Configures each dataset worker process. This is a helper function, called by Dataloader.
+
+    Calls WeatherBenchDataset.per_worker_init() on each dataset object.
+
+    Parameters
+    ----------
+    worker_id : int
+        Worker ID
+
+    Raises
+    ------
+    RuntimeError
+        If worker_info is None
+    """
+    worker_info = get_worker_info()  # information specific to each worker process
+    if worker_info is None:
+        LOGGER.error("worker_info is None! Set num_workers > 0 in your dataloader!")
+        raise RuntimeError
+    dataset_obj = (
+        worker_info.dataset
+    )  # the copy of the dataset held by this worker process.
+    dataset_obj.per_worker_init(
+        n_workers=worker_info.num_workers,
+        worker_id=worker_id,
+    )