feat(config)!: Improved configuration and data structures

graphs/src/anemoi/graphs/nodes/attributes.py

@@ -237,8 +237,6 @@ class NonmissingZarrVariable(BooleanBaseNodeAttribute):
     ----------
     variable : str
         Variable to read from the Zarr dataset.
-    norm : str
-        Normalization of the weights.
 
     Methods
     -------

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 = dataset
         self.idx_lon = idx_lon

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

@@ -51,7 +51,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.
@@ -133,3 +134,10 @@
 todo_include_todos = not read_the_docs_build
 
 autodoc_member_order = "bysource"  # Keep file order
+
+
+# https://autodoc-pydantic.readthedocs.io/en/stable/users/configuration.html
+
+autodoc_pydantic_model_show_json = True
+autodoc_pydantic_model_show_field_summary = False
+autodoc_pydantic_model_member_order = "bysource"

training/docs/dev/hydra.rst

@@ -2,9 +2,64 @@
  Configuration
 ###############
 
-Anemoi Training uses Hydra for configuration management, allowing for
-flexible and modular configuration of the training pipeline. This guide
-explains how to use Hydra effectively in the project.
+Anemoi Training uses Hydra and Pydantic for configuration management,
+allowing for flexible and modular configuration of the training pipeline
+while provide robustness through validation. This guide explains how to
+use Hydra and Pydantic effectively in the project.
+
+***************************************
+ Pydantic and Configuration Validation
+***************************************
+
+Pydantic is a package designed for data validation and settings
+management. It provides a simple way to define schemas which can be used
+to validate configuration files. For example, the following schema can
+be used to validate a training configuration:
+
+.. code:: python
+
+   from pydantic import BaseModel, Field, PositiveFloat, Literal
+
+   class TrainingSchema(BaseModel):
+      model: Literal{"AlexNet", "ResNet", "VGG"} = Field(default="AlexNet")
+      """Model architecture to use for training."""
+      learning_rate: PositiveFloat = Field(default=0.01)
+      """Learning rate."""
+      loss: str = Field(default="mse")
+      """Loss function."""
+
+To allow more complex configurations, Pydantic also supports nested
+schemas. For example, the following schema can be used to validate a
+configuration with a configurable model:
+
+.. code:: python
+
+   from pydantic import BaseModel, Field, PositiveFloat, Literal
+
+   from enum import StrEnum
+
+   class ActivationFunctions(StrEnum):
+       relu = "relu"
+       sigmoid = "sigmoid"
+       tanh = "tanh"
+
+   class ModelSchema(BaseModel):
+       num_layers: PositiveInt = Field(default=3)
+       """Number of layers in the model."""
+       activation: ActivationFunctions = Field(default="relu")
+       """Activation function to use."""
+
+   class TrainingSchema(BaseModel):
+       model: ModelSchema
+       """Model configuration."""
+       learning_rate: PositiveFloat = Field(default=0.01)
+       """Learning rate."""
+       loss: str = Field(default="mse")
+       """Loss function."""
+
+If your new feature requires a new configuration parameter, you should
+add it to the appropriate schemas and update the configuration files
+accordingly.
 
 **************
  Hydra Basics
@@ -48,13 +103,23 @@ Configuration in YAML:
      algorithm: SGD
      learning_rate: 0.01
 
+Pydantic schema:
+
+.. code:: python
+
+   from pydantic import BaseModel
+
+   class OptimizerSchema(BaseModel):
+       algorithm: str
+       learning_rate: float
+
 Instantiating in code:
 
 .. code:: python
 
    from hydra.utils import instantiate
 
-   optimizer = instantiate(config.optimizer)
+   optimizer = instantiate(config.optimizer.model_dump())
 
 ********************************************
  Configurable Components in Anemoi Training

training/docs/modules/schemas.rst

