feat(worker): opt-in container-image runtime for GPU apps (single-machine) (#156)

* feat(worker): opt-in container-image runtime for GPU apps (single-machine)

Add an opt-in path to run a Ray Serve replica inside a prebuilt container
image instead of the default pip/py_modules runtime. Pip stays the
recommended default; this is strictly gated behind the worker flag
--enable-container-runtime and the @bioengine.app(container_image=…) param.

- decorators: container_image= param drops a {"container": {"image": …}}
  marker into runtime_env; raises if combined with pip (Ray forbids mixing
  a container runtime_env with pip/py_modules).
- bootstrap: a container branch in build_and_run_application builds a
  minimal runtime_env ({container, env_vars} only — the sole keys Ray
  permits alongside container) and returns before the normal source/hook
  injection. The image is the sole source of code (bioengine + app source
  baked on PYTHONPATH). GPU wiring rides in podman run_options:
  --device nvidia.com/gpu=all + LD_LIBRARY_PATH to the driver libs.
- ray_cluster: when enabled, generate a podman-4.9.3-compatible CDI spec at
  startup (nvidia-ctk, --disable-hook update-ldcache, down-version to
  cdiVersion 0.6.0, strip additionalGids + its orphaned bare-int items) and
  export the driver-lib dir for the replica LD_LIBRARY_PATH.
- worker/__main__: --enable-container-runtime flag (Ray Cluster group,
  flows straight into RayCluster via ray_cluster_config).
- docker/worker-container-runtime.Dockerfile: dedicated rootful worker image
  (ubuntu:24.04 for podman 4.9.3, nvidia-container-toolkit, Ray keep-id
  patch). Kept a separate image from the lean python:3.11-slim prod worker
  on purpose — the podman pin forces the heavier base and its CVE surface
  should stay out of the production worker's scan report.

Single-machine mode only; KubeRay path unaffected.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* build(worker): add manual build script for container-runtime image

Manual, on-demand builder for the container-as-runtime worker image
(docker/worker-container-runtime.Dockerfile). Deliberately not wired into
docker-publish-worker.yml — the niche, heavier ubuntu:24.04 variant is
built only when needed and published as a separate GHCR package sharing
the worker version.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* chore(release): bump version to 0.14.0

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: nilsmechtel

bioengine/_app/bootstrap.py

@@ -532,6 +532,42 @@ def _with_pkg(cls: Any, spec_ray_opts: Dict[str, Any]) -> Any:
             if _rk in (spec_ray_opts or {}):
                 opts[_rk] = spec_ray_opts[_rk]
         runtime_env = dict(opts.get("runtime_env") or {})
+
+        container = runtime_env.get("container")
+        if container:
+            # Container-as-runtime (opt-in via @bioengine.app(container_image=…)).
+            # The image is the sole source of code: it bakes bioengine + the
+            # app source (on PYTHONPATH) + deps, so ``import bioengine`` and
+            # ``import <entry>`` resolve natively at cloudpickle.loads and the
+            # replica finder never needs to fetch source. Ray REJECTS any
+            # runtime_env that pairs ``container`` with py_modules/pip/
+            # worker_process_setup_hook (only env_vars/config are allowed),
+            # so this branch deliberately builds a minimal runtime_env and
+            # returns before the normal injection below.
+            #
+            # GPU + host wiring rides in podman ``run_options`` because Ray's
+            # plugin hardcodes its base flags and only appends run_options;
+            # ``--device nvidia.com/gpu=all`` is the CDI GPU injector and
+            # ``LD_LIBRARY_PATH`` points at the driver libs the disabled
+            # ``update-ldcache`` CDI hook would otherwise have wired up.
+            run_options = list(container.get("run_options") or [])
+            if opts.get("num_gpus"):
+                run_options += ["--device", "nvidia.com/gpu=all"]
+            ld_library_path = os.environ.get(
+                "BIOENGINE_CONTAINER_LD_LIBRARY_PATH", "/usr/lib64"
+            )
+            run_options += ["-e", f"LD_LIBRARY_PATH={ld_library_path}"]
+            opts["runtime_env"] = {
+                "container": {
+                    "image": container["image"],
+                    "run_options": run_options,
+                },
+                "env_vars": {
+                    **replica_env_vars,
+                    **(runtime_env.get("env_vars") or {}),
+                },
+            }
+            return cls.options(ray_actor_options=opts)
         # Ray Serve replicas do NOT inherit job-level py_modules (observed
         # empirically on KTH). The bioengine package has to ride in the
         # deployment's runtime_env to be on sys.path at

bioengine/_app/decorators.py

@@ -199,6 +199,7 @@ def app(
     num_gpus: float = 0,
     memory_mb: Optional[int] = None,
     pip: Optional[List[str]] = None,
+    container_image: Optional[str] = None,
     env_vars: Optional[Dict[str, str]] = None,
     max_ongoing_requests: int = 10,
     ray_actor_options: Optional[Dict[str, Any]] = None,
@@ -214,7 +215,15 @@ def app(
             convention internally.
         pip: Additional pip requirements for the replica's runtime_env on
             top of BioEngine's baseline (``hypha-rpc``, ``pydantic``,
-            ``httpx``, and the ``bioengine[worker]`` package).
+            ``httpx``, and the ``bioengine[worker]`` package). Mutually
+            exclusive with ``container_image``.
+        container_image: Opt-in alternative to pip: run the replica inside
+            this prebuilt container image instead of installing deps
+            per-actor. The image must be self-contained — Ray forbids
+            combining ``container`` with ``py_modules``/``pip``, so the
+            image has to bake ``bioengine`` plus the app source on
+            ``PYTHONPATH``. Requires the worker to run with
+            ``--enable-container-runtime`` (single-machine mode only).
         env_vars: Static env vars baked into the replica's runtime_env.
             Secrets passed at deploy time are layered on top.
         max_ongoing_requests: Concurrent request cap per replica.
@@ -234,6 +243,13 @@ def decorator(cls: type) -> Any:
                 f"{type(cls).__name__}"
             )
 
