Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow starting remotebmi containers #467

Merged
merged 7 commits into from
Nov 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)).

Expand Down
3 changes: 3 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ dependencies = [
# otherwise pip installation fails on esmpy
"Fiona",
"grpc4bmi>=0.4.0",
"remotebmi",
"hydrostats",
"matplotlib>=3.5.0",
"numpy",
Expand Down Expand Up @@ -205,6 +206,8 @@ ignore = [
"E501",
# Allow prints
"T201",
# Allow shadowing builtins
"A004",
]

[tool.ruff.lint.pydocstyle]
Expand Down
4 changes: 3 additions & 1 deletion src/ewatercycle/base/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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([])
Expand All @@ -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,
Expand Down
71 changes: 54 additions & 17 deletions src/ewatercycle/container.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -136,6 +137,7 @@
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.

Expand All @@ -152,6 +154,7 @@
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.
Expand Down Expand Up @@ -184,7 +187,13 @@

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(
Expand All @@ -193,6 +202,7 @@
input_dirs,
timeout,
delay,
protocol,
)
else:
msg = f"Unknown container technology: {CFG.container_engine}"
Expand All @@ -209,6 +219,7 @@
input_dirs: Iterable[str] = (),
timeout: int | None = None,
delay: int = 0,
protocol: Literal["grpc", "openapi"] = "grpc",
) -> Bmi:
"""Start Apptainer container with model inside.

Expand All @@ -220,6 +231,7 @@
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

Expand All @@ -234,13 +246,24 @@
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(

Check warning on line 258 in src/ewatercycle/container.py

View check run for this annotation

Codecov / codecov/patch

src/ewatercycle/container.py#L258

Added line #L258 was not covered by tests
image=image_fn,
work_dir=str(work_dir),
input_dirs=input_dirs,
delay=delay,
)
msg = f"Invalid protocol '{protocol}'!"
raise ValueError(msg)

Check warning on line 265 in src/ewatercycle/container.py

View check run for this annotation

Codecov / codecov/patch

src/ewatercycle/container.py#L264-L265

Added lines #L264 - L265 were not covered by tests

except FutureTimeoutError as exc:
msg = (
"Couldn't spawn container within allocated time limit "
Expand All @@ -258,6 +281,7 @@
image_port=55555,
timeout=None,
delay=0,
protocol: Literal["grpc", "openapi"] = "grpc",
):
"""Start Docker container with model inside.

Expand All @@ -268,6 +292,7 @@
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.
Expand All @@ -276,14 +301,26 @@
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(

Check warning on line 305 in src/ewatercycle/container.py

View check run for this annotation

Codecov / codecov/patch

src/ewatercycle/container.py#L305

Added line #L305 was not covered by tests
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(

Check warning on line 314 in src/ewatercycle/container.py

View check run for this annotation

Codecov / codecov/patch

src/ewatercycle/container.py#L314

Added line #L314 was not covered by tests
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)

Check warning on line 323 in src/ewatercycle/container.py

View check run for this annotation

Codecov / codecov/patch

src/ewatercycle/container.py#L322-L323

Added lines #L322 - L323 were not covered by tests
except FutureTimeoutError as exc:
# https://github.com/eWaterCycle/grpc4bmi/issues/95
# https://github.com/eWaterCycle/grpc4bmi/issues/100
Expand Down
19 changes: 19 additions & 0 deletions tests/src/base/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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),
Expand Down