@@ -0,0 +1,67 @@
+#########
+ Schemas
+#########
+
+This module defines pydantic schemas, which are used to validate the
+configuration before a training run is started. The top-level config
+yaml matches the BaseSchema.
+
+.. automodule:: anemoi.training.schemas.base_schema
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.data
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.dataloader
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.hardware
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.graphs.basegraph
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.graphs.node_schemas
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.graphs.edge_schemas
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.models.models
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.models.processor
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.models.encoder
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.models.decoder
+   :members:
+   :no-undoc-members:
+   :show-inheritance:
+
+.. automodule:: anemoi.training.schemas.training
+   :members:
+   :no-undoc-members:
+   :show-inheritance:

training/docs/start/hydra-intro.rst

@@ -86,3 +86,11 @@ The following missing config options which must be overridden by users:
 -  ``hardware.files.graph``: If you have pre-computed a specific graph,
    specify its filename here. Otherwise, a new graph will be constructed
    on the fly and written to the filename given.
+
+*********************************
+ Validation of the Configuration
+*********************************
+
+The configuration is validated using `Pydantic` before a training run
+starts. To turn this off, you can use the `--no-validation` flag in your
+top-level config.

training/docs/user-guide/configuring.rst

@@ -153,3 +153,138 @@ or combine everything together
 .. code:: bash
 
    anemoi-training train --config-name=debug.yaml model=transformer diagnostics.plot.enabled=False
