From b786752e3c4ebf54783396ef58eaba4231b930d6 Mon Sep 17 00:00:00 2001 From: nilsmechtel Date: Mon, 7 Sep 2026 14:14:38 +0200 Subject: [PATCH] feat(cli): add `bioengine worker start` to launch a worker container Starting a worker meant copying a ~10-line `docker run` block out of the deployment guide and adapting the GPU flag, the uid/gid and the bind mount by hand. `bioengine worker start` builds that command instead, picking the first of docker, podman and apptainer on PATH. The CLI is a launcher, not an in-process worker: it never imports ray, so `pip install "bioengine[cli]"` stays light and ray only ever runs inside the image. `--runtime native` is the escape hatch for environments that already have the worker extra installed. Worker arguments are forwarded verbatim after `--`, so the ~50 options on `python -m bioengine.worker` are available without being re-declared here and cannot drift out of sync. The image tag pins to the installed `bioengine` version so the CLI and the worker it starts cannot diverge. The Hypha token travels in the environment and is only ever named on the command line, never valued. Also adds `bioengine worker stop` / `logs`, and leads the deployment guide and README with the CLI while keeping the raw container commands as the documented fallback. --- README.md | 9 +- bioengine/cli/cli.py | 9 ++ bioengine/cli/worker.py | 305 +++++++++++++++++++++++++++++++++++ docs/deployment-guide.md | 36 ++++- tests/cli/test_worker_cli.py | 213 ++++++++++++++++++++++++ 5 files changed, 570 insertions(+), 2 deletions(-) create mode 100644 bioengine/cli/worker.py create mode 100644 tests/cli/test_worker_cli.py diff --git a/README.md b/README.md index ec270bf..50f5787 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,13 @@ Facility managers and system administrators can deploy a private worker on any h | **Kubernetes** | Production deployment with KubeRay | ```bash -# Docker — single machine quickstart +# Single machine quickstart — runs the worker image in a container +pip install "bioengine[cli]" +bioengine worker start -- --mode single-machine --head-num-cpus 4 +``` + +```bash +# Or from a clone, with docker compose git clone https://github.com/aicell-lab/bioengine.git cd bioengine mkdir -p .bioengine data @@ -131,6 +137,7 @@ pip install "bioengine[cli] @ git+https://github.com/aicell-lab/bioengine.git" bioengine call bioimage-io/bioengine-worker get_status bioengine apps list --worker bioimage-io/bioengine-worker +bioengine worker start -- --mode single-machine ``` ### Worker service API diff --git a/bioengine/cli/cli.py b/bioengine/cli/cli.py index 7fc002d..7965582 100644 --- a/bioengine/cli/cli.py +++ b/bioengine/cli/cli.py @@ -6,6 +6,7 @@ bioengine apps deploy ./my-app/ bioengine apps status bioengine cluster status + bioengine worker start -- --mode single-machine Environment variables: BIOENGINE_SERVER_URL Hypha server URL (default: https://hypha.aicell.io) @@ -18,6 +19,7 @@ from bioengine.cli.call import call_command from bioengine.cli.cluster import cluster_group from bioengine.cli.apps import apps_group +from bioengine.cli.worker import worker_group @click.group() @@ -43,6 +45,12 @@ def main(): Inspect cluster resources: bioengine cluster status + \b + Run a worker: + bioengine worker start -- --mode single-machine --head-num-cpus 4 + bioengine worker logs -f + bioengine worker stop + \b Environment variables: BIOENGINE_SERVER_URL Server URL (default: https://hypha.aicell.io) @@ -56,6 +64,7 @@ def main(): main.add_command(call_command) main.add_command(apps_group) main.add_command(cluster_group) +main.add_command(worker_group) if __name__ == "__main__": diff --git a/bioengine/cli/worker.py b/bioengine/cli/worker.py new file mode 100644 index 0000000..3056cec --- /dev/null +++ b/bioengine/cli/worker.py @@ -0,0 +1,305 @@ +""" +bioengine worker — start, stop and follow a BioEngine worker container. + +Wraps the container invocation from docs/deployment-guide.md. Worker arguments +are forwarded verbatim to ``python -m bioengine.worker`` inside the image, so +this module never has to know what they are. + +Examples: + bioengine worker start -- --mode single-machine --head-num-cpus 4 + bioengine worker start --dry-run -- --mode single-machine + bioengine worker logs -f + bioengine worker stop +""" +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path +from typing import List, Optional, Tuple + +import click + +from bioengine import __version__ +from bioengine.cli.utils import error_exit + +DEFAULT_IMAGE_REPO = "ghcr.io/aicell-lab/bioengine-worker" +DEFAULT_CONTAINER_NAME = "bioengine-worker" +DEFAULT_WORKSPACE_DIR = Path.home() / ".bioengine" +# The path the image mounts the workspace at — see docs/deployment-guide.md. +CONTAINER_WORKSPACE_DIR = "/.bioengine" +DEFAULT_SHM_SIZE = "8g" + +# Only the GPU flag differs between docker and podman (deployment-guide.md). +_GPU_FLAGS = { + "docker": ["--gpus=all"], + "podman": ["--device", "nvidia.com/gpu=all"], + "apptainer": ["--nv"], +} + +_RUNTIME_PREFERENCE = ("docker", "podman", "apptainer") + +_ENV_PASSTHROUGH = ("HYPHA_TOKEN", "BIOENGINE_SERVER_URL") + + +def _detect_runtime() -> Optional[str]: + for runtime in _RUNTIME_PREFERENCE: + if shutil.which(runtime): + return runtime + return None + + +def _has_gpu() -> bool: + """Whether to pass a GPU flag by default. + + Passing ``--gpus=all`` on a host without the NVIDIA container toolkit makes + the runtime refuse to start, so it cannot simply be on by default. + """ + return shutil.which("nvidia-smi") is not None + + +def build_command( + runtime: str, + image: str, + worker_args: Tuple[str, ...], + workspace_dir: Path, + container_name: str, + shm_size: str, + gpus: bool, + detach: bool, +) -> List[str]: + """Build the container invocation. Secrets travel in the environment, never argv.""" + entrypoint = ["python", "-m", "bioengine.worker", *worker_args] + + if runtime == "native": + return entrypoint + + if runtime == "apptainer": + command = ["apptainer", "exec"] + if gpus: + command += _GPU_FLAGS["apptainer"] + command += ["--bind", f"{workspace_dir}:{CONTAINER_WORKSPACE_DIR}"] + return command + [f"docker://{image}", *entrypoint] + + command = [runtime, "run", "--rm"] + command += ["--detach"] if detach else ["-it"] + command += ["--name", container_name] + command += ["--user", f"{os.getuid()}:{os.getgid()}"] + command += ["--shm-size", shm_size] + if gpus: + command += _GPU_FLAGS[runtime] + command += ["-v", f"{workspace_dir}:{CONTAINER_WORKSPACE_DIR}"] + for name in _ENV_PASSTHROUGH: + if os.environ.get(name): + command += ["-e", name] + return command + [image, *entrypoint] + + +def _subprocess_env(runtime: str, token: Optional[str], server_url: Optional[str]) -> dict: + env = dict(os.environ) + if token: + env["HYPHA_TOKEN"] = token + if server_url: + env["BIOENGINE_SERVER_URL"] = server_url + if runtime == "apptainer": + # Apptainer only forwards host variables it is told about explicitly. + for name in _ENV_PASSTHROUGH: + if env.get(name): + env[f"APPTAINERENV_{name}"] = env[name] + return env + + +def _resolve_runtime(runtime: str, require_available: bool = True) -> str: + """Resolve 'auto'. ``require_available`` is off for --dry-run, whose whole + point is producing a command to run somewhere else.""" + if runtime != "auto": + if require_available and runtime != "native" and not shutil.which(runtime): + error_exit( + f"Container runtime '{runtime}' is not on PATH.", + "Install it, or pass --runtime native to run the worker in this environment.", + ) + return runtime + + detected = _detect_runtime() + if not detected: + if not require_available: + return _RUNTIME_PREFERENCE[0] + error_exit( + "No container runtime found (looked for docker, podman, apptainer).", + "Install one, or pass --runtime native to run the worker in this environment.", + ) + return detected + + +@click.group("worker") +def worker_group(): + """Start, stop and follow a BioEngine worker container.""" + + +@worker_group.command( + "start", + context_settings={"ignore_unknown_options": True}, +) +@click.argument("worker_args", nargs=-1, type=click.UNPROCESSED) +@click.option( + "--runtime", + type=click.Choice(["auto", "docker", "podman", "apptainer", "native"]), + default="auto", + help="Container runtime. 'auto' picks the first of docker, podman, apptainer on PATH. " + "'native' runs the worker in this environment instead (requires the 'worker' extra).", +) +@click.option( + "--image", + default=None, + metavar="IMAGE", + help=f"Worker image (default: {DEFAULT_IMAGE_REPO}:).", +) +@click.option( + "--workspace-dir", + type=click.Path(file_okay=False, path_type=Path), + default=DEFAULT_WORKSPACE_DIR, + show_default=True, + help="Host directory mounted as the worker workspace.", +) +@click.option( + "--name", + "container_name", + default=DEFAULT_CONTAINER_NAME, + show_default=True, + help="Container name, used by 'bioengine worker stop' and 'logs'.", +) +@click.option( + "--shm-size", default=DEFAULT_SHM_SIZE, show_default=True, help="Shared memory size." +) +@click.option( + "--gpus/--no-gpus", + default=None, + help="Request GPUs. Defaults to on when nvidia-smi is present.", +) +@click.option("--detach", "-d", is_flag=True, help="Run the container in the background.") +@click.option( + "--token", + envvar=["HYPHA_TOKEN", "BIOENGINE_TOKEN"], + default=None, + metavar="TOKEN", + help="Hypha auth token (or HYPHA_TOKEN env var). Passed via the environment, not the command line.", +) +@click.option("--server-url", envvar="BIOENGINE_SERVER_URL", default=None, hidden=True) +@click.option("--dry-run", is_flag=True, help="Print the command instead of running it.") +def worker_start( + worker_args, + runtime, + image, + workspace_dir, + container_name, + shm_size, + gpus, + detach, + token, + server_url, + dry_run, +): + """ + Start a BioEngine worker. + + Everything after ``--`` is forwarded verbatim to ``python -m bioengine.worker`` + inside the container, so every worker option is available without this command + knowing about it. Run ``bioengine worker start -- --help`` to see them. + + \b + Examples: + bioengine worker start -- --mode single-machine --head-num-cpus 4 + bioengine worker start -d --no-gpus -- --mode single-machine + bioengine worker start --dry-run -- --mode single-machine + """ + runtime = _resolve_runtime(runtime, require_available=not dry_run) + + if gpus is None: + gpus = _has_gpu() + + # Pinned to the installed version so the CLI and the worker it starts cannot + # silently diverge. + image = image or f"{DEFAULT_IMAGE_REPO}:{__version__}" + + workspace_dir = workspace_dir.expanduser() + if runtime != "native" and not dry_run: + workspace_dir.mkdir(parents=True, exist_ok=True) + + command = build_command( + runtime=runtime, + image=image, + worker_args=worker_args, + workspace_dir=workspace_dir, + container_name=container_name, + shm_size=shm_size, + gpus=gpus, + detach=detach, + ) + + if dry_run: + click.echo(" ".join(command)) + return + + env = _subprocess_env(runtime, token, server_url) + try: + raise SystemExit(subprocess.call(command, env=env)) + except FileNotFoundError: + error_exit(f"Failed to execute '{command[0]}': not found on PATH.") + + +@worker_group.command("stop") +@click.option( + "--runtime", + type=click.Choice(["auto", "docker", "podman"]), + default="auto", + help="Container runtime. 'auto' picks the first of docker, podman on PATH.", +) +@click.option( + "--name", + "container_name", + default=DEFAULT_CONTAINER_NAME, + show_default=True, + help="Container name.", +) +def worker_stop(runtime, container_name): + """Stop a running BioEngine worker container.""" + runtime = _resolve_runtime(runtime) + if runtime not in ("docker", "podman"): + error_exit( + f"'{runtime}' has no named containers to stop.", + "Stop the worker process directly.", + ) + raise SystemExit(subprocess.call([runtime, "stop", container_name])) + + +@worker_group.command("logs") +@click.option( + "--runtime", + type=click.Choice(["auto", "docker", "podman"]), + default="auto", + help="Container runtime. 'auto' picks the first of docker, podman on PATH.", +) +@click.option( + "--name", + "container_name", + default=DEFAULT_CONTAINER_NAME, + show_default=True, + help="Container name.", +) +@click.option("--follow", "-f", is_flag=True, help="Follow log output.") +@click.option("--tail", default=None, metavar="N", help="Show only the last N lines.") +def worker_logs(runtime, container_name, follow, tail): + """Show the logs of a running BioEngine worker container.""" + runtime = _resolve_runtime(runtime) + if runtime not in ("docker", "podman"): + error_exit( + f"'{runtime}' has no named containers to read logs from.", + "Read the worker's own log file under the workspace directory instead.", + ) + command = [runtime, "logs"] + if follow: + command.append("-f") + if tail: + command += ["--tail", str(tail)] + raise SystemExit(subprocess.call(command + [container_name])) diff --git a/docs/deployment-guide.md b/docs/deployment-guide.md index c3f22b5..ce7ab42 100644 --- a/docs/deployment-guide.md +++ b/docs/deployment-guide.md @@ -15,7 +15,41 @@ BioEngine supports three deployment modes. The easiest way to generate deploymen Runs a local Ray cluster on one machine. Good for workstations, development, and small-scale analysis. -### Docker (recommended) +### The `bioengine` CLI (recommended) + +```bash +pip install "bioengine[cli]" + +bioengine worker start -- \ + --mode single-machine \ + --head-num-cpus 4 \ + --head-num-gpus 1 +``` + +This runs the worker image in a container. It picks the first of `docker`, `podman` and `apptainer` on your `PATH`, mounts `~/.bioengine` as the workspace, passes `HYPHA_TOKEN` through the environment, and pins the image tag to the installed `bioengine` version. Everything after `--` is forwarded verbatim to `python -m bioengine.worker` inside the container — see `bioengine worker start -- --help` for the full list. + +```bash +bioengine worker start --dry-run -- --mode single-machine # print the command, run nothing +bioengine worker start -d -- --mode single-machine # run in the background +bioengine worker logs -f +bioengine worker stop +``` + +| Option | Default | Description | +|---|---|---| +| `--runtime` | `auto` | `docker`, `podman`, `apptainer`, or `native` to run the worker in the current environment instead of a container | +| `--image` | `ghcr.io/aicell-lab/bioengine-worker:` | Worker image | +| `--workspace-dir` | `~/.bioengine` | Host directory mounted at `/.bioengine` | +| `--gpus` / `--no-gpus` | on when `nvidia-smi` is present | Whether to give the container GPUs | +| `--shm-size` | `8g` | Shared memory size | +| `--detach` / `-d` | off | Run in the background | +| `--dry-run` | off | Print the command instead of running it | + +`--gpus` only decides whether the *container* sees GPUs; tell Ray to use them with `--head-num-gpus` after the `--`. + +### Running the container directly + +The CLI is a thin wrapper — the underlying commands work on their own: ```bash docker run --rm -it \ diff --git a/tests/cli/test_worker_cli.py b/tests/cli/test_worker_cli.py new file mode 100644 index 0000000..4240956 --- /dev/null +++ b/tests/cli/test_worker_cli.py @@ -0,0 +1,213 @@ +""" +Contract for ``bioengine worker`` — the container launcher. + +The command's whole job is turning options into an argv list, so these tests +pin that argv: the shape documented in docs/deployment-guide.md for each +runtime, worker arguments forwarded verbatim, and the auth token never +appearing on a command line where `ps` would show it. +""" +import os +from pathlib import Path + +import pytest +from click.testing import CliRunner + +from bioengine import __version__ +from bioengine.cli.worker import ( + CONTAINER_WORKSPACE_DIR, + DEFAULT_IMAGE_REPO, + _subprocess_env, + build_command, + worker_group, +) + +WORKSPACE = Path("/home/someone/.bioengine") +IMAGE = f"{DEFAULT_IMAGE_REPO}:{__version__}" +WORKER_ARGS = ("--mode", "single-machine", "--head-num-cpus", "4") + + +def _build(runtime, **overrides): + kwargs = dict( + runtime=runtime, + image=IMAGE, + worker_args=WORKER_ARGS, + workspace_dir=WORKSPACE, + container_name="bioengine-worker", + shm_size="8g", + gpus=False, + detach=False, + ) + kwargs.update(overrides) + return build_command(**kwargs) + + +def _run(args, env=None): + base = {"HYPHA_TOKEN": "", "BIOENGINE_SERVER_URL": "", "BIOENGINE_TOKEN": ""} + base.update(env or {}) + return CliRunner().invoke(worker_group, args, env=base) + + +# ── The documented invocation, per runtime ──────────────────────────────────── + + +def test_the_docker_command_matches_the_deployment_guide(): + command = _build("docker", gpus=True) + assert command[:3] == ["docker", "run", "--rm"] + assert "--user" in command and f"{os.getuid()}:{os.getgid()}" in command + assert command[command.index("--shm-size") + 1] == "8g" + assert "--gpus=all" in command + assert f"{WORKSPACE}:{CONTAINER_WORKSPACE_DIR}" in command + assert command[-len(WORKER_ARGS) - 4 :] == [ + IMAGE, + "python", + "-m", + "bioengine.worker", + *WORKER_ARGS, + ] + + +def test_podman_uses_its_own_gpu_flag(): + command = _build("podman", gpus=True) + assert command[:2] == ["podman", "run"] + assert "--gpus=all" not in command + assert ["--device", "nvidia.com/gpu=all"] == command[ + command.index("--device") : command.index("--device") + 2 + ] + + +def test_apptainer_binds_instead_of_mounting(): + command = _build("apptainer", gpus=True) + assert command[:2] == ["apptainer", "exec"] + assert "--nv" in command + assert command[command.index("--bind") + 1] == f"{WORKSPACE}:{CONTAINER_WORKSPACE_DIR}" + assert f"docker://{IMAGE}" in command + # No container to name, detach or size — those flags belong to docker/podman. + for flag in ("--name", "--detach", "--shm-size", "--user"): + assert flag not in command + + +def test_native_runs_the_worker_without_a_container(): + command = _build("native", gpus=True) + assert command == ["python", "-m", "bioengine.worker", *WORKER_ARGS] + + +def test_the_gpu_flag_is_omitted_when_gpus_are_off(): + for runtime in ("docker", "podman", "apptainer"): + command = _build(runtime, gpus=False) + assert "--gpus=all" not in command + assert "--device" not in command + assert "--nv" not in command + + +def test_detaching_replaces_the_interactive_flags(): + assert "-it" in _build("docker") + detached = _build("docker", detach=True) + assert "--detach" in detached and "-it" not in detached + + +# ── The token must never reach argv ─────────────────────────────────────────── + + +def test_the_token_is_named_not_valued_in_the_container_command(monkeypatch): + monkeypatch.setenv("HYPHA_TOKEN", "secret-token-value") + for runtime in ("docker", "podman"): + command = _build(runtime) + assert "secret-token-value" not in command + assert command[command.index("-e") + 1] == "HYPHA_TOKEN" + + +def test_the_token_never_appears_in_any_runtimes_command(monkeypatch): + monkeypatch.setenv("HYPHA_TOKEN", "secret-token-value") + for runtime in ("docker", "podman", "apptainer", "native"): + assert "secret-token-value" not in " ".join(_build(runtime)) + + +def test_the_token_is_passed_through_the_environment(): + env = _subprocess_env("docker", token="secret-token-value", server_url=None) + assert env["HYPHA_TOKEN"] == "secret-token-value" + + +def test_apptainer_needs_the_prefixed_variable_to_forward_anything(): + env = _subprocess_env("apptainer", token="secret-token-value", server_url=None) + assert env["APPTAINERENV_HYPHA_TOKEN"] == "secret-token-value" + + +def test_an_unset_variable_is_not_forwarded(monkeypatch): + monkeypatch.delenv("HYPHA_TOKEN", raising=False) + assert "HYPHA_TOKEN" not in _build("docker") + assert "APPTAINERENV_HYPHA_TOKEN" not in _subprocess_env("apptainer", None, None) + + +# ── Worker arguments are forwarded, not interpreted ─────────────────────────── + + +def test_worker_arguments_are_forwarded_verbatim(): + args = ("--mode", "slurm", "--admin-users", "a@x.org,b@y.org", "--debug") + assert _build("native", worker_args=args)[3:] == list(args) + + +def test_an_option_the_cli_also_defines_still_reaches_the_worker(): + """``--workspace-dir`` after ``--`` configures the worker, not the container.""" + result = _run( + ["start", "--runtime", "native", "--dry-run", "--", "--workspace-dir", "/data/ws"] + ) + assert result.exit_code == 0, result.output + assert result.output.strip().endswith("--workspace-dir /data/ws") + + +def test_no_worker_arguments_still_starts_the_worker_module(): + assert _build("native", worker_args=())[-1] == "bioengine.worker" + + +# ── The CLI surface ─────────────────────────────────────────────────────────── + + +def test_the_image_is_pinned_to_the_installed_version(): + result = _run(["start", "--runtime", "docker", "--dry-run", "--", "--mode", "single-machine"]) + assert result.exit_code == 0, result.output + assert f"{DEFAULT_IMAGE_REPO}:{__version__}" in result.output + + +def test_a_dry_run_neither_creates_the_workspace_nor_needs_the_runtime(tmp_path): + workspace = tmp_path / "never-created" + result = _run( + [ + "start", + "--runtime", + "podman", + "--workspace-dir", + str(workspace), + "--dry-run", + "--", + "--mode", + "single-machine", + ] + ) + assert result.exit_code == 0, result.output + assert result.output.startswith("podman run") + assert not workspace.exists() + + +def test_a_missing_runtime_is_refused_when_actually_starting(monkeypatch): + monkeypatch.setattr("bioengine.cli.worker.shutil.which", lambda _: None) + result = _run(["start", "--runtime", "podman", "--", "--mode", "single-machine"]) + assert result.exit_code == 1 + assert "not on PATH" in result.output + + +def test_stop_and_logs_refuse_runtimes_without_named_containers(monkeypatch): + """On an apptainer-only host there is no container name to act on.""" + monkeypatch.setattr( + "bioengine.cli.worker.shutil.which", lambda name: name if name == "apptainer" else None + ) + for command in ("stop", "logs"): + result = _run([command]) + assert result.exit_code == 1 + assert "no named containers" in result.output + + +@pytest.mark.parametrize("command", ["start", "stop", "logs"]) +def test_every_subcommand_is_reachable(command): + result = _run([command, "--help"]) + assert result.exit_code == 0 + assert "bioengine-worker" in result.output or "worker" in result.output