1 feature improved configuration and data structures

graphs/src/anemoi/graphs/nodes/builders/from_file.py

@@ -68,15 +68,15 @@ class TextNodes(BaseNodeBuilder):
 
     Attributes
     ----------
-    dataset : str | DictConfig
-        The path to txt file containing the coordinates of the nodes.
+    dataset : str | Path
+        The path including filename to txt file containing the coordinates of the nodes.
     idx_lon : int
         The index of the longitude in the dataset.
     idx_lat : int
         The index of the latitude in the dataset.
     """
 
-    def __init__(self, dataset, name: str, idx_lon: int = 0, idx_lat: int = 1) -> None:
+    def __init__(self, dataset: str | Path, name: str, idx_lon: int = 0, idx_lat: int = 1) -> None:
         LOGGER.info("Reading the dataset from %s.", dataset)
         self.dataset = np.loadtxt(dataset)
         self.idx_lon = idx_lon

models/src/anemoi/models/models/encoder_processor_decoder.py

@@ -49,7 +49,7 @@ def __init__(
             Graph definition
         """
         super().__init__()
-
+        model_config = DotDict(model_config)
         self._graph_data = graph_data
         self._graph_name_data = model_config.graph.data
         self._graph_name_hidden = model_config.graph.hidden

training/docs/conf.py

@@ -70,6 +70,7 @@
     "sphinx.ext.napoleon",
     "sphinxarg.ext",
     "sphinx.ext.autosectionlabel",
+    "sphinxcontrib.autodoc_pydantic",
 ]
 
 # Add any paths that contain templates here, relative to this directory.

training/pyproject.toml

