Upgrade to BMI 2.0

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
       - id: trailing-whitespace
         args: [--markdown-linebreak-ext=md]
   - repo: https://github.com/adrienverge/yamllint
-    rev: "v1.26.0"
+    rev: "v1.29.0"
     hooks:
       - id: yamllint
   - repo: https://github.com/asottile/setup-cfg-fmt
@@ -23,7 +23,7 @@ repos:
     hooks:
       - id: black-jupyter
   - repo: https://github.com/PyCQA/isort
-    rev: "5.9.3"
+    rev: 5.12.0
     hooks:
       - id: isort
   # TODO renable when errors are fixed/ignored
@@ -78,7 +78,7 @@ repos:
     rev: 1.1.0
     hooks:
       - id: nbqa-isort
-        additional_dependencies: [isort==5.9.3]
+        additional_dependencies: [isort==5.11.2]
       - id: nbqa-mypy
         additional_dependencies: [mypy==0.910, types-python-dateutil]
       # TODO renable when errors are fixed/ignored

docs/conf.py

@@ -191,7 +191,7 @@ def setup(app):  # noqa: D103
 ]
 
 autodoc_mock_imports = [
-    "basic_modeling_interface",
+    "bmipy",
     "cftime",
     "dask",
     "esmvalcore",

setup.cfg

@@ -37,10 +37,10 @@ packages = find:
 install_requires =
     Fiona
     Shapely
-    basic_modeling_interface
+    bmipy
     cftime
     esmvaltool>=2.4.0
-    grpc4bmi>=0.2.12,<0.3
+    grpc4bmi@git+https://github.com/eWaterCycle/grpc4bmi@bmi2
     grpcio
     hydrostats
     matplotlib>=3.5.0

src/ewatercycle/config/_lisflood_versions.py

@@ -1,19 +1,10 @@
-"""
-Versions of Lisflood container images
-"""
-from pathlib import Path
+"""Versions of Lisflood container images."""
 
-version_images = {
+from ewatercycle.container import VersionImages
+
+version_images: VersionImages = {
     "20.10": {
         "docker": "ewatercycle/lisflood-grpc4bmi:20.10",
         "singularity": "ewatercycle-lisflood-grpc4bmi_20.10.sif",
     }
 }
-
-
-def get_docker_image(version):
-    return version_images[version]["docker"]
-
-
-def get_singularity_image(version, singularity_dir: Path):
-    return singularity_dir / version_images[version]["singularity"]

src/ewatercycle/container.py

@@ -0,0 +1,97 @@
+"""Container utilities."""
+from pathlib import Path
+from typing import Dict, Iterable, Literal, Mapping, Optional, Union
+
+from bmipy import Bmi
+from grpc import FutureTimeoutError
+from grpc4bmi.bmi_client_docker import BmiClientDocker
+from grpc4bmi.bmi_client_singularity import BmiClientSingularity
+from grpc4bmi.bmi_memoized import MemoizedBmi
+from grpc4bmi.bmi_optionaldest import OptionalDestBmi
+
+from ewatercycle import CFG
+
+ContainerEngines = Literal["docker", "singularity"]
+"""Supported container engines."""
+
+ImageForContainerEngines = Dict[ContainerEngines, str]
+"""Container image name for each container engine."""
+
+VersionImages = Mapping[str, ImageForContainerEngines]
+"""Dictionary of versions of a model.
+
+Each version has the image name for each container engine.
+"""
+
+
+def start_container(
+    work_dir: Union[str, Path],
+    image_engine: ImageForContainerEngines,
+    input_dirs: Optional[Iterable[str]] = None,
+    image_port=55555,
+    timeout=None,
+    delay=0,
+) -> Bmi:
+    """Start container with model inside.
+
+    The `ewatercycle.CFG['container_engine']` value determines
+    the engine used to start a container.
+
+    Args:
+        work_dir: Work directory
+        version_image: Image name for each container engine.
+        input_dirs: Additional directories to mount inside 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.
+
+    Raises:
+        ValueError: When unknown container technology is requested.
+        TimeoutError: When model inside container did not start quickly enough.
+
+    Returns:
+        _description_
+    """
+    engine: ContainerEngines = CFG["container_engine"]
+    image = image_engine[engine]
+    if input_dirs is None:
+        input_dirs = []
+    if engine == "docker":
+        try:
+            bmi = BmiClientDocker(
+                image=image,
+                image_port=image_port,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                timeout=timeout,
+                delay=delay,
+            )
+        except FutureTimeoutError as exc:
+            # https://github.com/eWaterCycle/grpc4bmi/issues/95
+            # https://github.com/eWaterCycle/grpc4bmi/issues/100
+            raise TimeoutError(
+                "Couldn't spawn container within allocated time limit "
+                f"({timeout} seconds). You may try pulling the docker image with"
+                f" `docker pull {image}` and then try again."
+            ) from exc
+    elif engine == "singularity":
+        image = str(CFG["singularity_dir"] / image)
+        try:
+            bmi = BmiClientSingularity(
+                image=image,
+                work_dir=str(work_dir),
+                input_dirs=input_dirs,
+                timeout=timeout,
+                delay=delay,
+            )
+        except FutureTimeoutError as exc:
+            docker_image = image_engine["docker"]
+            raise TimeoutError(
+                "Couldn't spawn container within allocated time limit "
+                f"({timeout} seconds). You may try pulling the docker image with"
+                f" `singularity build {image} "
+                f"docker://{docker_image}` and then try again."
+            ) from exc
+    else:
+        raise ValueError(f"Unknown container technology: {CFG['container_engine']}")
+    return OptionalDestBmi(MemoizedBmi(bmi))

src/ewatercycle/forcing/_lisvap.py

@@ -26,9 +26,10 @@
 from typing import Dict, Tuple
 
 from ewatercycle import CFG
+from ewatercycle.container import ContainerEngines
 from ewatercycle.parametersetdb.config import XmlConfig
 
-from ..config._lisflood_versions import get_docker_image, get_singularity_image
+from ..config._lisflood_versions import version_images
 from ..util import get_time
 
 
@@ -49,9 +50,10 @@ def lisvap(
         mask_map,
         forcing_dir,
     )
-
+    engine: ContainerEngines = CFG["container_engine"]
+    image = version_images[version][engine]
     if CFG["container_engine"].lower() == "singularity":
-        image = get_singularity_image(version, CFG["singularity_dir"])
+        image = CFG["singularity_dir"] / image
         args = [
             "singularity",
             "exec",
@@ -62,7 +64,6 @@ def lisvap(
             image,
         ]
     elif CFG["container_engine"].lower() == "docker":
-        image = get_docker_image(version)
         args = [
             "docker",
             "run",

src/ewatercycle/models/abstract.py

@@ -1,12 +1,13 @@
+"""Abstract class of a eWaterCycle model."""
 import logging
 import textwrap
 from abc import ABCMeta, abstractmethod
 from datetime import datetime
-from typing import Any, ClassVar, Generic, Iterable, Optional, Set, Tuple, TypeVar
+from typing import Any, ClassVar, Generic, Iterable, Optional, Tuple, TypeVar
 
 import numpy as np
 import xarray as xr
-from basic_modeling_interface import Bmi
+from bmipy import Bmi
 from cftime import num2date
 
 from ewatercycle.forcing import DefaultForcing
@@ -169,7 +170,7 @@ def get_value_as_xarray(self, name: str) -> xr.DataArray:
     @property
     @abstractmethod
     def parameters(self) -> Iterable[Tuple[str, Any]]:
-        """Default values for the setup() inputs"""
+        """Default values for the setup() inputs."""
 
     @property
     def start_time(self) -> float:

src/ewatercycle/models/hype.py

@@ -1,27 +1,24 @@
+"""eWaterCycle wrapper around Hype BMI."""
 import datetime
 import logging
 import shutil
 import types
 from typing import Any, Iterable, Optional, Tuple
 
-import numpy as np
 import xarray as xr
-from basic_modeling_interface import Bmi
-from cftime import num2date
 from dateutil.parser import parse
 from dateutil.tz import UTC
-from grpc4bmi.bmi_client_docker import BmiClientDocker
-from grpc4bmi.bmi_client_singularity import BmiClientSingularity
 
 from ewatercycle import CFG
+from ewatercycle.container import VersionImages, start_container
 from ewatercycle.forcing._hype import HypeForcing
 from ewatercycle.models.abstract import AbstractModel
 from ewatercycle.parameter_sets import ParameterSet
 from ewatercycle.util import geographical_distances, get_time, to_absolute_path
 
 logger = logging.getLogger(__name__)
 
-_version_images = {
+_version_images: VersionImages = {
     "feb2021": {
         "docker": "ewatercycle/hype-grpc4bmi:feb2021",
         "singularity": "ewatercycle-hype-grpc4bmi_feb2021.sif",
@@ -44,6 +41,7 @@ class Hype(AbstractModel[HypeForcing]):
     """
 
     available_versions = tuple(_version_images.keys())
+    """Show supported Hype versions in eWaterCycle"""
 
     def __init__(
         self,
@@ -52,7 +50,6 @@ def __init__(
         forcing: Optional[HypeForcing] = None,
     ):
         super().__init__(version, parameter_set, forcing)
-        assert version in _version_images
         self._setup_default_config()
 
     def _setup_default_config(self):
@@ -155,8 +152,10 @@ def setup(  # type: ignore
         cfg_file.write_text(self._cfg, encoding="cp437")
 
         # start container
-        work_dir = str(cfg_dir_as_path)
-        self.bmi = _start_container(self.version, work_dir)
+        self.bmi = start_container(
+            image_engine=_version_images[self.version],
+            work_dir=cfg_dir_as_path,
+        )
 
         since = self._start.strftime("%Y-%m-%dT%H:%M:%SZ")
 
@@ -167,7 +166,7 @@ def get_time_units(_self):
 
         self.bmi.get_time_units = types.MethodType(get_time_units, self.bmi)
 
-        return str(cfg_file), work_dir
+        return str(cfg_file), str(cfg_dir_as_path)
 
     @property
     def parameters(self) -> Iterable[Tuple[str, Any]]:
@@ -220,26 +219,6 @@ def _setup_cfg_dir(cfg_dir: Optional[str] = None):
     return work_dir
 
 
-def _start_container(version: str, work_dir: str):
-    if CFG["container_engine"].lower() == "singularity":
-        image = CFG["singularity_dir"] / _version_images[version]["singularity"]
-        return BmiClientSingularity(
-            image=str(image),
-            work_dir=work_dir,
-        )
-    elif CFG["container_engine"].lower() == "docker":
-        image = _version_images[version]["docker"]
-        return BmiClientDocker(
-            image=image,
-            image_port=55555,  # TODO needed?
-            work_dir=work_dir,
-        )
-    else:
-        raise ValueError(
-            f"Unknown container technology in CFG: {CFG['container_engine']}"
-        )
-
-
 def _get_code_in_cfg(content: str, code: str):
     lines = content.splitlines()
     for line in lines:
@@ -265,5 +244,5 @@ def _set_code_in_cfg(content: str, code: str, value: str) -> str:
 
 
 def _get_hype_time(value: str) -> datetime.datetime:
-    """Converts `yyyy-mm-dd [HH:MM]` string to datetime object"""
+    """Converts `yyyy-mm-dd [HH:MM]` string to datetime object."""
     return parse(value).replace(tzinfo=UTC)

src/ewatercycle/models/lisflood.py

@@ -8,15 +8,10 @@
 import numpy as np
 import xarray as xr
 from cftime import num2date
-from grpc4bmi.bmi_client_docker import BmiClientDocker
-from grpc4bmi.bmi_client_singularity import BmiClientSingularity
 
 from ewatercycle import CFG
-from ewatercycle.config._lisflood_versions import (
-    get_docker_image,
-    get_singularity_image,
-    version_images,
-)
+from ewatercycle.config._lisflood_versions import version_images
+from ewatercycle.container import start_container
 from ewatercycle.forcing._lisflood import LisfloodForcing
 from ewatercycle.models.abstract import AbstractModel
 from ewatercycle.parameter_sets import ParameterSet
@@ -115,27 +110,13 @@ def setup(  # type: ignore
                 # If not relative add dir
                 input_dirs.append(str(mask_map.parent))
 
-        if CFG["container_engine"].lower() == "singularity":
-            image = get_singularity_image(self.version, CFG["singularity_dir"])
-            self.bmi = BmiClientSingularity(
-                image=str(image),
-                input_dirs=input_dirs,
-                work_dir=str(cfg_dir_as_path),
-                timeout=300,
-            )
-        elif CFG["container_engine"].lower() == "docker":
-            image = get_docker_image(self.version)
-            self.bmi = BmiClientDocker(
-                image=image,
-                image_port=55555,
-                input_dirs=input_dirs,
-                work_dir=str(cfg_dir_as_path),
-                timeout=300,
-            )
-        else:
-            raise ValueError(
-                f"Unknown container technology in CFG: {CFG['container_engine']}"
-            )
+        self.bmi = start_container(
+            image_engine=version_images[self.version],
+            work_dir=cfg_dir_as_path,
+            input_dirs=input_dirs,
+            timeout=300,
+        )
+
         return str(config_file), str(cfg_dir_as_path)
 
     def _check_forcing(self, forcing):