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