+        if container_image and pip:
+            raise ValueError(
+                "@bioengine.app cannot combine 'container_image' with 'pip': "
+                "Ray forbids mixing a container runtime_env with pip. Bake "
+                "your dependencies into the image instead."
+            )
+
         _reject_reserved_names(cls)
         lifecycle, method_schemas = _scan_class(cls)
         composition_params = _extract_composition_params(cls)
@@ -253,6 +269,7 @@ def decorator(cls: type) -> Any:
             num_gpus=num_gpus,
             memory_mb=memory_mb,
             pip=pip,
+            container_image=container_image,
             env_vars=env_vars,
             extra=ray_actor_options,
         )
@@ -381,6 +398,7 @@ def _build_ray_actor_options(
     num_gpus: float,
     memory_mb: Optional[int],
     pip: Optional[List[str]],
+    container_image: Optional[str],
     env_vars: Optional[Dict[str, str]],
     extra: Optional[Dict[str, Any]],
 ) -> Dict[str, Any]:
@@ -390,6 +408,13 @@ def _build_ray_actor_options(
         opts["memory"] = int(memory_mb) * 1024 * 1024
 
     runtime_env: Dict[str, Any] = {}
+    if container_image:
+        # Marker only; the worker's bootstrap fills in the GPU/host
+        # ``run_options`` at deploy time (it knows whether the container
+        # runtime is enabled). Ray forbids ``container`` alongside pip/
+        # py_modules, so the container branch in bootstrap._with_pkg
+        # deliberately skips those.
+        runtime_env["container"] = {"image": container_image}
     if pip:
         runtime_env["pip"] = list(pip)
     if env_vars:

bioengine/_version.py

@@ -13,4 +13,4 @@
 Must stay in lock-step with ``pyproject.toml``'s ``version`` field.
 The ``version-check.yml`` CI workflow enforces the match.
 """
-__version__ = "0.13.0"
+__version__ = "0.14.0"

bioengine/cluster/ray_cluster.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import re
+import shutil
 import subprocess
 import time
 import uuid
@@ -93,6 +94,7 @@ def __init__(
         head_memory_in_gb: Optional[int] = None,
         runtime_env_pip_cache_size_gb: int = 30,  # Ray default is 10 GB
         force_clean_up: bool = True,
+        enable_container_runtime: bool = False,
         # SLURM Worker Configuration parameters
         image: str = f"ghcr.io/aicell-lab/bioengine-worker:{bioengine.__version__}",
         worker_workspace_dir: Optional[str] = None,
@@ -135,6 +137,10 @@ def __init__(
             head_memory_in_gb: Memory limit for head node in GB. If not set, Ray will auto-detect available memory.
             runtime_env_pip_cache_size_gb: Size of pip cache for runtime environments in GB. Default 30.
             force_clean_up: Force cleanup of previous Ray cluster on start. Default True.
+            enable_container_runtime: Opt-in container-as-runtime for apps
+                using ``@bioengine.app(container_image=…)`` (single-machine
+                mode). Generates a podman-compatible CDI spec at startup for
+                nested-GPU passthrough. Default False.
             image: Container image for workers (SLURM mode). Default bioengine-worker.
             worker_workspace_dir: Workspace directory mounted to worker containers (SLURM mode).
             default_num_gpus: Default GPU count per worker. Default 1.
@@ -174,6 +180,7 @@ def __init__(
                 "Supported modes are 'slurm', 'single-machine' and 'external-cluster'."
             )
         self.mode = mode
+        self.enable_container_runtime = bool(enable_container_runtime)
 
         # Initialize cluster state and monitoring attributes
         self.address = None
@@ -582,6 +589,105 @@ def _update_symlink(self, ray_temp_dir: Path) -> None:
             self.logger.error(f"Symlink '{symlink_path}' does not exist")
             raise FileNotFoundError(f"Symlink '{symlink_path}' does not exist")
 
+    async def _generate_cdi_spec(self) -> None:
+        """Emit a podman-4.9.3-compatible CDI spec for nested-GPU passthrough.
+
+        Container-as-runtime replicas run in a nested podman container; the
+        only GPU injector that works there is native podman CDI
+        (``--device nvidia.com/gpu=all``). Two host-specific fixups are baked
+        in because the recipe was validated on a double-nested cgroup-v1 host:
+
+        * ``--disable-hook update-ldcache`` — the ldcache hook runs ``ldconfig``
+          via a mount op that fails in the nested container; the driver libs
+          are reached via ``LD_LIBRARY_PATH`` instead (exported below and read
+          by ``bootstrap._with_pkg`` as a podman ``-e`` run option).
+        * down-version to ``cdiVersion: 0.6.0`` and strip the 0.7.0-only
+          ``additionalGids`` field — podman 4.9.3's CDI parser predates 0.7.0.
+
+        No-op with a warning if ``nvidia-ctk`` is absent or the process is not
+        root (a rootless worker can't write ``/etc/cdi`` nor create a usable
+        CUDA context anyway).
+        """
+        if os.geteuid() != 0:
+            self.logger.warning(
+                "enable_container_runtime is set but the worker is not running "
+                "as root; skipping CDI generation. GPU container apps will fail."
+            )
+            return
+
+        nvidia_ctk = shutil.which("nvidia-ctk")
+        if not nvidia_ctk:
+            self.logger.warning(
+                "enable_container_runtime is set but 'nvidia-ctk' was not found; "
+                "skipping CDI generation. GPU container apps will fail."
+            )
+            return
+
+        cdi_path = "/etc/cdi/nvidia.yaml"
+        await asyncio.to_thread(
+            lambda: Path(cdi_path).parent.mkdir(parents=True, exist_ok=True)
+        )
+
+        proc = await asyncio.create_subprocess_exec(
+            nvidia_ctk,
+            "cdi",
+            "generate",
+            f"--output={cdi_path}",
+            "--disable-hook",
+            "update-ldcache",
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        _, stderr = await proc.communicate()
+        if proc.returncode != 0:
+            raise subprocess.CalledProcessError(
+                proc.returncode,
+                "nvidia-ctk cdi generate",
+                stderr=stderr.decode() if stderr else "Unknown error",
+            )
+
+        def _downversion_spec() -> None:
+            lines = Path(cdi_path).read_text().splitlines()
+            out = []
+            for line in lines:
+                if "additionalgids" in line.lower():
+                    continue
+                # additionalGids' numeric list items (e.g. `- 44`) would otherwise
+                # orphan into the preceding deviceNodes list once the key is gone,
+                # producing a bare number where podman expects a DeviceNode. No
+                # other list entry in the spec is a bare integer, so drop them.
+                if re.fullmatch(r"\s*-\s*\d+\s*", line):
+                    continue
+                if line.startswith("cdiVersion:"):
+                    line = "cdiVersion: 0.6.0"
+                out.append(line)
+            Path(cdi_path).write_text("\n".join(out) + "\n")
+
+        await asyncio.to_thread(_downversion_spec)
+
+        # Point replicas' LD_LIBRARY_PATH at the host dir holding libcuda.so.1
+        # (CDI mounts driver libs at their host paths inside the container).
+        ld_dir = "/usr/lib64"
+        ldconfig = shutil.which("ldconfig")
+        if ldconfig:
+            proc = await asyncio.create_subprocess_exec(
+                ldconfig,
+                "-p",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.DEVNULL,
+            )
+            stdout, _ = await proc.communicate()
+            for line in stdout.decode().splitlines():
+                if "libcuda.so.1" in line and "=>" in line:
+                    ld_dir = str(Path(line.split("=>")[-1].strip()).parent)
+                    break
+        os.environ["BIOENGINE_CONTAINER_LD_LIBRARY_PATH"] = ld_dir
+
+        self.logger.info(
+            f"Generated CDI spec at {cdi_path} (cdiVersion 0.6.0); "
+            f"container GPU LD_LIBRARY_PATH={ld_dir}"
+        )
+
     async def _start_cluster(self) -> None:
         """Start Ray cluster head node with configured ports and resources.
 
@@ -614,6 +720,11 @@ async def _start_cluster(self) -> None:
             # Check and set cluster ports
             await asyncio.to_thread(self._set_cluster_ports)
 
+            # Opt-in container-as-runtime: emit a podman-compatible CDI spec so
+            # GPU apps can run inside their image (single-machine mode only).
+            if self.enable_container_runtime:
+                await self._generate_cdi_spec()
+
             # Start ray as the head node with the specified parameters
             args = [
                 "start",

bioengine/worker/__main__.py

@@ -319,6 +319,15 @@ def create_parser() -> argparse.ArgumentParser:
         help="Skip cleanup of previous Ray cluster processes and data. "
         "Use with caution as it may cause port conflicts or resource issues.",
     )
+    ray_cluster_group.add_argument(
+        "--enable-container-runtime",
+        action="store_true",
+        help="Opt-in: allow apps declared with @bioengine.app(container_image=…) "
+        "to run inside a prebuilt container image as their Ray Serve runtime "
+        "(single-machine mode only). Requires a rootful worker with podman + "
+        "nvidia-container-toolkit; the worker generates a podman-compatible CDI "
+        "spec at startup for GPU passthrough. Pip remains the default runtime.",
+    )
 
     # SLURM job configuration options (for HPC environments)
     slurm_job_group = parser.add_argument_group(

docker/worker-container-runtime.Dockerfile

@@ -0,0 +1,77 @@
+# PoC-only worker image for container-as-runtime GPU apps (single-machine,
+# ROOTFUL). Differs from worker.Dockerfile in three ways:
+#   1. Ubuntu 24.04 base — ships podman 4.9.3 via apt (podman 5.x rejects
+#      --pid=host on cgroup-v1 hosts, which Ray's container plugin requires).
+#   2. Bundles nvidia-container-toolkit (nvidia-ctk + CDI hooks) so the worker
+#      can emit a CDI spec at startup for nested-GPU passthrough.
+#   3. Patches the installed Ray to drop the hardcoded `--userns=keep-id`
+#      podman flag, which breaks /proc remount in a nested rootful container.
+#
+# This image is NOT a drop-in replacement for the production worker; it is the
+# rootful PoC target for @bioengine.app(container_image=…). Pip runtime stays
+# the default and is unaffected.
+FROM ubuntu:24.04
+
+ENV DEBIAN_FRONTEND=noninteractive \
+    PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1 \
+    SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt
+
+# System deps: build toolchain + git (framework), podman 4.9.3 + fuse-overlayfs
+# + uidmap (nested container runtime). Python 3.11 via deadsnakes to match the
+# app image's interpreter — the deployment class is cloudpickled from the worker
+# (head) into the app-image replica, so the two Python versions must agree.
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        git build-essential curl gnupg ca-certificates software-properties-common \
+        podman fuse-overlayfs uidmap \
+    && add-apt-repository -y ppa:deadsnakes/ppa \
+    && apt-get update && apt-get install -y --no-install-recommends \
+        python3.11 python3.11-venv python3.11-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# nvidia-container-toolkit: provides nvidia-ctk (CDI spec generation) and the
+# nvidia-cdi-hook binaries (create-symlinks etc.).
+RUN curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
+      | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
+ && curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
+      | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
+      > /etc/apt/sources.list.d/nvidia-container-toolkit.list \
+ && apt-get update && apt-get install -y --no-install-recommends nvidia-container-toolkit \
+ && rm -rf /var/lib/apt/lists/*
+
+# Isolated venv (Ubuntu 24.04 marks the system Python externally-managed).
+RUN python3.11 -m venv /opt/venv
+ENV PATH=/opt/venv/bin:$PATH
+
+WORKDIR /app
+
+# See worker.Dockerfile: requirements first (does NOT pin Ray), so the
+# RAY_VERSION build arg can change without invalidating this layer.
+COPY requirements-worker.txt /app/
+RUN pip install -U pip && \
+    pip install -r requirements-worker.txt
+
+COPY bioengine/ /app/bioengine/
+COPY pyproject.toml README.md LICENSE /app/
+RUN pip install --no-deps .
+
+ARG RAY_VERSION=2.55.1
+RUN pip install "ray[client,serve]==${RAY_VERSION}"
+ENV BIOENGINE_RAY_VERSION=${RAY_VERSION}
+
+# Drop Ray's hardcoded `--userns=keep-id` podman flag. keep-id is a rootless
+# volume-permission remap; combined with --pid=host in a nested ROOTFUL
+# container it fails to remount /proc ("mount proc to proc: Operation not
+# permitted"). The grep guard fails the build loudly if a Ray bump renames or
+# removes the anchor line, so we never silently ship an unpatched image.
+# The bare string `userns=keep-id` also appears in a doc comment (the example
+# command), so the post-patch guard must check the QUOTED code anchor is gone,
+# not the bare string.
+RUN f="$(python -c 'import ray._private.runtime_env.image_uri as m; print(m.__file__)')" \
+ && grep -q '"--userns=keep-id",' "$f" \
+ && sed -i '/"--userns=keep-id",/d' "$f" \
+ && ! grep -q '"--userns=keep-id",' "$f" \
+ && echo "Patched out --userns=keep-id in $f"
+
+CMD [ "/bin/bash" ]

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "bioengine"
-version = "0.13.0"
+version = "0.14.0"
 description = "BioEngine — CLI and SDK for deploying and calling AI model services on BioEngine workers"
 requires-python = ">=3.11"
 authors = [