+
+********************
+ Cconfig validation
+********************
+
+It is possible to validate your configuration before starting a training
+run using the following command:
+
+.. code:: bash
+
+   anemoi-training validate --name debug.yaml
+
+This will check that the configuration is valid and that all the
+required fields are present. If your config is correctly defined then
+the command will show an output similar to:
+
+.. code:: bash
+
+   2025-01-28 09:37:23 INFO Validating configs.
+   2025-01-28 09:37:23 INFO Prepending Anemoi Home (/home_path/.config/anemoi/training/config) to the search path.
+   2025-01-28 09:37:23 INFO Prepending current user directory (/repos_path/config_anemoi_core) to the search path.
+   2025-01-28 09:37:23 INFO Search path is now: [provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-home-searchpath-plugin, path=/home_path/.config/anemoi/training/config, provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands]
+   cfg = BaseSchema(**cfg)
+   2025-01-28 09:37:23 INFO Config files validated.
+
+Otherwise if there is an issue with some of your configuration fields,
+Pydantic will report an error message. See below example where we have a
+`debug.yaml` file with a field not correctly indented (in this case the
+`diagnostics.log` field):
+
+.. code:: yaml
+
+   defaults:
+   - data: zarr
+   - dataloader: native_grid
+   - diagnostics: evaluation
+   - hardware: example
+   - graph: multi_scale
+   - model: transformer # Change from default group
+   - training: default
+   - _self_
+
+
+   diagnostics:
+   log:
+   mlflow:
+      enabled: True
+      offline: True
+      experiment_name: 'test'
+      project_name: 'AIFS'
+      run_name: 'test_anemoi_core'
+      tracking_uri: 'https://mlflow-server.int'
+      authentication: True
+      terminal: True
+
+If we try to validate the above then the validate command will report
+the following error:
+
+.. code:: python
+
+   2025-01-28 09:37:23 INFO Validating configs.
+   2025-01-28 09:37:23 INFO Prepending Anemoi Home (/home_path/.config/anemoi/training/config) to the search path.
+   2025-01-28 09:37:23 INFO Prepending current user directory (/repos_path/config_anemoi_core) to the search path.
+   2025-01-28 09:37:23 INFO Search path is now: [provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-home-searchpath-plugin, path=/home_path/.config/anemoi/training/config, provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands]
+   pydantic_core._pydantic_core.ValidationError: 1 validation error for BaseSchema
+   diagnostics.log
+    Input should be a valid dictionary or instance of LoggingSchema [type=model_type, input_value=None, input_type=NoneType]
+      For further information visit https://errors.pydantic.dev/2.10/v/model_type
+   2025-01-28 09:54:08 ERROR
+   💣 1 validation error for BaseSchema
+   diagnostics.log
+   Input should be a valid dictionary or instance of LoggingSchema [type=model_type, input_value=None, input_type=NoneType]
+      For further information visit https://errors.pydantic.dev/2.10/v/model_type
+   2025-01-28 09:54:08 ERROR 💣 Exiting
+
+Which indicates that the `diagnostics.log` field is not correctly
+defined as it should be a dictionary or instance of `LoggingSchema`.
+Please note there might still be cases not captured by the current
+schemas, so it is always good to double check the configuration file
+before running the training. See below an example of a config with some
+typos that might still need to be fixed manually:
+
+.. code:: yaml
+
+   defaults:
+   - data: zarr
+   - dataloader: native_grid
+   - diagnostics: evaluation
+   - hardware: example
+   - graph: multi_scale
+   - model: transformer # Change from default group
+   - training: default
+   - _self_
+
+
+   diagnostics:
+   log:
+      mlflow:
+         enabled: True
+         ofline: True # this is a typo - should be offline
+         experiment_name: 'test'
+         project_name: 'AIFS'
+         run_name: 'test_anemoi_core'
+         tracking_uri: 'https://mlflow-server.int'
+         authentication: True
+         terminal: True
+
+In the example above, if there is a default already defined for
+`offline` under `diagnostics: evaluation` then the validation will be
+successful, and in the high-level config (ie `debug`) `ofline` it will
+just simply not be used, since it has a typo. Otherwise, if there is no
+default for `offline` then the validation will fail, with the following
+error:
+
+.. code:: python
+
+   2025-01-28 09:37:23 INFO Validating configs.
+   2025-01-28 09:37:23 INFO Prepending Anemoi Home (/home_path/.config/anemoi/training/config) to the search path.
+   2025-01-28 09:37:23 INFO Prepending current user directory (/repos_path/config_anemoi_core) to the search path.
+   2025-01-28 09:37:23 INFO Search path is now:  [provider=anemoi-cwd-searchpath-plugin, path=/repos_path/config_anemoi_core, provider=anemoi-home-searchpath-plugin, path=/home_path/.config/anemoi/training/config, provider=hydra, path=pkg://hydra.conf, provider=main, path=/repos_path/anemoi-core/training/src/anemoi/training/commands]
+   pydantic_core._pydantic_core.ValidationError: 1 validation error for BaseSchema
+   diagnostics.log.mlflow.offline
+   Field required [type=missing, input_value={'enabled': True, 'authen...onfig'], 'ofline': True}, input_type=DictConfig]
+      For further information visit https://errors.pydantic.dev/2.10/v/missing
+   2025-01-28 10:14:49 ERROR
+   💣 1 validation error for BaseSchema
+   diagnostics.log.mlflow.offline
+   Field required [type=missing, input_value={'enabled': True, 'authen...onfig'], 'ofline': True}, input_type=DictConfig]
+      For further information visit https://errors.pydantic.dev/2.10/v/missing
+   2025-01-28 10:14:49 ERROR 💣 Exiting
+
+That will indicate that the `offline` field is required and it is
+missing from the configuration file. If you identify any issues with the
+schemas or missing functionality, please raise an issue on the `Anemoi
+Core repository`.

training/pyproject.toml

@@ -51,6 +51,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",
@@ -68,6 +69,7 @@ optional-dependencies.dev = [
   "pre-commit>=3.3.3",
 ]
 optional-dependencies.docs = [
+  "autodoc-pydantic",
   "nbsphinx",
   "pandoc",
   "sphinx",
@@ -112,6 +114,9 @@ exclude = [ "docs/" ]
 preview = true
 lint.flake8-import-conventions.extend-aliases."pytorch_lightning" = "pl"
 lint.pydocstyle.convention = "numpy"
+# Preserve types, even if a file imports `from __future__ import annotations`.
+# Necessary for supporting PEP 604 and pydantic with python 3.9
+lint.pyupgrade.keep-runtime-typing = true
 select = [
   "A", # flake8 builtins
   # "ANN", # flake8-annotations