@@ -50,6 +50,7 @@ dependencies = [
   "matplotlib>=3.7.1",
   "mlflow>=2.11.1",
   "numpy<2",                         # Pinned until we can confirm it works with anemoi graphs
+  "pydantic>=2.9",
   "pynvml>=11.5",
   "pyshtools>=4.10.4",
   "pytorch-lightning>=2.1",
@@ -69,6 +70,7 @@ optional-dependencies.dev = [
 ]
 
 optional-dependencies.docs = [
+  "autodoc-pydantic",
   "nbsphinx",
   "pandoc",
   "sphinx",

training/src/anemoi/training/commands/config.py

@@ -16,7 +16,11 @@
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+from hydra import compose
+from hydra import initialize
+
 from anemoi.training.commands import Command
+from anemoi.training.schemas.base_schema import BaseSchema
 
 if TYPE_CHECKING:
     import argparse
@@ -48,24 +52,39 @@ def add_arguments(command_parser: argparse.ArgumentParser) -> None:
         )
         anemoi_training_home.add_argument("--overwrite", "-f", action="store_true")
 
+        help_msg = "Validate the Anemoi training configs."
+        validate = subparsers.add_parser("validate", help=help_msg, description=help_msg)
+
+        validate.add_argument("--name", help="Name of the primary config file")
+        validate.add_argument("--overwrite", "-f", action="store_true")
+
     def run(self, args: argparse.Namespace) -> None:
-        LOGGER.info(
-            "Generating configs, please wait.",
-        )
+
         self.overwrite = args.overwrite
         if args.subcommand == "generate":
-
+            LOGGER.info(
+                "Generating configs, please wait.",
+            )
             self.traverse_config(args.output)
 
             LOGGER.info("Inference checkpoint saved to %s", args.output)
             return
 
         if args.subcommand == "training-home":
             anemoi_home = Path.home() / ".config" / "anemoi" / "training" / "config"
+            LOGGER.info(
+                "Generating configs, please wait.",
+            )
             self.traverse_config(anemoi_home)
             LOGGER.info("Inference checkpoint saved to %s", anemoi_home)
             return
 
+        if args.subcommand == "validate":
+            LOGGER.info("Validating configs.")
+            self.validate_config(args.name)
+            LOGGER.info("Config files validated.")
+            return
+
     def traverse_config(self, destination_dir: Path | str) -> None:
         """Writes the given configuration data to the specified file path."""
         config_package = "anemoi.training.config"
@@ -97,5 +116,11 @@ def copy_file(item: Path, file_path: Path) -> None:
         except Exception:
             LOGGER.exception("Failed to copy %s", item.name)
 
+    def validate_config(self, name: Path | str) -> None:
+        """Validates the configuration files in the given directory."""
+        with initialize(version_base=None, config_path=""):
+            cfg = compose(config_name=name)
+            cfg = BaseSchema(**cfg)
+
 
 command = ConfigGenerator

training/src/anemoi/training/config/data/zarr.yaml

@@ -74,12 +74,12 @@ processors:
   #   config: ${data.imputer}
   normalizer:
     _target_: anemoi.models.preprocessing.normalizer.InputNormalizer
-    _convert_: all
+    #   _convert_: all # Is it still used ???
     config: ${data.normalizer}
   # remapper:
   #   _target_: anemoi.models.preprocessing.remapper.Remapper
   #   _convert_: all
   #   config: ${data.remapper}
 
-  # Values set in the code
+# Values set in the code
 num_features: null # number of features in the forecast state

training/src/anemoi/training/config/dataloader/native_grid.yaml

@@ -16,12 +16,10 @@ num_workers:
   training: 8
   validation: 8
   test: 8
-  predict: 8
 batch_size:
   training: 2
   validation: 4
   test: 4
-  predict: 4
 
 # ============
 # Default effective batch_size for training is 16
@@ -38,7 +36,6 @@ limit_batches:
   training: null
   validation: null
   test: 20
-  predict: 20
 
 # set a custom mask for grid points.
 # Useful for LAM (dropping unconnected nodes from forcing dataset)

training/src/anemoi/training/config/diagnostics/evaluation.yaml

@@ -52,6 +52,7 @@ log:
     on_resume_create_child: True
     expand_hyperparams: # Which keys in hyperparams to expand
       - config
+    http_max_retries: 35
   interval: 100
 
 enable_progress_bar: True

training/src/anemoi/training/config/training/default.yaml

@@ -3,6 +3,7 @@ run_id: null
 fork_run_id: null
 load_weights_only: null # only load model weights, do not restore optimiser states etc.
 transfer_learning: null # activate to perform transfer learning
+load_weights_only: False # only load model weights, do not restore optimiser states etc.
 
 # run in deterministic mode ; slows down
 deterministic: False

training/src/anemoi/training/data/datamodule.py

@@ -16,14 +16,14 @@
 
 import pytorch_lightning as pl
 from hydra.utils import instantiate
-from omegaconf import DictConfig
-from omegaconf import OmegaConf
 from torch.utils.data import DataLoader
 
 from anemoi.datasets.data import open_dataset
 from anemoi.models.data_indices.collection import IndexCollection
 from anemoi.training.data.dataset import NativeGridDataset
 from anemoi.training.data.dataset import worker_init_func
+from anemoi.training.schemas.base_schema import BaseSchema
+from anemoi.training.schemas.base_schema import convert_to_omegaconf
 from anemoi.utils.dates import frequency_to_seconds
 
 LOGGER = logging.getLogger(__name__)
@@ -37,12 +37,12 @@
 class AnemoiDatasetsDataModule(pl.LightningDataModule):
     """Anemoi Datasets data module for PyTorch Lightning."""
 
-    def __init__(self, config: DictConfig, graph_data: HeteroData) -> None:
+    def __init__(self, config: BaseSchema, graph_data: HeteroData) -> None:
         """Initialize Anemoi Datasets data module.
 
         Parameters
         ----------
-        config : DictConfig
+        config : BaseSchema
             Job configuration
 
         """
@@ -66,7 +66,7 @@ def __init__(self, config: DictConfig, graph_data: HeteroData) -> None:
             )
             self.config.dataloader.training.end = self.config.dataloader.validation.start - 1
 
