Apptainer socket poller and tests

.github/workflows/python.yml

@@ -90,6 +90,9 @@ jobs:
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
+      - uses: eWaterCycle/setup-apptainer@v2
+        with:
+          apptainer-version: 1.4.5
       - name: Setup Python
         uses: actions/setup-python@v5
         with:

.gitignore

@@ -34,3 +34,4 @@ openapi-generator-cli.jar
 openapitools.json
 PEQ_Hupsel.dat
 walrus.yml
+python/tests/client/fixtures/leakybucket.sif

python/src/remotebmi/__init__.py

@@ -2,4 +2,4 @@
 from remotebmi.client.client import RemoteBmiClient
 from remotebmi.client.docker import BmiClientDocker
 
-__all__ = ["RemoteBmiClient", "BmiClientApptainer", "BmiClientDocker"]
+__all__ = ["BmiClientApptainer", "BmiClientDocker", "RemoteBmiClient"]

python/src/remotebmi/client/apptainer.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import socket
 import subprocess
 import time
 from collections.abc import Iterable
@@ -9,15 +10,20 @@
 from remotebmi.client.client import RemoteBmiClient
 from remotebmi.client.utils import DeadContainerError, get_unique_port
 
+logger = logging.getLogger(__name__)
+
 
 class BmiClientApptainer(RemoteBmiClient):
     def __init__(
         self,
         image: str,
         work_dir: str,
         input_dirs: Iterable[str] = (),
-        delay: int = 0,
+        delay: float = 0,
         capture_logs: bool = True,
+        # wait up to 10 minutes for container to start, some models can be slow to start up
+        startup_timeout: float = 600.0,
+        startup_poll_interval: float = 0.1,
     ):
         if isinstance(input_dirs, str):
             msg = (
@@ -44,7 +50,7 @@ def __init__(
         # Change into working directory
         args += ["--pwd", self.work_dir]
         args.append(image)
-        logging.info(f"Running {image} apptainer container on port {port}")
+        logger.info(f"Running {image} apptainer container on port {port}")
         if capture_logs:
             self.logfile = SpooledTemporaryFile(  # noqa: SIM115 - file is closed in __del__
                 max_size=2**16,  # keep until 65Kb in memory if bigger write to disk
@@ -67,14 +73,64 @@ def __init__(
             msg = (
                 f"apptainer container {image} prematurely exited with code {returncode}"
             )
+            logger.error(msg)
             raise DeadContainerError(
                 msg,
                 returncode,
                 self.logs(),
             )
         url = f"http://{host}:{port}"
+
+        self._wait_until_connectable(
+            host=host,
+            port=port,
+            url=url,
+            image=image,
+            startup_timeout=startup_timeout,
+            startup_poll_interval=startup_poll_interval,
+        )
+
         super().__init__(url)
 
+    def _wait_until_connectable(
+        self,
+        host: str,
+        port: int,
+        url: str,
+        image: str,
+        startup_timeout: float,
+        startup_poll_interval: float,
+    ) -> None:
+        deadline = time.monotonic() + startup_timeout
+        while time.monotonic() < deadline:
+            returncode = self.container.poll()
+            if returncode is not None:
+                msg = f"apptainer container {image} prematurely exited with code {returncode}"
+                raise DeadContainerError(
+                    msg,
+                    returncode,
+                    self.logs(),
+                )
+
+            try:
+                with socket.create_connection(
+                    (host, port), timeout=startup_poll_interval
+                ):
+                    return
+            except OSError:
+                time.sleep(startup_poll_interval)
+
+        self.container.terminate()
+        self.container.wait()
+        logs = self.logs()
+        msg = (
+            f"Timed out after {startup_timeout} seconds waiting for apptainer "
+            f"container {image} to accept connections on {url}"
+        )
+        if logs:
+            msg += f". Container logs: {logs}"
+        raise TimeoutError(msg)
+
     def __del__(self) -> None:
         if hasattr(self, "container"):
             self.container.terminate()

python/src/remotebmi/client/docker.py

@@ -4,7 +4,7 @@
 from posixpath import abspath
 
 import docker
-from docker.models.containers import Container  # noqa: TCH002
+from docker.models.containers import Container  # noqa: TC002
 
 from remotebmi.client.client import RemoteBmiClient
 from remotebmi.client.utils import DeadContainerError, get_unique_port, getuser

python/src/remotebmi/server/__init__.py

@@ -47,7 +47,7 @@ async def lifespan_handler(app: ConnexionMiddleware) -> AsyncIterator:  # noqa:
 def main(**kwargs):  # type: ignore[no-untyped-def]
     model = from_env()
     app = make_app(model)
-    port = int(environ.get("BMI_PORT", 50051))
+    port = int(environ.get("BMI_PORT", "50051"))
     uvicorn.run(app, port=port, **kwargs)
 
 

python/tests/client/fixtures/leakybucket.def

@@ -0,0 +1,12 @@
+Bootstrap: docker
+From: python:3.14-slim
+
+%environment
+    export BMI_MODULE=leakybucket.leakybucket_bmi
+    export BMI_CLASS=LeakyBucketBmi
+
+%post
+    python3 -m pip install --no-cache-dir leakybucket remotebmi
+
+%runscript
+    exec run-bmi-server "$@"

python/tests/client/test_apptainer.py

@@ -0,0 +1,61 @@
+import subprocess
+from collections.abc import Generator
+from pathlib import Path
+
+import pytest
+
+from remotebmi.client.apptainer import BmiClientApptainer, DeadContainerError
+
+try:
+    _apptainer_available = (
+        subprocess.run(
+            ["apptainer", "--version"],  # noqa: S607
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        ).returncode
+        == 0
+    )
+except OSError:
+    _apptainer_available = False
+
+pytestmark = pytest.mark.skipif(
+    not _apptainer_available,
+    reason="apptainer is unavailable",
+)
+
+
+@pytest.fixture
+def leakybucket_def() -> Path:
+    return Path(__file__).parent / "fixtures" / "leakybucket.def"
+
+
+@pytest.fixture
+def leakybucket_image(leakybucket_def: Path) -> Path:
+    image = leakybucket_def.parent / "leakybucket.sif"
+    if image.exists():
+        return image
+    subprocess.run(["apptainer", "build", str(image), str(leakybucket_def)], check=True)  # noqa: S603, S607
+    return image
+
+
+@pytest.fixture
+def leakybucket_client(
+    leakybucket_image: Path, tmp_path: Path
+) -> Generator[BmiClientApptainer, None, None]:
+    client = BmiClientApptainer(
+        image=str(leakybucket_image),
+        work_dir=str(tmp_path),
+    )
+    yield client
+    del client
+
+
+def test_get_component_name(leakybucket_client: BmiClientApptainer) -> None:
+    assert leakybucket_client.get_component_name() == "leakybucket"
+
+
+def test_image_not_found(tmp_path: Path) -> None:
+    bad_image = tmp_path / "nonexistent_image.sif"
+    with pytest.raises(DeadContainerError):
+        BmiClientApptainer(image=str(bad_image), work_dir=str(tmp_path))