Allow starting remotebmi containers (#467)

* Allow starting remotebmi containers

* Correct formatting

* Ignore linter rule in notebooks

* Add remotebmi to dependencies

* Update changelog

* Update tests to reflect changes

* Add basic openapi container test

CHANGELOG.md

@@ -8,6 +8,7 @@ Formatted as described on [https://keepachangelog.com](https://keepachangelog.co
 
 ## Added
 
+- support for [Remote BMI](https://github.com/eWaterCycle/remotebmi), an OpenAPI based alternative for grpc4bmi ([#467](https://github.com/eWaterCycle/ewatercycle/pull/467)).
 - `.get_shape_area()` utility method to the ewatercycle Forcing objects. This returns the area of the shapefile in square meters, useful for converting the results of lumped models (e.g., from mm/day to m3/s) ([#464](https://github.com/eWaterCycle/ewatercycle/issues/464)).
 - `.plot_shape()` utility method to the ewatercycle Forcing objects. This allows plotting the shapefile in a single-line of code, or adds the shapefile to an existing plot ([#464](https://github.com/eWaterCycle/ewatercycle/issues/464)).
 

pyproject.toml

@@ -40,6 +40,7 @@ dependencies = [
     # otherwise pip installation fails on esmpy
     "Fiona",
     "grpc4bmi>=0.4.0",
+    "remotebmi",
     "hydrostats",
     "matplotlib>=3.5.0",
     "numpy",
@@ -205,6 +206,8 @@ ignore = [
     "E501",
     # Allow prints
     "T201",
+    # Allow shadowing builtins
+    "A004",
 ]
 
 [tool.ruff.lint.pydocstyle]

src/ewatercycle/base/model.py

@@ -8,7 +8,7 @@
 from contextlib import suppress
 from datetime import timezone
 from pathlib import Path
-from typing import Annotated, Any, cast
+from typing import Annotated, Any, Literal, cast
 
 import bmipy
 import numpy as np
@@ -439,6 +439,7 @@ class ContainerizedModel(eWaterCycleModel):
     """
 
     bmi_image: Annotated[ContainerImage, BeforeValidator(_parse_containerimage)]
+    protocol: Literal["grpc", "openapi"] = "grpc"
 
     # Create as empty list to allow models to append before bmi is made:
     _additional_input_dirs: list[str] = PrivateAttr([])
@@ -459,6 +460,7 @@ def _make_bmi_instance(self) -> OptionalDestBmi:
 
         return start_container(
             image=self.bmi_image,
+            protocol=self.protocol,
             work_dir=self._cfg_dir,
             input_dirs=self._additional_input_dirs,
             timeout=300,

src/ewatercycle/container.py

@@ -5,9 +5,10 @@
 import re
 from collections.abc import Iterable, Sequence
 from pathlib import Path
-from typing import Any, Protocol
+from typing import Any, Literal, Protocol
 
 import numpy as np
+import remotebmi
 from bmipy import Bmi
 from grpc import FutureTimeoutError
 from grpc4bmi.bmi_client_apptainer import BmiClientApptainer
@@ -136,6 +137,7 @@ def start_container(
     delay=0,
     # TODO replace Any type with Bmi + BmiFromOrigin
     wrappers: Sequence[type[Any]] = (MemoizedBmi, OptionalDestBmi),
+    protocol: Literal["grpc", "openapi"] = "grpc",
 ) -> OptionalDestBmi:
     """Start container with model inside.
 
@@ -152,6 +154,7 @@ def start_container(
         delay: Number of seconds to wait before connecting.
         wrappers: List of classes to wrap around the grcp4bmi object from container.
             Order is important. The first wrapper is the most inner wrapper.
+        protocol: Which protocol to use, grpc or openapi.
 
     Raises:
         ValueError: When unknown container technology is requested.
@@ -184,7 +187,13 @@ def start_container(
 
     if engine == "docker":
         bmi = start_docker_container(
-            work_dir, image, input_dirs, image_port, timeout, delay
+            work_dir,
+            image,
+            input_dirs,
+            image_port,
+            timeout,
+            delay,
+            protocol,
         )
     elif engine == "apptainer":
         bmi = start_apptainer_container(
@@ -193,6 +202,7 @@ def start_container(
             input_dirs,
             timeout,
             delay,
+            protocol,
         )
     else:
         msg = f"Unknown container technology: {CFG.container_engine}"
@@ -209,6 +219,7 @@ def start_apptainer_container(
     input_dirs: Iterable[str] = (),
     timeout: int | None = None,
     delay: int = 0,
+    protocol: Literal["grpc", "openapi"] = "grpc",
 ) -> Bmi:
     """Start Apptainer container with model inside.
 
@@ -220,6 +231,7 @@ def start_apptainer_container(
         input_dirs: Additional directories to mount inside container.
         timeout: Number of seconds to wait for grpc connection.
         delay: Number of seconds to wait before connecting.
+        protocol: Which protocol to use, grpc or openapi.
 
     .. _apptainer manual: https://apptainer.org/docs/user/latest/cli/apptainer_run.html
 
@@ -234,13 +246,24 @@ def start_apptainer_container(
         image_fn = str(CFG.apptainer_dir / image_fn)
 
     try:
-        return BmiClientApptainer(
-            image=image_fn,
-            work_dir=str(work_dir),
-            input_dirs=input_dirs,
-            timeout=timeout,
-            delay=delay,
-        )
+        if protocol == "grpc":
+            return BmiClientApptainer(
+                image=image_fn,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                timeout=timeout,
+                delay=delay,
+            )
+        if protocol == "openapi":
+            return remotebmi.BmiClientApptainer(
+                image=image_fn,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                delay=delay,
+            )
+        msg = f"Invalid protocol '{protocol}'!"
+        raise ValueError(msg)
+
     except FutureTimeoutError as exc:
         msg = (
             "Couldn't spawn container within allocated time limit "
@@ -258,6 +281,7 @@ def start_docker_container(
     image_port=55555,
     timeout=None,
     delay=0,
+    protocol: Literal["grpc", "openapi"] = "grpc",
 ):
     """Start Docker container with model inside.
 
@@ -268,6 +292,7 @@ def start_docker_container(
         image_port: Docker port inside container where grpc4bmi server is running.
         timeout: Number of seconds to wait for grpc connection.
         delay: Number of seconds to wait before connecting.
+        protocol: Which protocol to use, grpc or openapi.
 
     Raises:
         TimeoutError: When model inside container did not start quickly enough.
@@ -276,14 +301,26 @@ def start_docker_container(
         Bmi object which wraps the container.
     """
     try:
-        return BmiClientDocker(
-            image=image.docker_url,
-            image_port=image_port,
-            work_dir=str(work_dir),
-            input_dirs=input_dirs,
-            timeout=timeout,
-            delay=delay,
-        )
+        if protocol == "grpc":
+            return BmiClientDocker(
+                image=image.docker_url,
+                image_port=image_port,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                timeout=timeout,
+                delay=delay,
+            )
+        if protocol == "openapi":
+            return remotebmi.BmiClientDocker(
+                image=image.docker_url,
+                host="localhost",
+                image_port=50051,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                delay=delay,
+            )
+        msg = f"Invalid protocol '{protocol}'!"
+        raise ValueError(msg)
     except FutureTimeoutError as exc:
         # https://github.com/eWaterCycle/grpc4bmi/issues/95
         # https://github.com/eWaterCycle/grpc4bmi/issues/100

tests/src/base/test_model.py

@@ -312,6 +312,24 @@ def test_setup(self, mocked_start_container, tmp_path: Path):
 
         mocked_start_container.assert_called_once_with(
             image="ewatercycle/ewatercycle_dummy:latest",
+            protocol="grpc",
+            work_dir=tmp_path,
+            input_dirs=[],
+            timeout=300,
+        )
+
+    @patch("ewatercycle.base.model.start_container")
+    def test_remotebmi_setup(self, mocked_start_container, tmp_path: Path):
+        model = ContainerizedModel(
+            bmi_image="ewatercycle/ewatercycle_dummy:latest",
+            protocol="openapi",
+        )
+
+        model.setup(cfg_dir=str(tmp_path))
+
+        mocked_start_container.assert_called_once_with(
+            image="ewatercycle/ewatercycle_dummy:latest",
+            protocol="openapi",
             work_dir=tmp_path,
             input_dirs=[],
             timeout=300,
@@ -345,6 +363,7 @@ def test_setup_with_additional_input_dirs(
 
         mocked_start_container.assert_called_once_with(
             image="ewatercycle/ewatercycle_dummy:latest",
+            protocol="grpc",
             work_dir=tmp_path,
             input_dirs=[
                 str(parameter_set_dir),