feat(training)!: remove support for EDA (#651)

## Description
<!-- What issue or task does this change relate to? -->
This PR removes the EDA functionality, simplifying the logic of the
datasets and EnsForecast and paving the way for easier integration with
multiple datasets.

## What problem does this change solve?
<!-- Describe if it's a bugfix, new feature, doc update, or breaking
change -->

## What issue or task does this change relate to?
<!-- link to Issue Number -->

##  Additional notes ##
<!-- Include any additional information, caveats, or considerations that
the reviewer should be aware of. -->

***As a contributor to the Anemoi framework, please ensure that your
changes include unit tests, updates to any affected dependencies and
documentation, and have been tested in a parallel setting (i.e., with
multiple GPUs). As a reviewer, you are also responsible for verifying
these aspects and requesting changes if they are not adequately
addressed. For guidelines about those please refer to
https://anemoi.readthedocs.io/en/latest/***

By opening this pull request, I affirm that all authors agree to the
[Contributor License
Agreement.](https://github.com/ecmwf/codex/blob/main/Legal/contributor_license_agreement.md)


<!-- readthedocs-preview anemoi-training start -->
----
📚 Documentation preview 📚:
https://anemoi-training--651.org.readthedocs.build/en/651/

<!-- readthedocs-preview anemoi-training end -->

<!-- readthedocs-preview anemoi-graphs start -->
----
📚 Documentation preview 📚:
https://anemoi-graphs--651.org.readthedocs.build/en/651/

<!-- readthedocs-preview anemoi-graphs end -->

<!-- readthedocs-preview anemoi-models start -->
----
📚 Documentation preview 📚:
https://anemoi-models--651.org.readthedocs.build/en/651/

<!-- readthedocs-preview anemoi-models end -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Simon Lang <[email protected]>
Co-authored-by: Ana Prieto Nemesio <[email protected]>

Author: 4 people

models/src/anemoi/models/migrations/migrator.py

@@ -264,12 +264,12 @@ class MissingAttribute:
     """Placeholder type when encountering ImportError or AttributeError in Unpickler.find_class"""
 
 
-def _get_unpickler(replace_attrs: list[str] | bool = False):
+def _get_unpickler(replace_attrs: dict[str, list[str]] | bool = False):
     """Get the Unpickler
 
     Parameters
     ----------
-    replace_attrs : list[str] | bool, default False
+    replace_attrs : dict[str,list[str]] | bool, default False
         Replace the provided attrs by a ``MissingAttribute`` object. If False, Fill not
         try to replace attributes. If True, will replace every missing attribute. You can use
         * as a wildcard to be replaced by any attribute in a module.
@@ -289,11 +289,26 @@ def find_class(self, module_name: str, global_name: str, /) -> Any:
             try:
                 return super().find_class(module_name, global_name)
             except (ImportError, AttributeError) as e:
+
+                deleted_modules: list[str] = []
+                deleted_attributes: list[str] = []
+
+                # --- Normalize replace_attrs ---
+                if isinstance(replace_attrs, dict):
+                    deleted_modules = replace_attrs.get("deleted_modules", [])
+                    deleted_attributes = replace_attrs.get("deleted_attributes", [])
+
                 attr_name = f"{module_name}.{global_name}"
                 wild_name = f"{module_name}.*"
+
                 if replace_attrs is False:
                     raise e
-                if replace_attrs is True or attr_name in replace_attrs or wild_name in replace_attrs:
+                if (
+                    replace_attrs is True
+                    or attr_name in deleted_attributes
+                    or module_name in deleted_modules
+                    or wild_name in replace_attrs
+                ):
                     LOGGER.debug("Missing attribute %s.%s is checkpoint. Ignoring.", module_name, global_name)
                     return MissingAttribute
                 raise e
@@ -308,7 +323,7 @@ class UnpicklerWrapper:
     return UnpicklerWrapper
 
 
-def _load_ckpt(path: str | PathLike, replace_attrs: list[str] | bool = False) -> CkptType:
+def _load_ckpt(path: str | PathLike, replace_attrs: dict[str, list[str]] | bool = False) -> CkptType:
     """Loads a checkpoint
 
     Parameters
@@ -570,6 +585,11 @@ def _resolve_context(self, context: MigrationContext) -> None:
         context : MigrationContext
             The context object
         """
+        for module_path in getattr(context, "deleted_modules", []):
+            if module_path in sys.modules:
+                LOGGER.debug("Delete module %s.", module_path)
+                del sys.modules[module_path]
+
         for module_path_end, module_path_start in context.module_paths.items():
             LOGGER.debug("Move module %s to %s.", module_path_start, module_path_end)
             sys.modules[module_path_start] = sys.modules[module_path_end]
@@ -614,13 +634,14 @@ def sync(self, path: str | PathLike) -> tuple[CkptType, CkptType, list[BaseOp]]:
         compatible_migrations = self._grouped_migrations[-1]
         self._check_executed_migrations(ckpt, compatible_migrations)
         setups, ops = self._resolve_operations(ckpt, compatible_migrations)
-        replace_attrs: list[str] = []
+        replace_attrs: dict[str, list[str]] = {}
         if len(setups):
             context = MigrationContext()
             for setup in setups:
                 setup(context)
             self._resolve_context(context)
-            replace_attrs = context.deleted_attributes
+            replace_attrs["deleted_modules"] = context.deleted_modules
+            replace_attrs["deleted_attributes"] = context.deleted_attributes
         # Force reloading checkpoint without obfuscating import issues.
         ckpt = _load_ckpt(path, replace_attrs)
         ckpt["hyper_parameters"]["metadata"].setdefault("migrations", {}).setdefault("history", [])

models/src/anemoi/models/migrations/scripts/1762857427_deprecate_eda.py

@@ -0,0 +1,64 @@
+# (C) Copyright 2025 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+from anemoi.models.migrations import CkptType
+from anemoi.models.migrations import MigrationContext
+from anemoi.models.migrations import MigrationMetadata
+
+# DO NOT CHANGE -->
+metadata = MigrationMetadata(
+    versions={
+        "migration": "1.0.0",
+        "anemoi-models": "%NEXT_ANEMOI_MODELS_VERSION%",
+    },
+)
+# <-- END DO NOT CHANGE
+
+
+def migrate_setup(context: MigrationContext) -> None:
+    """Migrate setup callback to be run before loading the checkpoint.
+
+    Parameters
+    ----------
+    context : MigrationContext
+       A MigrationContext instance
+    """
+    context.delete_module("anemoi.training.schemas.datamodule")
+
+
+def migrate(ckpt: CkptType) -> CkptType:
+    """Migrate the checkpoint.
+
+    Parameters
+    ----------
+    ckpt : CkptType
+        The checkpoint dict.
+
+    Returns
+    -------
+    CkptType
+        The migrated checkpoint dict.
+    """
+    return ckpt
+
+
+def rollback(ckpt: CkptType) -> CkptType:
+    """Rollback the checkpoint.
+
+    Parameters
+    ----------
+    ckpt : CkptType
+        The checkpoint dict.
+
+    Returns
+    -------
+    CkptType
+        The rollbacked checkpoint dict.
+    """
+    return ckpt

models/src/anemoi/models/migrations/setup_context.py

@@ -29,6 +29,7 @@ def __init__(self) -> None:
         self.attribute_paths: dict[str, str] = {}
         self.module_paths: dict[str, str] = {}
         self.deleted_attributes: list[str] = []
+        self.deleted_modules: list[str] = []
 
     def delete_attribute(self, path: str) -> None:
         """Indicate that an attribute has been deleted. Any class referencing this module will
@@ -41,6 +42,10 @@ def delete_attribute(self, path: str) -> None:
         """
         self.deleted_attributes.append(path)
 
+    def delete_module(self, path: str) -> None:
+        """Mark a module for deletion."""
+        self.deleted_modules.append(path)
+
     def move_attribute(self, path_start: str, path_end: str) -> None:
         """Move and rename an attribute between modules.
 
@@ -76,6 +81,7 @@ class SerializedMigrationContext(TypedDict):
     attribute_paths: dict[str, str]
     module_paths: dict[str, str]
     deleted_attributes: list[str]
+    deleted_modules: list[str]
 
 
 def serialize_setup_callback(setup: Callable[[MigrationContext], None]) -> SerializedMigrationContext:
@@ -98,6 +104,7 @@ def serialize_setup_callback(setup: Callable[[MigrationContext], None]) -> Seria
         "attribute_paths": ctx.attribute_paths,
         "module_paths": ctx.module_paths,
         "deleted_attributes": ctx.deleted_attributes,
+        "deleted_modules": ctx.deleted_modules,
     }
 
 
@@ -112,6 +119,8 @@ def __call__(self, context: MigrationContext) -> None:
             context.delete_attribute(deleted_attribute)
         for path_end, path_start in self._ctx["attribute_paths"].items():
             context.move_attribute(path_start, path_end)
+        for deleted_module in self._ctx["deleted_modules"].items():
+            context.delete_module(deleted_module)
         for path_end, path_start in self._ctx["module_paths"].items():
             context.move_module(path_start, path_end)
 

models/src/anemoi/models/models/ens_encoder_processor_decoder.py

@@ -15,38 +15,18 @@
 import torch
 from hydra.utils import instantiate
 from torch.distributed.distributed_c10d import ProcessGroup
-from torch_geometric.data import HeteroData
 
 from anemoi.models.distributed.graph import shard_tensor
 from anemoi.models.distributed.shapes import get_or_apply_shard_shapes
 from anemoi.models.distributed.shapes import get_shard_shapes
 from anemoi.models.models import AnemoiModelEncProcDec
-from anemoi.utils.config import DotDict
 
 LOGGER = logging.getLogger(__name__)
 
 
 class AnemoiEnsModelEncProcDec(AnemoiModelEncProcDec):
     """Message passing graph neural network with ensemble functionality."""
 
-    def __init__(
-        self,
-        *,
-        model_config: DotDict,
-        data_indices: dict,
-        statistics: dict,
-        graph_data: HeteroData,
-        truncation_data: dict,
-    ) -> None:
-
-        super().__init__(
-            model_config=model_config,
-            data_indices=data_indices,
-            statistics=statistics,
-            graph_data=graph_data,
-            truncation_data=truncation_data,
-        )
-
     def _calculate_input_dim(self):
         base_input_dim = super()._calculate_input_dim()
         return base_input_dim + self.num_input_channels_prognostic + 1

training/docs/user-guide/configuring.rst

@@ -22,7 +22,6 @@ settings at the top as follows:
    - data: zarr
    - dataloader: native_grid
    - diagnostics: evaluation
-   - datamodule: single
    - hardware: example
    - graph: multi_scale
    - model: gnn
@@ -102,7 +101,6 @@ match the dataset you provide.
    - data: zarr
    - dataloader: native_grid
    - diagnostics: evaluation
-   - datamodule: single
    - hardware: example
    - graph: multi_scale
    - model: transformer # Change from default group

training/docs/user-guide/diffusion-set-up.rst

@@ -226,7 +226,6 @@ A minimal config file for standard diffusion training:
    - data: zarr
    - dataloader: native_grid
    - diagnostics: evaluation
-   - datamodule: single
    - hardware: example
    - graph: multi_scale
    - model: graphtransformer_diffusion  # Use diffusion model
@@ -247,7 +246,6 @@ For tendency-based diffusion, change the model config and model task:
    - data: zarr
    - dataloader: native_grid
    - diagnostics: evaluation
-   - datamodule: single
    - hardware: example
    - graph: multi_scale
    - model: graphtransformer_diffusiontend  # Use tendency diffusion model

training/docs/user-guide/yaml/example_crps_config.yaml

@@ -20,10 +20,6 @@ hardware:
   num_nodes: 1
   num_gpus_per_model: 1
 
-# Changes in datamodule
-datamodule:
-  _target_: anemoi.training.data.datamodule.AnemoiEnsDatasetsDataModule
-
 data:
   resolution: o96
 

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

@@ -2,7 +2,6 @@ defaults:
 - data: zarr
 - dataloader: native_grid
 - diagnostics: evaluation
-- datamodule: single
 - hardware: example
 - graph: multi_scale
 - model: gnn