diff --git a/.claude/skills/verify-warpforth/SKILL.md b/.claude/skills/verify-warpforth/SKILL.md index 13312e1..2fff263 100644 --- a/.claude/skills/verify-warpforth/SKILL.md +++ b/.claude/skills/verify-warpforth/SKILL.md @@ -1,14 +1,14 @@ --- name: verify-warpforth -description: Drive the real WarpForth toolchain to prove a change works. Compiles Forth with build/bin/warpforthc (and the warpforth-translate | warpforth-opt stages) locally, and executes kernels on a rented Vast.ai GPU via warpforth-runner. Use after changing the parser, dialect, conversions, pipeline, or CLI flags, when lit tests alone don't show the kernel actually computes the right numbers. +description: Drive the real WarpForth toolchain to prove a change works. Compiles Forth with build/bin/warpforthc (and the warpforth-translate | warpforth-opt stages) locally, and executes kernels on a rented Vast.ai GPU through the warpforth_gpu device worker. Use after changing the parser, dialect, conversions, pipeline, or CLI flags, when lit tests alone don't show the kernel actually computes the right numbers. --- # Verify WarpForth -WarpForth is a CLI compiler: the user writes a `.forth` kernel, runs `warpforthc` to get PTX, and runs that PTX on an NVIDIA GPU with `warpforth-runner`. There is no server or UI. This machine (macOS) has no CUDA GPU, no `nvcc`, no `ptxas`, so: +WarpForth is a CLI compiler: the user writes a `.forth` kernel, runs `warpforthc` to get PTX, and runs that PTX on an NVIDIA GPU through `warpforth_gpu.client.Device`, which talks to a stdlib-only Python worker (`warpforth_gpu/worker.py`) over newline-delimited JSON. There is no server or UI. This machine (macOS) has no CUDA GPU and no `ptxas`, so: - **Local surface** (free, always available): `warpforthc`, `warpforth-translate`, `warpforth-opt`. Proves parsing, diagnostics, lowering, and that PTX is emitted. -- **GPU surface** (costs money, needs `VASTAI_API_KEY`): rents one Vast.ai GPU (cap $0.50/hr, `nvidia/cuda:12.4.0-devel-ubuntu22.04`, host driver must support CUDA 12.4), builds `warpforth-runner` there with `nvcc`, executes PTX, reads back a buffer. This is the only proof that a kernel computes the right values. **Ask the user before each GPU run**: it spends their money. +- **GPU surface** (costs money, needs `VASTAI_API_KEY`): rents one Vast.ai GPU (cap $0.50/hr, `nvidia/cuda:12.4.0-devel-ubuntu22.04`, host driver must support CUDA 12.4), starts the device worker there over SSH (its source is passed with `python3 -c`, nothing is copied), executes PTX, reads back a buffer. This is the only proof that a kernel computes the right values. **Ask the user before each GPU run**: it spends their money. Other surfaces, not mapped yet: the GPT-2 demo (`demo/`, `docs/demo.md`) needs a PyTorch Vast.ai image and a manual scp. @@ -24,13 +24,13 @@ cmake --build build Ready when it exits 0 and `$S/doctor.sh` prints no `FAIL`. Each compile is its own short-lived process, so any number of agents can compile side by side. Teardown: nothing. -GPU launch happens inside `$S/gpu_run.py` (below): it searches offers, creates an instance labeled `warpforth-test--`, waits for sshd, compiles the runner with `nvcc`. Usually 2-7 minutes (the first proof run took 6: 5 to boot, 1 for nvcc and the jobs). Concurrent GPU runs are safe: each owns its own label and destroys only that. +GPU launch happens inside `$S/gpu_run.py` (below): it searches offers, creates an instance labeled `warpforth-test--`, waits for sshd, installs `python3` if the image lacks it. Usually 2-7 minutes, almost all of it instance boot. Concurrent GPU runs are safe: each owns its own label and destroys only that. ## Doctor ```bash $S/doctor.sh # local: binaries exist, `ninja -n` shows no pending build steps, smoke compile emits `.visible .entry main` -$S/doctor.sh --remote # plus: uv/ssh/scp on PATH, Vast.ai key accepted, lists any running warpforth-test-* instances +$S/doctor.sh --remote # plus: uv/ssh on PATH, Vast.ai key accepted, lists any running warpforth-test-* instances ``` Read-only (the smoke compile writes only to a mktemp dir). Exit 0 means worth driving. A stale build is the most common false result: a verification run against old binaries proves nothing, so rebuild when doctor says stale. `--remote` warnings about existing instances are informational; they may be another agent's live run. @@ -51,7 +51,7 @@ Runs exactly `build/bin/warpforthc --arch -o kernel.ptx` (the user uv run python $S/gpu_run.py ``` -One spec = one rented GPU for all its jobs. Spec format is in the script's docstring: per job a `.forth` source, `params` by header name, `grid`, `block`, `output_param` (index into the `\! param` list, must be an array), `output_count`, and optional `expect` (a prefix of the output; floats compared with `isclose`). It reuses `gpu_test/conftest.py`'s `VastSession` and `KernelRunner`, the same path as `uv run pytest -m gpu`. Exit 0 only if every `expect` matched and instance destruction was verified. A ready-made spec covering int arithmetic, multi-thread matmul (the README quickstart), and f64 math is at `.claude/skills/verify-warpforth/fixtures/smoke/spec.json`. For a change, write a new spec next to the kernels you want to prove, in scratch space, not inside the skill. +One spec = one rented GPU for all its jobs. Spec format is in the script's docstring: per job a `.forth` source, `params` by header name, `grid`, `block`, `output_param` (index into the `\! param` list, must be an array), `output_count`, and optional `expect` (a prefix of the output; floats compared with `isclose`). It reuses `warpforth_gpu.vast.VastSession` and `gpu_test/conftest.py`'s `KernelRunner`, the same path as `uv run pytest -m gpu`. The worker protocol itself is covered locally, without a GPU, by `uv run pytest gpu_test/test_worker.py`. Exit 0 only if every `expect` matched and instance destruction was verified. A ready-made spec covering int arithmetic, multi-thread matmul (the README quickstart), and f64 math is at `.claude/skills/verify-warpforth/fixtures/smoke/spec.json`. For a change, write a new spec next to the kernels you want to prove, in scratch space, not inside the skill. The full pytest GPU suite (`uv run pytest -v -m gpu`, ~100+ kernels) is the regression net; `gpu_run.py` is for proving the specific kernel your change affects with captured evidence. @@ -64,11 +64,11 @@ Everything lands in `build/verify-evidence//` (`build/` is gitignored; n Proof standards: -- Drive `warpforthc` / `warpforth-runner` as a user would. Do not call internal C++ APIs or hand-edit MLIR to make a case pass. +- Drive `warpforthc` and the `warpforth_gpu` client as a user would. Do not call internal C++ APIs or hand-edit MLIR to make a case pass. - A compile proof shows the source, the command's exit code, and the PTX (or the diagnostic for an error case, with its `file:line:col`). "It compiled" is not proof of correctness. Correct values need a GPU run. - A GPU proof shows the kernel source, the PTX that ran, the inputs (the spec), and the read-back buffer matching an expectation you computed independently, not one copied from the output. - Quote the `evidence:` path in your report. -- `--output-param` reads back one buffer. To prove a side effect in a second buffer, add a second job that outputs it. +- `output_param` reads back one buffer. To prove a side effect in a second buffer, add a second job that outputs it. ## Cleanup diff --git a/.claude/skills/verify-warpforth/features/gpu-execution.md b/.claude/skills/verify-warpforth/features/gpu-execution.md index 5df190e..4c51336 100644 --- a/.claude/skills/verify-warpforth/features/gpu-execution.md +++ b/.claude/skills/verify-warpforth/features/gpu-execution.md @@ -1,6 +1,6 @@ # GPU execution -A user runs compiled PTX on an NVIDIA GPU with `warpforth-runner`, passing array and scalar parameters, a grid and block shape, and reads back one array parameter as comma-separated values. This is the only journey that proves a kernel computes the right numbers. +A user runs compiled PTX on an NVIDIA GPU through `warpforth_gpu.client.Device`, uploading array buffers, passing scalar values, choosing a grid and block shape, and reading buffers back as numpy arrays. This is the only journey that proves a kernel computes the right numbers. ## Sub-features @@ -12,7 +12,7 @@ A user runs compiled PTX on an NVIDIA GPU with `warpforth-runner`, passing array ## How to get to it (user POV) -- Compile locally with `warpforthc`, then on a GPU host: `warpforth-runner kernel.ptx --kernel main --param 'i64[]:1,2,3' --param i64:42 --grid 6,1,1 --block 1,1,1 --output-param 0 --output-count 3`. +- Compile locally with `warpforthc`, then on a GPU host, from the repo root: `Device.spawn(worker_command())`, `device.load(ptx).kernel("main").launch(grid=..., block=..., args=[device.upload(array), 42])`, `buffer.read(np.int64)`. For a remote host, prefix the ssh argv: `Device.spawn([*ssh_argv, shlex.join(worker_command())])`. - `uv run pytest -v -m gpu` (regression suite, same harness). ## Driving it with gpu_run.py @@ -33,6 +33,7 @@ Preconditions: - Rental can fail with `No suitable GPU offers found` under the $0.50/hr cap; retry later or report it, don't raise the cap without the user. - `output_param` must index an array param; scalars cannot be read back. +- A CUDA error during a launch (for example an illegal address) is fatal to the worker; `KernelRunner` starts a fresh one for the next job, and that job's `result.json` shows the `DeviceError`. - `params` values are padded with zeros to the declared array size; `expect` is a prefix match, so trailing zeros need not be listed. - Higher `arch` in the spec (e.g. `sm_90`) narrows offers to newer GPUs and raises cost and wait. -- Instance boot plus `nvcc` takes minutes; run in the background and don't kill it mid-run. If you do, run `cleanup.py` on its evidence dir. +- Instance boot takes minutes; run in the background and don't kill it mid-run. If you do, run `cleanup.py` on its evidence dir. diff --git a/.claude/skills/verify-warpforth/scripts/cleanup.py b/.claude/skills/verify-warpforth/scripts/cleanup.py index 898aa5d..1bf96dd 100755 --- a/.claude/skills/verify-warpforth/scripts/cleanup.py +++ b/.claude/skills/verify-warpforth/scripts/cleanup.py @@ -26,7 +26,7 @@ def owned(sdk: VastAI, label: str) -> list[int]: instances = sdk.show_instances() if not isinstance(instances, list): msg = "Vast.ai returned an invalid instance listing" - raise RuntimeError(msg) + raise TypeError(msg) return [int(i["id"]) for i in instances if isinstance(i, dict) and i.get("label") == label] diff --git a/.claude/skills/verify-warpforth/scripts/doctor.sh b/.claude/skills/verify-warpforth/scripts/doctor.sh index f90426d..96d1b01 100755 --- a/.claude/skills/verify-warpforth/scripts/doctor.sh +++ b/.claude/skills/verify-warpforth/scripts/doctor.sh @@ -29,7 +29,7 @@ rm -rf "$smoke" if [ "${1:-}" = "--remote" ]; then command -v uv >/dev/null && ok "uv on PATH" || bad "uv missing" - command -v ssh >/dev/null && command -v scp >/dev/null && ok "ssh/scp on PATH" || bad "ssh/scp missing" + command -v ssh >/dev/null && ok "ssh on PATH" || bad "ssh missing" uv run --quiet python - <<'PY' || fail=1 import sys sys.path.insert(0, ".claude/skills/verify-warpforth/scripts") diff --git a/.claude/skills/verify-warpforth/scripts/gpu_run.py b/.claude/skills/verify-warpforth/scripts/gpu_run.py index a2d950e..a522132 100755 --- a/.claude/skills/verify-warpforth/scripts/gpu_run.py +++ b/.claude/skills/verify-warpforth/scripts/gpu_run.py @@ -33,13 +33,15 @@ import os import sys import traceback +from contextlib import closing from datetime import UTC, datetime from pathlib import Path REPO = Path(__file__).resolve().parents[4] sys.path.insert(0, str(REPO)) -from gpu_test.conftest import Compiler, KernelRunner, VastSession # noqa: E402 +from gpu_test.conftest import Compiler, KernelRunner, open_vast_device # noqa: E402 +from warpforth_gpu.vast import VastSession # noqa: E402 def load_api_key() -> str | None: @@ -99,9 +101,9 @@ def main() -> int: (out / "owned_label.txt").write_text(session.instance_label + "\n") print(f"evidence: {out}", flush=True) try: - with session: + runner = KernelRunner(lambda: open_vast_device(session), compiler) + with session, closing(runner): summary["instance_id"] = session.instance_id - runner = KernelRunner(session, compiler) for job in spec["jobs"]: job_dir = out / job["name"] job_dir.mkdir() diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d2ece84..727a504 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -86,7 +86,7 @@ jobs: run: cmake --build build --target check-format - name: Ruff check - run: uv run ruff check gpu_test/ && uv run ruff format --check gpu_test/ + run: uv run ruff check gpu_test/ warpforth_gpu/ && uv run ruff format --check gpu_test/ warpforth_gpu/ - name: Test GPU infrastructure run: uv run pytest -q -m "not gpu" diff --git a/AGENTS.md b/AGENTS.md index ff42db9..d7682d9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -38,7 +38,11 @@ CMake. **Tools:** - `tools/warpforth-translate/warpforth-translate.cpp` - Translation tool entry point - `tools/warpforth-opt/warpforth-opt.cpp` - Optimization tool entry point -- `tools/warpforth-runner/warpforth-runner.cpp` - PTX execution tool for GPU kernels + +**GPU execution:** +- `warpforth_gpu/worker.py` - Stdlib-only device worker; drives the CUDA driver via ctypes and speaks newline-delimited JSON on stdin/stdout +- `warpforth_gpu/client.py` - `Device`, `Module`, `Kernel`, `Buffer` client for the worker +- `warpforth_gpu/vast.py` - `VastSession`, rents a Vast.ai GPU and exposes `ssh_argv()` ## Tools Usage @@ -54,9 +58,21 @@ CMake. ./build/bin/warpforth-translate --forth-to-mlir test/example.forth | \ ./build/bin/warpforth-opt --warpforth-pipeline | \ ./build/bin/warpforth-translate --mlir-to-ptx > kernel.ptx +``` + +```python +# Execute PTX on a GPU (locally; prepend ssh argv and shlex.join the command for a remote host) +from pathlib import Path + +import numpy as np + +from warpforth_gpu.client import Device, worker_command -# Execute PTX on GPU -./warpforth-runner kernel.ptx --param i64[]:1,2,3 --param i64:42 --output-param 0 --output-count 3 +with Device.spawn(worker_command()) as device: + kernel = device.load(Path("kernel.ptx").read_text()).kernel("main") + data = device.upload(np.array([1, 2, 3], dtype=np.int64)) + kernel.launch(grid=(1, 1, 1), block=(3, 1, 1), args=[data, 42]) + print(data.read(np.int64)) ``` ## Adding New Operations @@ -79,15 +95,15 @@ Add corresponding conversion pattern in `lib/Conversion/ForthToMemRef/ForthToMem ## GPU Tests -End-to-end GPU execution tests live in `gpu_test/`. They compile Forth kernels locally, rent a GPU on Vast.ai, and verify output. +End-to-end GPU execution tests live in `gpu_test/`. They compile Forth kernels locally, rent a GPU on Vast.ai, run them through a device worker started over SSH, and verify output. `gpu_test/test_worker.py` covers the worker protocol locally with a fake driver. ```bash # Run GPU tests VASTAI_API_KEY=xxx uv run pytest -v -m gpu # Lint and format Python code -uv run ruff check gpu_test/ -uv run ruff format gpu_test/ +uv run ruff check gpu_test/ warpforth_gpu/ +uv run ruff format gpu_test/ warpforth_gpu/ ``` ## Architecture Notes diff --git a/README.md b/README.md index 846b739..061b954 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ See the [documentation](https://tetsuo-cpp.github.io/warpforth/). - LLVM/MLIR - CMake - C++17 compiler -- CUDA toolkit (for GPU execution) +- An NVIDIA driver and Python 3 on the GPU host (for GPU execution) - [uv](https://github.com/astral-sh/uv) (for Python test tooling) ## Building @@ -56,17 +56,30 @@ Compile to PTX: ./build/bin/warpforthc matmul.forth -o matmul.ptx --arch sm_80 ``` -Test on a GPU (A is 2x4 row-major, B is 4x3 row-major, C is 2x3 output): +Run it on a machine with an NVIDIA GPU, with the repository root on the Python +path (for example `PYTHONPATH=. uv run python run_matmul.py`). A is 2x4 +row-major, B is 4x3 row-major, C is 2x3 output: -```bash -./build/bin/warpforth-runner matmul.ptx \ - --param 'i64[]:1,2,3,4,5,6,7,8' \ - --param 'i64[]:1,2,3,4,5,6,7,8,9,10,11,12' \ - --param 'i64[]:0,0,0,0,0,0' \ - --grid 6,1,1 --block 1,1,1 \ - --output-param 2 --output-count 6 +```python +from pathlib import Path + +import numpy as np + +from warpforth_gpu.client import Device, worker_command + +with Device.spawn(worker_command()) as device: + kernel = device.load(Path("matmul.ptx").read_text()).kernel("main") + a = device.upload(np.arange(1, 9, dtype=np.int64)) + b = device.upload(np.arange(1, 13, dtype=np.int64)) + c = device.alloc(6 * 8) + kernel.launch(grid=(6, 1, 1), block=(1, 1, 1), args=[a, b, c]) + print(c.read(np.int64)) # [ 70 80 90 158 184 210] ``` +`worker_command()` runs the device worker from its inlined source, so the +same call works over SSH on a remote GPU host: +`Device.spawn(["ssh", host, shlex.join(worker_command())])`. + ## Toolchain | Tool | Description | @@ -74,7 +87,7 @@ Test on a GPU (A is 2x4 row-major, B is 4x3 row-major, C is 2x3 output): | `warpforthc` | Compiles Forth source to PTX | | `warpforth-translate` | Translates from Forth source to MLIR and MLIR to PTX assembly | | `warpforth-opt` | Runs individual MLIR passes or entire pipeline | -| `warpforth-runner` | Executes PTX kernels on a GPU for testing | +| `warpforth_gpu` | Python package that loads PTX and runs kernels on an NVIDIA GPU, locally or over SSH | These tools can be composed for debugging or inspecting intermediate stages: diff --git a/gpu_test/conftest.py b/gpu_test/conftest.py index 70f843b..37c4ff2 100644 --- a/gpu_test/conftest.py +++ b/gpu_test/conftest.py @@ -1,41 +1,29 @@ -"""GPU test infrastructure: Compiler, VastSession, KernelRunner, and fixtures.""" +"""GPU test infrastructure: Compiler, KernelRunner, and fixtures.""" from __future__ import annotations -import atexit -import logging import os +import shlex import subprocess import tempfile -import time from dataclasses import dataclass -from datetime import UTC, datetime from pathlib import Path from typing import TYPE_CHECKING -from uuid import uuid4 +import numpy as np import pytest -from requests import RequestException -from vastai import VastAI + +from warpforth_gpu.client import Device, worker_command +from warpforth_gpu.vast import DEFAULT_ARCHITECTURE, VastSession if TYPE_CHECKING: - from collections.abc import Generator - from typing import Self + from collections.abc import Callable, Generator -logger = logging.getLogger(__name__) + from warpforth_gpu.client import Buffer PROJECT_ROOT = Path(__file__).resolve().parent.parent WARPFORTHC = PROJECT_ROOT / "build" / "bin" / "warpforthc" -RUNNER_SRC = PROJECT_ROOT / "tools" / "warpforth-runner" / "warpforth-runner.cpp" - -MAX_COST_PER_HOUR = 0.50 -CUDA_VERSION = "12.4" -CUDA_IMAGE = f"nvidia/cuda:{CUDA_VERSION}.0-devel-ubuntu22.04" -POLL_INTERVAL_S = 10 -POLL_TIMEOUT_S = 300 -INSTANCE_LABEL_PREFIX = "warpforth-test-" -REMOTE_TMP = "/tmp" # noqa: S108 -DEFAULT_ARCHITECTURE = "sm_70" + SUPPORTED_ARCHITECTURES = ("sm_70", "sm_75", "sm_80", "sm_86", "sm_89", "sm_90") @@ -48,6 +36,10 @@ class ParamDecl: size: int # 0 for scalars base_type: str = "i64" # "i64" or "f64" + @property + def dtype(self) -> type[np.int64 | np.float64]: + return np.float64 if self.base_type == "f64" else np.int64 + class CompileError(Exception): """Raised when warpforthc fails to compile Forth source.""" @@ -85,447 +77,6 @@ def compile_source(self, forth_src: str) -> str: src_path.unlink() -class VastSession: - """Context manager that rents a GPU instance on Vast.ai. - - Handles searching for offers, launching an instance, waiting for SSH - readiness, and unconditional cleanup on exit. - """ - - def __init__(self, api_key: str, arch: str = DEFAULT_ARCHITECTURE) -> None: - self.sdk = VastAI(api_key, retry=1) - self.arch = arch - timestamp = datetime.now(tz=UTC).strftime("%Y%m%dT%H%M%SZ") - self.instance_label = f"{INSTANCE_LABEL_PREFIX}{timestamp}-{uuid4().hex}" - self.instance_id: int | None = None - self.ssh_host: str | None = None - self.ssh_port: int | None = None - self._atexit_registered = False - self._creation_ambiguous = False - self._ssh_key_dir: tempfile.TemporaryDirectory[str] | None = None - self.ssh_private_key: Path | None = None - - def __enter__(self) -> Self: - self._log_existing_instances() - logger.info("This GPU test session owns label=%s", self.instance_label) - atexit.register(self._atexit_cleanup) - self._atexit_registered = True - try: - self._prepare_ssh_key() - self._launch() - self._wait_for_ssh() - except BaseException: - try: - if self._destroy(): - self._unregister_atexit() - finally: - self._cleanup_ssh_key() - raise - else: - return self - - def __exit__(self, *_: object) -> None: - try: - if self._destroy(): - self._unregister_atexit() - finally: - self._cleanup_ssh_key() - - def _prepare_ssh_key(self) -> None: - """Generate an ephemeral SSH identity for this test session.""" - self._ssh_key_dir = tempfile.TemporaryDirectory(prefix="warpforth-ssh-") - self.ssh_private_key = Path(self._ssh_key_dir.name) / "id_ed25519" - subprocess.run( - [ - "ssh-keygen", - "-q", - "-t", - "ed25519", - "-N", - "", - "-f", - str(self.ssh_private_key), - ], - capture_output=True, - text=True, - check=True, - ) - - def _attach_ssh_key(self) -> None: - """Attach this session's ephemeral public key to its Vast instance.""" - if self.instance_id is None or self.ssh_private_key is None: - msg = "Cannot attach SSH key before instance and identity creation" - raise RuntimeError(msg) - - public_key = self.ssh_private_key.with_suffix(".pub").read_text().strip() - result = self.sdk.attach_ssh(instance_id=self.instance_id, ssh_key=public_key) - if not isinstance(result, dict) or result.get("success") is False: - msg = f"Failed to attach ephemeral SSH key to instance {self.instance_id}" - raise RuntimeError(msg) - logger.info("Attached ephemeral SSH key to instance %s", self.instance_id) - - def _cleanup_ssh_key(self) -> None: - """Remove this session's ephemeral local SSH key material.""" - if self._ssh_key_dir is not None: - self._ssh_key_dir.cleanup() - self._ssh_key_dir = None - self.ssh_private_key = None - - def _show_instances(self) -> list[dict] | None: - """Return the current account instances, or None if listing fails.""" - try: - instances = self.sdk.show_instances() - except Exception: - logger.exception("Failed to list Vast.ai instances") - return None - - if not isinstance(instances, list): - logger.warning("Vast.ai returned an invalid instance listing") - return None - return [inst for inst in instances if isinstance(inst, dict)] - - def _log_existing_instances(self) -> None: - """Log existing WarpForth instances without modifying them.""" - instances = self._show_instances() - if instances is None: - logger.warning("Could not check for existing WarpForth Vast.ai instances") - return - - matching = [ - inst - for inst in instances - if str(inst.get("label", "")).startswith(INSTANCE_LABEL_PREFIX) - ] - if not matching: - logger.info("No existing WarpForth Vast.ai instances found") - return - - total_cost = 0.0 - for inst in matching: - cost = float(inst.get("dph_total") or 0.0) - total_cost += cost - start_date = inst.get("start_date") - age = "unknown" - if isinstance(start_date, int | float): - age = f"{max(0.0, time.time() - start_date) / 60:.0f}m" - logger.warning( - "Existing WarpForth Vast.ai instance: id=%s label=%s status=%s " - "gpu=%s dph=$%.3f age=%s ownership=active run or orphan candidate", - inst.get("id"), - inst.get("label"), - inst.get("actual_status", inst.get("status", "unknown")), - inst.get("gpu_name", "unknown"), - cost, - age, - ) - logger.warning( - "Found %d existing WarpForth Vast.ai instance(s), total known cost=$%.3f/hr; " - "they may belong to other active test runs and will not be destroyed", - len(matching), - total_cost, - ) - - def _launch(self) -> None: - """Find the cheapest suitable offer and launch an instance.""" - minimum_compute_capability = int(self.arch.removeprefix("sm_")) * 10 - query = ( - f"num_gpus=1 rentable=True rented=False compute_cap>={minimum_compute_capability}" - f" reliability>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}" - f" cuda_max_good>={CUDA_VERSION}" - ) - offers = self.sdk.search_offers(query=query, order="dph", limit=5) - - if not offers: - msg = "No suitable GPU offers found on Vast.ai" - raise RuntimeError(msg) - - offer = offers[0] if isinstance(offers, list) else offers - details = offer if isinstance(offer, dict) else {"id": offer} - offer_id = int(details["id"]) - logger.info( - "Launching instance from offer %s: gpu=%s driver=%s cuda_max_good=%s", - offer_id, - details.get("gpu_name"), - details.get("driver_version"), - details.get("cuda_max_good"), - ) - self._creation_ambiguous = True - try: - result = self.sdk.create_instance( - id=offer_id, - image=CUDA_IMAGE, - disk=10.0, - runtype="ssh_direc ssh_proxy", - label=self.instance_label, - ) - except RequestException: - logger.warning( - "Instance creation result is ambiguous; reconciling label=%s", - self.instance_label, - ) - self._reconcile_instance() - return - - if ( - isinstance(result, dict) - and result.get("success") is True - and result.get("new_contract") - ): - self.instance_id = int(result["new_contract"]) - self._creation_ambiguous = False - logger.info( - "This GPU test session owns instance id=%s label=%s", - self.instance_id, - self.instance_label, - ) - return - - logger.warning( - "Instance creation response did not contain an ID; reconciling label=%s", - self.instance_label, - ) - self._reconcile_instance() - - def _reconcile_instance(self) -> None: - """Find this run's instance by its exact unique label.""" - deadline = time.monotonic() + POLL_TIMEOUT_S - while time.monotonic() < deadline: - instances = self._show_instances() - if instances is not None: - matches = [inst for inst in instances if inst.get("label") == self.instance_label] - if len(matches) > 1: - ids = [inst.get("id") for inst in matches] - msg = f"Multiple instances found for unique label {self.instance_label}: {ids}" - raise RuntimeError(msg) - if matches: - self.instance_id = int(matches[0]["id"]) - self._creation_ambiguous = False - logger.info( - "This GPU test session owns instance id=%s label=%s", - self.instance_id, - self.instance_label, - ) - return - time.sleep(POLL_INTERVAL_S) - - msg = f"Could not reconcile instance with label {self.instance_label} after creation" - raise TimeoutError(msg) - - def _wait_for_ssh(self) -> None: - """Poll until the instance is running and SSH is available.""" - deadline = time.monotonic() + POLL_TIMEOUT_S - - while time.monotonic() < deadline: - instances = self.sdk.show_instances() - if not isinstance(instances, list): - time.sleep(POLL_INTERVAL_S) - continue - - for inst in instances: - if not isinstance(inst, dict): - continue - if int(inst.get("id", -1)) != self.instance_id: - continue - - status = inst.get("actual_status", inst.get("status", "")) - if status == "running" and inst.get("ssh_host") and inst.get("ssh_port"): - self.ssh_host = inst["ssh_host"] - self.ssh_port = int(inst["ssh_port"]) - logger.info( - "Instance %s ready: %s:%s", - self.instance_id, - self.ssh_host, - self.ssh_port, - ) - self._attach_ssh_key() - self._wait_for_sshd() - self._compile_runner() - return - - time.sleep(POLL_INTERVAL_S) - - msg = f"Instance {self.instance_id} did not become SSH-ready within {POLL_TIMEOUT_S}s" - raise TimeoutError(msg) - - def _wait_for_sshd(self) -> None: - """Poll until sshd is actually accepting connections.""" - deadline = time.monotonic() + POLL_TIMEOUT_S - while time.monotonic() < deadline: - result = subprocess.run( - [*self._ssh_cmd(), "true"], - capture_output=True, - timeout=15, - check=False, - ) - if result.returncode == 0: - logger.info("sshd ready on %s:%s", self.ssh_host, self.ssh_port) - return - time.sleep(POLL_INTERVAL_S) - msg = ( - f"sshd on {self.ssh_host}:{self.ssh_port} did not become ready within {POLL_TIMEOUT_S}s" - ) - raise TimeoutError(msg) - - def _compile_runner(self) -> None: - """Upload warpforth-runner.cpp and compile it on the remote host.""" - self.scp_upload(RUNNER_SRC, f"{REMOTE_TMP}/warpforth-runner.cpp") - nvcc_cmd = ( - f"nvcc -o {REMOTE_TMP}/warpforth-runner" - f" {REMOTE_TMP}/warpforth-runner.cpp -lcuda -std=c++17" - ) - self.ssh_run(nvcc_cmd, timeout=60) - - def _record_labeled_instance_ids( - self, - instances: list[dict], - instance_ids: set[int], - ) -> None: - """Record instances owned by this session's exact unique label.""" - labeled_ids = { - int(inst["id"]) - for inst in instances - if inst.get("label") == self.instance_label and inst.get("id") is not None - } - instance_ids.update(labeled_ids) - if labeled_ids: - self._creation_ambiguous = False - - def _destroy(self) -> bool: - """Destroy only this session's instances and verify their absence.""" - instance_ids = {self.instance_id} if self.instance_id is not None else set() - retries = 3 - for attempt in range(retries): - instances = self._show_instances() - if instances is not None: - self._record_labeled_instance_ids(instances, instance_ids) - if not self._creation_ambiguous and not instance_ids: - logger.info( - "Verified cleanup for Vast.ai label=%s instance_ids=%s", - self.instance_label, - sorted(instance_ids), - ) - self.instance_id = None - return True - - for instance_id in instance_ids: - try: - self.sdk.destroy_instance(id=instance_id) - except Exception: - logger.exception( - "Destroy attempt %d/%d failed for instance %s", - attempt + 1, - retries, - instance_id, - ) - - instances = self._show_instances() - if instances is not None: - self._record_labeled_instance_ids(instances, instance_ids) - listed_ids = {int(inst["id"]) for inst in instances if inst.get("id") is not None} - if not self._creation_ambiguous and not instance_ids & listed_ids: - logger.info( - "Verified cleanup for Vast.ai label=%s instance_ids=%s", - self.instance_label, - sorted(instance_ids), - ) - self.instance_id = None - return True - - if attempt < retries - 1: - time.sleep(5) - - logger.error( - "Could not verify cleanup for Vast.ai label=%s instance_ids=%s; " - "check the Vast.ai account and destroy these instances manually", - self.instance_label, - sorted(instance_ids), - ) - return False - - def _unregister_atexit(self) -> None: - if self._atexit_registered: - atexit.unregister(self._atexit_cleanup) - self._atexit_registered = False - - def _atexit_cleanup(self) -> None: - """Last-resort cleanup registered before instance creation.""" - logger.warning( - "atexit: checking for Vast.ai instances with label=%s instance_id=%s", - self.instance_label, - self.instance_id, - ) - try: - self._destroy() - finally: - self._cleanup_ssh_key() - - def _ssh_options(self) -> list[str]: - if self.ssh_private_key is None: - msg = "Ephemeral SSH identity is not available" - raise RuntimeError(msg) - return [ - "-i", - str(self.ssh_private_key), - "-o", - "IdentitiesOnly=yes", - "-o", - "BatchMode=yes", - "-o", - "PreferredAuthentications=publickey", - "-o", - "StrictHostKeyChecking=no", - "-o", - "UserKnownHostsFile=/dev/null", - "-o", - "ConnectTimeout=10", - "-o", - "ConnectionAttempts=3", - "-o", - "LogLevel=ERROR", - ] - - def _ssh_cmd(self) -> list[str]: - """Build base SSH command with connection options.""" - return [ - "ssh", - *self._ssh_options(), - "-p", - str(self.ssh_port), - f"root@{self.ssh_host}", - ] - - def ssh_run(self, cmd: str, *, timeout: int = 120) -> str: - """Execute a command on the remote instance via SSH.""" - result = subprocess.run( - [*self._ssh_cmd(), cmd], - capture_output=True, - text=True, - timeout=timeout, - check=False, - ) - if result.returncode != 0: - msg = f"SSH command failed (rc={result.returncode}):\n{result.stderr}" - raise RuntimeError(msg) - return result.stdout - - def scp_upload(self, local_path: str | Path, remote_path: str) -> None: - """Upload a file to the remote instance via SCP.""" - subprocess.run( - [ - "scp", - *self._ssh_options(), - "-P", - str(self.ssh_port), - str(local_path), - f"root@{self.ssh_host}:{remote_path}", - ], - capture_output=True, - text=True, - timeout=60, - check=True, - ) - - def _parse_array_type(type_spec: str) -> tuple[str, int]: """Parse 'i64[256]' or 'f64[256]' into (base_type, size).""" if not type_spec.endswith("]"): @@ -598,13 +149,55 @@ def _parse_param_declarations(forth_source: str) -> list[ParamDecl]: return decls +def _initial_values( + decls: list[ParamDecl], + params: dict[str, list[int] | list[float] | int | float], +) -> list[np.ndarray | float]: + """Return each param's starting value: a zero-padded array, or a scalar.""" + values: list[np.ndarray | float] = [] + for decl in decls: + given = params.get(decl.name) + if decl.is_array: + if given is None: + given = [] + if not isinstance(given, list): + msg = f"Array param '{decl.name}' expects a list, got {type(given).__name__}" + raise TypeError(msg) + array = np.zeros(decl.size, dtype=decl.dtype) + array[: len(given)] = given + values.append(array) + else: + if isinstance(given, list): + msg = f"Scalar param '{decl.name}' expects a scalar, got list" + raise TypeError(msg) + given = given or 0 + values.append(float(given) if decl.base_type == "f64" else given) + return values + + +def open_vast_device(session: VastSession) -> Device: + """Start a device worker on the rented instance, streaming its source over SSH.""" + return Device.spawn([*session.ssh_argv(), shlex.join(worker_command())]) + + class KernelRunner: - """Orchestrates local compilation and remote GPU execution.""" + """Compiles Forth locally and runs it on one device worker, reopened after a fatal error.""" - def __init__(self, session: VastSession, compiler: Compiler) -> None: - self.session = session + def __init__(self, open_device: Callable[[], Device], compiler: Compiler) -> None: + self._open_device = open_device + self._device: Device | None = None self.compiler = compiler + def close(self) -> None: + if self._device is not None: + self._device.close() + self._device = None + + def _ready_device(self) -> Device: + if self._device is None or self._device.closed: + self._device = self._open_device() + return self._device + def run( self, forth_source: str, @@ -614,7 +207,7 @@ def run( output_param: int = 0, output_count: int | None = None, ) -> list[int] | list[float]: - """Compile Forth source locally, execute on remote GPU, return output values. + """Compile Forth source locally, execute on the GPU, return output values. Param buffer sizes are derived from the Forth source's 'param' declarations. The params dict maps param names to initial values: @@ -622,7 +215,6 @@ def run( - Scalar params: int or float Params not in the dict are zero-initialized. """ - # Parse kernel name and param declarations kernel_name = _parse_kernel_name(forth_source) decls = _parse_param_declarations(forth_source) if not decls: @@ -631,7 +223,6 @@ def run( params = params or {} - # Validate output_param if output_param < 0 or output_param >= len(decls): msg = f"output_param {output_param} out of range (have {len(decls)} params)" raise ValueError(msg) @@ -640,66 +231,24 @@ def run( msg = f"output_param {output_param} ('{name}') is a scalar and cannot be read back" raise ValueError(msg) - # Compile locally + initial = _initial_values(decls, params) ptx = self.compiler.compile_source(forth_source) + device = self._ready_device() + kernel = device.load(ptx).kernel(kernel_name) - # Write PTX to temp file and upload - with tempfile.NamedTemporaryFile(mode="w", suffix=".ptx", delete=False) as f: - f.write(ptx) - ptx_path = Path(f.name) - + buffers: dict[int, Buffer] = {} try: - self.session.scp_upload(ptx_path, f"{REMOTE_TMP}/kernel.ptx") + for index, value in enumerate(initial): + if isinstance(value, np.ndarray): + buffers[index] = device.upload(value) + args = [buffers.get(index, value) for index, value in enumerate(initial)] + kernel.launch(grid=grid, block=block, args=args) + output = buffers[output_param].read(decls[output_param].dtype) + return output[:output_count].tolist() finally: - ptx_path.unlink() - - # Build remote command - cmd_parts = [ - f"{REMOTE_TMP}/warpforth-runner", - f"{REMOTE_TMP}/kernel.ptx", - "--kernel", - kernel_name, - ] - - for decl in decls: - if decl.is_array: - values = params.get(decl.name, []) - if not isinstance(values, list): - msg = f"Array param '{decl.name}' expects a list, got {type(values).__name__}" - raise TypeError(msg) - zero = 0.0 if decl.base_type == "f64" else 0 - buf = [zero] * decl.size - for i, v in enumerate(values): - buf[i] = v - cmd_parts.extend(["--param", f"{decl.base_type}[]:{','.join(str(v) for v in buf)}"]) - else: - value = params.get(decl.name, 0.0 if decl.base_type == "f64" else 0) - if isinstance(value, list): - msg = f"Scalar param '{decl.name}' expects a scalar, got list" - raise TypeError(msg) - cmd_parts.extend(["--param", f"{decl.base_type}:{value}"]) - - cmd_parts.extend( - [ - "--grid", - f"{grid[0]},{grid[1]},{grid[2]}", - "--block", - f"{block[0]},{block[1]},{block[2]}", - "--output-param", - str(output_param), - ] - ) - - if output_count is not None: - cmd_parts.extend(["--output-count", str(output_count)]) - - cmd = " ".join(cmd_parts) - stdout = self.session.ssh_run(cmd, timeout=120) - - # Parse CSV output — type depends on the output param - out_type = decls[output_param].base_type - parse = float if out_type == "f64" else int - return [parse(v) for v in stdout.strip().split(",")] + if not device.closed: + for buffer in buffers.values(): + buffer.free() # --- Fixtures --- @@ -730,8 +279,10 @@ def gpu_session(pytestconfig: pytest.Config) -> Generator[VastSession]: @pytest.fixture(scope="session") -def kernel_runner(gpu_session: VastSession, compiler: Compiler) -> KernelRunner: - return KernelRunner(gpu_session, compiler) +def kernel_runner(gpu_session: VastSession, compiler: Compiler) -> Generator[KernelRunner]: + runner = KernelRunner(lambda: open_vast_device(gpu_session), compiler) + yield runner + runner.close() def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: diff --git a/gpu_test/fake_driver.py b/gpu_test/fake_driver.py new file mode 100644 index 0000000..5e72457 --- /dev/null +++ b/gpu_test/fake_driver.py @@ -0,0 +1,70 @@ +"""In-memory stand-in for the CUDA driver behind the device worker.""" + +from __future__ import annotations + +import itertools +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from warpforth_gpu.worker import CudaError + +if TYPE_CHECKING: + from warpforth_gpu.worker import Holder + + +@dataclass +class Launch: + kernel: str + grid: tuple[int, int, int] + block: tuple[int, int, int] + args: list[tuple[type, int | float]] + + +class FakeDriver: + """Allocations are zeroed bytearrays; launches are recorded instead of run.""" + + def __init__(self, *, load_error: str | None = None, launch_error: str | None = None) -> None: + self.load_error = load_error + self.launch_error = launch_error + self.memory: dict[int, bytearray] = {} + self.pointers: list[int] = [] + self.launches: list[Launch] = [] + self._handles = itertools.count(0x1000, 0x100) + self._functions: dict[int, str] = {} + + def load(self, _ptx: str) -> int: + if self.load_error is not None: + raise CudaError(self.load_error, "ptxas fatal: parse error") + return next(self._handles) + + def function(self, _module: int, name: str) -> int: + handle = next(self._handles) + self._functions[handle] = name + return handle + + def alloc(self, size: int) -> int: + pointer = next(self._handles) + self.memory[pointer] = bytearray(size) + self.pointers.append(pointer) + return pointer + + def write(self, pointer: int, data: bytes) -> None: + self.memory[pointer][: len(data)] = data + + def read(self, pointer: int, size: int) -> bytes: + return bytes(self.memory[pointer][:size]) + + def free(self, pointer: int) -> None: + del self.memory[pointer] + + def launch( + self, + function: int, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: list[Holder], + ) -> None: + if self.launch_error is not None: + raise CudaError(self.launch_error, "an illegal memory access was encountered") + packed = [(type(arg), arg.value) for arg in args] + self.launches.append(Launch(self._functions[function], grid, block, packed)) diff --git a/gpu_test/test_vast_session.py b/gpu_test/test_vast_session.py index 9db0a6d..85c8ad7 100644 --- a/gpu_test/test_vast_session.py +++ b/gpu_test/test_vast_session.py @@ -8,8 +8,8 @@ import pytest import requests -from gpu_test import conftest as gpu_conftest -from gpu_test.conftest import VastSession +from warpforth_gpu import vast +from warpforth_gpu.vast import VastSession if TYPE_CHECKING: from collections.abc import Callable @@ -96,8 +96,8 @@ def create_sdk(_api_key: str, **kwargs: object) -> FakeVastAI: sdk.constructor_kwargs = kwargs return sdk - monkeypatch.setattr(gpu_conftest, "VastAI", create_sdk) - monkeypatch.setattr(gpu_conftest.time, "sleep", lambda _seconds: None) + monkeypatch.setattr(vast, "VastAI", create_sdk) + monkeypatch.setattr(vast.time, "sleep", lambda _seconds: None) monkeypatch.setattr(VastSession, "_wait_for_ssh", wait_for_ssh or no_wait) @@ -115,14 +115,14 @@ def test_concurrent_sessions_destroy_only_their_own_instance( assert {instance["id"] for instance in sdk.instances} == {100, 101} assert [instance_id for instance_id, _key in sdk.attached_keys] == [100, 101] assert sdk.attached_keys[0][1] != sdk.attached_keys[1][1] - assert first._ssh_cmd()[1:5] == [ + assert first.ssh_argv()[1:5] == [ "-i", str(first_key_path), "-o", "IdentitiesOnly=yes", ] - assert "BatchMode=yes" in first._ssh_cmd() - assert "PreferredAuthentications=publickey" in first._ssh_cmd() + assert "BatchMode=yes" in first.ssh_argv() + assert "PreferredAuthentications=publickey" in first.ssh_argv() assert {instance["id"] for instance in sdk.instances} == {100} assert second_key_path is not None assert not second_key_path.exists() diff --git a/gpu_test/test_worker.py b/gpu_test/test_worker.py new file mode 100644 index 0000000..7fff022 --- /dev/null +++ b/gpu_test/test_worker.py @@ -0,0 +1,382 @@ +"""Local tests for the device worker protocol and its client, using a fake CUDA driver.""" + +from __future__ import annotations + +import base64 +import ctypes +import ctypes.util +import json +import queue +import sys +import threading +from pathlib import Path + +import numpy as np +import pytest + +from gpu_test.conftest import KernelRunner +from gpu_test.fake_driver import FakeDriver +from warpforth_gpu.client import Device, DeviceError, worker_command +from warpforth_gpu.worker import parse_entries, serve + +PROJECT_ROOT = Path(__file__).resolve().parent.parent + +# warpforthc --arch sm_70 output for a kernel with params A i64[4], B f64[4], N i64, X f64 +# and shared S i64[8]. +MIXED_PTX = """\ +// +// Generated by LLVM NVPTX Back-End +// + +.version 6.0 +.target sm_70 +.address_size 64 + + // .globl mixed +// __wg_mixed_0 has been demoted + +.visible .entry mixed( + .param .u64 .ptr .align 1 mixed_param_0, + .param .u64 .ptr .align 1 mixed_param_1, + .param .u64 mixed_param_2, + .param .f64 mixed_param_3 +) +{ + .reg .b32 %r<3>; + .reg .b64 %rd<9>; + // demoted variable + .shared .align 8 .b8 __wg_mixed_0[64]; + ld.param.b64 %rd1, [mixed_param_0]; + ld.param.b64 %rd2, [mixed_param_1]; + mov.u32 %r1, %tid.x; + shl.b32 %r2, %r1, 3; + ld.param.b64 %rd3, [mixed_param_2]; + cvt.u64.u32 %rd4, %r2; + ld.param.b64 %rd5, [mixed_param_3]; + mov.b64 %rd6, __wg_mixed_0; + add.s64 %rd7, %rd4, %rd6; + st.shared.b64 [%rd7], %rd3; + bar.sync 0; + ld.shared.b64 %rd8, [__wg_mixed_0]; + st.global.b64 [%rd1], %rd8; + st.global.b64 [%rd2], %rd5; + ret; + +} +""" + + +class Conversation: + """Drives serve() in a thread, one request and response at a time.""" + + def __init__(self, driver: FakeDriver) -> None: + self._requests: queue.Queue[str | None] = queue.Queue() + self._responses: queue.Queue[dict] = queue.Queue() + self._next_id = 0 + self.status: int | None = None + self._thread = threading.Thread(target=self._serve, args=(driver,), daemon=True) + self._thread.start() + + def _serve(self, driver: FakeDriver) -> None: + self.status = serve(driver, iter(self._requests.get, None), self) + + def write(self, text: str) -> None: + self._responses.put(json.loads(text)) + + def flush(self) -> None: + pass + + def send_line(self, line: str) -> dict: + self._requests.put(line + "\n") + return self._responses.get(timeout=5) + + def send(self, op: str, **fields: object) -> dict: + self._next_id += 1 + response = self.send_line(json.dumps({"id": self._next_id, "op": op, **fields})) + assert response["id"] == self._next_id + return response + + def ok(self, op: str, **fields: object) -> dict: + response = self.send(op, **fields) + assert response["ok"], response + return response["result"] + + def stopped(self) -> bool: + self._thread.join(timeout=5) + return not self._thread.is_alive() + + def finish(self) -> int | None: + self._requests.put(None) + self._thread.join(timeout=5) + return self.status + + +def b64(data: bytes) -> str: + return base64.b64encode(data).decode() + + +def test_parse_entries_reads_parameter_kinds_from_warpforthc_ptx() -> None: + assert parse_entries(MIXED_PTX) == {"mixed": ["ptr", "ptr", "i64", "f64"]} + + +def test_parse_entries_handles_several_entries_and_empty_parameter_lists() -> None: + ptx = ( + ".visible .entry first(\n\t.param .s64 first_param_0\n)\n{\n\tret;\n}\n" + ".visible .entry second()\n{\n\tret;\n}\n" + ) + assert parse_entries(ptx) == {"first": ["i64"], "second": []} + + +@pytest.mark.parametrize( + "param", + [".param .u32 k_param_0", ".param .align 8 .b8 k_param_0[16]", ".param .f32 k_param_0"], +) +def test_parse_entries_rejects_unsupported_parameters(param: str) -> None: + with pytest.raises(ValueError, match="unsupported kernel parameter"): + parse_entries(f".visible .entry k(\n\t{param}\n)\n{{\n}}\n") + + +def test_load_reports_kernels_and_their_parameter_kinds() -> None: + chat = Conversation(FakeDriver()) + result = chat.ok("load", ptx=MIXED_PTX) + assert result["kernels"] == {"mixed": ["ptr", "ptr", "i64", "f64"]} + assert isinstance(result["module"], int) + assert chat.finish() == 0 + + +def test_partial_write_then_read_returns_the_whole_zeroed_buffer() -> None: + chat = Conversation(FakeDriver()) + buffer = chat.ok("alloc", size=8)["buffer"] + assert base64.b64decode(chat.ok("read", buffer=buffer)["data"]) == bytes(8) + chat.ok("write", buffer=buffer, data=b64(b"\x01\x02\x03")) + assert base64.b64decode(chat.ok("read", buffer=buffer)["data"]) == b"\x01\x02\x03" + bytes(5) + assert chat.finish() == 0 + + +def test_launch_packs_arguments_in_ptx_order() -> None: + driver = FakeDriver() + chat = Conversation(driver) + module = chat.ok("load", ptx=MIXED_PTX)["module"] + a = chat.ok("alloc", size=32)["buffer"] + b = chat.ok("alloc", size=32)["buffer"] + chat.ok( + "launch", + module=module, + kernel="mixed", + grid=[2, 1, 1], + block=[64, 1, 1], + args=[{"buffer": a}, {"buffer": b}, {"value": -7}, {"value": 2}], + ) + pointer_a, pointer_b = driver.pointers + assert len(driver.launches) == 1 + launch = driver.launches[0] + assert (launch.kernel, launch.grid, launch.block) == ("mixed", (2, 1, 1), (64, 1, 1)) + assert launch.args == [ + (ctypes.c_uint64, pointer_a), + (ctypes.c_uint64, pointer_b), + (ctypes.c_int64, -7), + (ctypes.c_double, 2.0), + ] + assert chat.finish() == 0 + + +@pytest.mark.parametrize( + "args", + [ + [{"value": 1}, {"value": 2}, {"value": 3}], + [{"buffer": 999}, {"buffer": 999}, {"value": 1}, {"value": 1.0}], + ["not-a-buffer", "B", {"value": 1}, {"value": 1.0}], + ["A", "B", {"value": 1.5}, {"value": 1.0}], + ["A", "B", {"value": True}, {"value": 1.0}], + ["A", "B", {"value": 2**63}, {"value": 1.0}], + ["A", "B", {"value": 1}, {"value": "1.0"}], + ["A", {"value": 0}, {"value": 1}, {"value": 1.0}], + ], + ids=[ + "wrong-count", + "unknown-buffer", + "ptr-as-string", + "float-for-i64", + "bool-for-i64", + "i64-overflow", + "string-for-f64", + "value-for-ptr", + ], +) +def test_bad_launch_arguments_are_rejected_without_launching(args: list[object]) -> None: + driver = FakeDriver() + chat = Conversation(driver) + module = chat.ok("load", ptx=MIXED_PTX)["module"] + buffers = {"A": chat.ok("alloc", size=32)["buffer"], "B": chat.ok("alloc", size=32)["buffer"]} + resolved = [{"buffer": buffers[a]} if a in ("A", "B") else a for a in args] + response = chat.send( + "launch", module=module, kernel="mixed", grid=[1, 1, 1], block=[1, 1, 1], args=resolved + ) + assert response["ok"] is False + assert response["error"]["code"] == "BAD_REQUEST" + assert response["fatal"] is False + assert driver.launches == [] + chat.ok("read", buffer=buffers["A"]) + assert chat.finish() == 0 + + +@pytest.mark.parametrize( + ("line", "expected_id"), + [ + ("{not json", None), + ('{"op": "alloc", "size": 8}', None), + ('{"id": 4, "op": "explode"}', 4), + ('{"id": 5, "op": "read", "buffer": 12345}', 5), + ('{"id": 6, "op": "alloc", "size": 0}', 6), + ('{"id": 7, "op": "write", "buffer": 1, "data": "AAAAAAAAAAAA"}', 7), + ('{"id": 8, "op": "launch", "module": 77, "kernel": "mixed"}', 8), + ], + ids=[ + "malformed-json", + "missing-id", + "unknown-op", + "unknown-buffer", + "zero-size", + "write-past-end", + "unknown-module", + ], +) +def test_bad_requests_are_answered_and_the_worker_keeps_serving( + line: str, expected_id: int | None +) -> None: + chat = Conversation(FakeDriver()) + chat.ok("alloc", size=8) + response = chat.send_line(line) + assert response == { + "id": expected_id, + "ok": False, + "error": {"code": "BAD_REQUEST", "message": response["error"]["message"]}, + "fatal": False, + } + chat.ok("alloc", size=8) + assert chat.finish() == 0 + + +def test_load_driver_error_is_not_fatal() -> None: + chat = Conversation(FakeDriver(load_error="CUDA_ERROR_INVALID_PTX")) + response = chat.send("load", ptx=MIXED_PTX) + assert response["error"]["code"] == "CUDA_ERROR_INVALID_PTX" + assert "parse error" in response["error"]["message"] + assert response["fatal"] is False + chat.ok("alloc", size=8) + assert chat.finish() == 0 + + +def test_launch_driver_error_is_fatal_and_stops_the_worker() -> None: + chat = Conversation(FakeDriver(launch_error="CUDA_ERROR_ILLEGAL_ADDRESS")) + module = chat.ok("load", ptx=MIXED_PTX)["module"] + a = chat.ok("alloc", size=32)["buffer"] + response = chat.send( + "launch", + module=module, + kernel="mixed", + grid=[1, 1, 1], + block=[1, 1, 1], + args=[{"buffer": a}, {"buffer": a}, {"value": 1}, {"value": 1.0}], + ) + assert response["ok"] is False + assert response["error"]["code"] == "CUDA_ERROR_ILLEGAL_ADDRESS" + assert response["fatal"] is True + assert chat.stopped() + assert chat.status == 1 + + +def fake_worker_argv(**driver_options: str) -> list[str]: + code = ( + "import sys\n" + f"sys.path.insert(0, {str(PROJECT_ROOT)!r})\n" + "from gpu_test.fake_driver import FakeDriver\n" + "from warpforth_gpu.worker import serve\n" + f"sys.exit(serve(FakeDriver(**{driver_options!r}), sys.stdin, sys.stdout))\n" + ) + return [sys.executable, "-u", "-c", code] + + +def test_client_round_trips_buffers_through_a_worker_process() -> None: + with Device.spawn(fake_worker_argv()) as device: + module = device.load(MIXED_PTX) + assert module.kernel("mixed").kinds == ("ptr", "ptr", "i64", "f64") + values = np.array([1.5, -2.25, 3.0], dtype=np.float64) + buffer = device.upload(values) + assert buffer.size == 24 + np.testing.assert_array_equal(buffer.read(np.float64), values) + other = device.alloc(32) + module.kernel("mixed").launch(args=[buffer, other, np.int64(3), 0.5]) + buffer.free() + with pytest.raises(DeviceError) as excinfo: + buffer.read(np.float64) + assert excinfo.value.code == "BAD_REQUEST" + assert not device.closed + + +def test_client_closes_after_a_fatal_launch_error() -> None: + with Device.spawn(fake_worker_argv(launch_error="CUDA_ERROR_ILLEGAL_ADDRESS")) as device: + kernel = device.load(MIXED_PTX).kernel("mixed") + buffer = device.alloc(32) + with pytest.raises(DeviceError) as excinfo: + kernel.launch(args=[buffer, buffer, 1, 1.0]) + assert excinfo.value.code == "CUDA_ERROR_ILLEGAL_ADDRESS" + assert device.closed + with pytest.raises(DeviceError) as excinfo: + device.alloc(8) + assert excinfo.value.code == "CLOSED" + + +@pytest.mark.skipif(ctypes.util.find_library("cuda") is not None, reason="CUDA is installed") +def test_worker_without_cuda_exits_and_the_client_reports_its_stderr() -> None: + with Device.spawn(worker_command(sys.executable)) as device: + with pytest.raises(DeviceError) as excinfo: + device.alloc(8) + assert excinfo.value.code == "WORKER_EXITED" + assert "cannot initialize CUDA" in excinfo.value.message + assert device.closed + + +MIXED_SOURCE = """\\! kernel mixed +\\! param A i64[4] +\\! param B f64[4] +\\! param N i64 +\\! param X f64 +""" + + +class FixedPtxCompiler: + def compile_source(self, _forth_src: str) -> str: + return MIXED_PTX + + +def test_kernel_runner_reads_back_the_output_param_zero_padded_and_typed() -> None: + runner = KernelRunner(lambda: Device.spawn(fake_worker_argv()), FixedPtxCompiler()) + try: + params = {"A": [5, -6], "B": [0.5], "N": 3, "X": 2} + floats = runner.run(MIXED_SOURCE, params=params, output_param=1) + ints = runner.run(MIXED_SOURCE, params=params, output_param=0, output_count=3) + finally: + runner.close() + assert floats == [0.5, 0.0, 0.0, 0.0] + assert all(isinstance(v, float) for v in floats) + assert ints == [5, -6, 0] + assert all(isinstance(v, int) for v in ints) + + +def test_kernel_runner_opens_a_fresh_device_after_a_fatal_error() -> None: + worker_options = [{"launch_error": "CUDA_ERROR_ILLEGAL_ADDRESS"}, {}] + opened: list[Device] = [] + + def open_device() -> Device: + opened.append(Device.spawn(fake_worker_argv(**worker_options[len(opened)]))) + return opened[-1] + + runner = KernelRunner(open_device, FixedPtxCompiler()) + try: + with pytest.raises(DeviceError, match="CUDA_ERROR_ILLEGAL_ADDRESS"): + runner.run(MIXED_SOURCE, params={"A": [1]}) + assert runner.run(MIXED_SOURCE, params={"A": [1]}, output_count=1) == [1] + finally: + runner.close() + assert len(opened) == 2 diff --git a/pyproject.toml b/pyproject.toml index b1c1386..90187a4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,5 +36,7 @@ ignore = ["D", "COM812", "ISC001"] [tool.ruff.lint.per-file-ignores] "gpu_test/test_*.py" = ["S101", "PLR2004"] "gpu_test/test_vast_session.py" = ["SLF001"] -"gpu_test/conftest.py" = ["S603", "S607", "PLR0913"] +"gpu_test/conftest.py" = ["S603", "PLR0913"] +"warpforth_gpu/vast.py" = ["S603", "S607"] +".claude/skills/verify-warpforth/scripts/*.py" = ["T201", "PLR0915"] "demo/*.py" = ["T201", "INP001"] diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 7d4740f..f9e8857 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -1,4 +1,3 @@ add_subdirectory(warpforth-translate) add_subdirectory(warpforth-opt) add_subdirectory(warpforthc) -add_subdirectory(warpforth-runner) diff --git a/tools/warpforth-runner/CMakeLists.txt b/tools/warpforth-runner/CMakeLists.txt deleted file mode 100644 index a2786f2..0000000 --- a/tools/warpforth-runner/CMakeLists.txt +++ /dev/null @@ -1,9 +0,0 @@ -find_package(CUDAToolkit QUIET) - -if(CUDAToolkit_FOUND) - add_executable(warpforth-runner warpforth-runner.cpp) - target_link_libraries(warpforth-runner PRIVATE CUDA::cuda_driver) - target_compile_features(warpforth-runner PRIVATE cxx_std_17) -else() - message(STATUS "CUDAToolkit not found — skipping warpforth-runner") -endif() diff --git a/tools/warpforth-runner/warpforth-runner.cpp b/tools/warpforth-runner/warpforth-runner.cpp deleted file mode 100644 index 67e1d3f..0000000 --- a/tools/warpforth-runner/warpforth-runner.cpp +++ /dev/null @@ -1,347 +0,0 @@ -/// warpforth-runner: Execute PTX kernels via the CUDA Driver API. -/// -/// Single-file C++ program designed to be uploaded and compiled on a remote -/// GPU host with `nvcc -o warpforth-runner warpforth-runner.cpp -lcuda -/// -std=c++17`. -/// -/// Usage: -/// warpforth-runner kernel.ptx --param i64[]:1,2,3 --param f64:3.14 \ -/// --grid 4,1,1 --block 64,1,1 --kernel main \ -/// --output-param 0 --output-count 3 - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#define CHECK_CU(call) \ - do { \ - CUresult err = (call); \ - if (err != CUDA_SUCCESS) { \ - const char *errStr = nullptr; \ - cuGetErrorString(err, &errStr); \ - std::cerr << "CUDA error at " << __FILE__ << ":" << __LINE__ << ": " \ - << (errStr ? errStr : "unknown") << "\n"; \ - exit(1); \ - } \ - } while (0) - -template struct ArrayParam { - std::vector values; - CUdeviceptr devicePtr = 0; -}; - -template struct ScalarParam { - T value; -}; - -using Param = std::variant, ArrayParam, - ScalarParam, ScalarParam>; - -template static void allocDevice(ArrayParam &arr) { - size_t bytes = arr.values.size() * sizeof(T); - CHECK_CU(cuMemAlloc(&arr.devicePtr, bytes)); - CHECK_CU(cuMemcpyHtoD(arr.devicePtr, arr.values.data(), bytes)); -} - -template -static void printOutput(ArrayParam &arr, size_t count) { - std::vector output(arr.values.size()); - CHECK_CU(cuMemcpyDtoH(output.data(), arr.devicePtr, - arr.values.size() * sizeof(T))); - for (size_t i = 0; i < count; ++i) { - if (i > 0) - std::cout << ","; - if constexpr (std::is_floating_point_v) - std::cout << std::setprecision(17) << output[i]; - else - std::cout << output[i]; - } - std::cout << "\n"; -} - -static void *kernelArgPtr(Param &p) { - if (auto *a = std::get_if>(&p)) - return &a->devicePtr; - if (auto *a = std::get_if>(&p)) - return &a->devicePtr; - if (auto *s = std::get_if>(&p)) - return &s->value; - return &std::get>(p).value; -} - -static bool isScalar(const Param &p) { - return std::holds_alternative>(p) || - std::holds_alternative>(p); -} - -struct Dims { - unsigned x = 1, y = 1, z = 1; -}; - -static int parseIntArg(std::string_view s, std::string_view optName) { - int value = 0; - auto [ptr, ec] = std::from_chars(s.data(), s.data() + s.size(), value); - if (ec != std::errc{} || ptr != s.data() + s.size()) { - std::cerr << "Error: " << optName << " expects an integer, got: " << s - << "\n"; - exit(1); - } - return value; -} - -static Dims parseDims(std::string_view s) { - Dims d; - const char *p = s.data(); - const char *end = s.data() + s.size(); - - auto dimsErr = [&]() { - std::cerr << "Error: expected 3 comma-separated values, got: " << s << "\n"; - exit(1); - }; - - auto [p1, ec1] = std::from_chars(p, end, d.x); - if (ec1 != std::errc{} || p1 == end || *p1 != ',') - dimsErr(); - - auto [p2, ec2] = std::from_chars(p1 + 1, end, d.y); - if (ec2 != std::errc{} || p2 == end || *p2 != ',') - dimsErr(); - - auto [p3, ec3] = std::from_chars(p2 + 1, end, d.z); - if (ec3 != std::errc{} || p3 != end) - dimsErr(); - - return d; -} - -static Param parseParam(std::string_view s) { - std::string input(s); - - auto colonPos = input.find(':'); - if (colonPos == std::string::npos) { - std::cerr << "Error: --param requires type prefix (e.g. i64:42 or " - "f64[]:1.0,2.0), got: " - << s << "\n"; - exit(1); - } - - std::string typePrefix = input.substr(0, colonPos); - std::string valueStr = input.substr(colonPos + 1); - - if (valueStr.empty()) { - std::cerr << "Error: --param requires at least one value, got: " << s - << "\n"; - exit(1); - } - - // Parse comma-separated values into a typed vector - auto parseValues = [&](auto convert) { - using T = decltype(convert(std::string{})); - std::vector vals; - std::istringstream iss(valueStr); - std::string token; - while (std::getline(iss, token, ',')) - vals.push_back(convert(token)); - return vals; - }; - - auto toI64 = [&](const std::string &tok) -> int64_t { - try { - return std::stoll(tok); - } catch (const std::exception &) { - std::cerr << "Error: invalid integer value '" << tok << "' in --param " - << s << "\n"; - exit(1); - } - }; - auto toF64 = [&](const std::string &tok) -> double { - try { - return std::stod(tok); - } catch (const std::exception &) { - std::cerr << "Error: invalid float value '" << tok << "' in --param " << s - << "\n"; - exit(1); - } - }; - - if (typePrefix == "i64[]") - return Param{ArrayParam{parseValues(toI64)}}; - if (typePrefix == "f64[]") - return Param{ArrayParam{parseValues(toF64)}}; - - // Scalars — must be exactly one value - if (valueStr.find(',') != std::string::npos) { - std::cerr << "Error: scalar param expects exactly one value, got: " << s - << "\n"; - exit(1); - } - - if (typePrefix == "i64") - return Param{ScalarParam{toI64(valueStr)}}; - if (typePrefix == "f64") - return Param{ScalarParam{toF64(valueStr)}}; - - std::cerr << "Error: unsupported param type '" << typePrefix - << "' (expected i64, i64[], f64, or f64[]), got: " << s << "\n"; - exit(1); -} - -static std::string readFile(std::string_view path) { - std::ifstream f(std::string(path), std::ios::binary); - if (!f) { - std::cerr << "Error: cannot open " << path << "\n"; - exit(1); - } - std::ostringstream ss; - ss << f.rdbuf(); - return ss.str(); -} - -int main(int argc, char **argv) { - const char *ptxFile = nullptr; - const char *kernelName = nullptr; - std::vector params; - Dims grid, block; - int outputParam = 0; - int outputCount = -1; // -1 = all - - // Parse arguments - for (int i = 1; i < argc; ++i) { - std::string_view arg = argv[i]; - auto needsValue = [&](std::string_view opt) { - if (++i >= argc) { - std::cerr << "Error: " << opt << " requires a value\n"; - exit(1); - } - }; - if (arg == "--param") { - needsValue("--param"); - params.push_back(parseParam(argv[i])); - } else if (arg == "--grid") { - needsValue("--grid"); - grid = parseDims(argv[i]); - } else if (arg == "--block") { - needsValue("--block"); - block = parseDims(argv[i]); - } else if (arg == "--output-param") { - needsValue("--output-param"); - outputParam = parseIntArg(argv[i], "--output-param"); - } else if (arg == "--output-count") { - needsValue("--output-count"); - outputCount = parseIntArg(argv[i], "--output-count"); - } else if (arg == "--kernel") { - needsValue("--kernel"); - kernelName = argv[i]; - } else if (arg[0] == '-') { - std::cerr << "Error: unknown option " << arg << "\n"; - exit(1); - } else { - ptxFile = argv[i]; - } - } - - if (!ptxFile) { - std::cerr << "Usage: warpforth-runner kernel.ptx --kernel NAME " - "[--param i64[]:V,...] [--param f64[]:V,...] " - "[--param i64:V] [--param f64:V] [--grid X,Y,Z] " - "[--block X,Y,Z] [--output-param N] [--output-count N]\n"; - return 1; - } - - if (!kernelName) { - std::cerr << "Error: --kernel NAME is required\n"; - return 1; - } - - if (params.empty()) { - std::cerr << "Error: at least one --param is required\n"; - return 1; - } - - if (outputParam < 0 || outputParam >= static_cast(params.size())) { - std::cerr << "Error: output-param " << outputParam << " out of range (have " - << params.size() << " params)\n"; - return 1; - } - - if (isScalar(params[outputParam])) { - std::cerr << "Error: output-param " << outputParam - << " is a scalar (cannot read back)\n"; - return 1; - } - - // Read PTX - std::string ptx = readFile(ptxFile); - - // Initialize CUDA - CHECK_CU(cuInit(0)); - - CUdevice device; - CHECK_CU(cuDeviceGet(&device, 0)); - - CUcontext ctx; - CHECK_CU(cuCtxCreate(&ctx, 0, device)); - - // Load PTX module - CUmodule module; - CHECK_CU(cuModuleLoadData(&module, ptx.c_str())); - - CUfunction func; - CHECK_CU(cuModuleGetFunction(&func, module, kernelName)); - - // Allocate device buffers for array params - for (auto &p : params) { - if (auto *a = std::get_if>(&p)) - allocDevice(*a); - else if (auto *a = std::get_if>(&p)) - allocDevice(*a); - } - - // Set up kernel parameters — Driver API expects array of pointers to args - std::vector kernelArgs(params.size()); - for (size_t i = 0; i < params.size(); ++i) - kernelArgs[i] = kernelArgPtr(params[i]); - - // Launch kernel - CHECK_CU(cuLaunchKernel(func, grid.x, grid.y, grid.z, block.x, block.y, - block.z, 0, nullptr, kernelArgs.data(), nullptr)); - - CHECK_CU(cuCtxSynchronize()); - - // Copy back and print output param - size_t count = outputCount >= 0 ? static_cast(outputCount) : 0; - if (auto *iArr = std::get_if>(¶ms[outputParam])) { - if (outputCount < 0) - count = iArr->values.size(); - printOutput(*iArr, count); - } else { - auto &fArr = std::get>(params[outputParam]); - if (outputCount < 0) - count = fArr.values.size(); - printOutput(fArr, count); - } - - // Cleanup — only free device memory for array params - for (auto &p : params) { - if (auto *a = std::get_if>(&p)) - cuMemFree(a->devicePtr); - else if (auto *a = std::get_if>(&p)) - cuMemFree(a->devicePtr); - } - cuModuleUnload(module); - cuCtxDestroy(ctx); - - return 0; -} diff --git a/warpforth_gpu/__init__.py b/warpforth_gpu/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/warpforth_gpu/client.py b/warpforth_gpu/client.py new file mode 100644 index 0000000..4c0997c --- /dev/null +++ b/warpforth_gpu/client.py @@ -0,0 +1,214 @@ +"""Client for the device worker: load PTX, move buffers, and launch kernels on a GPU.""" + +from __future__ import annotations + +import base64 +import contextlib +import itertools +import json +import subprocess +import tempfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import IO, TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + from collections.abc import Sequence + from typing import Self + + from numpy.typing import DTypeLike + +WORKER_SOURCE = Path(__file__).with_name("worker.py") +_CLOSE_TIMEOUT_S = 30 + +CLOSED = "CLOSED" +PROTOCOL = "PROTOCOL" +WORKER_EXITED = "WORKER_EXITED" + + +def worker_command(python: str = "python3") -> list[str]: + """Return argv that runs the worker from its inlined source, so no file must be copied.""" + return [python, "-u", "-c", WORKER_SOURCE.read_text()] + + +class DeviceError(Exception): + """The worker rejected a request, hit a CUDA error, or exited.""" + + def __init__(self, code: str, message: str) -> None: + super().__init__(f"{code}: {message}") + self.code = code + self.message = message + + +class Device: + """One worker process and the GPU context it owns. Requests are synchronous.""" + + def __init__(self, process: subprocess.Popen[bytes], stderr: IO[bytes]) -> None: + if process.stdin is None or process.stdout is None: + msg = "the worker process needs piped stdin and stdout" + raise ValueError(msg) + self._process = process + self._stdin = process.stdin + self._stdout = process.stdout + self._stderr = stderr + self._stderr_text = "" + self._ids = itertools.count(1) + self._closed = False + + @classmethod + def spawn(cls, argv: Sequence[str]) -> Device: + """Start any command that speaks the worker protocol on its stdin and stdout.""" + # A file rather than a pipe, so a chatty worker cannot block on a full stderr pipe. + stderr = tempfile.TemporaryFile() # noqa: SIM115 + process = subprocess.Popen( # noqa: S603 + list(argv), stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=stderr + ) + return cls(process, stderr) + + @property + def closed(self) -> bool: + return self._closed + + def __enter__(self) -> Self: + return self + + def __exit__(self, *_: object) -> None: + self.close() + + def close(self) -> None: + """Close the worker's stdin and wait for it to exit.""" + if self._closed: + return + self._closed = True + with contextlib.suppress(BrokenPipeError): + self._stdin.close() + try: + self._process.wait(timeout=_CLOSE_TIMEOUT_S) + except subprocess.TimeoutExpired: + self._process.kill() + self._process.wait() + self._stdout.close() + self._stderr.seek(0) + self._stderr_text = self._stderr.read().decode(errors="replace").strip() + self._stderr.close() + + def _fail(self, code: str, message: str) -> DeviceError: + """Close the device and build an error that carries the worker's stderr.""" + self.close() + if self._stderr_text: + message += f"\nworker stderr:\n{self._stderr_text}" + return DeviceError(code, message) + + def _exited(self) -> DeviceError: + self.close() # the exit status is known only after the process is reaped + return self._fail(WORKER_EXITED, f"worker exited with status {self._process.returncode}") + + def request(self, op: str, **fields: object) -> dict: + """Send one request and return its result, raising DeviceError on any failure.""" + if self._closed: + msg = "device is closed" + raise DeviceError(CLOSED, msg) + request_id = next(self._ids) + line = json.dumps({"id": request_id, "op": op, **fields}).encode() + b"\n" + try: + self._stdin.write(line) + self._stdin.flush() + except BrokenPipeError: + raise self._exited() from None + reply = self._stdout.readline() + if not reply: + raise self._exited() + try: + response = json.loads(reply) + except json.JSONDecodeError: + raise self._fail(PROTOCOL, f"worker sent a line that is not JSON: {reply!r}") from None + if not isinstance(response, dict) or response.get("id") != request_id: + raise self._fail(PROTOCOL, f"expected a reply to request {request_id}, got {reply!r}") + if response["ok"]: + return response["result"] + error = response["error"] + if response.get("fatal"): + raise self._fail(error["code"], error["message"]) + raise DeviceError(error["code"], error["message"]) + + def load(self, ptx: str) -> Module: + result = self.request("load", ptx=ptx) + kernels = {name: tuple(kinds) for name, kinds in result["kernels"].items()} + return Module(self, result["module"], kernels) + + def alloc(self, size: int) -> Buffer: + """Allocate ``size`` bytes of zeroed device memory.""" + return Buffer(self, self.request("alloc", size=size)["buffer"], size) + + def upload(self, array: np.ndarray) -> Buffer: + """Allocate a buffer the size of ``array`` and copy it in.""" + buffer = self.alloc(array.nbytes) + buffer.write(array) + return buffer + + +@dataclass(frozen=True) +class Buffer: + device: Device = field(repr=False) + id: int + size: int + + def write(self, data: bytes | np.ndarray) -> None: + """Copy ``data`` to the start of the buffer. It may be shorter than the buffer.""" + raw = np.ascontiguousarray(data).tobytes() if isinstance(data, np.ndarray) else data + self.device.request("write", buffer=self.id, data=base64.b64encode(raw).decode()) + + def read(self, dtype: DTypeLike) -> np.ndarray: + """Copy the whole buffer back as an array of ``dtype``.""" + data = base64.b64decode(self.device.request("read", buffer=self.id)["data"]) + return np.frombuffer(bytearray(data), dtype=dtype) + + def free(self) -> None: + self.device.request("free", buffer=self.id) + + +@dataclass(frozen=True) +class Module: + device: Device = field(repr=False) + id: int + kernels: dict[str, tuple[str, ...]] + + def kernel(self, name: str) -> Kernel: + if name not in self.kernels: + msg = f"module has no kernel {name!r}; it has {sorted(self.kernels)}" + raise KeyError(msg) + return Kernel(self, name, self.kernels[name]) + + +@dataclass(frozen=True) +class Kernel: + module: Module + name: str + kinds: tuple[str, ...] + + def launch( + self, + *, + grid: Sequence[int] = (1, 1, 1), + block: Sequence[int] = (1, 1, 1), + args: Sequence[Buffer | float | np.generic] = (), + ) -> None: + """Launch the kernel and wait for it to finish.""" + self.module.device.request( + "launch", + module=self.module.id, + kernel=self.name, + grid=[int(n) for n in grid], + block=[int(n) for n in block], + args=[_encode_arg(arg) for arg in args], + ) + + +def _encode_arg(arg: Buffer | float | np.generic) -> dict[str, object]: + if isinstance(arg, Buffer): + return {"buffer": arg.id} + if isinstance(arg, np.generic): + return {"value": arg.item()} + return {"value": arg} diff --git a/warpforth_gpu/vast.py b/warpforth_gpu/vast.py new file mode 100644 index 0000000..06cada2 --- /dev/null +++ b/warpforth_gpu/vast.py @@ -0,0 +1,452 @@ +"""Rent one GPU instance on Vast.ai for the lifetime of a context manager.""" + +from __future__ import annotations + +import atexit +import logging +import subprocess +import tempfile +import time +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING +from uuid import uuid4 + +from requests import RequestException +from vastai import VastAI + +if TYPE_CHECKING: + from typing import Self + +logger = logging.getLogger(__name__) + +MAX_COST_PER_HOUR = 0.50 +CUDA_VERSION = "12.4" +CUDA_IMAGE = f"nvidia/cuda:{CUDA_VERSION}.0-devel-ubuntu22.04" +POLL_INTERVAL_S = 10 +POLL_TIMEOUT_S = 300 +INSTANCE_LABEL_PREFIX = "warpforth-test-" +DEFAULT_ARCHITECTURE = "sm_70" + + +class VastSession: + """Context manager that rents a GPU instance on Vast.ai. + + Handles searching for offers, launching an instance, waiting for SSH + readiness, and unconditional cleanup on exit. + """ + + def __init__(self, api_key: str, arch: str = DEFAULT_ARCHITECTURE) -> None: + self.sdk = VastAI(api_key, retry=1) + self.arch = arch + timestamp = datetime.now(tz=UTC).strftime("%Y%m%dT%H%M%SZ") + self.instance_label = f"{INSTANCE_LABEL_PREFIX}{timestamp}-{uuid4().hex}" + self.instance_id: int | None = None + self.ssh_host: str | None = None + self.ssh_port: int | None = None + self._atexit_registered = False + self._creation_ambiguous = False + self._ssh_key_dir: tempfile.TemporaryDirectory[str] | None = None + self.ssh_private_key: Path | None = None + + def __enter__(self) -> Self: + self._log_existing_instances() + logger.info("This GPU test session owns label=%s", self.instance_label) + atexit.register(self._atexit_cleanup) + self._atexit_registered = True + try: + self._prepare_ssh_key() + self._launch() + self._wait_for_ssh() + except BaseException: + try: + if self._destroy(): + self._unregister_atexit() + finally: + self._cleanup_ssh_key() + raise + else: + return self + + def __exit__(self, *_: object) -> None: + try: + if self._destroy(): + self._unregister_atexit() + finally: + self._cleanup_ssh_key() + + def _prepare_ssh_key(self) -> None: + """Generate an ephemeral SSH identity for this test session.""" + self._ssh_key_dir = tempfile.TemporaryDirectory(prefix="warpforth-ssh-") + self.ssh_private_key = Path(self._ssh_key_dir.name) / "id_ed25519" + subprocess.run( + [ + "ssh-keygen", + "-q", + "-t", + "ed25519", + "-N", + "", + "-f", + str(self.ssh_private_key), + ], + capture_output=True, + text=True, + check=True, + ) + + def _attach_ssh_key(self) -> None: + """Attach this session's ephemeral public key to its Vast instance.""" + if self.instance_id is None or self.ssh_private_key is None: + msg = "Cannot attach SSH key before instance and identity creation" + raise RuntimeError(msg) + + public_key = self.ssh_private_key.with_suffix(".pub").read_text().strip() + result = self.sdk.attach_ssh(instance_id=self.instance_id, ssh_key=public_key) + if not isinstance(result, dict) or result.get("success") is False: + msg = f"Failed to attach ephemeral SSH key to instance {self.instance_id}" + raise RuntimeError(msg) + logger.info("Attached ephemeral SSH key to instance %s", self.instance_id) + + def _cleanup_ssh_key(self) -> None: + """Remove this session's ephemeral local SSH key material.""" + if self._ssh_key_dir is not None: + self._ssh_key_dir.cleanup() + self._ssh_key_dir = None + self.ssh_private_key = None + + def _show_instances(self) -> list[dict] | None: + """Return the current account instances, or None if listing fails.""" + try: + instances = self.sdk.show_instances() + except Exception: + logger.exception("Failed to list Vast.ai instances") + return None + + if not isinstance(instances, list): + logger.warning("Vast.ai returned an invalid instance listing") + return None + return [inst for inst in instances if isinstance(inst, dict)] + + def _log_existing_instances(self) -> None: + """Log existing WarpForth instances without modifying them.""" + instances = self._show_instances() + if instances is None: + logger.warning("Could not check for existing WarpForth Vast.ai instances") + return + + matching = [ + inst + for inst in instances + if str(inst.get("label", "")).startswith(INSTANCE_LABEL_PREFIX) + ] + if not matching: + logger.info("No existing WarpForth Vast.ai instances found") + return + + total_cost = 0.0 + for inst in matching: + cost = float(inst.get("dph_total") or 0.0) + total_cost += cost + start_date = inst.get("start_date") + age = "unknown" + if isinstance(start_date, int | float): + age = f"{max(0.0, time.time() - start_date) / 60:.0f}m" + logger.warning( + "Existing WarpForth Vast.ai instance: id=%s label=%s status=%s " + "gpu=%s dph=$%.3f age=%s ownership=active run or orphan candidate", + inst.get("id"), + inst.get("label"), + inst.get("actual_status", inst.get("status", "unknown")), + inst.get("gpu_name", "unknown"), + cost, + age, + ) + logger.warning( + "Found %d existing WarpForth Vast.ai instance(s), total known cost=$%.3f/hr; " + "they may belong to other active test runs and will not be destroyed", + len(matching), + total_cost, + ) + + def _launch(self) -> None: + """Find the cheapest suitable offer and launch an instance.""" + minimum_compute_capability = int(self.arch.removeprefix("sm_")) * 10 + query = ( + f"num_gpus=1 rentable=True rented=False compute_cap>={minimum_compute_capability}" + f" reliability>=0.95 inet_up>=100 dph<={MAX_COST_PER_HOUR}" + f" cuda_max_good>={CUDA_VERSION}" + ) + offers = self.sdk.search_offers(query=query, order="dph", limit=5) + + if not offers: + msg = "No suitable GPU offers found on Vast.ai" + raise RuntimeError(msg) + + offer = offers[0] if isinstance(offers, list) else offers + details = offer if isinstance(offer, dict) else {"id": offer} + offer_id = int(details["id"]) + logger.info( + "Launching instance from offer %s: gpu=%s driver=%s cuda_max_good=%s", + offer_id, + details.get("gpu_name"), + details.get("driver_version"), + details.get("cuda_max_good"), + ) + self._creation_ambiguous = True + try: + result = self.sdk.create_instance( + id=offer_id, + image=CUDA_IMAGE, + disk=10.0, + runtype="ssh_direc ssh_proxy", + label=self.instance_label, + ) + except RequestException: + logger.warning( + "Instance creation result is ambiguous; reconciling label=%s", + self.instance_label, + ) + self._reconcile_instance() + return + + if ( + isinstance(result, dict) + and result.get("success") is True + and result.get("new_contract") + ): + self.instance_id = int(result["new_contract"]) + self._creation_ambiguous = False + logger.info( + "This GPU test session owns instance id=%s label=%s", + self.instance_id, + self.instance_label, + ) + return + + logger.warning( + "Instance creation response did not contain an ID; reconciling label=%s", + self.instance_label, + ) + self._reconcile_instance() + + def _reconcile_instance(self) -> None: + """Find this run's instance by its exact unique label.""" + deadline = time.monotonic() + POLL_TIMEOUT_S + while time.monotonic() < deadline: + instances = self._show_instances() + if instances is not None: + matches = [inst for inst in instances if inst.get("label") == self.instance_label] + if len(matches) > 1: + ids = [inst.get("id") for inst in matches] + msg = f"Multiple instances found for unique label {self.instance_label}: {ids}" + raise RuntimeError(msg) + if matches: + self.instance_id = int(matches[0]["id"]) + self._creation_ambiguous = False + logger.info( + "This GPU test session owns instance id=%s label=%s", + self.instance_id, + self.instance_label, + ) + return + time.sleep(POLL_INTERVAL_S) + + msg = f"Could not reconcile instance with label {self.instance_label} after creation" + raise TimeoutError(msg) + + def _wait_for_ssh(self) -> None: + """Poll until the instance is running and SSH is available.""" + deadline = time.monotonic() + POLL_TIMEOUT_S + + while time.monotonic() < deadline: + instances = self.sdk.show_instances() + if not isinstance(instances, list): + time.sleep(POLL_INTERVAL_S) + continue + + for inst in instances: + if not isinstance(inst, dict): + continue + if int(inst.get("id", -1)) != self.instance_id: + continue + + status = inst.get("actual_status", inst.get("status", "")) + if status == "running" and inst.get("ssh_host") and inst.get("ssh_port"): + self.ssh_host = inst["ssh_host"] + self.ssh_port = int(inst["ssh_port"]) + logger.info( + "Instance %s ready: %s:%s", + self.instance_id, + self.ssh_host, + self.ssh_port, + ) + self._attach_ssh_key() + self._wait_for_sshd() + self._ensure_python() + return + + time.sleep(POLL_INTERVAL_S) + + msg = f"Instance {self.instance_id} did not become SSH-ready within {POLL_TIMEOUT_S}s" + raise TimeoutError(msg) + + def _wait_for_sshd(self) -> None: + """Poll until sshd is actually accepting connections.""" + deadline = time.monotonic() + POLL_TIMEOUT_S + while time.monotonic() < deadline: + result = subprocess.run( + [*self.ssh_argv(), "true"], + capture_output=True, + timeout=15, + check=False, + ) + if result.returncode == 0: + logger.info("sshd ready on %s:%s", self.ssh_host, self.ssh_port) + return + time.sleep(POLL_INTERVAL_S) + msg = ( + f"sshd on {self.ssh_host}:{self.ssh_port} did not become ready within {POLL_TIMEOUT_S}s" + ) + raise TimeoutError(msg) + + def _ensure_python(self) -> None: + """Install python3 for the device worker if the image lacks it.""" + self.ssh_run( + "command -v python3 >/dev/null || (apt-get update -qq && " + "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3)", + timeout=600, + ) + + def _record_labeled_instance_ids( + self, + instances: list[dict], + instance_ids: set[int], + ) -> None: + """Record instances owned by this session's exact unique label.""" + labeled_ids = { + int(inst["id"]) + for inst in instances + if inst.get("label") == self.instance_label and inst.get("id") is not None + } + instance_ids.update(labeled_ids) + if labeled_ids: + self._creation_ambiguous = False + + def _destroy(self) -> bool: + """Destroy only this session's instances and verify their absence.""" + instance_ids = {self.instance_id} if self.instance_id is not None else set() + retries = 3 + for attempt in range(retries): + instances = self._show_instances() + if instances is not None: + self._record_labeled_instance_ids(instances, instance_ids) + if not self._creation_ambiguous and not instance_ids: + logger.info( + "Verified cleanup for Vast.ai label=%s instance_ids=%s", + self.instance_label, + sorted(instance_ids), + ) + self.instance_id = None + return True + + for instance_id in instance_ids: + try: + self.sdk.destroy_instance(id=instance_id) + except Exception: + logger.exception( + "Destroy attempt %d/%d failed for instance %s", + attempt + 1, + retries, + instance_id, + ) + + instances = self._show_instances() + if instances is not None: + self._record_labeled_instance_ids(instances, instance_ids) + listed_ids = {int(inst["id"]) for inst in instances if inst.get("id") is not None} + if not self._creation_ambiguous and not instance_ids & listed_ids: + logger.info( + "Verified cleanup for Vast.ai label=%s instance_ids=%s", + self.instance_label, + sorted(instance_ids), + ) + self.instance_id = None + return True + + if attempt < retries - 1: + time.sleep(5) + + logger.error( + "Could not verify cleanup for Vast.ai label=%s instance_ids=%s; " + "check the Vast.ai account and destroy these instances manually", + self.instance_label, + sorted(instance_ids), + ) + return False + + def _unregister_atexit(self) -> None: + if self._atexit_registered: + atexit.unregister(self._atexit_cleanup) + self._atexit_registered = False + + def _atexit_cleanup(self) -> None: + """Last-resort cleanup registered before instance creation.""" + logger.warning( + "atexit: checking for Vast.ai instances with label=%s instance_id=%s", + self.instance_label, + self.instance_id, + ) + try: + self._destroy() + finally: + self._cleanup_ssh_key() + + def _ssh_options(self) -> list[str]: + if self.ssh_private_key is None: + msg = "Ephemeral SSH identity is not available" + raise RuntimeError(msg) + return [ + "-i", + str(self.ssh_private_key), + "-o", + "IdentitiesOnly=yes", + "-o", + "BatchMode=yes", + "-o", + "PreferredAuthentications=publickey", + "-o", + "StrictHostKeyChecking=no", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "ConnectTimeout=10", + "-o", + "ConnectionAttempts=3", + "-o", + "LogLevel=ERROR", + ] + + def ssh_argv(self) -> list[str]: + """Return the ssh argv for this instance; append a remote command to run it.""" + return [ + "ssh", + *self._ssh_options(), + "-p", + str(self.ssh_port), + f"root@{self.ssh_host}", + ] + + def ssh_run(self, cmd: str, *, timeout: int = 120) -> str: + """Execute a command on the remote instance via SSH.""" + result = subprocess.run( + [*self.ssh_argv(), cmd], + capture_output=True, + text=True, + timeout=timeout, + check=False, + ) + if result.returncode != 0: + msg = f"SSH command failed (rc={result.returncode}):\n{result.stderr}" + raise RuntimeError(msg) + return result.stdout diff --git a/warpforth_gpu/worker.py b/warpforth_gpu/worker.py new file mode 100644 index 0000000..f28d93a --- /dev/null +++ b/warpforth_gpu/worker.py @@ -0,0 +1,401 @@ +"""Device worker: loads PTX and launches kernels for a client over newline-delimited JSON. + +Stdlib only, because it is sent to a bare GPU host as ``python3 -u -c ``. +Each request is one JSON object per line on stdin; each response is one line on stdout. +""" + +from __future__ import annotations + +import base64 +import binascii +import ctypes +import itertools +import json +import re +import sys +from typing import TYPE_CHECKING, Literal, Protocol + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + from typing import IO + + Holder = ctypes.c_uint64 | ctypes.c_int64 | ctypes.c_double + +Kind = Literal["ptr", "i64", "f64"] + +_ENTRY = re.compile(r"\.entry\s+([A-Za-z_$%][\w$]*)\s*(?:\(([^)]*)\))?") +_I64_MIN = -(2**63) +_I64_MAX = 2**63 - 1 + +# A driver error during launch or synchronize leaves the context unusable, so +# the worker reports it and exits instead of serving further requests. +_STICKY_OPS = frozenset({"launch"}) + + +class BadRequestError(Exception): + """The request is malformed or refers to something that does not exist.""" + + +class CudaError(Exception): + """A CUDA driver call failed.""" + + def __init__(self, code: str, message: str) -> None: + super().__init__(f"{code}: {message}") + self.code = code + self.message = message + + +def _param_kind(decl: str) -> Kind: + match decl.split(): + case [".param", ".u64", *modifiers, _name] if ".ptr" in modifiers: + return "ptr" + case [".param", ".u64" | ".s64" | ".b64", _name]: + return "i64" + case [".param", ".f64", _name]: + return "f64" + msg = f"unsupported kernel parameter: {decl.strip()}" + raise ValueError(msg) + + +def parse_entries(ptx: str) -> dict[str, list[Kind]]: + """Map each ``.entry`` kernel in ``ptx`` to the kinds of its parameters, in order.""" + entries: dict[str, list[Kind]] = {} + for match in _ENTRY.finditer(ptx): + params = (match.group(2) or "").strip() + entries[match.group(1)] = [_param_kind(p) for p in params.split(",")] if params else [] + return entries + + +class Driver(Protocol): + def load(self, ptx: str) -> int: ... + def function(self, module: int, name: str) -> int: ... + def alloc(self, size: int) -> int: ... + def write(self, pointer: int, data: bytes) -> None: ... + def read(self, pointer: int, size: int) -> bytes: ... + def free(self, pointer: int) -> None: ... + def launch( + self, + function: int, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: list[Holder], + ) -> None: ... + + +_CU_JIT_ERROR_LOG_BUFFER = 5 +_CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES = 6 +_JIT_LOG_BYTES = 16384 + +_SIGNATURES: dict[str, list[object]] = { + "cuInit": [ctypes.c_uint], + "cuDeviceGet": [ctypes.POINTER(ctypes.c_int), ctypes.c_int], + "cuDevicePrimaryCtxRetain": [ctypes.POINTER(ctypes.c_void_p), ctypes.c_int], + "cuCtxSetCurrent": [ctypes.c_void_p], + "cuModuleLoadDataEx": [ + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_char_p, + ctypes.c_uint, + ctypes.POINTER(ctypes.c_int), + ctypes.POINTER(ctypes.c_void_p), + ], + "cuModuleGetFunction": [ctypes.POINTER(ctypes.c_void_p), ctypes.c_void_p, ctypes.c_char_p], + "cuMemAlloc_v2": [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t], + "cuMemsetD8_v2": [ctypes.c_uint64, ctypes.c_ubyte, ctypes.c_size_t], + "cuMemcpyHtoD_v2": [ctypes.c_uint64, ctypes.c_char_p, ctypes.c_size_t], + "cuMemcpyDtoH_v2": [ctypes.c_void_p, ctypes.c_uint64, ctypes.c_size_t], + "cuMemFree_v2": [ctypes.c_uint64], + "cuLaunchKernel": [ + ctypes.c_void_p, + *[ctypes.c_uint] * 7, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ctypes.POINTER(ctypes.c_void_p), + ], + "cuCtxSynchronize": [], + "cuGetErrorName": [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)], + "cuGetErrorString": [ctypes.c_int, ctypes.POINTER(ctypes.c_char_p)], +} + + +class CudaDriver: + """The CUDA driver API on device 0's primary context, through ctypes.""" + + def __init__(self, library: str = "libcuda.so.1") -> None: + self._cu = ctypes.CDLL(library) + for name, argtypes in _SIGNATURES.items(): + fn = getattr(self._cu, name) + fn.argtypes = argtypes + fn.restype = ctypes.c_int + self._check(self._cu.cuInit(0)) + device = ctypes.c_int() + self._check(self._cu.cuDeviceGet(ctypes.byref(device), 0)) + context = ctypes.c_void_p() + self._check(self._cu.cuDevicePrimaryCtxRetain(ctypes.byref(context), device)) + self._check(self._cu.cuCtxSetCurrent(context)) + + def _check(self, result: int, detail: str = "") -> None: + if result == 0: + return + name = ctypes.c_char_p() + text = ctypes.c_char_p() + self._cu.cuGetErrorName(result, ctypes.byref(name)) + self._cu.cuGetErrorString(result, ctypes.byref(text)) + code = name.value.decode() if name.value else f"CUDA_ERROR_{result}" + message = text.value.decode() if text.value else "unknown CUDA error" + raise CudaError(code, f"{message}\n{detail}" if detail else message) + + def load(self, ptx: str) -> int: + log = ctypes.create_string_buffer(_JIT_LOG_BYTES) + options = (ctypes.c_int * 2)(_CU_JIT_ERROR_LOG_BUFFER, _CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES) + values = (ctypes.c_void_p * 2)(ctypes.addressof(log), _JIT_LOG_BYTES) + module = ctypes.c_void_p() + result = self._cu.cuModuleLoadDataEx(ctypes.byref(module), ptx.encode(), 2, options, values) + self._check(result, log.value.decode(errors="replace")) + return module.value or 0 + + def function(self, module: int, name: str) -> int: + function = ctypes.c_void_p() + self._check(self._cu.cuModuleGetFunction(ctypes.byref(function), module, name.encode())) + return function.value or 0 + + def alloc(self, size: int) -> int: + pointer = ctypes.c_uint64() + self._check(self._cu.cuMemAlloc_v2(ctypes.byref(pointer), size)) + try: + self._check(self._cu.cuMemsetD8_v2(pointer, 0, size)) + except CudaError: + self._cu.cuMemFree_v2(pointer) + raise + return pointer.value + + def write(self, pointer: int, data: bytes) -> None: + self._check(self._cu.cuMemcpyHtoD_v2(pointer, data, len(data))) + + def read(self, pointer: int, size: int) -> bytes: + out = ctypes.create_string_buffer(size) + self._check(self._cu.cuMemcpyDtoH_v2(ctypes.addressof(out), pointer, size)) + return out.raw + + def free(self, pointer: int) -> None: + self._check(self._cu.cuMemFree_v2(pointer)) + + def launch( + self, + function: int, + grid: tuple[int, int, int], + block: tuple[int, int, int], + args: list[Holder], + ) -> None: + params = (ctypes.c_void_p * len(args))(*(ctypes.addressof(a) for a in args)) + self._check(self._cu.cuLaunchKernel(function, *grid, *block, 0, None, params, None)) + self._check(self._cu.cuCtxSynchronize()) + + +def _field(request: dict, name: str) -> object: + if name not in request: + msg = f"missing field {name!r}" + raise BadRequestError(msg) + return request[name] + + +def _is_int(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _positive_int(request: dict, name: str) -> int: + value = _field(request, name) + if not _is_int(value) or value <= 0: + msg = f"{name} must be a positive integer, got {value!r}" + raise BadRequestError(msg) + return value + + +def _dims(request: dict, name: str) -> tuple[int, int, int]: + match _field(request, name): + case [x, y, z] if all(_is_int(v) and v > 0 for v in (x, y, z)): + return x, y, z + case value: + msg = f"{name} must be three positive integers, got {value!r}" + raise BadRequestError(msg) + + +class Worker: + """Owns the modules and buffers created through one driver and serves requests.""" + + def __init__(self, driver: Driver) -> None: + self._driver = driver + self._ids = itertools.count(1) + self._kernels: dict[int, dict[str, tuple[int, list[Kind]]]] = {} + self._buffers: dict[int, tuple[int, int]] = {} + self._ops: dict[str, Callable[[dict], dict]] = { + "load": self._load, + "alloc": self._alloc, + "write": self._write, + "read": self._read, + "free": self._free, + "launch": self._launch, + } + + def handle(self, request: dict) -> dict: + op = request.get("op") + handler = self._ops.get(op) if isinstance(op, str) else None + if handler is None: + msg = f"unknown op {op!r}" + raise BadRequestError(msg) + return handler(request) + + def _buffer(self, request: dict, name: str = "buffer") -> tuple[int, int]: + handle = _field(request, name) + if not _is_int(handle) or handle not in self._buffers: + msg = f"unknown buffer {handle!r}" + raise BadRequestError(msg) + return self._buffers[handle] + + def _load(self, request: dict) -> dict: + ptx = _field(request, "ptx") + if not isinstance(ptx, str): + msg = "ptx must be a string" + raise BadRequestError(msg) + try: + entries = parse_entries(ptx) + except ValueError as e: + raise BadRequestError(str(e)) from e + module = self._driver.load(ptx) + handle = next(self._ids) + self._kernels[handle] = { + name: (self._driver.function(module, name), kinds) for name, kinds in entries.items() + } + return {"module": handle, "kernels": entries} + + def _alloc(self, request: dict) -> dict: + size = _positive_int(request, "size") + pointer = self._driver.alloc(size) + handle = next(self._ids) + self._buffers[handle] = (pointer, size) + return {"buffer": handle} + + def _write(self, request: dict) -> dict: + pointer, size = self._buffer(request) + encoded = _field(request, "data") + if not isinstance(encoded, str): + msg = "data must be a base64 string" + raise BadRequestError(msg) + try: + data = base64.b64decode(encoded, validate=True) + except binascii.Error as e: + msg = f"data is not valid base64: {e}" + raise BadRequestError(msg) from e + if len(data) > size: + msg = f"data is {len(data)} bytes but the buffer holds {size}" + raise BadRequestError(msg) + self._driver.write(pointer, data) + return {} + + def _read(self, request: dict) -> dict: + pointer, size = self._buffer(request) + return {"data": base64.b64encode(self._driver.read(pointer, size)).decode()} + + def _free(self, request: dict) -> dict: + pointer, _size = self._buffer(request) + self._driver.free(pointer) + del self._buffers[request["buffer"]] + return {} + + def _launch(self, request: dict) -> dict: + module = _field(request, "module") + if not _is_int(module) or module not in self._kernels: + msg = f"unknown module {module!r}" + raise BadRequestError(msg) + name = _field(request, "kernel") + kernels = self._kernels[module] + if not isinstance(name, str) or name not in kernels: + msg = f"module {module} has no kernel {name!r}; it has {sorted(kernels)}" + raise BadRequestError(msg) + function, kinds = kernels[name] + grid = _dims(request, "grid") + block = _dims(request, "block") + args = _field(request, "args") + if not isinstance(args, list) or len(args) != len(kinds): + count = len(args) if isinstance(args, list) else args + msg = f"kernel {name!r} takes {len(kinds)} arguments ({kinds}), got {count!r}" + raise BadRequestError(msg) + holders = [ + self._pack(i, kind, arg) for i, (kind, arg) in enumerate(zip(kinds, args, strict=True)) + ] + self._driver.launch(function, grid, block, holders) + return {} + + def _pack(self, index: int, kind: Kind, arg: object) -> Holder: + if kind == "ptr": + if isinstance(arg, dict) and "buffer" in arg: + return ctypes.c_uint64(self._buffer(arg)[0]) + elif isinstance(arg, dict) and "value" in arg: + value = arg["value"] + if kind == "i64" and _is_int(value) and _I64_MIN <= value <= _I64_MAX: + return ctypes.c_int64(value) + if kind == "f64" and (_is_int(value) or isinstance(value, float)): + try: + return ctypes.c_double(float(value)) + except OverflowError: + pass + expected = {"ptr": '{"buffer": id}', "i64": "a 64-bit integer", "f64": "a number"}[kind] + msg = f"argument {index} is {kind} and needs {expected}, got {arg!r}" + raise BadRequestError(msg) + + +def _error(request_id: object, code: str, message: str, *, fatal: bool) -> dict: + return { + "id": request_id, + "ok": False, + "error": {"code": code, "message": message}, + "fatal": fatal, + } + + +def _parse_request(line: str) -> dict: + try: + request = json.loads(line) + except json.JSONDecodeError as e: + msg = f"malformed JSON: {e}" + raise BadRequestError(msg) from e + if not isinstance(request, dict) or "id" not in request: + msg = "request must be a JSON object with an id" + raise BadRequestError(msg) + return request + + +def serve(driver: Driver, requests: Iterable[str], responses: IO[str]) -> int: + """Answer each request line until input ends or a sticky error; return the exit status.""" + worker = Worker(driver) + for line in requests: + if not line.strip(): + continue + request_id = None + fatal = False + try: + request = _parse_request(line) + request_id = request["id"] + response = {"id": request_id, "ok": True, "result": worker.handle(request)} + except BadRequestError as e: + response = _error(request_id, "BAD_REQUEST", str(e), fatal=False) + except CudaError as e: + fatal = request.get("op") in _STICKY_OPS + response = _error(request_id, e.code, e.message, fatal=fatal) + responses.write(json.dumps(response) + "\n") + responses.flush() + if fatal: + return 1 + return 0 + + +def main() -> int: + try: + driver = CudaDriver() + except (OSError, CudaError) as e: + sys.stderr.write(f"warpforth worker: cannot initialize CUDA: {e}\n") + return 2 + return serve(driver, sys.stdin, sys.stdout) + + +if __name__ == "__main__": + sys.exit(main())