-        if not self.config.dataloader.get("pin_memory", True):
+        if not self.config.dataloader.pin_memory:
             LOGGER.info("Data loader memory pinning disabled.")
 
     @cached_property
@@ -83,12 +83,16 @@ def supporting_arrays(self) -> dict:
 
     @cached_property
     def data_indices(self) -> IndexCollection:
-        return IndexCollection(self.config, self.ds_train.name_to_index)
+        return IndexCollection(convert_to_omegaconf(self.config), self.ds_train.name_to_index)
 
     @cached_property
     def grid_indices(self) -> type[BaseGridIndices]:
-        reader_group_size = self.config.dataloader.get("read_group_size", self.config.hardware.num_gpus_per_model)
-        grid_indices = instantiate(self.config.dataloader.grid_indices, reader_group_size=reader_group_size)
+        reader_group_size = self.config.dataloader.read_group_size
+
+        grid_indices = instantiate(
+            self.config.dataloader.grid_indices.model_dump(by_alias=True),
+            reader_group_size=reader_group_size,
+        )
         grid_indices.setup(self.graph_data)
         return grid_indices
 
@@ -123,13 +127,14 @@ def timeincrement(self) -> int:
     @cached_property
     def ds_train(self) -> NativeGridDataset:
         return self._get_dataset(
-            open_dataset(OmegaConf.to_container(self.config.dataloader.training, resolve=True)),
+            open_dataset(self.config.dataloader.training.model_dump()),
             label="train",
         )
 
     @cached_property
     def ds_valid(self) -> NativeGridDataset:
-        r = max(self.rollout, self.config.dataloader.get("validation_rollout", 1))
+        r = self.rollout
+        r = max(r, self.config.dataloader.validation_rollout)
 
         if not self.config.dataloader.training.end < self.config.dataloader.validation.start:
             LOGGER.warning(
@@ -138,7 +143,7 @@ def ds_valid(self) -> NativeGridDataset:
                 self.config.dataloader.validation.start,
             )
         return self._get_dataset(
-            open_dataset(OmegaConf.to_container(self.config.dataloader.validation, resolve=True)),
+            open_dataset(self.config.dataloader.validation.model_dump()),
             shuffle=False,
             rollout=r,
             label="validation",
@@ -155,7 +160,7 @@ def ds_test(self) -> NativeGridDataset:
             f"test start date {self.config.dataloader.test.start}"
         )
         return self._get_dataset(
-            open_dataset(OmegaConf.to_container(self.config.dataloader.test, resolve=True)),
+            open_dataset(self.config.dataloader.test.model_dump()),
             shuffle=False,
             label="test",
         )
