feat: Resolve config with omegaconf (ecmwf#252)

## Description
Uses omegaconf to resolve the configuration. This allows to use
interpolations in the configuration, similar to the anemoi-core
packages.

## 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?
Fixes ecmwf#243.

##  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/***


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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Gert Mertes <13658335+gmertes@users.noreply.github.com>

Author: 3 people

docs/inference/apis/level3.rst

@@ -26,6 +26,11 @@ You can also override values by providing them on the command line:
 .. literalinclude:: code/level3_2.sh
    :language: bash
 
+Overrides are parsed as an `OmegaConf
+<https://omegaconf.readthedocs.io/en/2.2_branch/usage.html#from-a-dot-list>`_
+dotlist, so list items can be accessed with ``list.index`` or
+``list[index]``.
+
 You can also run entirely from the command line without a config file,
 by passing all required options as an override:
 
@@ -39,9 +44,9 @@ that was used to train the model, by setting ``dataset`` entry to
 .. literalinclude:: code/level3_2.yaml
    :language: yaml
 
-It is also possible to override list entries and append to lists on the
-command line by using the list indices as key. Running inference with
-following command:
+Below is an example of how to override list entries and append to lists
+on the command line by using the dotlist notation. Running inference
+with following command:
 
 .. literalinclude:: code/level3_4.sh
    :language: bash

docs/inference/configs/introduction.rst

@@ -8,10 +8,13 @@ This document provides an overview of the configuration to provide to
 the :ref:`anemoi-inference run <run_command>` command line tool.
 
 The configuration file is a YAML file that specifies various options. It
-is composed of :ref:`top level <top-level>` options which are usually
-simple values such as strings, number or booleans. The configuration
-also provide ways to specify which internal classes to use for the
-:ref:`inputs <inputs>` and :ref:`outputs <outputs>`, and how to
+is extended by `OmegaConf <https://github.com/omry/omegaconf>`_ such
+that `interpolations
+<https://omegaconf.readthedocs.io/en/2.2_branch/usage.html#variable-interpolation>`_
+can be used. It is composed of :ref:`top level <top-level>` options
+which are usually simple values such as strings, number or booleans. The
+configuration also provide ways to specify which internal classes to use
+for the :ref:`inputs <inputs>` and :ref:`outputs <outputs>`, and how to
 configure them.
 
 In that case, the general format is shown below. The first entry

pyproject.toml

@@ -51,6 +51,7 @@ dependencies = [
   "earthkit-data>=0.12.4",
   "eccodes>=2.38.3",
   "numpy",
+  "omegaconf>=2.2,<2.4",
   "packaging",
   "pydantic",
   "pyyaml",

src/anemoi/inference/config/__init__.py

@@ -11,7 +11,7 @@
 
 import logging
 import os
-from copy import deepcopy
+from datetime import datetime
 from typing import Any
 from typing import Dict
 from typing import List
@@ -20,9 +20,13 @@
 from typing import TypeVar
 from typing import Union
 
-import yaml
+from earthkit.data.utils.dates import to_datetime
+from omegaconf import DictConfig
+from omegaconf import ListConfig
+from omegaconf import OmegaConf
 from pydantic import BaseModel
 from pydantic import ConfigDict
+from pydantic import field_validator
 
 LOG = logging.getLogger(__name__)
 
@@ -34,6 +38,15 @@ class Configuration(BaseModel):
 
     model_config = ConfigDict(extra="forbid")
 
+    date: Union[datetime, None] = None
+    """The starting date for the forecast. If not provided, the date will depend on the selected Input object. If a string, it is parsed by :func:`earthkit.data.utils.dates`."""
+
+    @field_validator("date", mode="before")
+    @classmethod
+    def to_datetime(cls, date: Union[str, int, datetime, None]) -> Optional[datetime]:
+        if date is not None:
+            return to_datetime(date)
+
     @classmethod
     def load(
         cls: Type[T],
@@ -58,56 +71,47 @@ def load(
             The loaded configuration.
         """
 
-        config = {}
+        configs: List[Union[DictConfig, ListConfig]] = []
 
         # Set default values
         if defaults is not None:
             if not isinstance(defaults, list):
                 defaults = [defaults]
             for d in defaults:
                 if isinstance(d, str):
-                    with open(d) as f:
-                        d = yaml.safe_load(f)
-                config.update(d)
+                    configs.append(OmegaConf.load(d))
+                    continue
+                configs.append(OmegaConf.create(d))
 
         # Load the user configuration
         if isinstance(path, dict):
-            user_config = deepcopy(path)
+            configs.append(OmegaConf.create(path))
         else:
-            with open(path) as f:
-                user_config = yaml.safe_load(f)
-
-        cls._merge_configs(config, user_config)
+            configs.append(OmegaConf.load(path))
 
-        # Apply overrides
         if not isinstance(overrides, list):
             overrides = [overrides]
 
+        # unsafe merge should be fine as we don't re-use the original configs
+        oc_config = OmegaConf.unsafe_merge(*configs)
+
         for override in overrides:
             if isinstance(override, dict):
-                cls._merge_configs(config, override)
+                oc_config = OmegaConf.unsafe_merge(oc_config, OmegaConf.create(override))
             else:
-                path = config
-                key, value = override.split("=")
-                keys = key.split(".")
-                for key in keys[:-1]:
-                    if key.isdigit() and isinstance(path, list):
-                        index = int(key)
-                        if index < len(path):
-                            LOG.debug(f"key {key} is used as list index in list{path}")
-                            path = path[index]
-                        elif index == len(path):
-                            LOG.debug(f"key {key} is used to append to list {path}")
-                            path.append({})
-                            path = path[index]
-                        else:
-                            raise IndexError(f"Index {index} out of range for list {path} of length {len(path)}")
-                    else:
-                        path = path.setdefault(key, {})
-                path[keys[-1]] = value
-
-        # Validate the configuration
-        config = cls(**config)
+                # use from_dotlist to use OmegaConf split
+                # which allows for "param.val" or "param[val]".
+                override_conf = OmegaConf.from_dotlist([override])
+                # We can't directly merge reconstructed with the config because
+                # omegaconf prefers parsing digits (like 0 in key.0) into dict keys
+                # rather than lists.
+                # Instead, we provide a reference config and we try to merge the override
+                # into the reference and keep types provided by the reference.
+                oc_config = OmegaConf.unsafe_merge(_merge_configs(oc_config, override_conf))
+
+        resolved_config = OmegaConf.to_container(oc_config, resolve=True)
+
+        config = cls.model_validate(resolved_config)
 
         # Set environment variables found in the configuration
         # as soon as possible
@@ -116,19 +120,47 @@ def load(
 
         return config
 
-    @classmethod
-    def _merge_configs(cls, a: Dict[Any, Any], b: Dict[Any, Any]) -> None:
-        """Merge two configurations.
 
-        Parameters
-        ----------
-        a : Dict[Any, Any]
-            The first configuration.
-        b : Dict[Any, Any]
-            The second configuration.
-        """
-        for key, value in b.items():
-            if key in a and isinstance(a[key], dict) and isinstance(value, dict):
-                cls._merge_configs(a[key], value)
-            else:
-                a[key] = value
+def _merge_configs(ref_conf: Any, new_conf: Any) -> Any:
+    """Recursively merges a new OmegaConf object into a reference OmegaConf object
+
+    Parameters
+    ----------
+    ref_conf : Any
+        reference OmegaConf object. Should be a DictConfig or ListConfig.
+    new_conf : Any
+        new OmegaConf object.
+
+    Returns
+    -------
+    Any
+        The merged OmegaConf config
+    """
+    if isinstance(new_conf, DictConfig) and len(new_conf):
+        key, rest = next(iter(new_conf.items()))
+        key = str(key)
+    elif isinstance(new_conf, ListConfig) and len(new_conf):
+        key, rest = 0, new_conf[0]
+    else:
+        return new_conf
+    if isinstance(ref_conf, ListConfig):
+        if isinstance(key, str) and not key.isdigit():
+            raise ValueError(f"Expected int key, got {key}")
+        index = int(key)
+        if index < len(ref_conf):
+            LOG.debug(f"key {key} is used as list key in list{ref_conf}")
+            ref_conf[index] = _merge_configs(ref_conf[index], rest)
+        elif index == len(ref_conf):
+            LOG.debug(f"key {key} is used to append to list {ref_conf}")
+            ref_conf.append(rest)
+        else:
+            raise IndexError(f"key {key} out of range for list {ref_conf} of length {len(ref_conf)}")
+        return ref_conf
+    elif isinstance(ref_conf, DictConfig) and key in ref_conf:
+        ref_conf[key] = _merge_configs(ref_conf[key], rest)
+        return ref_conf
+    elif isinstance(ref_conf, DictConfig):
+        ref_conf[key] = rest
+        return ref_conf
+    else:
+        raise ValueError(f"ref is of unexpected type {type(ref_conf)}. Should be ListConfig or DictConfig")

src/anemoi/inference/config/couple.py

@@ -18,7 +18,6 @@
 from typing import Tuple
 
 from anemoi.inference.config import Configuration
-from anemoi.inference.types import Date
 
 LOG = logging.getLogger(__name__)
 
@@ -28,9 +27,6 @@ class CoupleConfiguration(Configuration):
 
     description: Optional[str] = None
 
-    date: Optional[Date] = None
-    """The starting date for the forecast. If not provided, the date will depend on the selected Input object. If a string, it is parsed by :func:`anemoi.utils.dates.as_datetime`."""
-
     lead_time: Optional[Tuple[str, int, datetime.timedelta]] = None
     """The lead time for the forecast. This can be a string, an integer or a timedelta object.
     If an integer, it represents a number of hours. Otherwise, it is parsed by :func:`anemoi.utils.dates.as_timedelta`.

src/anemoi/inference/config/run.py

@@ -34,9 +34,6 @@ class RunConfiguration(Configuration):
     runner: Union[str, Dict[str, Any]] = "default"
     """The runner to use."""
 
-    date: Union[str, int, datetime.datetime, None] = None
-    """The starting date for the forecast. If not provided, the date will depend on the selected Input object. If a string, it is parsed by :func:`anemoi.utils.dates.as_datetime`."""
-
     lead_time: Union[str, int, datetime.timedelta] = "10d"
     """The lead time for the forecast. This can be a string, an integer or a timedelta object.
     If an integer, it represents a number of hours. Otherwise, it is parsed by :func:`anemoi.utils.dates.as_timedelta`.

src/anemoi/inference/inputs/cds.py

@@ -151,8 +151,6 @@ def create_input_state(self, *, date: Optional[Date]) -> State:
             date = to_datetime(-1)
             LOG.warning("CDSInput: `date` parameter not provided, using yesterday's date: %s", date)
 
-        date = to_datetime(date)
-
         return self._create_input_state(
             self.retrieve(
                 self.variables,

src/anemoi/inference/inputs/dataset.py

@@ -18,7 +18,6 @@
 from typing import Tuple
 
 import numpy as np
-from earthkit.data.utils.dates import to_datetime
 
 from anemoi.inference.context import Context
 from anemoi.inference.types import Date
@@ -103,7 +102,6 @@ def create_input_state(self, *, date: Optional[Date] = None) -> State:
         if date is None:
             raise ValueError("`date` must be provided")
 
-        date = to_datetime(date)
         latitudes = self.ds.latitudes
         longitudes = self.ds.longitudes
 

src/anemoi/inference/inputs/ekd.py

@@ -280,7 +280,6 @@ def _create_state(
 
         n_points = fields[0].to_numpy(dtype=dtype, flatten=flatten).size
         for field in fields:
-
             name, valid_datetime = field.metadata("name"), field.metadata("valid_datetime")
             if name not in state_fields:
                 state_fields[name] = np.full(
@@ -368,8 +367,6 @@ def _create_input_state(
                 "%s: `date` not provided, using the most recent date: %s", self.__class__.__name__, date.isoformat()
             )
 
-        # TODO: where we do this might change in the future
-        date = to_datetime(date)
         dates = [date + h for h in self.checkpoint.lagged]
 
         return self._create_state(

src/anemoi/inference/inputs/mars.py

@@ -247,8 +247,6 @@ def create_input_state(self, *, date: Optional[Date]) -> State:
             date = to_datetime(-1)
             LOG.warning("MarsInput: `date` parameter not provided, using yesterday's date: %s", date)
 
-        date = to_datetime(date)
-
         return self._create_input_state(
             self.retrieve(
                 self.variables,