Merge pull request #226 from metno/feat/multiple_datasets

Feat/multiple datasets

Author: havardhhaugen

.github/workflows/run_test.yml

@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.10"]
+        python-version: ["3.11"]
 
     steps:
       - uses: actions/checkout@v4
@@ -61,7 +61,7 @@ jobs:
 
       - name: Run tox unit tests (6G)
         run: |
-          tox -e py310
+          tox -e py311
 
       - name: Run tox security check (<50M)
         run: |

bris/__main__.py

@@ -15,6 +15,10 @@
     LOGGER,
     create_config,
     get_all_leadtimes,
+    get_dataset_config,
+    get_interpolator_timestep_seconds,
+    get_model_multistep_input,
+    get_model_timestep,
     parse_args,
     set_base_seed,
     set_encoder_decoder_num_chunks,
@@ -27,6 +31,7 @@ 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())
@@ -46,28 +51,20 @@ def main(arg_list: list[str] | None = None):
     set_base_seed()
 
     # Compute timestep_seconds for each checkpoint
-    config.checkpoints.forecaster.timestep = checkpoints[
-        "forecaster"
-    ].config.data.timestep
+    config.checkpoints.forecaster.timestep = get_model_timestep(
+        checkpoints["forecaster"]
+    )
     config.checkpoints.forecaster.timestep_seconds = frequency_to_seconds(
         config.checkpoints.forecaster.timestep
     )
 
     if "interpolator" in checkpoints:
-        target_times = checkpoints[
-            "interpolator"
-        ].metadata.config.training.explicit_times.target
-        input_times = checkpoints[
-            "interpolator"
-        ].metadata.config.training.explicit_times.input
-        if target_times[-1] == input_times[-1]:
-            config.checkpoints.interpolator.timestep_seconds = int(
-                config.checkpoints.forecaster.timestep_seconds / len(target_times)
-            )
-        else:
-            config.checkpoints.interpolator.timestep_seconds = int(
-                config.checkpoints.forecaster.timestep_seconds / (len(target_times) + 1)
+        config.checkpoints.interpolator.timestep_seconds = (
+            get_interpolator_timestep_seconds(
+                checkpoints["interpolator"],
+                config.checkpoints.forecaster.timestep_seconds,
             )
+        )
 
     num_members = config["hardware"].get("num_members", 1)
 
@@ -89,11 +86,7 @@ def main(arg_list: list[str] | None = None):
         num_members_in_parallel = num_members
 
     # Get multistep. A default of 2 to ignore multistep in start_date calculation if not set.
-    multistep = 2
-    try:
-        multistep = checkpoints["forecaster"].config.training.multistep_input
-    except KeyError:
-        LOGGER.debug("Multistep not found in checkpoint")
+    multistep = get_model_multistep_input(checkpoints["forecaster"])
 
     # If no start_date given, calculate as end_date-((multistep-1)*timestep)
     if "start_date" not in config or config.start_date is None:
@@ -116,13 +109,8 @@ def main(arg_list: list[str] | None = None):
             ),
             "%Y-%m-%dT%H:%M:%S",
         )
-
-    config.dataset = {
-        "dataset": config.dataset,
-        "start": config.start_date,
-        "end": config.end_date,
-        "frequency": config.frequency,
-    }
+    # Get dataset config with backwards comapatibility for single dataset config setup
+    config.datasets = get_dataset_config(config)
 
     datamodule = DataModule(
         config=config,

bris/checkpoint.py

@@ -2,7 +2,6 @@
 import os
 from copy import deepcopy
 from functools import cached_property
-from typing import Optional
 
 import torch
 from anemoi.utils.checkpoints import load_metadata
@@ -48,7 +47,7 @@ class Metadata(DotDict):
 class Checkpoint:
     """This class makes accessible various information stored in Anemoi checkpoints."""
 
-    def __init__(self, path: str, graph: Optional[str] = None):
+    def __init__(self, path: str, graph: str | None = None):
         assert os.path.exists(path), f"The given checkpoint {path} does not exist!"
 
         self.path = path
@@ -133,8 +132,50 @@ def _load_model(self) -> torch.nn.Module:
                     "check module versions."
                 ) from e
             raise e
+        if not torch.cuda.is_available():
+            self._apply_triton_cpu_fallback(inst)
         return inst
 
+    def _apply_triton_cpu_fallback(self, model: torch.nn.Module) -> None:
+        """Replace Triton graph attention with the PyG backend when running on CPU.
+
+        anemoi-models checks is_triton_available() at model construction time and falls
+        back to the PyG backend automatically. However, when a model is loaded from a
+        checkpoint via torch.load() (weights_only=False), __init__ is not called —
+        pickle restores __dict__ directly — so the Triton function reference is
+        preserved even when no GPU is available. This method applies the same fallback
+        after loading.
+
+        GraphTransformerConv has no trainable parameters, so the swap is safe.
+        """
+        try:
+            from anemoi.models.layers.block import GraphTransformerBaseBlock
+            from anemoi.models.layers.conv import GraphTransformerConv
+        except ImportError:
+            LOGGER.warning(
+                "Could not import anemoi.models layers to apply Triton->PyG CPU fallback."
+            )
+            return
+
+        patched = 0
+        for module in model.modules():
+            if (
+                isinstance(module, GraphTransformerBaseBlock)
+                and module.graph_attention_backend == "triton"
+            ):
+                module.graph_attention_backend = "pyg"
+                module.conv = GraphTransformerConv(
+                    out_channels=module.out_channels_conv
+                )
+                patched += 1
+
+        if patched:
+            LOGGER.warning(
+                "Checkpoint was saved with the Triton graph attention backend but no GPU "
+                "is available. Fell back to the PyG backend for %d block(s).",
+                patched,
+            )
+
     @property
     def graph(self) -> HeteroData:
         """
@@ -180,7 +221,7 @@ def graph(self) -> HeteroData:
     #     _model_params = self._model_instance.named_parameters()
     #     return deepcopy(dict(_model_params))
 
-    def update_graph(self, path: Optional[str] = None) -> HeteroData:
+    def update_graph(self, path: str | None = None) -> HeteroData:
         """
         Replaces existing graph object within model instance.
         The new graph is either provided as an torch file or