@@ -172,7 +177,7 @@ def _get_dataset(
 
         # Compute effective batch size
         effective_bs = (
-            self.config.dataloader.batch_size["training"]
+            self.config.dataloader.batch_size.training
             * self.config.hardware.num_gpus_per_node
             * self.config.hardware.num_nodes
             // self.config.hardware.num_gpus_per_model
@@ -193,12 +198,12 @@ def _get_dataloader(self, ds: NativeGridDataset, stage: str) -> DataLoader:
         assert stage in {"training", "validation", "test"}
         return DataLoader(
             ds,
-            batch_size=self.config.dataloader.batch_size[stage],
+            batch_size=self.config.dataloader.batch_size.model_dump()[stage],
             # number of worker processes
-            num_workers=self.config.dataloader.num_workers[stage],
+            num_workers=self.config.dataloader.num_workers.model_dump()[stage],
             # use of pinned memory can speed up CPU-to-GPU data transfers
             # see https://pytorch.org/docs/stable/notes/cuda.html#cuda-memory-pinning
-            pin_memory=self.config.dataloader.get("pin_memory", True),
+            pin_memory=self.config.dataloader.pin_memory,
             # worker initializer
             worker_init_fn=worker_init_func,
             # prefetch batches

training/src/anemoi/training/diagnostics/callbacks/__init__.py

@@ -18,6 +18,7 @@
 
 from hydra.utils import instantiate
 from omegaconf import DictConfig
+from pydantic import BaseModel
 
 from anemoi.training.diagnostics.callbacks.checkpoint import AnemoiCheckpoint
 from anemoi.training.diagnostics.callbacks.optimiser import LearningRateMonitor
@@ -28,21 +29,23 @@
 if TYPE_CHECKING:
     from pytorch_lightning.callbacks import Callback
 
+    from anemoi.training.schemas.base_schema import BaseSchema
+
 LOGGER = logging.getLogger(__name__)
 
 
-def nestedget(conf: DictConfig, key: str, default: Any) -> Any:
+def nestedget(config: DictConfig, key: str, default: Any) -> Any:
     """Get a nested key from a DictConfig object.
 
     E.g.
     >>> nestedget(config, "diagnostics.log.wandb.enabled", False)
     """
     keys = key.split(".")
     for k in keys:
-        conf = conf.get(k, default)
-        if not isinstance(conf, (dict, DictConfig)):
+        config = getattr(config, k, default)
+        if not isinstance(config, (BaseModel, dict, DictConfig)):
             break
-    return conf
+    return config
 
 
 # Callbacks to add according to flags in the config
@@ -57,9 +60,9 @@ def nestedget(conf: DictConfig, key: str, default: Any) -> Any:
 ]
 
 
-def _get_checkpoint_callback(config: DictConfig) -> list[AnemoiCheckpoint]:
+def _get_checkpoint_callback(config: BaseSchema) -> list[AnemoiCheckpoint]:
     """Get checkpointing callbacks."""
-    if not config.diagnostics.get("enable_checkpointing", True):
+    if not config.diagnostics.enable_checkpointing:
         return []
 
     checkpoint_settings = {
@@ -77,11 +80,11 @@ def _get_checkpoint_callback(config: DictConfig) -> list[AnemoiCheckpoint]:
     ckpt_frequency_save_dict = {}
 
     for key, frequency_dict in config.diagnostics.checkpoint.items():
-        frequency = frequency_dict["save_frequency"]
-        n_saved = frequency_dict["num_models_saved"]
-        if key == "every_n_minutes" and frequency_dict["save_frequency"] is not None:
+        frequency = frequency_dict.save_frequency
+        n_saved = frequency_dict.num_models_saved
+        if key == "every_n_minutes" and frequency_dict.save_frequency is not None:
             target = "train_time_interval"
-            frequency = timedelta(minutes=frequency_dict["save_frequency"])
+            frequency = timedelta(minutes=frequency_dict.save_frequency)
         else:
             target = key
         ckpt_frequency_save_dict[target] = (
@@ -143,7 +146,7 @@ def check_key(config: dict, key: str | Iterable[str] | Callable[[DictConfig], bo
     return callbacks
 
 
-def get_callbacks(config: DictConfig) -> list[Callback]:
+def get_callbacks(config: BaseSchema) -> list[Callback]:
     """Setup callbacks for PyTorch Lightning trainer.
 
     Set `config.diagnostics.callbacks` to a list of callback configurations
@@ -183,14 +186,11 @@ def get_callbacks(config: DictConfig) -> list[Callback]:
     trainer_callbacks.extend(_get_checkpoint_callback(config))
 
     # Base callbacks
-    trainer_callbacks.extend(
-        instantiate(callback, config) for callback in config.diagnostics.get("callbacks", None) or []
-    )
+    trainer_callbacks.extend(instantiate(callback, config) for callback in config.diagnostics.callbacks)
 
     # Plotting callbacks
-
     trainer_callbacks.extend(
-        instantiate(callback, config) for callback in config.diagnostics.plot.get("callbacks", None) or []
+        instantiate(callback.model_dump(by_alias=True), config) for callback in config.diagnostics.plot.callbacks
     )
 
     # Extend with config enabled callbacks