feat: plugin support (#187)

Support plugins in inference

Co-authored-by: Harrison Cook <[email protected]>

Author: b8raoult

docs/conf.py

@@ -112,6 +112,10 @@
         "https://anemoi-transform.readthedocs.io/en/latest/",
         ("../../anemoi-transform/docs/_build/html/objects.inv", None),
     ),
+    "anemoi-plugins": (
+        "https://anemoi-plugins.readthedocs.io/en/latest/",
+        ("../../anemoi-plugins/docs/_build/html/objects.inv", None),
+    ),
 }
 
 
@@ -148,3 +152,10 @@
 autodoc_pydantic_model_show_json = True
 autodoc_pydantic_model_show_field_summary = False
 autodoc_pydantic_model_member_order = "bysource"
+
+html_context = {
+    "display_github": True,
+    "github_user": "ecmwf",
+    "github_repo": "anemoi-inference",
+    "github_version": "main/docs/",
+}

docs/index.rst

@@ -84,9 +84,9 @@ You may also have to install pandoc on MacOS:
 
    brew install pandoc
 
-*****************
- Anemoi packages
-*****************
+***********************
+ Other Anemoi packages
+***********************
 
 -  :ref:`anemoi-utils <anemoi-utils:index-page>`
 -  :ref:`anemoi-transform <anemoi-transform:index-page>`
@@ -96,6 +96,7 @@ You may also have to install pandoc on MacOS:
 -  :ref:`anemoi-training <anemoi-training:index-page>`
 -  :ref:`anemoi-inference <anemoi-inference:index-page>`
 -  :ref:`anemoi-registry <anemoi-registry:index-page>`
+-  :ref:`anemoi-plugins <anemoi-plugins:index-page>`
 
 *********
  License

pyproject.toml

@@ -8,6 +8,8 @@
 # nor does it submit to any jurisdiction.
 
 [build-system]
+build-backend = "setuptools.build_meta"
+
 requires = [ "setuptools>=60", "setuptools-scm>=8" ]
 
 [project]
@@ -42,7 +44,7 @@ classifiers = [
 
 dynamic = [ "version" ]
 dependencies = [
-  "anemoi-transform>=0.1.3",
+  "anemoi-transform>=0.1.4",
   "anemoi-utils[text,provenance]>=0.4.16",
   "aniso8601",
   "anytree",
@@ -86,9 +88,18 @@ entry-points."ai_models.model".anemoi = "anemoi.inference.plugin:AIModelPlugin"
 [tool.setuptools.package-data]
 "anemoi.inference.grib.templates" = [ "*.yaml" ]
 
+[tool.setuptools.packages.find]
+where = [ "src" ]
+
 [tool.setuptools_scm]
 version_file = "src/anemoi/inference/_version.py"
 
+[tool.pytest.ini_options]
+markers = [
+  "skip_on_hpc: mark a test that should not be run on HPC",
+]
+testpaths = "tests"
+
 [tool.mypy]
 exclude = [ "docs/" ]
 strict = false

src/anemoi/inference/post_processors/backward_transform_filter.py

@@ -14,15 +14,19 @@
 from anemoi.transform.filters import filter_registry
 
 from anemoi.inference.context import Context
+from anemoi.inference.decorators import main_argument
 from anemoi.inference.types import State
 
 from ..processor import Processor
+from . import post_processor_registry
 from .earthkit_state import unwrap_state
 from .earthkit_state import wrap_state
 
 LOG = logging.getLogger(__name__)
 
 
+@post_processor_registry.register("backward_transform_filter")
+@main_argument("filter")
 class BackwardTransformFilter(Processor):
     """A processor that applies a backward transform filter to a given state.
 

src/anemoi/inference/post_processors/cos_sin_mean_wave_direction.py

@@ -14,11 +14,16 @@
 import earthkit.data as ekd
 from anemoi.transform.filters import filter_registry
 
+from anemoi.inference.decorators import main_argument
+
 from ..processor import Processor
+from . import pre_processor_registry
 
 LOG = logging.getLogger(__name__)
 
 
+@pre_processor_registry.register("forward_transform_filter")
+@main_argument("filter")
 class ForwardTransformFilter(Processor):
     """A processor that applies a forward transform filter to the given fields.
 

src/anemoi/inference/pre_processors/cos_sin_mean_wave_direction.py

@@ -100,3 +100,9 @@ def files_for_tests(name: str) -> str:
     bits.append("tests")
     bits.append(name)
     return os.path.sep.join(bits)
+
+
+class TestingContext:
+    """A context for testing plugins."""
+
+    pass

src/anemoi/inference/pre_processors/forward_transform_filter.py

@@ -6,10 +6,11 @@ lead_time: 48
 report_error: true
 
 pre_processors:
-- cos_sin_mean_wave_direction
+- forward_transform_filter: cos_sin_mean_wave_direction
 - no_missing_values
 
 post_processors:
-- cos_sin_mean_wave_direction
+- backward_transform_filter: cos_sin_mean_wave_direction
+
 
 use_grib_paramid: true

src/anemoi/inference/testing/__init__.py

@@ -47,6 +47,8 @@ def test_inference_mwd() -> None:
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
+    test_inference_mwd()
+    exit(0)
     for name, obj in list(globals().items()):
         if name.startswith("test_") and callable(obj):
             print(f"Running {name}...")