@@ -219,129 +260,10 @@ def update_graph(self, path: Optional[str] = None) -> HeteroData:
         self._model_instance.load_state_dict(new_state_dict)
         return self._model_instance.graph_data
 
-    @property
-    def name_to_index(self) -> tuple[dict[str, int], ...]:
-        """
-        Mapping between name and their corresponding variable index.
-        Returns a tuple. If the model is a multiencoder-decoder model
-        the tuple will contain two dicts, one for each decoder. If not
-        the tuple will contain a single dict.
-        """
-        _data_indices = self._model_instance.data_indices
-        if isinstance(_data_indices, (tuple, list)) and len(_data_indices) >= 2:
-            return tuple(
-                _data_indices[k].name_to_index for k in range(len(_data_indices))
-            )
-
-        return (_data_indices.name_to_index,)
-
-    @property
-    def index_to_name(self) -> tuple[dict[int, str], ...]:
-        """
-        Mapping between index and their corresponding variable name.
-        Returns a tuple. If the model is a multiencoder-decoder model
-        the tuple will contain two dicts, one for each decoder. If not
-        the tuple will contain a single dict.
-        """
-        _data_indices = self._model_instance.data_indices
-        if isinstance(_data_indices, (tuple, list)) and len(_data_indices) >= 2:
-            return tuple(
-                {
-                    index: var
-                    for (var, index) in self.name_to_index[decoder_index].items()
-                }
-                for decoder_index in range(len(self.name_to_index))
-            )
-        return ({index: name for name, index in _data_indices.name_to_index.items()},)
-
-    def _make_indices_mapping(self, indices_from, indices_to):
-        """
-        Creates a mapping for a given model and data output
-        or model input and data input, and vice versa.
-
-        args:
-            indices_from (dict)
-            indices_to (dict)
-        return
-            a mapping between indices_from and indices_to
-        """
-        assert len(indices_from) == len(indices_to)
-        return dict(zip(indices_from, indices_to))
-
-    @property
-    def model_output_index_to_name(self) -> tuple[dict[int, str], ...]:
-        """
-        A mapping from model output to data output. This
-        dict returns index and name pairs according to model.output.full to
-        data.output.full
-
-        args:
-            None
-        return:
-            tuple of dicts where the tuple index represents decoder index. Each
-            tuple dicts contains a dict with a mapping between model.output.full and
-            data.output.full
-        Example:
-        -> if the parameter skt has index 27 in name_to_index or data.output.full it will
-        have index 17 in model.output.full. This dict will yield that index 17 is skt
-        """
-        if (
-            isinstance(self._metadata.data_indices, (tuple, list))
-            and len(self._metadata.data_indices) >= 2
-        ):
-            mapping = {
-                grid_index: self._make_indices_mapping(
-                    self._metadata.data_indices[grid_index].model.output.full,
-                    self._metadata.data_indices[grid_index].data.output.full,
-                )
-                for grid_index in range(len(self._metadata.data_indices))
-            }
-            return tuple(
-                {name: self.index_to_name[k][index] for (name, index) in v.items()}
-                for (k, v) in mapping.items()
-            )
-
-        mapping = self._make_indices_mapping(
-            self._metadata.data_indices.model.output.full,
-            self._metadata.data_indices.data.output.full,
-        )
-        return ({k: self._metadata.dataset.variables[v] for k, v in mapping.items()},)
-
-    @property
-    def model_output_name_to_index(self) -> tuple[dict[str, int], ...]:
-        """
-        A mapping from model output to data output. This
-        dict returns name and index pairs according to model.output.full to
-        data.output.full
-
-        args:
-            None
-        return:
-            Identical tuple[dict] from model_output_index_to_name but key
-            value pairs are switched.
-        """
-        if (
-            isinstance(self._metadata.data_indices, (tuple, list))
-            and len(self._metadata.data_indices) >= 2
-        ):
-            return tuple(
-                {name: index for (index, name) in v.items()}
-                for (k, v) in enumerate(self.model_output_index_to_name)
-            )
-
-        return (
-            {name: index for index, name in self.model_output_index_to_name[0].items()},
-        )
-
     @cached_property
-    def data_indices(self) -> tuple[IndexCollection, ...]:
-        """
-        Wrapper for model.data_indices. Returns a tuple of dict and None or two dicts.
-        """
-
-        # If Multiencdec checkpoint
-        if isinstance(self._model_instance.data_indices, (tuple, list)):
-            return tuple(self._model_instance.data_indices)
-
-        # If simple checkpoint
-        return (self._model_instance.data_indices,)
+    def data_indices(self) -> dict[str, IndexCollection]:
+        _data_indices = self._model_instance.data_indices
+        if isinstance(_data_indices, IndexCollection):  # Backwards compatibility
+            return {"data": _data_indices}
+        else:
+            return _data_indices