From 6608824cb0aa343716b0c845ed781418a23db47e Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Wed, 1 Jul 2026 12:44:25 +0400 Subject: [PATCH] feat(mvp-chat): local e2e chat on cuda gpu Stand up an interactive end-to-end chat over a CUDA GPU, provisioning a Dockerized node that loads a GGUF model and serves prompts over TCP. - prompt_rpc: add the newline-JSON prompt protocol (`SubmitPrompt` + `PromptEvent::{TextDelta,Done,Fault}`) carried over TCP - mvp_chat: add an interactive REPL client connecting to the prompt RPC port (default 127.0.0.1:19777) - mvp_orch_one_node / mvp_one_node_chat: add the single-node orchestrator that provisions a `LocalDockerPlugin` node, loads `bartowski/Llama-3.2-1B-Instruct-GGUF` (Q4_K_M), and exposes the prompt RPC listener with boot/route/weight timeouts - mvp_node: add the GPU worker binary that spawns `tinygrad_worker.py` (default device CUDA) and ships runtime telemetry via a `ClusterFrameSink` - vastai_provisioning / bootstrap_datastream: add the vast.ai provider adapter (`VastAiProvisioningConfig`, `VastAiLeaseClient`) wrapping `swactor_vastai`, plus a bridge that folds provision stdout onto a per-node datastream - apps/mvp-node: add CUDA base/runtime Dockerfiles (nvidia/cuda 12.6.3, tinygrad 0.12.0, sshd), `mvp_entrypoint.sh` (sshd + mvp-node, held for postmortem), `local_docker_e2e.sh`, the GGUF tinygrad worker, and one-node-chat/bootstrap/vastai guarantee tests Signed-off-by: Zachery Aaron Shores-Chmielewski --- .cargo/config.toml | 2 + .dockerignore | 11 + Cargo.lock | 2 + apps/mvp-node/Dockerfile | 15 + apps/mvp-node/Dockerfile.base | 68 ++ apps/mvp-node/local_docker_e2e.sh | 54 + apps/mvp-node/mvp_entrypoint.sh | 34 + apps/mvp-node/tinygrad_worker.py | 299 ++++++ crates/dashboard/src/lib.rs | 2 +- crates/mvp-system/Cargo.toml | 19 + crates/mvp-system/src/actors/mod.rs | 1 + crates/mvp-system/src/actors/node_agent.rs | 53 +- crates/mvp-system/src/bin/local_e2e.rs | 136 ++- .../mvp-system/src/bin/local_e2e_cluster.rs | 135 ++- crates/mvp-system/src/bin/mvp_chat.rs | 174 ++++ crates/mvp-system/src/bin/mvp_node.rs | 794 ++++++++++++++ .../mvp-system/src/bin/mvp_one_node_chat.rs | 512 +++++++++ .../mvp-system/src/bin/mvp_orch_one_node.rs | 983 ++++++++++++++++++ crates/mvp-system/src/bootstrap_datastream.rs | 183 ++++ crates/mvp-system/src/lib.rs | 3 + crates/mvp-system/src/prompt_rpc.rs | 136 +++ crates/mvp-system/src/provisioning.rs | 80 +- crates/mvp-system/src/run_plan.rs | 6 +- crates/mvp-system/src/stage_controller.rs | 30 +- .../tests/bootstrap_datastream_guarantees.rs | 115 ++ .../src/tests/local_mock/environment.rs | 9 +- crates/mvp-system/src/tests/mod.rs | 2 + .../src/tests/node_provisioning_guarantees.rs | 86 ++ .../src/tests/run_plan_guarantees.rs | 1 + .../src/tests/stage_controller_guarantees.rs | 2 +- .../tests/vastai_provisioning_guarantees.rs | 351 +++++++ crates/mvp-system/src/vastai_provisioning.rs | 619 +++++++++++ .../local_e2e_cluster/tinygrad_cpu_worker.py | 17 + crates/mvp-system/tests/one_node_chat_e2e.rs | 362 +++++++ mvp-node-image-gguf-shape.md | 167 +++ vastai-datastream-plugin-shape.md | 110 ++ 36 files changed, 5445 insertions(+), 128 deletions(-) create mode 100644 apps/mvp-node/Dockerfile create mode 100644 apps/mvp-node/Dockerfile.base create mode 100755 apps/mvp-node/local_docker_e2e.sh create mode 100755 apps/mvp-node/mvp_entrypoint.sh create mode 100755 apps/mvp-node/tinygrad_worker.py create mode 100644 crates/mvp-system/src/bin/mvp_chat.rs create mode 100644 crates/mvp-system/src/bin/mvp_node.rs create mode 100644 crates/mvp-system/src/bin/mvp_one_node_chat.rs create mode 100644 crates/mvp-system/src/bin/mvp_orch_one_node.rs create mode 100644 crates/mvp-system/src/bootstrap_datastream.rs create mode 100644 crates/mvp-system/src/prompt_rpc.rs create mode 100644 crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs create mode 100644 crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs create mode 100644 crates/mvp-system/src/vastai_provisioning.rs create mode 100644 crates/mvp-system/tests/one_node_chat_e2e.rs create mode 100644 mvp-node-image-gguf-shape.md create mode 100644 vastai-datastream-plugin-shape.md diff --git a/.cargo/config.toml b/.cargo/config.toml index 35049cb..65e7c3b 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,2 +1,4 @@ [alias] xtask = "run --package xtask --" +mvp-chat = "run -p mvp-system --features local-e2e --bin mvp-one-node-chat --" +mvp-chat-test = "test -p mvp-system --features local-e2e --test one_node_chat_e2e -- --nocapture" diff --git a/.dockerignore b/.dockerignore index 72e8ffc..7d5f374 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1 +1,12 @@ * +!apps/ +!apps/mvp-node/ +!apps/mvp-node/** +!target/ +target/* +!target/release/ +target/release/* +!target/release/mvp-node +!target/debug/ +target/debug/* +!target/debug/mvp-node diff --git a/Cargo.lock b/Cargo.lock index 38682ed..9e9f761 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2477,10 +2477,12 @@ dependencies = [ "iroh", "iroh-driver", "libc", + "parking_lot", "serde", "serde_json", "swactor", "swactor-transport", + "swactor-vastai", "tokio", ] diff --git a/apps/mvp-node/Dockerfile b/apps/mvp-node/Dockerfile new file mode 100644 index 0000000..99d8692 --- /dev/null +++ b/apps/mvp-node/Dockerfile @@ -0,0 +1,15 @@ +# MVP node CODE image — thin layer over the CUDA/tinygrad base. +# Build from workspace root after compiling the Rust binary: +# cargo build --release -p mvp-system --bin mvp-node +# docker build -f apps/mvp-node/Dockerfile.base -t swactor-mvp-node-base:cuda12.6 . +# docker build -f apps/mvp-node/Dockerfile --build-arg BASE_IMAGE=swactor-mvp-node-base:cuda12.6 -t swactor-mvp-node:latest . + +ARG BASE_IMAGE=swactor-mvp-node-base:cuda12.6 +ARG MVP_NODE_BIN=target/release/mvp-node +FROM ${BASE_IMAGE} +ARG MVP_NODE_BIN=target/release/mvp-node + +COPY ${MVP_NODE_BIN} /usr/local/bin/mvp-node +COPY apps/mvp-node/tinygrad_worker.py /usr/local/share/mvp/tinygrad_worker.py + +RUN chmod +x /usr/local/bin/mvp-node /usr/local/share/mvp/tinygrad_worker.py diff --git a/apps/mvp-node/Dockerfile.base b/apps/mvp-node/Dockerfile.base new file mode 100644 index 0000000..411276f --- /dev/null +++ b/apps/mvp-node/Dockerfile.base @@ -0,0 +1,68 @@ +# MVP node BASE image — CUDA/Python/tinygrad/sshd/PID-1 foundation. +# Build from workspace root: +# docker build -f apps/mvp-node/Dockerfile.base -t swactor-mvp-node-base:cuda12.6 . + +FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04 AS builder + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 \ + python3-pip \ + ca-certificates \ + cuda-cudart-dev-12-6 && \ + rm -rf /var/lib/apt/lists/* + +RUN python3 -m pip install --no-cache-dir --break-system-packages \ + --target=/opt/mvp-pydeps \ + tinygrad==0.12.0 numpy && \ + find /opt/mvp-pydeps -depth -type d \ + \( -name '__pycache__' -o -name 'tests' -o -name 'test' \) \ + -exec rm -rf {} + && \ + find /opt/mvp-pydeps -name '*.pyc' -delete && \ + find /opt/mvp-pydeps -type d -name '*.dist-info' -exec rm -rf {} + + +RUN mkdir -p /opt/mvp-nvrtc-include && \ + cp -rL /usr/local/cuda/include/. /opt/mvp-nvrtc-include/ && \ + test -f /opt/mvp-nvrtc-include/vector_types.h + +FROM nvidia/cuda:12.6.3-base-ubuntu24.04 AS runtime + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + python3 \ + cuda-cudart-12-6 \ + cuda-nvrtc-12-6 \ + ca-certificates \ + procps \ + openssh-server && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* \ + /var/cache/apt/archives/* \ + /var/cache/apt/*.bin \ + /var/log/apt/* \ + /var/log/dpkg.log \ + /tmp/* \ + /var/tmp/* \ + /usr/share/doc/* \ + /usr/share/man/* \ + /usr/share/info/* && \ + find /usr -depth -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null || true && \ + find /usr -type f -name '*.pyc' -delete 2>/dev/null || true && \ + mkdir -p /usr/local/share/mvp /var/cache/mvp-models /var/log + +COPY --from=builder /opt/mvp-nvrtc-include/ /usr/local/cuda/include/ +COPY --from=builder /opt/mvp-pydeps /opt/mvp-pydeps +ENV PYTHONPATH=/opt/mvp-pydeps +ENV PYTHONDONTWRITEBYTECODE=1 +ENV CUDA=1 +ENV DEV=CUDA +ENV MVP_MODEL_CACHE_DIR=/var/cache/mvp-models +ENV MVP_TINYGRAD_WORKER=/usr/local/share/mvp/tinygrad_worker.py + +COPY apps/mvp-node/mvp_entrypoint.sh /usr/local/bin/mvp_entrypoint.sh +RUN chmod +x /usr/local/bin/mvp_entrypoint.sh + +ARG DEPLOY_PUBKEY="" +RUN if [ -n "$DEPLOY_PUBKEY" ]; then printf '%s\n' "$DEPLOY_PUBKEY" > /etc/mvp_deploy_key.pub; fi + +ENTRYPOINT ["/usr/local/bin/mvp_entrypoint.sh"] diff --git a/apps/mvp-node/local_docker_e2e.sh b/apps/mvp-node/local_docker_e2e.sh new file mode 100755 index 0000000..71ff5cd --- /dev/null +++ b/apps/mvp-node/local_docker_e2e.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +set -euo pipefail + +BASE_IMAGE=${BASE_IMAGE:-swactor-mvp-node-base:cuda12.6} +IMAGE=${IMAGE:-swactor-mvp-node:latest} +CONTAINER=${CONTAINER:-swactor-mvp-node-e2e-$$} +GPUS=${MVP_CUDA_GPUS:-all} +PROMPT=${MVP_NODE_SELF_TEST_PROMPT:-ping} +TIMEOUT_SECS=${MVP_NODE_E2E_TIMEOUT_SECS:-1800} +FRAME_LOG=${MVP_DATASTREAM_FRAME_LOG:-/var/log/mvp-datastream.ndjson} + +cleanup() { + docker rm -f "$CONTAINER" >/dev/null 2>&1 || true +} +trap cleanup EXIT + +cargo build --release -p mvp-system --bin mvp-node +docker build -f apps/mvp-node/Dockerfile.base -t "$BASE_IMAGE" . +docker build -f apps/mvp-node/Dockerfile --build-arg BASE_IMAGE="$BASE_IMAGE" -t "$IMAGE" . + +docker run -d \ + --name "$CONTAINER" \ + --gpus "$GPUS" \ + -e MVP_NODE_SELF_TEST_PROMPT="$PROMPT" \ + -e MVP_NODE_MAX_RUNTIME_SECS=1 \ + -e MVP_SELF_TEST_MAX_TOKENS="${MVP_SELF_TEST_MAX_TOKENS:-1}" \ + -e MVP_MODEL_CACHE_DIR=/var/cache/mvp-models \ + -e MVP_DATASTREAM_FRAME_LOG="$FRAME_LOG" \ + ${HF_TOKEN:+-e HF_TOKEN="$HF_TOKEN"} \ + "$IMAGE" >/dev/null + +deadline=$((SECONDS + TIMEOUT_SECS)) +while (( SECONDS < deadline )); do + logs=$(docker logs "$CONTAINER" 2>&1 || true) + if grep -q '"type":"ready"' <<<"$logs" && grep -q '"type":"self_test_completed"' <<<"$logs"; then + frames=$(docker exec "$CONTAINER" cat "$FRAME_LOG" 2>/dev/null || true) + if grep -q '"channel":"mvp.node.ready"' <<<"$frames" && + grep -q '"channel":"mvp.worker.weights"' <<<"$frames" && + grep -q '"channel":"mvp.worker.prompt"' <<<"$frames"; then + printf '%s\n' "$logs" + printf '%s\n' "$frames" + exit 0 + fi + fi + if grep -q 'WorkerFatal\|mvp-node: .*failed\|ModelLoadFailed\|GgufDownloadFailed' <<<"$logs"; then + printf '%s\n' "$logs" >&2 + exit 1 + fi + sleep 5 +done + +docker logs "$CONTAINER" 2>&1 || true +echo "mvp-node Docker E2E timed out after ${TIMEOUT_SECS}s" >&2 +exit 1 diff --git a/apps/mvp-node/mvp_entrypoint.sh b/apps/mvp-node/mvp_entrypoint.sh new file mode 100755 index 0000000..5fd6af5 --- /dev/null +++ b/apps/mvp-node/mvp_entrypoint.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +set -u + +if ! mkdir /run/mvp_entrypoint.lock 2>/dev/null; then + echo "mvp-entrypoint: already running (lock held); parking" >&2 + exec sleep infinity +fi + +mkdir -p /root/.ssh && chmod 700 /root/.ssh +: > /root/.ssh/authorized_keys +if [ -n "${PUBLIC_KEY:-}" ]; then printf '%s\n' "$PUBLIC_KEY" >> /root/.ssh/authorized_keys; fi +if [ -n "${SSH_PUBLIC_KEY:-}" ]; then printf '%s\n' "$SSH_PUBLIC_KEY" >> /root/.ssh/authorized_keys; fi +if [ -f /etc/mvp_deploy_key.pub ]; then cat /etc/mvp_deploy_key.pub >> /root/.ssh/authorized_keys; fi +chmod 600 /root/.ssh/authorized_keys + +if [ ! -s /root/.ssh/authorized_keys ]; then + echo "mvp-entrypoint: WARNING no SSH public key found (PUBLIC_KEY/SSH_PUBLIC_KEY unset, no baked key); SSH will reject logins" >&2 +fi + +mkdir -p /run/sshd /var/log /var/cache/mvp-models +ssh-keygen -A 2>/dev/null || true +/usr/sbin/sshd -e +echo "mvp-entrypoint: sshd up on :22" >&2 + +NODE_LOG=/var/log/mvp-node.log +echo "mvp-entrypoint: launching mvp-node (log -> $NODE_LOG)" >&2 +set -o pipefail +/usr/local/bin/mvp-node "$@" 2>&1 | tee "$NODE_LOG" +code=${PIPESTATUS[0]} + +echo "mvp-entrypoint: mvp-node exited with code $code; NOT restarting (node held for postmortem)" >&2 +echo "mvp-entrypoint: --- last 80 lines of $NODE_LOG ---" >&2 +tail -n 80 "$NODE_LOG" >&2 || true +exec sleep infinity diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py new file mode 100755 index 0000000..e685cbc --- /dev/null +++ b/apps/mvp-node/tinygrad_worker.py @@ -0,0 +1,299 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import hashlib +import json +import os +import sys +import time +import traceback +import urllib.parse +import urllib.request +from pathlib import Path +from typing import Any + +Tensor: Any = None +dtypes: Any = None +model: Any = None +tokenizer: Any = None +role: dict[str, Any] = {} +loaded: dict[str, Any] = {} + + +def control(**event: Any) -> None: + print(json.dumps(event, separators=(",", ":")), flush=True) + + +def log(message: str) -> None: + print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True) + + +def fatal(reason: str, **fields: Any) -> None: + control(type="WorkerFatal", reason=reason, **fields) + raise SystemExit(1) + + +def test_mode() -> bool: + return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"} + + + +def initialize(cmd: dict[str, Any]) -> None: + global Tensor, dtypes + if int(cmd.get("helper_abi_version", 1)) != 1: + fatal("UnsupportedHelperAbi", helper_abi_version=cmd.get("helper_abi_version")) + device = str(cmd.get("backend", {}).get("device") or os.environ.get("DEV") or "CUDA") + os.environ["DEV"] = device + started = time.monotonic() + control(type="TinygradImportStarted", device=device) + from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes + + Tensor = TinyTensor + dtypes = tiny_dtypes + value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() + control( + type="WorkerReady", + pid=os.getpid(), + backend={"device": device}, + cuda_probe=value, + elapsed_ms=int((time.monotonic() - started) * 1000), + ) + + +def configure_role(cmd: dict[str, Any]) -> None: + config = cmd.get("config", {}) + role.clear() + role.update( + role_id=int(cmd.get("role_id", 1)), + run_id=int(config.get("run_id", 1)), + stage_index=int(config.get("stage_index", 0)), + layer_start=int(config.get("layer_start", 0)), + layer_end_exclusive=int(config.get("layer_end_exclusive", 0)), + ) + control(type="RoleConfigured", role_id=role["role_id"], stage_index=role["stage_index"]) + + +def cache_root() -> Path: + raw = os.environ.get("MVP_MODEL_CACHE_DIR", "").strip() + root = Path(raw).expanduser() if raw else Path.home() / ".cache" / "mvp-node" + root.mkdir(parents=True, exist_ok=True) + return root + + +def hf_url(repo: str, file: str, revision: str | None) -> str: + encoded_file = "/".join(urllib.parse.quote(part) for part in file.split("/")) + return f"https://huggingface.co/{repo}/resolve/{revision or 'main'}/{encoded_file}" + + +def source_url(source: dict[str, Any]) -> str | None: + if "HuggingFaceGguf" not in source: + return None + hf = source["HuggingFaceGguf"] + return hf_url(str(hf["repo"]), str(hf["file"]), hf.get("revision")) + + +def source_path(source: dict[str, Any]) -> Path | None: + if "LocalPath" not in source: + return None + return Path(str(source["LocalPath"])).expanduser() + + +def cache_path_for(url: str) -> Path: + parsed = urllib.parse.urlparse(url) + basename = Path(parsed.path).name or "model.gguf" + digest = hashlib.sha256(url.encode("utf-8")).hexdigest()[:16] + return cache_root() / f"{digest}-{basename}" + + +def request_headers() -> dict[str, str]: + headers = {"User-Agent": "swactor-mvp-node/0.1"} + token = os.environ.get("HF_TOKEN", "").strip() + if token: + headers["Authorization"] = f"Bearer {token}" + return headers + + +def fetch_whole(source: dict[str, Any]) -> Path: + local = source_path(source) + if local is not None: + if not local.is_file(): + fatal("GgufLocalPathMissing", path=str(local)) + control(type="GgufCacheReady", path=str(local), cache_hit=True, source="local") + return local + + url = source_url(source) + if not url: + fatal("UnsupportedGgufSource", source=source) + target = cache_path_for(url) + if target.is_file() and target.stat().st_size > 0: + control(type="GgufCacheReady", path=str(target), bytes=target.stat().st_size, cache_hit=True, url=url) + return target + + partial = target.with_name(target.name + ".partial") + started = time.monotonic() + req = urllib.request.Request(url, headers=request_headers()) + control(type="GgufDownloadStarted", url=url, path=str(target)) + try: + with urllib.request.urlopen(req, timeout=60) as response, partial.open("wb") as out: + total = int(response.headers.get("Content-Length") or 0) + done = 0 + last_event = 0.0 + while True: + chunk = response.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + done += len(chunk) + now = time.monotonic() + if now - last_event >= float(os.environ.get("MVP_DOWNLOAD_PROGRESS_SECS", "5")): + control( + type="GgufDownloadProgress", + bytes_done=done, + bytes_total=total, + elapsed_ms=int((now - started) * 1000), + ) + last_event = now + partial.replace(target) + except Exception as exc: + try: + partial.unlink(missing_ok=True) + except Exception: + pass + fatal("GgufDownloadFailed", url=url, error=str(exc)) + control( + type="GgufCacheReady", + path=str(target), + bytes=target.stat().st_size, + cache_hit=False, + elapsed_ms=int((time.monotonic() - started) * 1000), + url=url, + ) + return target + + +def require_tinygrad() -> Any: + if Tensor is None: + fatal("BackendNotInitialized") + return Tensor + + +def load_weights(cmd: dict[str, Any]) -> None: + global model, tokenizer + TensorCls = require_tinygrad() + started = time.monotonic() + model_id = str(cmd["model_id"]) + if test_mode(): + model = {"test_mode": True} + tokenizer = {"test_mode": True} + loaded.clear() + loaded.update( + model_id=model_id, + path="mvp-tinygrad-test-mode", + layer_start=int(cmd.get("layer_start", 0)), + layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), + ) + control( + type="WeightsLoaded", + model_id=model_id, + path=loaded["path"], + test_mode=True, + elapsed_ms=int((time.monotonic() - started) * 1000), + ) + return + source = cmd["gguf_source"] + path = fetch_whole(source) + try: + from tinygrad.apps.llm import SimpleTokenizer, Transformer + + max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512") + max_context = int(max_context_raw) if max_context_raw else 512 + model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True) + tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None}) + if "EmbeddedGguf" in tok_src: + tokenizer = SimpleTokenizer.from_gguf_kv(kv) + else: + fatal("UnsupportedTokenizerSource", tokenizer=tok_src) + except SystemExit: + raise + except Exception as exc: + tb = traceback.format_exc() + print(tb, file=sys.stderr, flush=True) + fatal("ModelLoadFailed", error=str(exc), traceback=tb) + loaded.clear() + loaded.update( + model_id=model_id, + path=str(path), + layer_start=int(cmd.get("layer_start", 0)), + layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)), + ) + control( + type="WeightsLoaded", + model_id=model_id, + path=str(path), + elapsed_ms=int((time.monotonic() - started) * 1000), + ) + + +def infer_prompt(cmd: dict[str, Any]) -> None: + if model is None or tokenizer is None: + fatal("WeightsNotLoaded") + prompt = str(cmd.get("prompt", "")) + max_tokens = int(cmd.get("max_tokens", 1)) + started = time.monotonic() + if test_mode(): + text = f"mvp-test response: {prompt}" + control( + type="PromptCompleted", + model_id=loaded.get("model_id"), + prompt_tokens=[], + generated_tokens=list(range(min(max_tokens, 3))), + text=text, + test_mode=True, + elapsed_ms=int((time.monotonic() - started) * 1000), + ) + return + prompt_tokens = tokenizer.encode(prompt) + generated: list[int] = [] + for token in model.generate(list(prompt_tokens)): + generated.append(int(token)) + if len(generated) >= max_tokens: + break + text = tokenizer.decode(generated) if generated else "" + control( + type="PromptCompleted", + model_id=loaded.get("model_id"), + prompt_tokens=prompt_tokens, + generated_tokens=generated, + text=text, + elapsed_ms=int((time.monotonic() - started) * 1000), + ) + + +def shutdown_worker(_: dict[str, Any]) -> None: + control(type="WorkerStopped", reason="Graceful") + raise SystemExit(0) + + +HANDLERS = { + "InitializeWorker": initialize, + "ConfigureRole": configure_role, + "LoadWeights": load_weights, + "InferPrompt": infer_prompt, + "ShutdownWorker": shutdown_worker, +} + +for raw in sys.stdin: + if not raw.strip(): + continue + try: + command = json.loads(raw) + handler = HANDLERS.get(command.get("type")) + if handler is None: + fatal("UnknownCommand", command=command.get("type")) + handler(command) + except SystemExit: + raise + except Exception as exc: + tb = traceback.format_exc() + print(tb, file=sys.stderr, flush=True) + fatal("UnhandledWorkerException", error=str(exc), traceback=tb) diff --git a/crates/dashboard/src/lib.rs b/crates/dashboard/src/lib.rs index 581381b..98228ba 100644 --- a/crates/dashboard/src/lib.rs +++ b/crates/dashboard/src/lib.rs @@ -1,6 +1,6 @@ +mod live_explorer; mod server; mod store; -mod live_explorer; pub mod swactor; pub mod view; diff --git a/crates/mvp-system/Cargo.toml b/crates/mvp-system/Cargo.toml index 2612fb8..727900e 100644 --- a/crates/mvp-system/Cargo.toml +++ b/crates/mvp-system/Cargo.toml @@ -19,10 +19,29 @@ distribution = { path = "../distribution" } iroh-driver = { path = "../iroh-driver" } iroh = "0.98" tokio = { version = "1", features = ["rt-multi-thread", "macros", "process", "io-util", "sync", "time"] } +swactor-vastai = { path = "../../tools/vastai" } +parking_lot = "0.12" [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2" +[[bin]] +name = "mvp-node" +path = "src/bin/mvp_node.rs" + +[[bin]] +name = "mvp-orch-one-node" +path = "src/bin/mvp_orch_one_node.rs" + +[[bin]] +name = "mvp-chat" +path = "src/bin/mvp_chat.rs" + +[[bin]] +name = "mvp-one-node-chat" +path = "src/bin/mvp_one_node_chat.rs" +required-features = ["local-e2e"] + [[bin]] name = "mvp-local-e2e" path = "src/bin/local_e2e.rs" diff --git a/crates/mvp-system/src/actors/mod.rs b/crates/mvp-system/src/actors/mod.rs index 210aeaa..77baa32 100644 --- a/crates/mvp-system/src/actors/mod.rs +++ b/crates/mvp-system/src/actors/mod.rs @@ -14,4 +14,5 @@ pub fn register_mvp_actor_codecs(registry: &mut swactor_transport::CodecRegistry node_agent::register_codecs(registry); orchestrator::register_codecs(registry); provisioner::register_codecs(registry); + crate::prompt_rpc::register_codecs(registry); } diff --git a/crates/mvp-system/src/actors/node_agent.rs b/crates/mvp-system/src/actors/node_agent.rs index d70b612..6ce510b 100644 --- a/crates/mvp-system/src/actors/node_agent.rs +++ b/crates/mvp-system/src/actors/node_agent.rs @@ -3,7 +3,7 @@ use swactor::actor::{ActorAddress, ActorInterface}; use swactor::runtime::Ctx; use swactor_transport::{CodecRegistry, NetworkMessage}; -use crate::stage_controller as stage; +use crate::{run_plan, stage_controller as stage}; use super::codec::JsonCodec; use super::orchestrator::OrchestratorMsg; @@ -19,7 +19,9 @@ pub struct StageProvisionWire { pub layer_end_exclusive: u32, pub inbound_edge_id: u64, pub outbound_edge_id: u64, - pub weight_artifact: String, + pub model_id: String, + pub gguf_source: run_plan::GgufSource, + pub tokenizer: run_plan::TokenizerSource, } impl StageProvisionWire { @@ -36,7 +38,11 @@ impl StageProvisionWire { }, inbound: stage::EdgeProvision::inbound(stage::EdgeId(self.inbound_edge_id)), outbound: stage::EdgeProvision::outbound(stage::EdgeId(self.outbound_edge_id)), - weight_source: stage::WeightSource::TestArtifact(self.weight_artifact.clone()), + weight_source: stage::WeightSource::new( + self.model_id.clone(), + self.gguf_source.clone(), + self.tokenizer.clone(), + ), } } } @@ -81,6 +87,12 @@ pub enum NodeAgentMsg { Snapshot { reply_to: ActorAddress, }, + InferPrompt { + request_id: u64, + prompt: String, + max_tokens: u32, + reply_to: ActorAddress, + }, } impl NetworkMessage for NodeAgentMsg { @@ -104,7 +116,9 @@ pub enum StageCommandWire { layer_end_exclusive: u32, }, LoadWeights { - artifact: String, + model_id: String, + gguf_source: run_plan::GgufSource, + tokenizer: run_plan::TokenizerSource, layer_start: u32, layer_end_exclusive: u32, }, @@ -156,6 +170,12 @@ pub enum StageLifecycleWire { pub enum NodeAgentReport { Command(StageCommandWire), Lifecycle(StageLifecycleWire), + PromptRequested { + request_id: u64, + prompt: String, + max_tokens: u32, + reply_to: ActorAddress, + }, Snapshot { commands: Vec, events: Vec, @@ -255,6 +275,25 @@ impl NodeAgentActor { run_id: stage::RunId(run_id), }) } + NodeAgentMsg::InferPrompt { + request_id, + prompt, + max_tokens, + reply_to, + } => { + if let Some(report_to) = self.report_to { + let _ = ctx.send( + report_to, + NodeAgentReport::PromptRequested { + request_id, + prompt, + max_tokens, + reply_to, + }, + ); + } + return; + } NodeAgentMsg::Snapshot { reply_to } => { let _ = ctx.send( reply_to, @@ -365,9 +404,9 @@ impl From<&stage::StageCommand> for StageCommandWire { layer_end_exclusive: layer_range.end_exclusive, }, stage::StageCommand::LoadWeights { source, range } => Self::LoadWeights { - artifact: match source { - stage::WeightSource::TestArtifact(value) => value.clone(), - }, + model_id: source.model_id.clone(), + gguf_source: source.gguf_source.clone(), + tokenizer: source.tokenizer.clone(), layer_start: range.start, layer_end_exclusive: range.end_exclusive, }, diff --git a/crates/mvp-system/src/bin/local_e2e.rs b/crates/mvp-system/src/bin/local_e2e.rs index 24c0a75..4411236 100644 --- a/crates/mvp-system/src/bin/local_e2e.rs +++ b/crates/mvp-system/src/bin/local_e2e.rs @@ -3,6 +3,7 @@ use std::io::{BufRead, BufReader, Write}; use std::net::{SocketAddr, TcpListener, TcpStream}; use std::path::PathBuf; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc::{self, Receiver}; use std::thread; use std::time::{Duration, Instant}; @@ -36,7 +37,10 @@ const NODE0_LOGICAL_ID: u64 = 11; const NODE1_LOGICAL_ID: u64 = 12; const MAX_TOKENS: u64 = 1; +static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); + fn main() -> ExitCode { + install_signal_handlers(); let args = std::env::args().collect::>(); let result = if args.iter().any(|arg| arg == "--role=node") { run_node_role(&args) @@ -228,7 +232,7 @@ fn run_supervisor_once(run_id: u64, print_summary: bool) -> Result<(), String> { let mut token_in_object_allocator = edge_actor::ObjectIdAllocator::new(edge_actor::EdgeId(stage0.inbound_edge.0)); - while started_at.elapsed() < Duration::from_secs(30) { + while !STOP_REQUESTED.load(Ordering::SeqCst) && started_at.elapsed() < Duration::from_secs(30) { pump_network(&mut driver, &stack); stage_ready_count += drain_node_stdout(&node0.stdout_rx); stage_ready_count += drain_node_stdout(&node1.stdout_rx); @@ -380,6 +384,9 @@ fn run_supervisor_once(run_id: u64, print_summary: bool) -> Result<(), String> { shutdown_node(&mut node0); shutdown_node(&mut node1); + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Err("interrupted".to_owned()); + } Err(format!( "timed out: injected={injected} completed={completed} torn_down={torn_down} stop0={sent_stop_to_node0} stop1={sent_stop_to_node1}" )) @@ -480,7 +487,7 @@ fn run_node_role(args: &[String]) -> Result<(), String> { let mut outbound_object_allocator: Option = None; let started_at = Instant::now(); - loop { + while !STOP_REQUESTED.load(Ordering::SeqCst) { if let Ok(value) = stdin_rx.try_recv() { if value.get("type").and_then(|value| value.as_str()) == Some("shutdown") { let _ = worker.shutdown(); @@ -504,7 +511,7 @@ fn run_node_role(args: &[String]) -> Result<(), String> { ); let _ = std::io::stdout().flush(); } - NodeAgentReport::Snapshot { .. } => {} + NodeAgentReport::PromptRequested { .. } | NodeAgentReport::Snapshot { .. } => {} } } @@ -533,10 +540,15 @@ fn run_node_role(args: &[String]) -> Result<(), String> { pending_commands = deferred; if started_at.elapsed() > Duration::from_secs(60) { - return Err("node role timed out".to_owned()); + break; } thread::sleep(Duration::from_millis(10)); } + let _ = worker.shutdown(); + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Err("interrupted".to_owned()); + } + Err("node role timed out".to_owned()) } fn new_driver(handle: tokio::runtime::Handle) -> Result { @@ -673,7 +685,9 @@ fn stage_provision_wire( layer_end_exclusive: provision.layer_end_exclusive, inbound_edge_id: provision.inbound.edge_id.0, outbound_edge_id: provision.outbound.edge_id.0, - weight_artifact: "local-e2e-fixture".to_owned(), + model_id: provision.model.model_id, + gguf_source: provision.gguf_source, + tokenizer: provision.tokenizer, }) } @@ -875,6 +889,7 @@ struct WorkerProc { child: Child, stdin: ChildStdin, stdout: BufReader, + cleaned: bool, } impl WorkerProc { @@ -885,18 +900,25 @@ impl WorkerProc { .stderr(Stdio::null()) .spawn() .map_err(|e| format!("spawn dumb worker: {e}"))?; - let stdin = child - .stdin - .take() - .ok_or_else(|| "worker stdin missing".to_owned())?; - let stdout = child - .stdout - .take() - .ok_or_else(|| "worker stdout missing".to_owned())?; + let stdin = match child.stdin.take() { + Some(stdin) => stdin, + None => { + kill_child(&mut child); + return Err("worker stdin missing".to_owned()); + } + }; + let stdout = match child.stdout.take() { + Some(stdout) => stdout, + None => { + kill_child(&mut child); + return Err("worker stdout missing".to_owned()); + } + }; Ok(Self { child, stdin, stdout: BufReader::new(stdout), + cleaned: false, }) } @@ -915,7 +937,17 @@ impl WorkerProc { } fn shutdown(&mut self) -> Result<(), String> { - self.command(json!({"type":"ShutdownWorker"}), "WorkerStopped") + let result = self.command(json!({"type":"ShutdownWorker"}), "WorkerStopped"); + self.terminate(); + result + } + + fn terminate(&mut self) { + if self.cleaned { + return; + } + self.cleaned = true; + kill_child(&mut self.child); } fn command(&mut self, command: serde_json::Value, expected: &str) -> Result<(), String> { @@ -939,7 +971,7 @@ impl WorkerProc { impl Drop for WorkerProc { fn drop(&mut self) { - let _ = self.child.try_wait(); + self.terminate(); } } @@ -948,6 +980,7 @@ struct NodeChild { stdin: ChildStdin, stdout_rx: Receiver, ready: NodeReady, + cleaned: bool, } fn spawn_node_process( @@ -978,22 +1011,35 @@ fn spawn_node_process( .spawn() .map_err(|e| format!("spawn node {stage_index}: {e}"))?; - let stdin = child - .stdin - .take() - .ok_or_else(|| "node stdin missing".to_owned())?; - let stdout = child - .stdout - .take() - .ok_or_else(|| "node stdout missing".to_owned())?; + let stdin = match child.stdin.take() { + Some(stdin) => stdin, + None => { + kill_child(&mut child); + return Err("node stdin missing".to_owned()); + } + }; + let stdout = match child.stdout.take() { + Some(stdout) => stdout, + None => { + kill_child(&mut child); + return Err("node stdout missing".to_owned()); + } + }; let mut reader = BufReader::new(stdout); let mut ready_line = String::new(); - reader - .read_line(&mut ready_line) - .map_err(|e| format!("read node ready: {e}"))?; - let ready: NodeReady = serde_json::from_str(&ready_line) - .map_err(|e| format!("parse node ready {ready_line:?}: {e}"))?; + if let Err(error) = reader.read_line(&mut ready_line) { + kill_child(&mut child); + return Err(format!("read node ready: {error}")); + } + let ready: NodeReady = match serde_json::from_str(&ready_line) { + Ok(ready) => ready, + Err(error) => { + kill_child(&mut child); + return Err(format!("parse node ready {ready_line:?}: {error}")); + } + }; if ready.kind != "ready" { + kill_child(&mut child); return Err(format!("node first line was not ready: {ready_line}")); } @@ -1011,10 +1057,21 @@ fn spawn_node_process( stdin, stdout_rx: rx, ready, + cleaned: false, }) } +impl Drop for NodeChild { + fn drop(&mut self) { + shutdown_node(self); + } +} + fn shutdown_node(node: &mut NodeChild) { + if node.cleaned { + return; + } + node.cleaned = true; let _ = writeln!(node.stdin, "{}", json!({"type":"shutdown"})); let _ = node.stdin.flush(); let started = Instant::now(); @@ -1024,8 +1081,7 @@ fn shutdown_node(node: &mut NodeChild) { } thread::sleep(Duration::from_millis(20)); } - let _ = node.child.kill(); - let _ = node.child.wait(); + kill_child(&mut node.child); } fn drain_node_stdout(rx: &Receiver) -> usize { @@ -1044,6 +1100,26 @@ fn drain_node_stdout(rx: &Receiver) -> usize { stage_ready_count } +fn kill_child(child: &mut Child) { + if !matches!(child.try_wait(), Ok(Some(_))) { + let _ = child.kill(); + } + let _ = child.wait(); +} + +extern "C" fn request_stop(_: libc::c_int) { + STOP_REQUESTED.store(true, Ordering::SeqCst); +} + +fn install_signal_handlers() { + #[cfg(target_os = "linux")] + unsafe { + libc::signal(libc::SIGINT, request_stop as *const () as usize); + libc::signal(libc::SIGTERM, request_stop as *const () as usize); + libc::signal(libc::SIGHUP, request_stop as *const () as usize); + } +} + fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> { let index = args .iter() diff --git a/crates/mvp-system/src/bin/local_e2e_cluster.rs b/crates/mvp-system/src/bin/local_e2e_cluster.rs index 23b194d..fb3dbcb 100644 --- a/crates/mvp-system/src/bin/local_e2e_cluster.rs +++ b/crates/mvp-system/src/bin/local_e2e_cluster.rs @@ -4,6 +4,7 @@ use std::path::PathBuf; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; use std::sync::{ Arc, + atomic::{AtomicBool, Ordering}, mpsc::{self, Receiver, Sender}, }; use std::thread; @@ -60,6 +61,7 @@ const ARENA_BYTES: usize = 16 * 1024; const RING_BYTES: usize = 4096; const EDGE_READY_TIMEOUT: Duration = Duration::from_secs(20); const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500); +static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); struct MvpDashboard { url: String, @@ -185,6 +187,19 @@ impl MvpDashboard { } } +extern "C" fn request_stop(_: libc::c_int) { + STOP_REQUESTED.store(true, Ordering::SeqCst); +} + +fn install_signal_handlers() { + #[cfg(target_os = "linux")] + unsafe { + libc::signal(libc::SIGINT, request_stop as *const () as usize); + libc::signal(libc::SIGTERM, request_stop as *const () as usize); + libc::signal(libc::SIGHUP, request_stop as *const () as usize); + } +} + fn runtime_snapshot_interval_from_env() -> Result { let Some(value) = std::env::var_os("MVP_RUNTIME_SNAPSHOT_MS") else { return Ok(DEFAULT_RUNTIME_SNAPSHOT_INTERVAL); @@ -197,6 +212,7 @@ fn runtime_snapshot_interval_from_env() -> Result { } fn main() -> ExitCode { + install_signal_handlers(); let args = std::env::args().collect::>(); let result = if args.iter().any(|arg| arg == "--role=node") { run_node_role(&args) @@ -243,6 +259,26 @@ struct ProvisionedDockerNode { provider_process_id: Option, provisioner: LocalDockerNodeProvisioner, events: Receiver, + cleaned: bool, +} + +impl ProvisionedDockerNode { + fn stop(&mut self) -> Result<(), String> { + if self.cleaned { + return Ok(()); + } + self.provisioner + .stop() + .map_err(|e| format!("stop node {}: {e:?}", self.node_id))?; + self.cleaned = true; + Ok(()) + } +} + +impl Drop for ProvisionedDockerNode { + fn drop(&mut self) { + let _ = self.stop(); + } } #[derive(Default)] @@ -420,7 +456,7 @@ fn run_supervisor_dashboard_loop() -> Result<(), String> { "mvp-local-e2e-cluster: MVP_DASHBOARD=1, repeating Docker cluster scenario until Ctrl+C" ); let mut run_id = RUN_ID; - loop { + while !STOP_REQUESTED.load(Ordering::SeqCst) { match run_supervisor_once( run_id, Some(&mut dashboard), @@ -433,6 +469,7 @@ fn run_supervisor_dashboard_loop() -> Result<(), String> { run_id = run_id.saturating_add(1); thread::sleep(Duration::from_secs(1)); } + Ok(()) } fn run_supervisor_once( @@ -628,7 +665,7 @@ fn run_supervisor_once( let mut edge_stream_count = 0usize; let mut token_out_streams = HashMap::>::new(); - while started_at.elapsed() < Duration::from_secs(30) { + while !STOP_REQUESTED.load(Ordering::SeqCst) && started_at.elapsed() < Duration::from_secs(30) { pump_network(&mut driver, &stack); driver_runtime.poll_iroh(&driver); while let Some(event) = driver_runtime.try_recv() { @@ -878,6 +915,17 @@ fn run_supervisor_once( thread::sleep(Duration::from_millis(10)); } + if STOP_REQUESTED.load(Ordering::SeqCst) { + record_dashboard_event(&mut dashboard, run_fault_event(run_id)); + let _ = stop_provisioned_nodes( + &mut [&mut node0, &mut node1], + &mut driver, + &stack, + &mut dashboard, + ); + return Err("interrupted".to_owned()); + } + record_dashboard_event(&mut dashboard, run_fault_event(run_id)); let _ = stop_provisioned_nodes( &mut [&mut node0, &mut node1], @@ -992,7 +1040,7 @@ fn run_node_role(args: &[String]) -> Result<(), String> { let mut inbound_ring_id = None; let mut outbound_ring_id = None; - loop { + while !STOP_REQUESTED.load(Ordering::SeqCst) { pump_network(&mut driver, &stack); driver_runtime.poll_iroh(&driver); while let Some(event) = driver_runtime.try_recv() { @@ -1102,7 +1150,7 @@ fn run_node_role(args: &[String]) -> Result<(), String> { .flush() .map_err(|e| format!("flush lifecycle stdout: {e}"))?; } - NodeAgentReport::Snapshot { .. } => {} + NodeAgentReport::PromptRequested { .. } | NodeAgentReport::Snapshot { .. } => {} } } @@ -1161,7 +1209,7 @@ fn run_node_role(args: &[String]) -> Result<(), String> { )?; } - if shutdown_rx.try_recv().is_ok() { + if shutdown_rx.try_recv().is_ok() || STOP_REQUESTED.load(Ordering::SeqCst) { worker.shutdown().ok(); return Ok(()); } @@ -1171,6 +1219,8 @@ fn run_node_role(args: &[String]) -> Result<(), String> { } thread::sleep(Duration::from_millis(10)); } + worker.shutdown().ok(); + Ok(()) } fn new_driver(handle: tokio::runtime::Handle) -> Result { @@ -1308,7 +1358,9 @@ fn stage_provision_wire( layer_end_exclusive: provision.layer_end_exclusive, inbound_edge_id: provision.inbound.edge_id.0, outbound_edge_id: provision.outbound.edge_id.0, - weight_artifact: "local-e2e-cluster-tinygrad-cpu-fixture".to_owned(), + model_id: provision.model.model_id, + gguf_source: provision.gguf_source, + tokenizer: provision.tokenizer, }) } @@ -1427,9 +1479,24 @@ fn handle_node_command( .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) .map_err(|e| format!("mark worker ready: {e}")) } - StageCommandWire::LoadWeights { .. } => runtime - .send_to(node_actor, NodeAgentMsg::MarkWeightsReady) - .map_err(|e| format!("mark weights ready: {e}")), + StageCommandWire::LoadWeights { + model_id, + gguf_source, + tokenizer, + layer_start, + layer_end_exclusive, + } => { + worker.load_weights( + model_id, + gguf_source, + tokenizer, + layer_start, + layer_end_exclusive, + )?; + runtime + .send_to(node_actor, NodeAgentMsg::MarkWeightsReady) + .map_err(|e| format!("mark weights ready: {e}")) + } StageCommandWire::ExecuteStep { step_id, input_edge_id, @@ -1971,8 +2038,6 @@ impl GpuWorkerRuntime { "stage_index":stage_index, "layer_start":layer_start, "layer_end_exclusive":layer_end_exclusive, - "model_id":"local-e2e-cluster-tinygrad-cpu-fixture", - "gguf_source":"docker-cpu://local-e2e-cluster-tinygrad-cpu-fixture", } }), "RoleConfigured", @@ -1980,6 +2045,28 @@ impl GpuWorkerRuntime { .map(|_| ()) } + fn load_weights( + &mut self, + model_id: String, + gguf_source: plan::GgufSource, + tokenizer: plan::TokenizerSource, + layer_start: u32, + layer_end_exclusive: u32, + ) -> Result<(), String> { + self.command( + json!({ + "type":"LoadWeights", + "model_id":model_id, + "gguf_source":gguf_source, + "tokenizer":tokenizer, + "layer_start":layer_start, + "layer_end_exclusive":layer_end_exclusive, + }), + "WeightsLoaded", + ) + .map(|_| ()) + } + fn ring_readable(&mut self, ring_id: u64, edge_id: u64) -> Result { let value = self.command( json!({"type":"RingReadable","ring_id":ring_id}), @@ -2515,7 +2602,7 @@ fn provision_local_docker_node( let provider_process_id = provisioner.provider().cli().provider_process_id(); let started = Instant::now(); - while started.elapsed() < Duration::from_secs(30) { + while !STOP_REQUESTED.load(Ordering::SeqCst) && started.elapsed() < Duration::from_secs(30) { pump_network(driver, stack); drain_dashboard(dashboard); publish_runtime_snapshot_throttled(dashboard, stack); @@ -2548,6 +2635,7 @@ fn provision_local_docker_node( provider_process_id, provisioner, events: events_rx, + cleaned: false, }); } } @@ -2569,7 +2657,11 @@ fn provision_local_docker_node( } thread::sleep(Duration::from_millis(10)); } - Err(format!("timed out provisioning node {expected_node_id}")) + if STOP_REQUESTED.load(Ordering::SeqCst) { + Err(format!("interrupted provisioning node {expected_node_id}")) + } else { + Err(format!("timed out provisioning node {expected_node_id}")) + } } fn ready_node_from_stdout( @@ -2690,15 +2782,24 @@ fn stop_provisioned_nodes( stack: &DistributionRuntimeStack, dashboard: &mut Option<&mut MvpDashboard>, ) -> Result<(), String> { + let mut first_error = None; for node in nodes.iter_mut().rev() { - node.provisioner - .stop() - .map_err(|e| format!("stop node {}: {e:?}", node.node_id))?; + if first_error.is_none() { + if let Err(error) = node.stop() { + first_error = Some(error); + } + } else { + let _ = node.stop(); + } } pump_network(driver, stack); drain_dashboard(dashboard); publish_runtime_snapshot(dashboard, stack); - Ok(()) + if let Some(error) = first_error { + Err(error) + } else { + Ok(()) + } } fn spawn_send_pump( diff --git a/crates/mvp-system/src/bin/mvp_chat.rs b/crates/mvp-system/src/bin/mvp_chat.rs new file mode 100644 index 0000000..154e11f --- /dev/null +++ b/crates/mvp-system/src/bin/mvp_chat.rs @@ -0,0 +1,174 @@ +use std::io::{self, BufRead, BufReader, Write}; +use std::net::TcpStream; +use std::process::ExitCode; + +use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, write_json_line}; + +const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; +const DEFAULT_MAX_TOKENS: u32 = 64; +const DEFAULT_TIMEOUT_MS: u64 = 120_000; + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("mvp-chat: {error}"); + ExitCode::from(1) + } + } +} + +fn run() -> Result<(), String> { + let config = Config::from_env_and_args()?; + let mut stream = TcpStream::connect(&config.addr) + .map_err(|e| format!("connect prompt RPC {}: {e}", config.addr))?; + let mut reader = BufReader::new( + stream + .try_clone() + .map_err(|e| format!("clone prompt RPC stream: {e}"))?, + ); + let stdin = io::stdin(); + let mut next_request_id = 1_u64; + + eprintln!("mvp-chat: connected to {}", config.addr); + loop { + print!("> "); + io::stdout() + .flush() + .map_err(|e| format!("flush prompt: {e}"))?; + let mut prompt = String::new(); + let n = stdin + .read_line(&mut prompt) + .map_err(|e| format!("read stdin: {e}"))?; + if n == 0 { + return Ok(()); + } + let prompt = prompt.trim_end().to_owned(); + if prompt.eq_ignore_ascii_case("/quit") || prompt.eq_ignore_ascii_case("/exit") { + return Ok(()); + } + if prompt.trim().is_empty() { + continue; + } + + let request_id = next_request_id; + next_request_id = next_request_id.wrapping_add(1).max(1); + let request = SubmitPrompt { + request_id, + prompt_text: prompt, + max_tokens: config.max_tokens, + timeout_ms: config.timeout_ms, + }; + write_json_line(&mut stream, &request)?; + + loop { + let mut line = String::new(); + let n = reader + .read_line(&mut line) + .map_err(|e| format!("read prompt event: {e}"))?; + if n == 0 { + return Err("prompt RPC closed".to_owned()); + } + let event = serde_json::from_str::(&line) + .map_err(|e| format!("parse prompt event: {e}"))?; + match event { + PromptEvent::TextDelta { + request_id: seen, + text, + } if seen == request_id => { + print!("{text}"); + io::stdout() + .flush() + .map_err(|e| format!("flush text delta: {e}"))?; + } + PromptEvent::Done { + request_id: seen, + tokens_generated, + elapsed_ms, + .. + } if seen == request_id => { + println!(); + eprintln!( + "mvp-chat: done request={} tokens={} elapsed_ms={}", + seen, tokens_generated, elapsed_ms + ); + break; + } + PromptEvent::Fault { + request_id: seen, + error, + } if seen == request_id => { + eprintln!("mvp-chat: fault request={seen}: {error}"); + break; + } + _ => {} + } + } + } +} + +struct Config { + addr: String, + max_tokens: u32, + timeout_ms: u64, +} + +impl Config { + fn from_env_and_args() -> Result { + let mut config = Self { + addr: std::env::var("MVP_PROMPT_RPC_ADDR") + .unwrap_or_else(|_| DEFAULT_RPC_ADDR.to_owned()), + max_tokens: env_u32("MVP_PROMPT_MAX_TOKENS", DEFAULT_MAX_TOKENS)?, + timeout_ms: env_u64("MVP_PROMPT_TIMEOUT_MS", DEFAULT_TIMEOUT_MS)?, + }; + let mut args = std::env::args().skip(1); + while let Some(arg) = args.next() { + match arg.as_str() { + "--addr" => { + config.addr = args + .next() + .ok_or_else(|| "missing value after --addr".to_owned())? + } + "--max-tokens" => { + config.max_tokens = parse_next(&mut args, "--max-tokens")?; + } + "--timeout-ms" => { + config.timeout_ms = parse_next(&mut args, "--timeout-ms")?; + } + other => return Err(format!("unknown argument {other:?}")), + } + } + Ok(config) + } +} + +fn env_u64(name: &str, default: u64) -> Result { + match std::env::var(name).ok().filter(|value| !value.is_empty()) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn env_u32(name: &str, default: u32) -> Result { + match std::env::var(name).ok().filter(|value| !value.is_empty()) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn parse_next(args: &mut impl Iterator, name: &str) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + let value = args + .next() + .ok_or_else(|| format!("missing value after {name}"))?; + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) +} diff --git a/crates/mvp-system/src/bin/mvp_node.rs b/crates/mvp-system/src/bin/mvp_node.rs new file mode 100644 index 0000000..58910a2 --- /dev/null +++ b/crates/mvp-system/src/bin/mvp_node.rs @@ -0,0 +1,794 @@ +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; +use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; +use std::sync::{ + Arc, OnceLock, + mpsc::{self, Receiver}, +}; +use std::thread; +use std::time::{Duration, Instant}; + +use datastream::emit::{ClusterFrameSink, DatastreamEmitter, EmitterConfig, FrameSink, NoopSink}; +use datastream::{ChannelId, DATASTREAM_SINK_NAME}; + +use distribution::node::DistributedNodeConfig; +use iroh::EndpointAddr; +use iroh_driver::{IrohDriver, IrohDriverConfig}; +use mvp_system::actors::node_agent::{ + NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, +}; +use mvp_system::actors::register_mvp_actor_codecs; +use mvp_system::distribution_stack::DistributionRuntimeStack; +use mvp_system::prompt_rpc::PromptEvent; +use mvp_system::run_plan::{GgufSource, TokenizerSource}; +use mvp_system::stage_controller as stage; +use serde_json::{Value, json}; +use swactor::actor::ActorAddress; + +const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py"; +const DEFAULT_DEVICE: &str = "CUDA"; +const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF"; +const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; +const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; +const PUMP_INTERVAL: Duration = Duration::from_millis(10); + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("mvp-node: {error}"); + ExitCode::from(1) + } + } +} + +fn run() -> Result<(), String> { + let config = DeploymentConfig::from_env()?; + eprintln!( + "mvp-node: boot run={} logical_node={} stage={} worker={} device={}", + config.run_id, + config.logical_node_id, + config.stage_index, + config.worker_script, + config.device + ); + + let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; + let mut driver = IrohDriver::with_handle( + tokio.handle().clone(), + IrohDriverConfig { + secret_key: None, + relay_mode: config.relay_mode.clone(), + node: DistributedNodeConfig::default(), + peer_auth: None, + additional_alpns: vec![], + }, + ) + .map_err(|e| format!("create iroh driver: {e}"))?; + if let Some(coordinator) = &config.coordinator_endpoint { + eprintln!("mvp-node: joining coordinator {coordinator:?}"); + driver.join(std::slice::from_ref(coordinator)); + } else { + eprintln!("mvp-node: no MVP_COORDINATOR_ENDPOINT set; running standalone until joined"); + } + + let stack = DistributionRuntimeStack::new_with_codecs( + driver.node_id(), + DistributedNodeConfig::default(), + |registry| { + register_mvp_actor_codecs(registry); + datastream::wire::register_datastream_codec(registry); + }, + ); + driver.enable_actor_bridge( + stack.runtime.clone(), + stack.codec.clone(), + stack.actor_bridge_routes(), + stack.actors.swim, + stack.relay_mirror.clone(), + stack.route_view.clone(), + ); + + let mut datastream = node_datastream(&config, &stack); + + let reports = stack + .runtime + .new_inbox::() + .map_err(|e| format!("node report inbox: {e}"))?; + let orchestrator = config + .orchestrator_actor + .unwrap_or_else(ActorAddress::new_random); + let node_actor = stack + .runtime + .spawn(NodeAgentActor::new( + stage::NodeId(config.logical_node_id), + orchestrator, + Some(*reports.addr()), + )) + .map_err(|e| format!("spawn node agent: {e}"))?; + stack.register_local_actor(driver.register_actor(node_actor, 1)); + + let mut worker = TinygradWorker::spawn(&config)?; + let mut initial_pump = || {}; + worker.initialize(&config.device, &mut datastream, &mut initial_pump)?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) + .map_err(|e| format!("mark initialized worker ready: {e}"))?; + + let ready = json!({ + "type":"ready", + "role":"node", + "endpoint": driver.endpoint_addr(), + "node_actor": node_actor, + "logical_node_id": config.logical_node_id, + "stage_index": config.stage_index, + }); + println!("{ready}"); + datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string()); + std::io::stdout() + .flush() + .map_err(|e| format!("flush ready line: {e}"))?; + + if let Some(prompt) = &config.self_test_prompt { + run_self_test(&mut worker, &config, prompt, &mut datastream)?; + } + + let shutdown_rx = spawn_stdin_shutdown_listener(); + let started = Instant::now(); + loop { + pump_network(&mut driver, &stack); + datastream.tick(); + while let Some(report) = reports.try_recv() { + handle_node_report( + report, + &stack, + &mut driver, + node_actor, + &mut worker, + &mut datastream, + )?; + } + if shutdown_rx.try_recv().is_ok() { + eprintln!("mvp-node: shutdown requested on stdin"); + let mut pump = || pump_network(&mut driver, &stack); + let _ = worker.shutdown(&mut datastream, &mut pump); + return Ok(()); + } + if let Some(status) = worker.try_wait()? { + let _ = stack + .runtime + .send_to(node_actor, NodeAgentMsg::WorkerCrashed); + return Err(format!("tinygrad helper exited with {status}")); + } + if started.elapsed() > config.max_runtime && config.max_runtime != Duration::ZERO { + eprintln!("mvp-node: max runtime elapsed; shutting down cleanly"); + let mut pump = || pump_network(&mut driver, &stack); + let _ = worker.shutdown(&mut datastream, &mut pump); + return Ok(()); + } + thread::sleep(PUMP_INTERVAL); + } +} + +fn pump_network(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) { + stack.tick_protocol_actors(Instant::now()); + driver.pump_inbound_to_actors(); + stack.pump_runtime_once(); + driver.drain_outbox(&stack.outbox); +} + +fn node_datastream( + config: &DeploymentConfig, + stack: &DistributionRuntimeStack, +) -> DatastreamEmitter { + let mut sinks: Vec> = Vec::new(); + if let Some(actor) = config.datastream_sink_actor { + let sink_addr = Arc::new(OnceLock::new()); + let _ = sink_addr.set(actor); + eprintln!("mvp-node: native datastream targeting {DATASTREAM_SINK_NAME} actor {actor:?}"); + sinks.push(Box::new(ClusterFrameSink::new( + stack.runtime.clone(), + sink_addr, + ))); + } + if let Some(path) = &config.datastream_frame_log { + match JsonlFrameSink::open(path) { + Ok(sink) => { + eprintln!("mvp-node: native datastream frame log {path}"); + sinks.push(Box::new(sink)); + } + Err(error) => { + eprintln!("mvp-node: failed to open MVP_DATASTREAM_FRAME_LOG={path:?}: {error}"); + } + } + } + if sinks.is_empty() { + eprintln!("mvp-node: no datastream sink configured; native datastream drops locally"); + sinks.push(Box::new(NoopSink)); + } + let sink: Box = if sinks.len() == 1 { + sinks.pop().expect("one sink") + } else { + Box::new(TeeFrameSink { sinks }) + }; + DatastreamEmitter::new( + EmitterConfig { + node_hex: config.logical_node_id.to_string(), + life: config.run_id, + mux_capacity: 256, + }, + sink, + ) +} + +struct TeeFrameSink { + sinks: Vec>, +} + +impl FrameSink for TeeFrameSink { + fn ship(&mut self, stream: &datastream::StreamId, frame: &datastream::Frame) { + for sink in &mut self.sinks { + sink.ship(stream, frame); + } + } +} + +struct JsonlFrameSink { + file: File, +} + +impl JsonlFrameSink { + fn open(path: &str) -> std::io::Result { + Ok(Self { + file: OpenOptions::new().create(true).append(true).open(path)?, + }) + } +} + +impl FrameSink for JsonlFrameSink { + fn ship(&mut self, stream: &datastream::StreamId, frame: &datastream::Frame) { + let record = json!({ + "stream":stream.to_string(), + "channel":frame.channel.as_str(), + "position":frame.position.0, + "payload":String::from_utf8_lossy(&frame.payload), + }); + let _ = serde_json::to_writer(&mut self.file, &record); + let _ = writeln!(self.file); + let _ = self.file.flush(); + } +} + +fn handle_node_report( + report: NodeAgentReport, + stack: &DistributionRuntimeStack, + driver: &mut IrohDriver, + node_actor: ActorAddress, + worker: &mut TinygradWorker, + datastream: &mut DatastreamEmitter, +) -> Result<(), String> { + match report { + NodeAgentReport::Command(command) => { + handle_stage_command(command, stack, driver, node_actor, worker, datastream) + } + NodeAgentReport::Lifecycle(event) => { + let record = json!({"type":"node_lifecycle","event":format!("{event:?}")}); + println!("{record}"); + datastream.submit_text(ChannelId::new("mvp.node.lifecycle"), record.to_string()); + std::io::stdout() + .flush() + .map_err(|e| format!("flush lifecycle: {e}")) + } + NodeAgentReport::PromptRequested { + request_id, + prompt, + max_tokens, + reply_to, + } => handle_prompt_request( + request_id, prompt, max_tokens, reply_to, stack, driver, worker, datastream, + ), + NodeAgentReport::Snapshot { .. } => Ok(()), + } +} + +fn handle_prompt_request( + request_id: u64, + prompt: String, + max_tokens: u32, + reply_to: ActorAddress, + stack: &DistributionRuntimeStack, + driver: &mut IrohDriver, + worker: &mut TinygradWorker, + datastream: &mut DatastreamEmitter, +) -> Result<(), String> { + let started = Instant::now(); + let mut pump = || pump_network(driver, stack); + match worker.infer_prompt(&prompt, max_tokens, datastream, &mut pump) { + Ok(result) => { + let text = result + .get("text") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(); + let tokens_generated = result + .get("generated_tokens") + .and_then(Value::as_array) + .map_or(0, |tokens| tokens.len() as u32); + if !text.is_empty() { + stack + .runtime + .send_to( + reply_to, + PromptEvent::TextDelta { + request_id, + text: text.clone(), + }, + ) + .map_err(|e| format!("send prompt text delta: {e}"))?; + } + stack + .runtime + .send_to( + reply_to, + PromptEvent::Done { + request_id, + final_text: text, + tokens_generated, + elapsed_ms: result + .get("elapsed_ms") + .and_then(Value::as_u64) + .unwrap_or_else(|| started.elapsed().as_millis() as u64), + }, + ) + .map_err(|e| format!("send prompt done: {e}")) + } + Err(error) => stack + .runtime + .send_to(reply_to, PromptEvent::Fault { request_id, error }) + .map_err(|e| format!("send prompt fault: {e}")), + } +} + +fn handle_stage_command( + command: StageCommandWire, + stack: &DistributionRuntimeStack, + driver: &mut IrohDriver, + node_actor: ActorAddress, + worker: &mut TinygradWorker, + datastream: &mut DatastreamEmitter, +) -> Result<(), String> { + match command { + StageCommandWire::ConfigureWorkerRole { + run_id, + stage_index, + layer_start, + layer_end_exclusive, + } => { + let mut pump = || pump_network(driver, stack); + worker.configure_role( + run_id, + stage_index, + layer_start, + layer_end_exclusive, + datastream, + &mut pump, + )?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) + .map_err(|e| format!("mark worker ready: {e}")) + } + StageCommandWire::LoadWeights { + model_id, + gguf_source, + tokenizer, + layer_start, + layer_end_exclusive, + } => { + let mut pump = || pump_network(driver, stack); + worker.load_weights( + model_id, + gguf_source, + tokenizer, + layer_start, + layer_end_exclusive, + datastream, + &mut pump, + )?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkWeightsReady) + .map_err(|e| format!("mark weights ready: {e}")) + } + StageCommandWire::StopLocalEdges { run_id } => { + stack + .runtime + .send_to(node_actor, NodeAgentMsg::LocalEdgesStopped { run_id }) + .map_err(|e| format!("mark local edges stopped: {e}"))?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id }) + .map_err(|e| format!("mark worker rings quiesced: {e}")) + } + StageCommandWire::ReleaseRunDeviceObjects { run_id } => { + stack + .runtime + .send_to(node_actor, NodeAgentMsg::DeviceObjectsReleased { run_id }) + .map_err(|e| format!("mark device objects released: {e}"))?; + stack + .runtime + .send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id }) + .map_err(|e| format!("mark worker role reset: {e}")) + } + StageCommandWire::EstablishInboundEdge { edge_id } => stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id }) + .map_err(|e| format!("mark inbound edge ready: {e}")), + StageCommandWire::EstablishOutboundEdge { edge_id } => stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id }) + .map_err(|e| format!("mark outbound edge ready: {e}")), + StageCommandWire::RewireEdge { .. } | StageCommandWire::ReleaseInputHandle { .. } => Ok(()), + StageCommandWire::ExecuteStep { .. } => Err(format!( + "mvp-node image path does not carry ring payload commands yet: {command:?}" + )), + } +} + +fn run_self_test( + worker: &mut TinygradWorker, + config: &DeploymentConfig, + prompt: &str, + datastream: &mut DatastreamEmitter, +) -> Result<(), String> { + eprintln!("mvp-node: running self-test prompt"); + let mut pump = || {}; + worker.configure_role( + config.run_id, + config.stage_index, + 0, + config.self_test_layer_end, + datastream, + &mut pump, + )?; + worker.load_weights( + config.model_id.clone(), + config.gguf_source.clone(), + config.tokenizer.clone(), + 0, + config.self_test_layer_end, + datastream, + &mut pump, + )?; + let result = worker.infer_prompt(prompt, config.self_test_max_tokens, datastream, &mut pump)?; + let record = json!({ + "type":"self_test_completed", + "prompt":prompt, + "result":result, + }); + println!("{record}"); + datastream.submit_text(ChannelId::new("mvp.node.self_test"), record.to_string()); + std::io::stdout() + .flush() + .map_err(|e| format!("flush self-test: {e}")) +} + +#[derive(Clone)] +struct DeploymentConfig { + run_id: u64, + logical_node_id: u64, + stage_index: u32, + coordinator_endpoint: Option, + orchestrator_actor: Option, + datastream_sink_actor: Option, + datastream_frame_log: Option, + relay_mode: iroh::RelayMode, + worker_script: String, + device: String, + model_id: String, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + self_test_prompt: Option, + self_test_layer_end: u32, + self_test_max_tokens: u32, + max_runtime: Duration, +} + +impl DeploymentConfig { + fn from_env() -> Result { + Ok(Self { + run_id: env_u64("MVP_RUN_ID", 1)?, + logical_node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, + stage_index: env_u32("MVP_STAGE_INDEX", 0)?, + coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, + orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, + datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?, + datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), + relay_mode: relay_mode_from_env()?, + worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), + device: env_string("DEV", DEFAULT_DEVICE), + model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), + gguf_source: gguf_source_from_env(), + tokenizer: tokenizer_from_env(), + self_test_prompt: env_optional("MVP_NODE_SELF_TEST_PROMPT"), + self_test_layer_end: env_u32("MVP_SELF_TEST_LAYER_END", 16)?, + self_test_max_tokens: env_u32("MVP_SELF_TEST_MAX_TOKENS", 1)?, + max_runtime: Duration::from_secs(env_u64("MVP_NODE_MAX_RUNTIME_SECS", 0)?), + }) + } +} + +struct TinygradWorker { + child: Child, + stdin: ChildStdin, + stdout: BufReader, +} + +impl TinygradWorker { + fn spawn(config: &DeploymentConfig) -> Result { + let mut child = Command::new("python3") + .arg(&config.worker_script) + .env("DEV", &config.device) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .map_err(|e| format!("spawn tinygrad helper {}: {e}", config.worker_script))?; + let stdin = child + .stdin + .take() + .ok_or_else(|| "tinygrad helper stdin missing".to_owned())?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "tinygrad helper stdout missing".to_owned())?; + Ok(Self { + child, + stdin, + stdout: BufReader::new(stdout), + }) + } + + fn initialize( + &mut self, + device: &str, + datastream: &mut DatastreamEmitter, + pump: &mut dyn FnMut(), + ) -> Result<(), String> { + self.command( + json!({"type":"InitializeWorker","helper_abi_version":1,"backend":{"device":device}}), + "WorkerReady", + datastream, + ChannelId::new("mvp.worker.initialize"), + pump, + ) + .map(|_| ()) + } + + fn configure_role( + &mut self, + run_id: u64, + stage_index: u32, + layer_start: u32, + layer_end_exclusive: u32, + datastream: &mut DatastreamEmitter, + pump: &mut dyn FnMut(), + ) -> Result<(), String> { + self.command( + json!({ + "type":"ConfigureRole", + "role_id":stage_index + 1, + "config":{ + "run_id":run_id, + "stage_index":stage_index, + "layer_start":layer_start, + "layer_end_exclusive":layer_end_exclusive, + } + }), + "RoleConfigured", + datastream, + ChannelId::new("mvp.worker.role"), + pump, + ) + .map(|_| ()) + } + + fn load_weights( + &mut self, + model_id: String, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + layer_start: u32, + layer_end_exclusive: u32, + datastream: &mut DatastreamEmitter, + pump: &mut dyn FnMut(), + ) -> Result<(), String> { + self.command( + json!({ + "type":"LoadWeights", + "model_id":model_id, + "gguf_source":gguf_source, + "tokenizer":tokenizer, + "layer_start":layer_start, + "layer_end_exclusive":layer_end_exclusive, + }), + "WeightsLoaded", + datastream, + ChannelId::new("mvp.worker.weights"), + pump, + ) + .map(|_| ()) + } + + fn infer_prompt( + &mut self, + prompt: &str, + max_tokens: u32, + datastream: &mut DatastreamEmitter, + pump: &mut dyn FnMut(), + ) -> Result { + self.command( + json!({"type":"InferPrompt","prompt":prompt,"max_tokens":max_tokens}), + "PromptCompleted", + datastream, + ChannelId::new("mvp.worker.prompt"), + pump, + ) + } + + fn shutdown( + &mut self, + datastream: &mut DatastreamEmitter, + pump: &mut dyn FnMut(), + ) -> Result<(), String> { + self.command( + json!({"type":"ShutdownWorker"}), + "WorkerStopped", + datastream, + ChannelId::new("mvp.worker.shutdown"), + pump, + ) + .map(|_| ()) + } + + fn try_wait(&mut self) -> Result, String> { + self.child + .try_wait() + .map_err(|e| format!("poll tinygrad helper: {e}")) + } + + fn command( + &mut self, + command: Value, + expected: &str, + datastream: &mut DatastreamEmitter, + channel: ChannelId, + pump: &mut dyn FnMut(), + ) -> Result { + writeln!(self.stdin, "{command}").map_err(|e| format!("write helper command: {e}"))?; + self.stdin + .flush() + .map_err(|e| format!("flush helper command: {e}"))?; + self.expect_event(expected, datastream, channel, pump) + } + + fn expect_event( + &mut self, + expected: &str, + datastream: &mut DatastreamEmitter, + channel: ChannelId, + pump: &mut dyn FnMut(), + ) -> Result { + loop { + let mut line = String::new(); + let n = self + .stdout + .read_line(&mut line) + .map_err(|e| format!("read helper stdout: {e}"))?; + if n == 0 { + return Err(format!( + "tinygrad helper stdout closed while waiting for {expected}" + )); + } + let value: Value = serde_json::from_str(&line) + .map_err(|e| format!("parse helper stdout {line:?}: {e}"))?; + datastream.submit_text(channel.clone(), value.to_string()); + datastream.tick(); + pump(); + if value.get("type").and_then(Value::as_str) == Some(expected) { + return Ok(value); + } + println!("{}", json!({"type":"tinygrad_event","event":value})); + std::io::stdout() + .flush() + .map_err(|e| format!("flush helper passthrough: {e}"))?; + } + } +} + +impl Drop for TinygradWorker { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +fn spawn_stdin_shutdown_listener() -> Receiver<()> { + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + let stdin = std::io::stdin(); + for line in stdin.lock().lines().map_while(Result::ok) { + if line.trim().eq_ignore_ascii_case("shutdown") { + let _ = tx.send(()); + break; + } + } + }); + rx +} + +fn env_string(name: &str, default: &str) -> String { + env_optional(name).unwrap_or_else(|| default.to_owned()) +} + +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn env_u64(name: &str, default: u64) -> Result { + match env_optional(name) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn env_u32(name: &str, default: u32) -> Result { + match env_optional(name) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn env_json(name: &str) -> Result, String> +where + T: serde::de::DeserializeOwned, +{ + env_optional(name) + .map(|value| serde_json::from_str(&value).map_err(|e| format!("invalid {name} JSON: {e}"))) + .transpose() +} + +fn relay_mode_from_env() -> Result { + match env_string("MVP_IROH_RELAY_MODE", "disabled").as_str() { + "disabled" => Ok(iroh::RelayMode::Disabled), + "default" => Ok(iroh::RelayMode::Default), + other => Err(format!( + "unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default" + )), + } +} + +fn gguf_source_from_env() -> GgufSource { + if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { + return GgufSource::LocalPath(path); + } + GgufSource::HuggingFaceGguf { + repo: env_string("MVP_GGUF_REPO", DEFAULT_HF_REPO), + file: env_string("MVP_GGUF_FILE", DEFAULT_HF_FILE), + revision: env_optional("MVP_GGUF_REVISION"), + } +} + +fn tokenizer_from_env() -> TokenizerSource { + env_optional("MVP_TOKENIZER_LOCAL_PATH") + .map(TokenizerSource::LocalPath) + .unwrap_or(TokenizerSource::EmbeddedGguf) +} diff --git a/crates/mvp-system/src/bin/mvp_one_node_chat.rs b/crates/mvp-system/src/bin/mvp_one_node_chat.rs new file mode 100644 index 0000000..11848fe --- /dev/null +++ b/crates/mvp-system/src/bin/mvp_one_node_chat.rs @@ -0,0 +1,512 @@ +use std::io::{self, BufRead, BufReader, Write}; +use std::net::TcpStream; +use std::path::PathBuf; +use std::process::{Child, ChildStdin, Command, ExitCode, Stdio}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::mpsc; +use std::thread; +use std::time::{Duration, Instant}; + +use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, write_json_line}; +use serde_json::Value; + +#[cfg(target_os = "linux")] +use std::os::unix::process::CommandExt; + +const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; +const DEFAULT_NODE_IMAGE: &str = "swactor-mvp-node:latest"; +const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; +const DEFAULT_MAX_TOKENS: u32 = 64; +const DEFAULT_TIMEOUT_MS: u64 = 120_000; +const ORCH_READY_TIMEOUT: Duration = Duration::from_secs(1_200); +const ORCH_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30); +const CHAT_READ_TIMEOUT: Duration = Duration::from_millis(100); + +static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); + +fn main() -> ExitCode { + install_signal_handlers(); + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("mvp-one-node-chat: {error}"); + ExitCode::from(1) + } + } +} + +fn run() -> Result<(), String> { + let config = Config::from_args()?; + prepare_runtime(&config)?; + let mut orch = OrchChild::spawn(&config)?; + let rpc_addr = orch.wait_ready(config.rpc_addr.clone())?; + eprintln!("mvp-one-node-chat: prompt loop ready at {rpc_addr}"); + let result = run_chat_loop(&rpc_addr, config.max_tokens, config.timeout_ms); + orch.shutdown(); + result +} + +struct Config { + orch_bin: PathBuf, + orch_args: Vec, + rpc_addr: String, + node_image: String, + max_tokens: u32, + timeout_ms: u64, + dashboard: bool, + build_image: bool, +} + +impl Config { + fn from_args() -> Result { + let mut config = Self { + orch_bin: default_orch_bin()?, + orch_args: Vec::new(), + rpc_addr: std::env::var("MVP_PROMPT_RPC_ADDR") + .or_else(|_| std::env::var("MVP_PROMPT_RPC_BIND")) + .unwrap_or_else(|_| DEFAULT_RPC_ADDR.to_owned()), + node_image: std::env::var("MVP_NODE_IMAGE") + .unwrap_or_else(|_| DEFAULT_NODE_IMAGE.to_owned()), + max_tokens: env_u32("MVP_PROMPT_MAX_TOKENS", DEFAULT_MAX_TOKENS)?, + timeout_ms: env_u64("MVP_PROMPT_TIMEOUT_MS", DEFAULT_TIMEOUT_MS)?, + dashboard: env_bool("MVP_DASHBOARD", true)?, + build_image: env_bool("MVP_BUILD_NODE_IMAGE", true)?, + }; + + let mut args = std::env::args().skip(1).peekable(); + while let Some(arg) = args.next() { + match arg.as_str() { + "--orch-bin" => config.orch_bin = PathBuf::from(next_arg(&mut args, "--orch-bin")?), + "--addr" => config.rpc_addr = next_arg(&mut args, "--addr")?, + "--image" => { + let image = next_arg(&mut args, "--image")?; + config.node_image = image.clone(); + config.orch_args.push("--image".to_owned()); + config.orch_args.push(image); + } + "--max-tokens" => config.max_tokens = parse_next(&mut args, "--max-tokens")?, + "--timeout-ms" => config.timeout_ms = parse_next(&mut args, "--timeout-ms")?, + "--dashboard" => config.dashboard = true, + "--no-dashboard" => config.dashboard = false, + "--no-build-image" => config.build_image = false, + "--" => { + config.orch_args.extend(args); + break; + } + other => config.orch_args.push(other.to_owned()), + } + } + Ok(config) + } +} + +struct OrchChild { + child: Child, + stdin: Option, + ready_rx: mpsc::Receiver, + cleaned: bool, +} + +impl OrchChild { + fn spawn(config: &Config) -> Result { + let mut command = Command::new(&config.orch_bin); + command + .args(&config.orch_args) + .env("MVP_NODE_IMAGE", &config.node_image) + .env("MVP_PROMPT_RPC_BIND", &config.rpc_addr) + .env("MVP_PROMPT_RPC_ADDR", &config.rpc_addr) + .env("MVP_PROMPT_MAX_TOKENS", config.max_tokens.to_string()) + .env("MVP_PROMPT_TIMEOUT_MS", config.timeout_ms.to_string()) + .env("MVP_DASHBOARD", if config.dashboard { "1" } else { "0" }) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()); + #[cfg(target_os = "linux")] + unsafe { + command.pre_exec(|| { + if libc::setpgid(0, 0) == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } + }); + } + let mut child = command + .spawn() + .map_err(|e| format!("spawn {}: {e}", config.orch_bin.display()))?; + let stdin = child.stdin.take(); + let stdout = child + .stdout + .take() + .ok_or_else(|| "orchestrator stdout missing".to_owned())?; + let (ready_tx, ready_rx) = mpsc::channel(); + thread::spawn(move || { + for line in BufReader::new(stdout).lines().map_while(Result::ok) { + eprintln!("mvp-orch-one-node: {line}"); + if let Ok(value) = serde_json::from_str::(&line) + && value.get("type").and_then(Value::as_str) == Some("prompt_loop_ready") + && let Some(addr) = value.get("addr").and_then(Value::as_str) + { + let _ = ready_tx.send(addr.to_owned()); + } + } + }); + Ok(Self { + child, + stdin, + ready_rx, + cleaned: false, + }) + } + + fn wait_ready(&mut self, fallback_addr: String) -> Result { + let start = Instant::now(); + loop { + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Err("interrupted before orchestrator became ready".to_owned()); + } + if let Ok(addr) = self.ready_rx.recv_timeout(Duration::from_millis(100)) { + return Ok(addr); + } + if let Some(status) = self + .child + .try_wait() + .map_err(|e| format!("poll orchestrator: {e}"))? + { + return Err(format!( + "orchestrator exited before prompt loop ready: {status}" + )); + } + if start.elapsed() > ORCH_READY_TIMEOUT { + return Err(format!( + "timed out waiting for orchestrator prompt loop; try connecting to {fallback_addr} if it is still booting" + )); + } + } + } + + fn shutdown(&mut self) { + if self.cleaned { + return; + } + self.cleaned = true; + if let Some(stdin) = self.stdin.as_mut() { + let _ = writeln!(stdin, "shutdown"); + let _ = stdin.flush(); + } + let start = Instant::now(); + while start.elapsed() < ORCH_SHUTDOWN_TIMEOUT { + match self.child.try_wait() { + Ok(Some(_)) => return, + Ok(None) => thread::sleep(Duration::from_millis(100)), + Err(_) => break, + } + } + terminate_process_group(self.child.id()); + let _ = self.child.wait(); + } +} + +impl Drop for OrchChild { + fn drop(&mut self) { + self.shutdown(); + } +} + +fn prepare_runtime(config: &Config) -> Result<(), String> { + run_status( + cargo_command(), + &[ + "build", + "-p", + "mvp-system", + "--features", + "local-e2e", + "--bin", + "mvp-orch-one-node", + ], + "build mvp-orch-one-node", + )?; + if !config.build_image { + eprintln!("mvp-one-node-chat: skipping node image rebuild (--no-build-image)"); + return Ok(()); + } + run_status( + cargo_command(), + &["build", "-p", "mvp-system", "--bin", "mvp-node"], + "build mvp-node", + )?; + if !docker_image_exists(BASE_NODE_IMAGE) { + run_status( + "docker", + &[ + "build", + "-f", + "apps/mvp-node/Dockerfile.base", + "-t", + BASE_NODE_IMAGE, + ".", + ], + "build mvp node base image", + )?; + } + let node_bin = node_bin_for_current_profile()?; + run_status( + "docker", + &[ + "build", + "-f", + "apps/mvp-node/Dockerfile", + "--build-arg", + &format!("BASE_IMAGE={BASE_NODE_IMAGE}"), + "--build-arg", + &format!("MVP_NODE_BIN={}", node_bin.display()), + "-t", + &config.node_image, + ".", + ], + "build mvp node image", + ) +} + +fn run_chat_loop(addr: &str, max_tokens: u32, timeout_ms: u64) -> Result<(), String> { + let mut stream = + TcpStream::connect(addr).map_err(|e| format!("connect prompt RPC {addr}: {e}"))?; + stream + .set_read_timeout(Some(CHAT_READ_TIMEOUT)) + .map_err(|e| format!("set prompt RPC read timeout: {e}"))?; + let mut reader = BufReader::new( + stream + .try_clone() + .map_err(|e| format!("clone prompt RPC stream: {e}"))?, + ); + let (input_tx, input_rx) = mpsc::channel::(); + thread::spawn(move || { + let stdin = io::stdin(); + for line in stdin.lock().lines().map_while(Result::ok) { + if input_tx.send(line).is_err() { + break; + } + } + }); + let mut next_request_id = 1_u64; + eprintln!( + "mvp-one-node-chat: Ctrl-C cleans up the orchestrator and Docker node; /exit exits cleanly" + ); + + loop { + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Ok(()); + } + print!("> "); + io::stdout() + .flush() + .map_err(|e| format!("flush prompt: {e}"))?; + let prompt = loop { + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Ok(()); + } + match input_rx.recv_timeout(Duration::from_millis(100)) { + Ok(line) => break line.trim_end().to_owned(), + Err(mpsc::RecvTimeoutError::Timeout) => continue, + Err(mpsc::RecvTimeoutError::Disconnected) => return Ok(()), + } + }; + if prompt.eq_ignore_ascii_case("/quit") || prompt.eq_ignore_ascii_case("/exit") { + return Ok(()); + } + if prompt.trim().is_empty() { + continue; + } + + let request_id = next_request_id; + next_request_id = next_request_id.wrapping_add(1).max(1); + write_json_line( + &mut stream, + &SubmitPrompt { + request_id, + prompt_text: prompt, + max_tokens, + timeout_ms, + }, + )?; + + loop { + if STOP_REQUESTED.load(Ordering::SeqCst) { + return Ok(()); + } + let mut line = String::new(); + match reader.read_line(&mut line) { + Ok(0) => return Err("prompt RPC closed".to_owned()), + Ok(_) => {} + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::WouldBlock | std::io::ErrorKind::TimedOut + ) => + { + continue; + } + Err(error) => return Err(format!("read prompt event: {error}")), + } + let event = serde_json::from_str::(&line) + .map_err(|e| format!("parse prompt event: {e}"))?; + match event { + PromptEvent::TextDelta { + request_id: seen, + text, + } if seen == request_id => { + print!("{text}"); + io::stdout() + .flush() + .map_err(|e| format!("flush response text: {e}"))?; + } + PromptEvent::Done { + request_id: seen, + tokens_generated, + elapsed_ms, + .. + } if seen == request_id => { + println!(); + eprintln!( + "mvp-one-node-chat: done request={} tokens={} elapsed_ms={}", + seen, tokens_generated, elapsed_ms + ); + break; + } + PromptEvent::Fault { + request_id: seen, + error, + } if seen == request_id => { + eprintln!("mvp-one-node-chat: fault request={seen}: {error}"); + break; + } + _ => {} + } + } + } +} + +fn default_orch_bin() -> Result { + if let Some(path) = std::env::var_os("MVP_ORCH_BIN") { + return Ok(PathBuf::from(path)); + } + let mut path = std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; + path.set_file_name("mvp-orch-one-node"); + Ok(path) +} + +fn node_bin_for_current_profile() -> Result { + let mut path = std::env::current_exe().map_err(|e| format!("current exe: {e}"))?; + path.set_file_name("mvp-node"); + let cwd = std::env::current_dir().map_err(|e| format!("current dir: {e}"))?; + if let Ok(relative) = path.strip_prefix(&cwd) { + Ok(relative.to_path_buf()) + } else { + Ok(path) + } +} + +fn cargo_command() -> &'static str { + "cargo" +} + +fn run_status(program: &str, args: &[&str], label: &str) -> Result<(), String> { + eprintln!("mvp-one-node-chat: {label}"); + let status = Command::new(program) + .args(args) + .stdin(Stdio::null()) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .map_err(|e| format!("run {label}: {e}"))?; + if status.success() { + Ok(()) + } else { + Err(format!("{label} failed with {status}")) + } +} + +fn docker_image_exists(image: &str) -> bool { + Command::new("docker") + .args(["image", "inspect", image]) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(false) +} + +extern "C" fn request_stop(_: libc::c_int) { + STOP_REQUESTED.store(true, Ordering::SeqCst); +} + +fn install_signal_handlers() { + #[cfg(target_os = "linux")] + unsafe { + libc::signal(libc::SIGINT, request_stop as *const () as usize); + libc::signal(libc::SIGTERM, request_stop as *const () as usize); + } +} + +fn terminate_process_group(pid: u32) { + #[cfg(target_os = "linux")] + unsafe { + let pgid = -(pid as libc::pid_t); + let _ = libc::kill(pgid, libc::SIGTERM); + thread::sleep(Duration::from_secs(2)); + let _ = libc::kill(pgid, libc::SIGKILL); + } + #[cfg(not(target_os = "linux"))] + { + let _ = pid; + } +} + +fn env_bool(name: &str, default: bool) -> Result { + match std::env::var(name).ok().filter(|value| !value.is_empty()) { + None => Ok(default), + Some(value) => match value.to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Ok(true), + "0" | "false" | "no" | "off" => Ok(false), + _ => Err(format!( + "invalid {name}={value:?}; use 1/0, true/false, yes/no, or on/off" + )), + }, + } +} + +fn env_u64(name: &str, default: u64) -> Result { + match std::env::var(name).ok().filter(|value| !value.is_empty()) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn env_u32(name: &str, default: u32) -> Result { + match std::env::var(name).ok().filter(|value| !value.is_empty()) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn next_arg(args: &mut impl Iterator, name: &str) -> Result { + args.next() + .ok_or_else(|| format!("missing value after {name}")) +} + +fn parse_next(args: &mut impl Iterator, name: &str) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + let value = next_arg(args, name)?; + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) +} diff --git a/crates/mvp-system/src/bin/mvp_orch_one_node.rs b/crates/mvp-system/src/bin/mvp_orch_one_node.rs new file mode 100644 index 0000000..f145b2a --- /dev/null +++ b/crates/mvp-system/src/bin/mvp_orch_one_node.rs @@ -0,0 +1,983 @@ +use std::io::{BufRead, BufReader}; +use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::process::ExitCode; +use std::sync::{Arc, mpsc}; +use std::thread; +use std::time::{Duration, Instant}; + +use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, NodeId, Position, StreamId}; +use distribution::node::DistributedNodeConfig; +use iroh::EndpointAddr; +use iroh_driver::{IrohDriver, IrohDriverConfig}; +use mvp_system::actors::node_agent::{NodeAgentMsg, StageProvisionWire}; +use mvp_system::actors::register_mvp_actor_codecs; +#[cfg(feature = "local-e2e")] +use mvp_system::dashboard_view::MvpClusterDashboardView; +use mvp_system::distribution_stack::DistributionRuntimeStack; +use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, read_submit_prompt, write_json_line}; +use mvp_system::provisioning::{ + LocalDockerPlugin, NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, + ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, ProvisionPlugin, +}; +use mvp_system::run_plan::{GgufSource, TokenizerSource}; +use mvp_system::telemetry::{ + MVP_PROVISIONING_EVENTS, MvpProvisionEventRecord, MvpProvisionLogRecord, + mvp_provision_log_channel, +}; +use parking_lot::Mutex; +use serde_json::{Value, json}; +use swactor::actor::ActorAddress; + +const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest"; +const DEFAULT_RPC_BIND: &str = "127.0.0.1:19777"; +const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF"; +const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; +const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; +const DEFAULT_MAX_TOKENS: u32 = 64; +const DEFAULT_TIMEOUT_MS: u64 = 120_000; +const PUMP_INTERVAL: Duration = Duration::from_millis(10); +const BOOT_TIMEOUT: Duration = Duration::from_secs(180); +const ROUTE_TIMEOUT: Duration = Duration::from_secs(30); +const WEIGHT_TIMEOUT: Duration = Duration::from_secs(900); + +fn main() -> ExitCode { + match run() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("mvp-orch-one-node: {error}"); + ExitCode::from(1) + } + } +} + +fn run() -> Result<(), String> { + let config = Config::from_env_and_args()?; + let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; + let mut driver = IrohDriver::with_handle( + tokio.handle().clone(), + IrohDriverConfig { + secret_key: None, + relay_mode: config.relay_mode.clone(), + node: DistributedNodeConfig::default(), + peer_auth: None, + additional_alpns: vec![], + }, + ) + .map_err(|e| format!("create iroh driver: {e}"))?; + let stack = DistributionRuntimeStack::new_with_codecs( + driver.node_id(), + DistributedNodeConfig::default(), + |registry| { + register_mvp_actor_codecs(registry); + datastream::wire::register_datastream_codec(registry); + }, + ); + driver.enable_actor_bridge( + stack.runtime.clone(), + stack.codec.clone(), + stack.actor_bridge_routes(), + stack.actors.swim, + stack.relay_mirror.clone(), + stack.route_view.clone(), + ); + + let (frame_tx, frame_rx) = mpsc::channel::<(StreamId, Frame)>(); + let datastream_sink = stack + .runtime + .spawn(DatastreamSink::new(move |stream, frame| { + let _ = frame_tx.send((stream, frame)); + })) + .map_err(|e| format!("spawn datastream sink: {e}"))?; + stack.register_local_actor(driver.register_actor(datastream_sink, 1)); + let dashboard = DashboardSupport::start_from_env()?; + let mut orch_datastream = OrchDatastream::new(config.run_id); + + let prompt_events = stack + .runtime + .new_inbox::() + .map_err(|e| format!("prompt event inbox: {e}"))?; + let prompt_reply_actor = *prompt_events.addr(); + stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); + + let (work_tx, work_rx) = mpsc::channel::(); + let rpc_addr = spawn_prompt_rpc( + config.rpc_bind, + work_tx, + config.default_max_tokens, + config.default_timeout_ms, + )?; + println!( + "{}", + json!({"type":"prompt_rpc_ready","addr":rpc_addr.to_string()}) + ); + + let mut docker = LocalDockerPlugin::new("mvp-orch-one-node"); + let (obs_tx, obs_rx) = mpsc::channel::(); + let sink = PluginSink::new(Arc::new(ChannelObservationSink { + tx: Mutex::new(obs_tx), + })); + orch_datastream.emit_event( + dashboard.as_ref(), + ProvisionEvent { + run_id: config.run_id, + node_id: config.node_id, + kind: ProvisionEventKind::ProvisionStart, + message: Some(format!("starting Docker image {}", config.image)), + }, + ); + let handle = docker.start_node( + config.node_spec(driver.endpoint_addr(), datastream_sink)?, + sink, + )?; + + let ready = wait_for_runtime_ready( + &mut driver, + &stack, + &obs_rx, + &frame_rx, + dashboard.as_ref(), + &mut orch_datastream, + )?; + driver.join(std::slice::from_ref(&ready.endpoint)); + wait_for_route(&mut driver, &stack, ready.node_actor)?; + provision_stage(&stack, ready.node_actor, &config)?; + wait_for_weights_loaded( + &mut driver, + &stack, + &obs_rx, + &frame_rx, + dashboard.as_ref(), + &mut orch_datastream, + )?; + + println!( + "{}", + json!({"type":"prompt_loop_ready","addr":rpc_addr.to_string(),"node_actor":ready.node_actor}) + ); + + let stop_rx = spawn_stop_listener(); + let result = serve_prompts( + &mut driver, + &stack, + &obs_rx, + &frame_rx, + &work_rx, + &prompt_events, + &stop_rx, + dashboard.as_ref(), + &mut orch_datastream, + ready.node_actor, + prompt_reply_actor, + ); + let stop_result = docker.stop_node(&handle); + result.and(stop_result) +} + +#[derive(Clone)] +struct Config { + image: String, + docker_gpus: String, + rpc_bind: SocketAddr, + run_id: u64, + node_id: u64, + stage_index: u32, + layer_end_exclusive: u32, + model_id: String, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + default_max_tokens: u32, + default_timeout_ms: u64, + relay_mode: iroh::RelayMode, +} + +impl Config { + fn from_env_and_args() -> Result { + let mut args = std::env::args().skip(1); + let mut config = Self { + image: env_string("MVP_NODE_IMAGE", DEFAULT_IMAGE), + docker_gpus: env_string("MVP_DOCKER_GPUS", "all"), + rpc_bind: env_string("MVP_PROMPT_RPC_BIND", DEFAULT_RPC_BIND) + .parse() + .map_err(|e| format!("invalid MVP_PROMPT_RPC_BIND: {e}"))?, + run_id: env_u64("MVP_RUN_ID", 1)?, + node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, + stage_index: env_u32("MVP_STAGE_INDEX", 0)?, + layer_end_exclusive: env_u32("MVP_LAYER_END_EXCLUSIVE", 16)?, + model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), + gguf_source: gguf_source_from_env(), + tokenizer: tokenizer_from_env(), + default_max_tokens: env_u32("MVP_PROMPT_MAX_TOKENS", DEFAULT_MAX_TOKENS)?, + default_timeout_ms: env_u64("MVP_PROMPT_TIMEOUT_MS", DEFAULT_TIMEOUT_MS)?, + relay_mode: relay_mode_from_env()?, + }; + + while let Some(arg) = args.next() { + match arg.as_str() { + "--image" => config.image = next_arg(&mut args, "--image")?, + "--gpus" => config.docker_gpus = next_arg(&mut args, "--gpus")?, + "--rpc-bind" => { + config.rpc_bind = next_arg(&mut args, "--rpc-bind")? + .parse() + .map_err(|e| format!("invalid --rpc-bind: {e}"))? + } + "--run-id" => config.run_id = parse_next(&mut args, "--run-id")?, + "--node-id" => config.node_id = parse_next(&mut args, "--node-id")?, + "--max-tokens" => { + config.default_max_tokens = parse_next(&mut args, "--max-tokens")? + } + "--timeout-ms" => { + config.default_timeout_ms = parse_next(&mut args, "--timeout-ms")? + } + "--model-id" => config.model_id = next_arg(&mut args, "--model-id")?, + "--gguf-local-path" => { + config.gguf_source = + GgufSource::LocalPath(next_arg(&mut args, "--gguf-local-path")?) + } + "--gguf-repo" => { + let repo = next_arg(&mut args, "--gguf-repo")?; + config.gguf_source = match config.gguf_source { + GgufSource::HuggingFaceGguf { file, revision, .. } => { + GgufSource::HuggingFaceGguf { + repo, + file, + revision, + } + } + GgufSource::LocalPath(_) => GgufSource::HuggingFaceGguf { + repo, + file: env_string("MVP_GGUF_FILE", DEFAULT_HF_FILE), + revision: env_optional("MVP_GGUF_REVISION"), + }, + }; + } + "--gguf-file" => { + let file = next_arg(&mut args, "--gguf-file")?; + config.gguf_source = match config.gguf_source { + GgufSource::HuggingFaceGguf { repo, revision, .. } => { + GgufSource::HuggingFaceGguf { + repo, + file, + revision, + } + } + GgufSource::LocalPath(_) => GgufSource::HuggingFaceGguf { + repo: env_string("MVP_GGUF_REPO", DEFAULT_HF_REPO), + file, + revision: env_optional("MVP_GGUF_REVISION"), + }, + }; + } + other => return Err(format!("unknown argument {other:?}")), + } + } + Ok(config) + } + + fn node_spec( + &self, + coordinator: EndpointAddr, + datastream_sink: ActorAddress, + ) -> Result { + let mut env = vec![ + ("MVP_RUN_ID".to_owned(), self.run_id.to_string()), + ("MVP_LOGICAL_NODE_ID".to_owned(), self.node_id.to_string()), + ("MVP_STAGE_INDEX".to_owned(), self.stage_index.to_string()), + ( + "MVP_COORDINATOR_ENDPOINT".to_owned(), + serde_json::to_string(&coordinator) + .map_err(|e| format!("serialize coordinator endpoint: {e}"))?, + ), + ( + "MVP_DATASTREAM_SINK_ACTOR".to_owned(), + serde_json::to_string(&datastream_sink) + .map_err(|e| format!("serialize datastream sink actor: {e}"))?, + ), + ("MVP_MODEL_ID".to_owned(), self.model_id.clone()), + ("MVP_NODE_MAX_RUNTIME_SECS".to_owned(), "0".to_owned()), + ("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone()), + ]; + env.extend(optional_env("MVP_TINYGRAD_TEST_MODE")); + env.extend(optional_env("MVP_MODEL_CACHE_DIR")); + env.extend(optional_env("HF_TOKEN")); + match &self.gguf_source { + GgufSource::LocalPath(path) => { + env.push(("MVP_GGUF_LOCAL_PATH".to_owned(), path.clone())) + } + GgufSource::HuggingFaceGguf { + repo, + file, + revision, + } => { + env.push(("MVP_GGUF_REPO".to_owned(), repo.clone())); + env.push(("MVP_GGUF_FILE".to_owned(), file.clone())); + if let Some(revision) = revision { + env.push(("MVP_GGUF_REVISION".to_owned(), revision.clone())); + } + } + } + if let TokenizerSource::LocalPath(path) = &self.tokenizer { + env.push(("MVP_TOKENIZER_LOCAL_PATH".to_owned(), path.clone())); + } + Ok(NodeProvisionSpec { + run_id: self.run_id, + node_id: self.node_id, + stage_index: Some(self.stage_index), + image: self.image.clone(), + env, + args: Vec::new(), + }) + } +} + +#[derive(Clone)] +struct RuntimeReady { + endpoint: EndpointAddr, + node_actor: ActorAddress, +} + +struct PromptWork { + request: SubmitPrompt, + events: mpsc::Sender, +} + +struct ActivePrompt { + request: SubmitPrompt, + events: mpsc::Sender, + deadline: Instant, +} + +struct OrchDatastream { + stream: StreamId, + next_position: u64, +} + +impl OrchDatastream { + fn new(run_id: u64) -> Self { + Self { + stream: StreamId::new(NodeId::new("mvp-orchestrator"), Lifetime(run_id)), + next_position: 0, + } + } + + fn emit_event(&mut self, dashboard: Option<&DashboardSupport>, event: ProvisionEvent) { + let payload = serde_json::to_vec(&MvpProvisionEventRecord::new(event)) + .expect("serialize provisioning event"); + self.emit_bytes(dashboard, ChannelId::new(MVP_PROVISIONING_EVENTS), payload); + } + + fn emit_log(&mut self, dashboard: Option<&DashboardSupport>, line: ProvisionLogLine) { + let channel = mvp_provision_log_channel(line.node_id, line.stream); + let payload = + serde_json::to_vec(&MvpProvisionLogRecord::new(line)).expect("serialize provision log"); + self.emit_bytes(dashboard, channel, payload); + } + + fn emit_bytes( + &mut self, + dashboard: Option<&DashboardSupport>, + channel: ChannelId, + payload: Vec, + ) { + let frame = Frame::new(channel, Position(self.next_position), payload); + self.next_position += 1; + ingest_dashboard_frame(dashboard, &self.stream, &frame); + eprintln!( + "mvp-orch-one-node: datastream {} {}", + frame.channel.as_str(), + String::from_utf8_lossy(&frame.payload) + ); + } +} + +#[cfg(feature = "local-e2e")] +struct DashboardSupport { + handle: dashboard::DashboardHandle, +} + +#[cfg(feature = "local-e2e")] +impl DashboardSupport { + fn start_from_env() -> Result, String> { + if !env_bool("MVP_DASHBOARD", false)? { + return Ok(None); + } + let mut config = dashboard::DashboardConfig::default(); + if let Some(port) = env_optional("MVP_DASHBOARD_PORT") { + config.port = port + .parse::() + .map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?; + } + let url = format!("http://127.0.0.1:{}/view/datastream/live", config.port); + let handle = dashboard::start_dashboard(config); + handle.register_view(Arc::new(MvpClusterDashboardView::new())); + handle.start_http_standalone(); + println!("{}", json!({"type":"dashboard_ready","url":url})); + Ok(Some(Self { handle })) + } + + fn ingest(&self, stream: &StreamId, frame: &Frame) { + self.handle.ingest(stream, frame); + } +} + +#[cfg(not(feature = "local-e2e"))] +struct DashboardSupport; + +#[cfg(not(feature = "local-e2e"))] +impl DashboardSupport { + fn start_from_env() -> Result, String> { + if env_bool("MVP_DASHBOARD", false)? { + return Err( + "MVP_DASHBOARD requires building mvp-system with feature local-e2e".to_owned(), + ); + } + Ok(None) + } + + fn ingest(&self, _stream: &StreamId, _frame: &Frame) {} +} + +struct ChannelObservationSink { + tx: Mutex>, +} + +impl PluginObservationSink for ChannelObservationSink { + fn observe(&self, observation: PluginObservation) { + let _ = self.tx.lock().send(observation); + } +} + +fn spawn_prompt_rpc( + bind: SocketAddr, + work_tx: mpsc::Sender, + default_max_tokens: u32, + default_timeout_ms: u64, +) -> Result { + let listener = TcpListener::bind(bind).map_err(|e| format!("bind prompt RPC {bind}: {e}"))?; + let addr = listener + .local_addr() + .map_err(|e| format!("read prompt RPC addr: {e}"))?; + thread::spawn(move || { + for accepted in listener.incoming() { + match accepted { + Ok(stream) => { + let tx = work_tx.clone(); + thread::spawn(move || { + if let Err(error) = handle_prompt_connection( + stream, + tx, + default_max_tokens, + default_timeout_ms, + ) { + eprintln!("mvp-orch-one-node: prompt connection closed: {error}"); + } + }); + } + Err(error) => eprintln!("mvp-orch-one-node: accept prompt RPC: {error}"), + } + } + }); + Ok(addr) +} + +fn handle_prompt_connection( + stream: TcpStream, + work_tx: mpsc::Sender, + default_max_tokens: u32, + default_timeout_ms: u64, +) -> Result<(), String> { + let mut reader = BufReader::new( + stream + .try_clone() + .map_err(|e| format!("clone prompt stream: {e}"))?, + ); + let mut writer = stream; + loop { + let request = match read_submit_prompt(&mut reader) { + Ok(Some(request)) => request, + Ok(None) => break, + Err(error) if error.contains("expected value at line 1 column 1") => break, + Err(error) => return Err(error), + }; + let request = request.with_defaults(default_max_tokens, default_timeout_ms); + let (event_tx, event_rx) = mpsc::channel(); + work_tx + .send(PromptWork { + request, + events: event_tx, + }) + .map_err(|_| "prompt loop stopped".to_owned())?; + for event in event_rx { + let terminal = event.is_terminal(); + write_json_line(&mut writer, &event)?; + if terminal { + break; + } + } + } + Ok(()) +} + +fn wait_for_runtime_ready( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver<(StreamId, Frame)>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, +) -> Result { + let start = Instant::now(); + loop { + pump(driver, stack); + drain_frames(frame_rx, dashboard); + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, &observation); + match observation { + PluginObservation::RuntimeReady { + endpoint, + node_actor, + .. + } => { + return Ok(RuntimeReady { + endpoint, + node_actor, + }); + } + PluginObservation::ProviderLine { line, .. } + | PluginObservation::StdoutLine { line, .. } + | PluginObservation::StderrLine { line, .. } => { + eprintln!("mvp-orch-one-node: node: {line}"); + } + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited before ready: {status:?}")); + } + } + } + if start.elapsed() > BOOT_TIMEOUT { + return Err("timed out waiting for node ready".to_owned()); + } + thread::sleep(PUMP_INTERVAL); + } +} + +fn wait_for_route( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + actor: ActorAddress, +) -> Result<(), String> { + let start = Instant::now(); + while start.elapsed() <= ROUTE_TIMEOUT { + pump(driver, stack); + if stack + .route_view + .read() + .expect("route view poisoned") + .contains_key(&actor) + { + return Ok(()); + } + thread::sleep(PUMP_INTERVAL); + } + Err(format!( + "timed out waiting for route to node actor {actor:?}" + )) +} + +fn provision_stage( + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + config: &Config, +) -> Result<(), String> { + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::ProvisionStage(StageProvisionWire { + run_id: config.run_id, + authorized_orchestrator: 0, + node_id: config.node_id, + stage_index: config.stage_index, + stage_count: 1, + layer_start: 0, + layer_end_exclusive: config.layer_end_exclusive, + inbound_edge_id: 1, + outbound_edge_id: 2, + model_id: config.model_id.clone(), + gguf_source: config.gguf_source.clone(), + tokenizer: config.tokenizer.clone(), + }), + ) + .map_err(|e| format!("send stage provision: {e}")) +} + +fn wait_for_weights_loaded( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver<(StreamId, Frame)>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, +) -> Result<(), String> { + let start = Instant::now(); + loop { + pump(driver, stack); + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, &observation); + match observation { + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited while loading weights: {status:?}")); + } + PluginObservation::ProviderLine { line, .. } + | PluginObservation::StdoutLine { line, .. } + | PluginObservation::StderrLine { line, .. } => { + eprintln!("mvp-orch-one-node: node: {line}"); + } + PluginObservation::RuntimeReady { .. } => {} + } + } + while let Ok((stream, frame)) = frame_rx.try_recv() { + ingest_dashboard_frame(dashboard, &stream, &frame); + let payload = String::from_utf8_lossy(&frame.payload); + eprintln!( + "mvp-orch-one-node: datastream {} {}", + frame.channel.as_str(), + payload + ); + if frame.channel == ChannelId::new("mvp.worker.weights") + && json_type_is(&payload, "WeightsLoaded") + { + return Ok(()); + } + if json_type_is(&payload, "WorkerFatal") { + return Err(format!("worker fatal while loading weights: {payload}")); + } + } + if start.elapsed() > WEIGHT_TIMEOUT { + return Err("timed out waiting for weights loaded".to_owned()); + } + thread::sleep(PUMP_INTERVAL); + } +} + +fn serve_prompts( + driver: &mut IrohDriver, + stack: &DistributionRuntimeStack, + obs_rx: &mpsc::Receiver, + frame_rx: &mpsc::Receiver<(StreamId, Frame)>, + work_rx: &mpsc::Receiver, + prompt_events: &swactor::runtime::Inbox, + stop_rx: &mpsc::Receiver<()>, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, + node_actor: ActorAddress, + reply_to: ActorAddress, +) -> Result<(), String> { + let mut active: Option = None; + loop { + pump(driver, stack); + drain_observations(obs_rx, dashboard, orch_datastream)?; + drain_frames(frame_rx, dashboard); + if stop_rx.try_recv().is_ok() { + eprintln!("mvp-orch-one-node: stop requested"); + return Ok(()); + } + + if active.is_none() { + if let Ok(work) = work_rx.try_recv() { + let request = work.request; + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::InferPrompt { + request_id: request.request_id, + prompt: request.prompt_text.clone(), + max_tokens: request.max_tokens, + reply_to, + }, + ) + .map_err(|e| format!("send prompt request: {e}"))?; + active = Some(ActivePrompt { + deadline: Instant::now() + Duration::from_millis(request.timeout_ms), + request, + events: work.events, + }); + } + } + + while let Some(event) = prompt_events.try_recv() { + if let Some(current) = active.as_ref() { + if event.request_id() == current.request.request_id { + let terminal = event.is_terminal(); + let _ = current.events.send(event); + if terminal { + active = None; + } + } + } + } + + if let Some(current) = active.as_ref() + && Instant::now() >= current.deadline + { + let _ = current.events.send(PromptEvent::Fault { + request_id: current.request.request_id, + error: "prompt timed out".to_owned(), + }); + active = None; + } + + thread::sleep(PUMP_INTERVAL); + } +} + +fn spawn_stop_listener() -> mpsc::Receiver<()> { + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + let stdin = std::io::stdin(); + for line in stdin.lock().lines().map_while(Result::ok) { + let trimmed = line.trim(); + if trimmed.eq_ignore_ascii_case("stop") + || trimmed.eq_ignore_ascii_case("shutdown") + || trimmed.eq_ignore_ascii_case("quit") + { + let _ = tx.send(()); + break; + } + } + }); + rx +} + +fn drain_observations( + obs_rx: &mpsc::Receiver, + dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, +) -> Result<(), String> { + while let Ok(observation) = obs_rx.try_recv() { + emit_plugin_observation(orch_datastream, dashboard, &observation); + match observation { + PluginObservation::Failed { reason, .. } => return Err(reason), + PluginObservation::Exited { status, .. } => { + return Err(format!("node exited: {status:?}")); + } + PluginObservation::ProviderLine { line, .. } + | PluginObservation::StdoutLine { line, .. } + | PluginObservation::StderrLine { line, .. } => { + eprintln!("mvp-orch-one-node: node: {line}"); + } + PluginObservation::RuntimeReady { .. } => {} + } + } + Ok(()) +} + +fn emit_plugin_observation( + orch_datastream: &mut OrchDatastream, + dashboard: Option<&DashboardSupport>, + observation: &PluginObservation, +) { + match observation { + PluginObservation::StdoutLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Stdout, + line: line.clone(), + }, + ), + PluginObservation::StderrLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Stderr, + line: line.clone(), + }, + ), + PluginObservation::ProviderLine { + run_id, + node_id, + line, + } => orch_datastream.emit_log( + dashboard, + ProvisionLogLine { + run_id: *run_id, + node_id: *node_id, + stream: ProvisionLogStream::Provider, + line: line.clone(), + }, + ), + PluginObservation::RuntimeReady { + run_id, node_id, .. + } => orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: *run_id, + node_id: *node_id, + kind: ProvisionEventKind::NodeLive, + message: None, + }, + ), + PluginObservation::Exited { + run_id, + node_id, + status, + } => orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: *run_id, + node_id: *node_id, + kind: ProvisionEventKind::NodeStopped, + message: Some(format!("node process exited with {status:?}")), + }, + ), + PluginObservation::Failed { + run_id, + node_id, + reason, + } => orch_datastream.emit_event( + dashboard, + ProvisionEvent { + run_id: *run_id, + node_id: *node_id, + kind: ProvisionEventKind::ProvisionFailed, + message: Some(reason.clone()), + }, + ), + } +} + +fn ingest_dashboard_frame(dashboard: Option<&DashboardSupport>, stream: &StreamId, frame: &Frame) { + if let Some(dashboard) = dashboard { + dashboard.ingest(stream, frame); + } +} + +fn drain_frames( + frame_rx: &mpsc::Receiver<(StreamId, Frame)>, + dashboard: Option<&DashboardSupport>, +) { + while let Ok((stream, frame)) = frame_rx.try_recv() { + ingest_dashboard_frame(dashboard, &stream, &frame); + eprintln!( + "mvp-orch-one-node: datastream {} {}", + frame.channel.as_str(), + String::from_utf8_lossy(&frame.payload) + ); + } +} + +fn pump(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) { + stack.tick_protocol_actors(Instant::now()); + driver.pump_inbound_to_actors(); + stack.pump_runtime_once(); + driver.drain_outbox(&stack.outbox); +} + +fn json_type_is(payload: &str, expected: &str) -> bool { + serde_json::from_str::(payload) + .ok() + .and_then(|value| value.get("type").and_then(Value::as_str).map(str::to_owned)) + .as_deref() + == Some(expected) +} + +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn optional_env(name: &str) -> Option<(String, String)> { + env_optional(name).map(|value| (name.to_owned(), value)) +} + +fn env_string(name: &str, default: &str) -> String { + env_optional(name).unwrap_or_else(|| default.to_owned()) +} + +fn env_bool(name: &str, default: bool) -> Result { + match env_optional(name) { + None => Ok(default), + Some(value) => match value.to_ascii_lowercase().as_str() { + "1" | "true" | "yes" | "on" => Ok(true), + "0" | "false" | "no" | "off" => Ok(false), + _ => Err(format!( + "invalid {name}={value:?}; use 1/0, true/false, yes/no, or on/off" + )), + }, + } +} + +fn env_u64(name: &str, default: u64) -> Result { + match env_optional(name) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn env_u32(name: &str, default: u32) -> Result { + match env_optional(name) { + Some(value) => value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")), + None => Ok(default), + } +} + +fn gguf_source_from_env() -> GgufSource { + if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { + return GgufSource::LocalPath(path); + } + GgufSource::HuggingFaceGguf { + repo: env_string("MVP_GGUF_REPO", DEFAULT_HF_REPO), + file: env_string("MVP_GGUF_FILE", DEFAULT_HF_FILE), + revision: env_optional("MVP_GGUF_REVISION"), + } +} + +fn tokenizer_from_env() -> TokenizerSource { + env_optional("MVP_TOKENIZER_LOCAL_PATH") + .map(TokenizerSource::LocalPath) + .unwrap_or(TokenizerSource::EmbeddedGguf) +} + +fn relay_mode_from_env() -> Result { + match env_string("MVP_IROH_RELAY_MODE", "disabled").as_str() { + "disabled" => Ok(iroh::RelayMode::Disabled), + "default" => Ok(iroh::RelayMode::Default), + other => Err(format!( + "unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default" + )), + } +} + +fn next_arg(args: &mut impl Iterator, name: &str) -> Result { + args.next() + .ok_or_else(|| format!("missing value after {name}")) +} + +fn parse_next(args: &mut impl Iterator, name: &str) -> Result +where + T: std::str::FromStr, + T::Err: std::fmt::Display, +{ + let value = next_arg(args, name)?; + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) +} diff --git a/crates/mvp-system/src/bootstrap_datastream.rs b/crates/mvp-system/src/bootstrap_datastream.rs new file mode 100644 index 0000000..4f8114b --- /dev/null +++ b/crates/mvp-system/src/bootstrap_datastream.rs @@ -0,0 +1,183 @@ +use std::io::{BufRead, BufReader, Read}; +use std::thread::{self, JoinHandle}; + +use datastream::{ChannelId, DatastreamProducer, Lifetime, NodeId, StreamId}; +use iroh::EndpointAddr; +use serde::Deserialize; +use swactor::actor::ActorAddress; + +use crate::provisioning::{ + NodeProvisionSpec, PluginObservation, PluginSink, ProvisionLogLine, ProvisionLogStream, +}; +use crate::telemetry::{MvpProvisionLogRecord, mvp_provision_log_channel}; + +pub fn node_datastream_id(node_id: u64) -> String { + node_id.to_string() +} + +pub fn node_stream_id(run_id: u64, node_id: u64) -> StreamId { + StreamId::new(NodeId::new(&node_datastream_id(node_id)), Lifetime(run_id)) +} + +#[derive(Clone)] +pub struct BootstrapDatastreamBridge { + spec: NodeProvisionSpec, + sink: PluginSink, + producer: Option, +} + +impl BootstrapDatastreamBridge { + pub fn new( + spec: NodeProvisionSpec, + sink: PluginSink, + producer: Option, + ) -> Self { + Self { + spec, + sink, + producer, + } + } + + pub fn spec(&self) -> &NodeProvisionSpec { + &self.spec + } + + pub fn stream_id(&self) -> StreamId { + node_stream_id(self.spec.run_id, self.spec.node_id) + } + + pub fn observe_stdout_line(&self, line: impl Into) { + let line = line.into(); + self.submit_log(ProvisionLogStream::Stdout, &line); + self.sink.observe(PluginObservation::StdoutLine { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + line: line.clone(), + }); + if let Some(ready) = parse_runtime_ready(&self.spec, &line) { + self.sink.observe(ready); + } + } + + pub fn observe_stderr_line(&self, line: impl Into) { + let line = line.into(); + self.submit_log(ProvisionLogStream::Stderr, &line); + self.sink.observe(PluginObservation::StderrLine { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + line, + }); + } + + pub fn observe_provider_line(&self, line: impl Into) { + let line = line.into(); + self.submit_log(ProvisionLogStream::Provider, &line); + self.sink.observe(PluginObservation::ProviderLine { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + line, + }); + } + + pub fn spawn_stdout_reader(&self, stdout: R) -> JoinHandle<()> + where + R: Read + Send + 'static, + { + let bridge = self.clone(); + thread::spawn(move || bridge.read_stdout(stdout)) + } + + pub fn spawn_stderr_reader(&self, stderr: R) -> JoinHandle<()> + where + R: Read + Send + 'static, + { + let bridge = self.clone(); + thread::spawn(move || bridge.read_stderr(stderr)) + } + + fn read_stdout(&self, stdout: R) + where + R: Read, + { + let reader = BufReader::new(stdout); + for next in reader.lines() { + match next { + Ok(line) => self.observe_stdout_line(line), + Err(error) => { + self.sink.observe(PluginObservation::Failed { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + reason: format!("read stdout: {error}"), + }); + break; + } + } + } + } + + fn read_stderr(&self, stderr: R) + where + R: Read, + { + let reader = BufReader::new(stderr); + for next in reader.lines() { + match next { + Ok(line) => self.observe_stderr_line(line), + Err(error) => { + self.sink.observe(PluginObservation::Failed { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + reason: format!("read stderr: {error}"), + }); + break; + } + } + } + } + + fn submit_log(&self, stream: ProvisionLogStream, line: &str) { + let Some(producer) = &self.producer else { + return; + }; + let record = MvpProvisionLogRecord::new(ProvisionLogLine { + run_id: self.spec.run_id, + node_id: self.spec.node_id, + stream, + line: line.to_owned(), + }); + let payload = serde_json::to_vec(&record).expect("serialize bootstrap log record"); + producer.submit_bytes( + mvp_provision_log_channel(self.spec.node_id, stream), + payload, + ); + } +} + +#[derive(Deserialize)] +struct RuntimeReadyLine { + #[serde(rename = "type")] + kind: String, + endpoint: EndpointAddr, + node_actor: ActorAddress, + logical_node_id: u64, + stage_index: u32, +} + +pub fn parse_runtime_ready(spec: &NodeProvisionSpec, line: &str) -> Option { + let ready = serde_json::from_str::(line).ok()?; + if ready.kind != "ready" || ready.logical_node_id != spec.node_id { + return None; + } + Some(PluginObservation::RuntimeReady { + run_id: spec.run_id, + node_id: spec.node_id, + stage_index: Some(ready.stage_index), + endpoint: ready.endpoint, + node_actor: ready.node_actor, + }) +} + +pub fn bootstrap_log_channel(node_id: u64, stream: ProvisionLogStream) -> ChannelId { + mvp_provision_log_channel(node_id, stream) +} diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 0a8f4d8..4600e8e 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -3,6 +3,7 @@ extern crate self as mvp_system; pub mod actors; pub mod arena_manager; +pub mod bootstrap_datastream; #[cfg(feature = "local-e2e")] pub mod dashboard_view; pub mod device_bridge; @@ -21,6 +22,7 @@ pub mod node_provisioning; pub mod observability_surface; pub mod orchestrator_run_fsm; pub mod orchestrator_token_endpoint; +pub mod prompt_rpc; pub mod provisioning; pub mod resource_inventory; pub mod run_plan; @@ -30,6 +32,7 @@ pub mod shared_ring_helper_abi; pub mod stage_controller; pub mod telemetry; pub mod tx_rx_edge_actor; +pub mod vastai_provisioning; pub mod weight_lifecycle; pub mod weight_shards; diff --git a/crates/mvp-system/src/prompt_rpc.rs b/crates/mvp-system/src/prompt_rpc.rs new file mode 100644 index 0000000..82ae451 --- /dev/null +++ b/crates/mvp-system/src/prompt_rpc.rs @@ -0,0 +1,136 @@ +use std::io::{BufRead, Write}; + +use serde::{Deserialize, Serialize}; +use swactor_transport::{CodecRegistry, NetworkMessage}; + +use crate::actors::codec::JsonCodec; + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SubmitPrompt { + pub request_id: u64, + pub prompt_text: String, + pub max_tokens: u32, + pub timeout_ms: u64, +} + +impl SubmitPrompt { + pub fn with_defaults(mut self, max_tokens: u32, timeout_ms: u64) -> Self { + if self.max_tokens == 0 { + self.max_tokens = max_tokens; + } + if self.timeout_ms == 0 { + self.timeout_ms = timeout_ms; + } + self + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum PromptEvent { + TextDelta { + request_id: u64, + text: String, + }, + Done { + request_id: u64, + final_text: String, + tokens_generated: u32, + elapsed_ms: u64, + }, + Fault { + request_id: u64, + error: String, + }, +} + +impl PromptEvent { + pub fn request_id(&self) -> u64 { + match self { + Self::TextDelta { request_id, .. } + | Self::Done { request_id, .. } + | Self::Fault { request_id, .. } => *request_id, + } + } + + pub fn is_terminal(&self) -> bool { + matches!(self, Self::Done { .. } | Self::Fault { .. }) + } +} + +impl NetworkMessage for PromptEvent { + fn type_tag() -> &'static str { + "mvp_system::PromptEvent" + } +} + +pub fn register_codecs(registry: &mut CodecRegistry) { + registry.register::(JsonCodec::::default()); +} + +pub fn write_json_line(writer: &mut impl Write, value: &T) -> Result<(), String> { + serde_json::to_writer(&mut *writer, value).map_err(|e| format!("serialize JSON line: {e}"))?; + writer + .write_all(b"\n") + .map_err(|e| format!("write JSON line: {e}"))?; + writer.flush().map_err(|e| format!("flush JSON line: {e}")) +} + +pub fn read_submit_prompt(reader: &mut impl BufRead) -> Result, String> { + let mut line = String::new(); + let n = reader + .read_line(&mut line) + .map_err(|e| format!("read prompt request: {e}"))?; + if n == 0 { + return Ok(None); + } + serde_json::from_str::(&line) + .map(Some) + .map_err(|e| format!("parse prompt request: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn zero_request_limits_take_loop_defaults() { + let request = SubmitPrompt { + request_id: 7, + prompt_text: "hello".to_owned(), + max_tokens: 0, + timeout_ms: 0, + } + .with_defaults(32, 1_000); + + assert_eq!(request.max_tokens, 32); + assert_eq!(request.timeout_ms, 1_000); + } + + #[test] + fn event_terminal_state_is_explicit() { + assert!( + !PromptEvent::TextDelta { + request_id: 1, + text: "a".to_owned(), + } + .is_terminal() + ); + assert!( + PromptEvent::Done { + request_id: 1, + final_text: "a".to_owned(), + tokens_generated: 1, + elapsed_ms: 2, + } + .is_terminal() + ); + assert!( + PromptEvent::Fault { + request_id: 1, + error: "boom".to_owned(), + } + .is_terminal() + ); + } +} diff --git a/crates/mvp-system/src/provisioning.rs b/crates/mvp-system/src/provisioning.rs index a070897..82f8262 100644 --- a/crates/mvp-system/src/provisioning.rs +++ b/crates/mvp-system/src/provisioning.rs @@ -4,7 +4,7 @@ //! report observations back to the provisioner actor through [`PluginSink`]. use std::collections::BTreeMap; -use std::io::{BufRead, BufReader, Write}; +use std::io::Write; use std::process::{ChildStdin, Command, Stdio}; use std::sync::Arc; use std::thread; @@ -13,6 +13,8 @@ use iroh::EndpointAddr; use serde::{Deserialize, Serialize}; use swactor::actor::ActorAddress; +use crate::bootstrap_datastream::BootstrapDatastreamBridge; + #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct NodeProvisionSpec { pub run_id: u64, @@ -165,6 +167,17 @@ impl ProvisionPlugin for LocalDockerPlugin { .arg("--name") .arg(&container_name) .arg("-i"); + let spec_gpus = spec + .env + .iter() + .find(|(key, _)| key == "MVP_DOCKER_GPUS") + .map(|(_, value)| value.clone()); + if let Some(gpus) = spec_gpus + .or_else(|| std::env::var("MVP_DOCKER_GPUS").ok()) + .filter(|value| !value.trim().is_empty()) + { + command.arg("--gpus").arg(gpus); + } for (key, value) in &spec.env { command.arg("-e").arg(format!("{key}={value}")); } @@ -250,54 +263,12 @@ impl ProvisionPlugin for LocalDockerPlugin { } } -#[derive(Deserialize)] -struct RuntimeReadyLine { - #[serde(rename = "type")] - kind: String, - endpoint: EndpointAddr, - node_actor: ActorAddress, - logical_node_id: u64, - stage_index: u32, -} - fn spawn_stdout_reader( spec: NodeProvisionSpec, sink: PluginSink, stdout: impl std::io::Read + Send + 'static, ) { - thread::spawn(move || { - let reader = BufReader::new(stdout); - for next in reader.lines() { - match next { - Ok(line) => { - sink.observe(PluginObservation::StdoutLine { - run_id: spec.run_id, - node_id: spec.node_id, - line: line.clone(), - }); - if let Ok(ready) = serde_json::from_str::(&line) { - if ready.kind == "ready" && ready.logical_node_id == spec.node_id { - sink.observe(PluginObservation::RuntimeReady { - run_id: spec.run_id, - node_id: spec.node_id, - stage_index: Some(ready.stage_index), - endpoint: ready.endpoint, - node_actor: ready.node_actor, - }); - } - } - } - Err(error) => { - sink.observe(PluginObservation::Failed { - run_id: spec.run_id, - node_id: spec.node_id, - reason: format!("read stdout: {error}"), - }); - break; - } - } - } - }); + BootstrapDatastreamBridge::new(spec, sink, None).spawn_stdout_reader(stdout); } fn spawn_stderr_reader( @@ -305,24 +276,5 @@ fn spawn_stderr_reader( sink: PluginSink, stderr: impl std::io::Read + Send + 'static, ) { - thread::spawn(move || { - let reader = BufReader::new(stderr); - for next in reader.lines() { - match next { - Ok(line) => sink.observe(PluginObservation::StderrLine { - run_id: spec.run_id, - node_id: spec.node_id, - line, - }), - Err(error) => { - sink.observe(PluginObservation::Failed { - run_id: spec.run_id, - node_id: spec.node_id, - reason: format!("read stderr: {error}"), - }); - break; - } - } - } - }); + BootstrapDatastreamBridge::new(spec, sink, None).spawn_stderr_reader(stderr); } diff --git a/crates/mvp-system/src/run_plan.rs b/crates/mvp-system/src/run_plan.rs index c99176a..30a5986 100644 --- a/crates/mvp-system/src/run_plan.rs +++ b/crates/mvp-system/src/run_plan.rs @@ -54,7 +54,7 @@ pub struct ModelFacts { pub tokenizer: TokenizerSource, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum GgufSource { LocalPath(String), HuggingFaceGguf { @@ -64,7 +64,7 @@ pub enum GgufSource { }, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum TokenizerSource { EmbeddedGguf, LocalPath(String), @@ -345,6 +345,7 @@ pub struct ProvisionStage { pub stage_index: u32, pub stage_count: u32, pub gguf_source: GgufSource, + pub tokenizer: TokenizerSource, pub layer_start: u32, pub layer_end_exclusive: u32, pub inbound: InboundEdgeProvision, @@ -547,6 +548,7 @@ pub fn derive_stage_provision( stage_index, stage_count: stage.stage_count, gguf_source: stage.gguf_source.clone(), + tokenizer: plan.model.tokenizer.clone(), layer_start: stage.layer_start, layer_end_exclusive: stage.layer_end_exclusive, inbound: InboundEdgeProvision { diff --git a/crates/mvp-system/src/stage_controller.rs b/crates/mvp-system/src/stage_controller.rs index e0602a1..abe8d37 100644 --- a/crates/mvp-system/src/stage_controller.rs +++ b/crates/mvp-system/src/stage_controller.rs @@ -1,3 +1,5 @@ +use crate::run_plan::{GgufSource, TokenizerSource}; + #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct RunId(pub u64); #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] @@ -28,8 +30,32 @@ pub struct LayerRange { } #[derive(Clone, Debug, PartialEq, Eq)] -pub enum WeightSource { - TestArtifact(String), +pub struct WeightSource { + pub model_id: String, + pub gguf_source: GgufSource, + pub tokenizer: TokenizerSource, +} + +impl WeightSource { + pub fn new( + model_id: impl Into, + gguf_source: GgufSource, + tokenizer: TokenizerSource, + ) -> Self { + Self { + model_id: model_id.into(), + gguf_source, + tokenizer, + } + } + + pub fn embedded_gguf(model_id: impl Into, path: impl Into) -> Self { + Self::new( + model_id, + GgufSource::LocalPath(path.into()), + TokenizerSource::EmbeddedGguf, + ) + } } #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs new file mode 100644 index 0000000..17c949d --- /dev/null +++ b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs @@ -0,0 +1,115 @@ +use std::sync::Arc; + +use datastream::{DatastreamEndpoint, Record}; +use iroh::{EndpointAddr, SecretKey}; +use mvp_system::bootstrap_datastream::{BootstrapDatastreamBridge, node_stream_id}; +use mvp_system::provisioning::{ + NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, ProvisionLogStream, +}; +use mvp_system::telemetry::MvpProvisionLogRecord; +use parking_lot::Mutex; +use serde_json::json; +use swactor::actor::ActorAddress; + +#[derive(Default)] +struct RecordingSink { + observations: Mutex>, +} + +impl RecordingSink { + fn observations(&self) -> Vec { + self.observations.lock().clone() + } +} + +impl PluginObservationSink for RecordingSink { + fn observe(&self, observation: PluginObservation) { + self.observations.lock().push(observation); + } +} + +fn recording_sink() -> (Arc, PluginSink) { + let recording = Arc::new(RecordingSink::default()); + (recording.clone(), PluginSink::new(recording)) +} + +fn spec() -> NodeProvisionSpec { + NodeProvisionSpec { + run_id: 7, + node_id: 42, + stage_index: Some(3), + image: "worker:latest".to_owned(), + env: Vec::new(), + args: Vec::new(), + } +} + +#[test] +fn bootstrap_bridge_writes_node_stream_and_forwards_plugin_observations() { + let endpoint = DatastreamEndpoint::new(node_stream_id(7, 42)); + let subscription = endpoint.subscribe_all("test"); + let (recording, sink) = recording_sink(); + let bridge = BootstrapDatastreamBridge::new(spec(), sink, Some(endpoint.producer())); + + bridge.observe_stdout_line("boot entered"); + bridge.observe_stderr_line("warning"); + endpoint.tick(); + + let observations = recording.observations(); + assert_eq!( + observations, + vec![ + PluginObservation::StdoutLine { + run_id: 7, + node_id: 42, + line: "boot entered".to_owned(), + }, + PluginObservation::StderrLine { + run_id: 7, + node_id: 42, + line: "warning".to_owned(), + }, + ] + ); + + let deliveries = subscription.drain_available(); + assert_eq!(deliveries.len(), 2); + assert_eq!(deliveries[0].stream, node_stream_id(7, 42)); + assert_eq!(deliveries[1].stream, node_stream_id(7, 42)); + + let stdout = MvpProvisionLogRecord::decode(&deliveries[0].frame.payload).unwrap(); + let stderr = MvpProvisionLogRecord::decode(&deliveries[1].frame.payload).unwrap(); + assert_eq!(stdout.line.stream, ProvisionLogStream::Stdout); + assert_eq!(stdout.line.line, "boot entered"); + assert_eq!(stderr.line.stream, ProvisionLogStream::Stderr); + assert_eq!(stderr.line.line, "warning"); +} + +#[test] +fn ready_json_on_stdout_emits_runtime_ready_through_plugin_sink() { + let (recording, sink) = recording_sink(); + let bridge = BootstrapDatastreamBridge::new(spec(), sink, None); + let endpoint = EndpointAddr::new(SecretKey::from_bytes(&[7; 32]).public()); + let node_actor = ActorAddress::new_random(); + let line = serde_json::to_string(&json!({ + "type": "ready", + "endpoint": endpoint, + "node_actor": node_actor, + "logical_node_id": 42, + "stage_index": 3, + })) + .unwrap(); + + bridge.observe_stdout_line(line); + + assert!(recording.observations().iter().any(|observation| matches!( + observation, + PluginObservation::RuntimeReady { + run_id: 7, + node_id: 42, + stage_index: Some(3), + endpoint: observed_endpoint, + node_actor: observed_actor, + } if observed_endpoint == &endpoint && observed_actor == &node_actor + ))); +} diff --git a/crates/mvp-system/src/tests/local_mock/environment.rs b/crates/mvp-system/src/tests/local_mock/environment.rs index 74c99ec..3a1e468 100644 --- a/crates/mvp-system/src/tests/local_mock/environment.rs +++ b/crates/mvp-system/src/tests/local_mock/environment.rs @@ -797,10 +797,11 @@ impl LocalMockCluster { }, inbound: stage::EdgeProvision::inbound(stage::EdgeId(provision.inbound.edge_id.0)), outbound: stage::EdgeProvision::outbound(stage::EdgeId(provision.outbound.edge_id.0)), - weight_source: stage::WeightSource::TestArtifact(format!( - "mock-stage-{}", - provision.stage_index - )), + weight_source: stage::WeightSource::new( + provision.model.model_id, + provision.gguf_source, + provision.tokenizer, + ), } } diff --git a/crates/mvp-system/src/tests/mod.rs b/crates/mvp-system/src/tests/mod.rs index 6481daa..a006607 100644 --- a/crates/mvp-system/src/tests/mod.rs +++ b/crates/mvp-system/src/tests/mod.rs @@ -1,4 +1,5 @@ mod arena_manager_guarantees; +mod bootstrap_datastream_guarantees; mod device_bridge_guarantees; mod docker_cluster_provisioning_guarantees; mod edge_establisher_guarantees; @@ -21,4 +22,5 @@ mod shared_ring_helper_abi_guarantees; mod stage_controller_guarantees; mod telemetry_guarantees; mod tx_rx_edge_actor_guarantees; +mod vastai_provisioning_guarantees; mod weight_lifecycle_guarantees; diff --git a/crates/mvp-system/src/tests/node_provisioning_guarantees.rs b/crates/mvp-system/src/tests/node_provisioning_guarantees.rs index cc379de..ff04554 100644 --- a/crates/mvp-system/src/tests/node_provisioning_guarantees.rs +++ b/crates/mvp-system/src/tests/node_provisioning_guarantees.rs @@ -342,6 +342,92 @@ fn bootstrap_session_streams_logs_flushes_and_closes_on_convergence() { )); } +#[test] +fn node_manager_closes_bootstrap_only_after_expected_swactor_convergence() { + let spec = one_logical_node(); + let logical_node_id = spec.logical_node_id.clone(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let bootstrap_spec = start_bootstrap_command(commands); + let mut session = provision::BootstrapSession::new(bootstrap_spec); + let mut datastream = provision::InMemoryBootstrapDatastream::default(); + let events = session.start( + &provision::MockBootstrapScript::successful(vec![( + provision::BootstrapLogStream::Stdout, + "boot entered".into(), + )]), + &mut datastream, + ); + for event in events { + if let provision::BootstrapSessionEvent::Observed(observation) = event { + manager + .handle(provision::NodeManagerMsg::BootstrapObserved(observation)) + .expect("bootstrap observation accepted"); + } + } + + let wrong_join = manager.handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id: provision::LogicalNodeId("workers-99".into()), + swactor_id: provision::SwactorId("swactor-wrong".into()), + }); + assert!(wrong_join.is_err()); + assert_eq!( + manager.active_bootstrap(), + Some(provision::BootstrapSessionId(1)) + ); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::BootstrapRunning + ); + assert!(!manager.is_ready()); + + let commands = manager + .handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id, + swactor_id: provision::SwactorId("swactor-a".into()), + }) + .expect("expected swactor join accepted"); + assert!(matches!( + commands.as_slice(), + [provision::NodeManagerCommand::BootstrapConvergenceObserved { + session_id: provision::BootstrapSessionId(1), + swactor_id + }] if swactor_id == &provision::SwactorId("swactor-a".into()) + )); + assert!(!manager.is_ready(), "join alone must not close bootstrap"); + + let events = + session.convergence_observed(provision::SwactorId("swactor-a".into()), &mut datastream); + assert_eq!(datastream.flush_count(), 1); + for event in events { + match event { + provision::BootstrapSessionEvent::Observed(observation) => { + manager + .handle(provision::NodeManagerMsg::BootstrapObserved(observation)) + .expect("convergence observation accepted"); + } + provision::BootstrapSessionEvent::Closed => { + manager + .handle(provision::NodeManagerMsg::BootstrapClosed) + .expect("bootstrap close accepted"); + } + provision::BootstrapSessionEvent::Failed(reason) => { + panic!("convergence must not fail: {reason}"); + } + } + } + + let record = manager.record().expect("record exists"); + assert_eq!(record.stage, provision::NodeStage::Dormant); + assert!(record.ready); + assert_eq!(manager.active_bootstrap(), None); + assert!(session.is_closed()); +} + #[test] fn bootstrap_session_reports_boot_check_failure_without_handoff() { let spec = one_logical_node(); diff --git a/crates/mvp-system/src/tests/run_plan_guarantees.rs b/crates/mvp-system/src/tests/run_plan_guarantees.rs index b0d3a63..ad8020a 100644 --- a/crates/mvp-system/src/tests/run_plan_guarantees.rs +++ b/crates/mvp-system/src/tests/run_plan_guarantees.rs @@ -403,6 +403,7 @@ fn provision_stage_projection_is_deterministic_and_stage_local() { assert_eq!(first.layer_start, stage.layer_start); assert_eq!(first.layer_end_exclusive, stage.layer_end_exclusive); assert_eq!(first.gguf_source, stage.gguf_source); + assert_eq!(first.tokenizer, plan.model.tokenizer); assert_eq!(first.model.model_id, plan.model.model_id); assert_eq!(first.model.hidden_dim, plan.model.hidden_dim); assert_eq!(first.model.dtype_family, plan.model.dtype_family); diff --git a/crates/mvp-system/src/tests/stage_controller_guarantees.rs b/crates/mvp-system/src/tests/stage_controller_guarantees.rs index 9412511..62907a4 100644 --- a/crates/mvp-system/src/tests/stage_controller_guarantees.rs +++ b/crates/mvp-system/src/tests/stage_controller_guarantees.rs @@ -26,7 +26,7 @@ fn valid_provision() -> stage::ProvisionStage { }, inbound: stage::EdgeProvision::inbound(stage::EdgeId(7001)), outbound: stage::EdgeProvision::outbound(stage::EdgeId(7002)), - weight_source: stage::WeightSource::TestArtifact("model.gguf".into()), + weight_source: stage::WeightSource::embedded_gguf("model", "model.gguf"), } } diff --git a/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs new file mode 100644 index 0000000..4c2e07d --- /dev/null +++ b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs @@ -0,0 +1,351 @@ +use std::collections::VecDeque; +use std::sync::Arc; + +use mvp_system::node_provisioning as provision; +use mvp_system::node_provisioning::ProviderPlugin; +use mvp_system::provisioning::{ + NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, ProvisionPlugin, +}; +use mvp_system::vastai_provisioning::{ + VastAiBootstrapLauncher, VastAiLeaseClient, VastAiProviderPlugin, VastAiProvisioningConfig, + VastAiProvisioningPlugin, VastAiSshEndpoint, +}; +use parking_lot::Mutex; +use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy}; + +#[derive(Default)] +struct RecordingSink { + observations: Mutex>, +} + +impl PluginObservationSink for RecordingSink { + fn observe(&self, observation: PluginObservation) { + self.observations.lock().push(observation); + } +} + +fn sink() -> PluginSink { + PluginSink::new(Arc::new(RecordingSink::default())) +} + +#[derive(Default)] +struct FakeLeaseClient { + requests: Vec, + endpoint_lookups: Vec<(u64, String, String)>, + endpoint_results: VecDeque>, + destroyed: Vec, + destroy_result: Option>, + next_contract_id: u64, +} + +impl FakeLeaseClient { + fn with_contract(mut self, contract_id: u64) -> Self { + self.next_contract_id = contract_id; + self + } +} + +impl VastAiLeaseClient for FakeLeaseClient { + fn provision_one(&mut self, request: ProvisionRequest) -> Result { + self.requests.push(request); + let contract_id = self.next_contract_id; + self.next_contract_id = self.next_contract_id.wrapping_add(1).max(1); + Ok(ProvisionedInstance { + index: 0, + contract_id, + offer_id: 55, + host_id: Some(77), + gpu_name: "RTX 4090".to_owned(), + gpu_ram: Some(24_000.0), + dph_total: 0.42, + }) + } + + fn ssh_endpoint( + &mut self, + contract_id: u64, + label: &str, + _lifecycle: &LifecyclePolicy, + ssh_user: &str, + ) -> Result { + self.endpoint_lookups + .push((contract_id, label.to_owned(), ssh_user.to_owned())); + self.endpoint_results.pop_front().unwrap_or_else(|| { + Ok(VastAiSshEndpoint { + host: "ssh5.vast.ai".to_owned(), + port: 22017, + user: ssh_user.to_owned(), + }) + }) + } + + fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> { + self.destroyed.push(contract_id); + self.destroy_result.clone().unwrap_or(Ok(())) + } +} + +#[derive(Default)] +struct FakeBootstrap { + starts: Vec<(NodeProvisionSpec, VastAiSshEndpoint)>, + stops: Vec, + fail: Option, + next_handle: usize, +} + +impl VastAiBootstrapLauncher for FakeBootstrap { + type Handle = usize; + + fn start_bootstrap( + &mut self, + spec: NodeProvisionSpec, + endpoint: VastAiSshEndpoint, + _sink: PluginSink, + _producer: Option, + ) -> Result { + if let Some(reason) = self.fail.clone() { + return Err(reason); + } + self.starts.push((spec, endpoint)); + self.next_handle += 1; + Ok(self.next_handle) + } + + fn stop_bootstrap(&mut self, handle: &mut Self::Handle) { + self.stops.push(*handle); + } +} + +fn spec() -> NodeProvisionSpec { + NodeProvisionSpec { + run_id: 9, + node_id: 11, + stage_index: Some(2), + image: "registry.example/mvp-worker:latest".to_owned(), + env: vec![("EXISTING".to_owned(), "1".to_owned())], + args: vec!["python".to_owned(), "worker.py".to_owned()], + } +} + +fn config() -> VastAiProvisioningConfig { + VastAiProvisioningConfig { + label_prefix: "test-mvp".to_owned(), + disk_gb: 80, + ssh_user: "ubuntu".to_owned(), + selection: SelectionPolicy::default(), + lifecycle: LifecyclePolicy::default(), + confirm_lease: false, + onstart: None, + } +} + +fn logical_spec() -> provision::LogicalNodeSpec { + let logical_node_id = provision::LogicalNodeId("workers-0".into()); + provision::LogicalNodeSpec { + run_id: provision::RunId(9), + logical_node_id: logical_node_id.clone(), + group_id: provision::NodeGroupId("workers".into()), + role: provision::RoleId("worker".into()), + provider: provision::ProviderKind::VastAi, + shape: provision::DesiredNodeShape { + image: "registry.example/mvp-worker:latest".to_owned(), + disk_gb: 80, + gpu_name: Some("RTX 4090".to_owned()), + min_gpu_ram_mb: Some(20_000), + min_down_mbps: Some(150.0), + min_up_mbps: Some(25.0), + min_reliability: Some(0.98), + require_verified: true, + provider_labels: [("system".to_owned(), "mvp".to_owned())] + .into_iter() + .collect(), + }, + boot: provision::BootSpec { + ssh_user: "ubuntu".to_owned(), + verify_commands: vec!["test -x /opt/mvp/swactor".to_owned()], + start_swactor_command: "/opt/mvp/swactor-node --join".to_owned(), + stdout_sources: vec!["/var/log/mvp/stdout.log".to_owned()], + stderr_sources: vec!["/var/log/mvp/stderr.log".to_owned()], + timeout_policy: provision::BootstrapTimeoutPolicy { + ssh_connect_secs: 10, + boot_check_secs: 20, + swactor_join_secs: 30, + }, + }, + swarm_join: provision::SwarmJoinSpec { + orch_swactor_addr: "quic://orch.example:9443".to_owned(), + join_token_ref: "secret://join".to_owned(), + expected_logical_node_id: logical_node_id, + }, + } +} + +#[test] +fn vastai_plugin_builds_one_node_request_and_starts_bootstrap() { + let mut plugin = VastAiProvisioningPlugin::new( + FakeLeaseClient::default().with_contract(100), + FakeBootstrap::default(), + config(), + ); + + let handle = plugin.start_node(spec(), sink()).unwrap(); + + assert_eq!(handle.id, 1); + assert_eq!(handle.provider_process_id, None); + assert_eq!(plugin.active_contract_count(), 1); + + let request = &plugin.client().requests[0]; + assert_eq!(request.count, 1); + assert_eq!(request.image, "registry.example/mvp-worker:latest"); + assert_eq!(request.label.as_deref(), Some("test-mvp-9-11")); + assert_eq!(request.disk_gb, 80); + assert_eq!(request.onstart.as_deref(), Some("python worker.py")); + assert_eq!(request.env.get("EXISTING").map(String::as_str), Some("1")); + + assert_eq!( + plugin.client().endpoint_lookups, + vec![(100, "test-mvp-9-11".to_owned(), "ubuntu".to_owned())] + ); + assert_eq!(plugin.bootstrap().starts.len(), 1); + assert_eq!(plugin.bootstrap().starts[0].0.env, spec().env); + assert_eq!( + plugin.bootstrap().starts[0].1, + VastAiSshEndpoint { + host: "ssh5.vast.ai".to_owned(), + port: 22017, + user: "ubuntu".to_owned(), + } + ); +} + +#[test] +fn stop_destroys_known_vastai_contract_exactly_once() { + let mut plugin = VastAiProvisioningPlugin::new( + FakeLeaseClient::default().with_contract(100), + FakeBootstrap::default(), + config(), + ); + let handle = plugin.start_node(spec(), sink()).unwrap(); + + plugin.stop_node(&handle).unwrap(); + plugin.stop_node(&handle).unwrap(); + + assert_eq!(plugin.client().destroyed, vec![100]); + assert_eq!(plugin.bootstrap().stops, vec![1]); + assert_eq!(plugin.active_contract_count(), 0); +} + +#[test] +fn start_failure_after_contract_creation_destroys_contract_once() { + let mut client = FakeLeaseClient::default().with_contract(100); + client + .endpoint_results + .push_back(Err("ssh missing".to_owned())); + let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config()); + + let error = plugin.start_node(spec(), sink()).unwrap_err(); + + assert!(error.contains("vastai SSH endpoint node 11: ssh missing")); + assert_eq!(plugin.client().destroyed, vec![100]); + assert_eq!(plugin.active_contract_count(), 0); +} + +#[test] +fn start_failure_reports_original_error_and_cleanup_failure() { + let mut client = FakeLeaseClient::default().with_contract(100); + client + .endpoint_results + .push_back(Err("ssh missing".to_owned())); + client.destroy_result = Some(Err("destroy refused".to_owned())); + let mut plugin = VastAiProvisioningPlugin::new(client, FakeBootstrap::default(), config()); + + let error = plugin.start_node(spec(), sink()).unwrap_err(); + + assert!(error.contains("vastai SSH endpoint node 11: ssh missing")); + assert!(error.contains("cleanup destroy 100 failed: destroy refused")); + assert_eq!(plugin.client().destroyed, vec![100]); + assert_eq!(plugin.active_contract_count(), 0); +} + +#[test] +fn vastai_provider_plugin_maps_contracts_into_node_manager_lease_model() { + let mut provider = + VastAiProviderPlugin::new(FakeLeaseClient::default().with_contract(100), config()); + let spec = logical_spec(); + + let result = provider + .create_lease(provision::CreateLeaseRequest { spec: spec.clone() }) + .expect("vastai lease succeeds"); + + assert_eq!(result.endpoint, None); + assert_eq!(result.lease.provider, provision::ProviderKind::VastAi); + assert_eq!( + result.lease.lease_id, + provision::ProviderLeaseId("vastai:100".into()) + ); + assert_eq!(result.lease.provider_contract_id, "100"); + assert_eq!( + result.lease.destroy_handle, + provision::DestroyHandle { + provider: provision::ProviderKind::VastAi, + lease_id: provision::ProviderLeaseId("vastai:100".into()), + provider_contract_id: "100".into(), + } + ); + assert_eq!( + result + .lease + .provider_metadata + .get("label") + .map(String::as_str), + Some("test-mvp-9-workers-0") + ); + assert_eq!( + result + .lease + .provider_metadata + .get("ssh_user") + .map(String::as_str), + Some("ubuntu") + ); + + let request = &provider.client().requests[0]; + assert_eq!(request.count, 1); + assert_eq!(request.image, spec.shape.image); + assert_eq!(request.disk_gb, 80); + assert_eq!(request.label.as_deref(), Some("test-mvp-9-workers-0")); + assert_eq!(request.env.get("MVP_RUN_ID").map(String::as_str), Some("9")); + assert_eq!( + request.env.get("MVP_LOGICAL_NODE_ID").map(String::as_str), + Some("workers-0") + ); + assert_eq!(request.selection.gpu_name.as_deref(), Some("RTX 4090")); + assert_eq!(request.selection.min_gpu_ram_mb, Some(20_000)); + assert_eq!(request.selection.min_down_mbps, 150.0); + assert_eq!(request.selection.min_up_mbps, Some(25.0)); + assert_eq!(request.selection.min_reliability, 0.98); + assert!(request.selection.require_verified); + + let endpoint = provider + .lookup_endpoint(&result.lease) + .expect("lookup succeeds") + .expect("vastai lookup yields endpoint"); + assert_eq!( + endpoint, + provision::SshEndpoint { + host: "ssh5.vast.ai".into(), + port: 22017, + user: "ubuntu".into(), + auth_ref: "vastai:100:ssh".into(), + } + ); + assert_eq!( + provider.client().endpoint_lookups, + vec![(100, "test-mvp-9-workers-0".to_owned(), "ubuntu".to_owned())] + ); + + provider + .destroy_lease(&result.lease.destroy_handle) + .expect("destroy succeeds"); + assert_eq!(provider.client().destroyed, vec![100]); +} diff --git a/crates/mvp-system/src/vastai_provisioning.rs b/crates/mvp-system/src/vastai_provisioning.rs new file mode 100644 index 0000000..dfc2676 --- /dev/null +++ b/crates/mvp-system/src/vastai_provisioning.rs @@ -0,0 +1,619 @@ +use parking_lot::Mutex; +use std::collections::BTreeMap; +use std::process::{Child, Command, Stdio}; +use std::sync::Arc; + +use datastream::DatastreamProducer; +use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy}; + +use crate::bootstrap_datastream::{BootstrapDatastreamBridge, node_stream_id}; +use crate::node_provisioning::{ + CreateLeaseRequest, CreateLeaseResult, DestroyHandle, LeaseFacts, LogicalNodeSpec, + ProviderError, ProviderKind, ProviderLeaseId, ProviderPlugin, SshEndpoint, +}; +use crate::provisioning::{ + NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginSink, ProvisionPlugin, +}; + +#[derive(Clone, Debug)] +pub struct VastAiProvisioningConfig { + pub label_prefix: String, + pub disk_gb: u32, + pub ssh_user: String, + pub selection: SelectionPolicy, + pub lifecycle: LifecyclePolicy, + pub confirm_lease: bool, + pub onstart: Option, +} + +impl Default for VastAiProvisioningConfig { + fn default() -> Self { + Self { + label_prefix: "mvp".to_owned(), + disk_gb: 50, + ssh_user: "root".to_owned(), + selection: SelectionPolicy::default(), + lifecycle: LifecyclePolicy::default(), + confirm_lease: false, + onstart: None, + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct VastAiSshEndpoint { + pub host: String, + pub port: u16, + pub user: String, +} + +pub trait VastAiLeaseClient: Send { + fn provision_one(&mut self, request: ProvisionRequest) -> Result; + + fn ssh_endpoint( + &mut self, + contract_id: u64, + label: &str, + lifecycle: &LifecyclePolicy, + ssh_user: &str, + ) -> Result; + + fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String>; +} + +pub struct ToolsVastAiLeaseClient { + client: swactor_vastai::VastClient, + runtime: tokio::runtime::Runtime, +} + +impl ToolsVastAiLeaseClient { + pub fn new(client: swactor_vastai::VastClient) -> Result { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| format!("vastai tokio runtime: {e}"))?; + Ok(Self { client, runtime }) + } + + pub fn from_api_key(api_key: impl Into) -> Result { + Self::new(swactor_vastai::VastClient::new(api_key)) + } + + pub fn client(&self) -> &swactor_vastai::VastClient { + &self.client + } +} + +impl VastAiLeaseClient for ToolsVastAiLeaseClient { + fn provision_one(&mut self, request: ProvisionRequest) -> Result { + let fleet = self.runtime.block_on(self.client.provision(request))?; + let mut instances = fleet.instances; + if instances.len() != 1 { + return Err(format!( + "vastai provision expected one instance, got {}", + instances.len() + )); + } + Ok(instances.remove(0)) + } + + fn ssh_endpoint( + &mut self, + contract_id: u64, + label: &str, + lifecycle: &LifecyclePolicy, + ssh_user: &str, + ) -> Result { + self.runtime.block_on(async { + let instances = self.client.list_by_label(label).await?; + if let Some(instance) = instances + .into_iter() + .find(|instance| instance.contract_id == contract_id) + { + let host = if instance.ssh_host.is_empty() { + instance.public_ipaddr + } else { + instance.ssh_host + }; + return endpoint_from_parts(contract_id, host, instance.ssh_port, ssh_user); + } + + let running = self.client.wait_for_running(contract_id, lifecycle).await?; + endpoint_from_parts(contract_id, running.ip, running.port, ssh_user) + }) + } + + fn destroy_contract(&mut self, contract_id: u64) -> Result<(), String> { + self.runtime + .block_on(swactor_vastai::destroy_instance_with_retry( + self.client.http(), + self.client.base_url(), + self.client.api_key(), + contract_id, + )) + } +} + +fn endpoint_from_parts( + contract_id: u64, + host: String, + port: u16, + ssh_user: &str, +) -> Result { + if host.is_empty() || host == "unknown" { + return Err(format!("vastai contract {contract_id} has no SSH host")); + } + if port == 0 { + return Err(format!("vastai contract {contract_id} has no SSH port")); + } + Ok(VastAiSshEndpoint { + host, + port, + user: ssh_user.to_owned(), + }) +} + +#[derive(Clone, Debug)] +pub struct VastAiProviderPlugin +where + C: VastAiLeaseClient, +{ + client: C, + config: VastAiProvisioningConfig, +} + +impl VastAiProviderPlugin +where + C: VastAiLeaseClient, +{ + pub fn new(client: C, config: VastAiProvisioningConfig) -> Self { + Self { client, config } + } + + pub fn client(&self) -> &C { + &self.client + } + + pub fn client_mut(&mut self) -> &mut C { + &mut self.client + } + + pub fn config(&self) -> &VastAiProvisioningConfig { + &self.config + } + + fn label_for(&self, spec: &LogicalNodeSpec) -> String { + format!( + "{}-{}-{}", + self.config.label_prefix, spec.run_id.0, spec.logical_node_id.0 + ) + } + + fn selection_for(&self, spec: &LogicalNodeSpec) -> SelectionPolicy { + let mut selection = self.config.selection.clone(); + if let Some(gpu_name) = &spec.shape.gpu_name { + selection.gpu_name = Some(gpu_name.clone()); + } + if let Some(min_gpu_ram_mb) = spec.shape.min_gpu_ram_mb { + selection.min_gpu_ram_mb = Some(min_gpu_ram_mb); + } + if let Some(min_down_mbps) = spec.shape.min_down_mbps { + selection.min_down_mbps = min_down_mbps; + } + if let Some(min_up_mbps) = spec.shape.min_up_mbps { + selection.min_up_mbps = Some(min_up_mbps); + } + if let Some(min_reliability) = spec.shape.min_reliability { + selection.min_reliability = min_reliability; + } + selection.require_verified = spec.shape.require_verified; + selection + } + + fn build_request(&self, spec: &LogicalNodeSpec, label: String) -> ProvisionRequest { + let mut env = BTreeMap::new(); + env.insert("MVP_RUN_ID".into(), spec.run_id.0.to_string()); + env.insert("MVP_LOGICAL_NODE_ID".into(), spec.logical_node_id.0.clone()); + env.insert( + "MVP_ORCH_SWACTOR_ADDR".into(), + spec.swarm_join.orch_swactor_addr.clone(), + ); + env.insert( + "MVP_JOIN_TOKEN_REF".into(), + spec.swarm_join.join_token_ref.clone(), + ); + + ProvisionRequest { + count: 1, + image: spec.shape.image.clone(), + label: Some(label), + disk_gb: spec.shape.disk_gb, + env, + per_instance_env: vec![BTreeMap::new()], + onstart: self.config.onstart.clone(), + selection: self.selection_for(spec), + lifecycle: self.config.lifecycle.clone(), + confirm_lease: self.config.confirm_lease, + } + } + + fn lease_from_instance( + spec: &LogicalNodeSpec, + label: &str, + instance: ProvisionedInstance, + ) -> LeaseFacts { + let contract_id = instance.contract_id.to_string(); + let lease_id = ProviderLeaseId(format!("vastai:{contract_id}")); + let mut provider_metadata = BTreeMap::new(); + provider_metadata.insert("contract_id".into(), contract_id.clone()); + provider_metadata.insert("label".into(), label.to_owned()); + provider_metadata.insert("image".into(), spec.shape.image.clone()); + provider_metadata.insert("logical_node_id".into(), spec.logical_node_id.0.clone()); + provider_metadata.insert("ssh_user".into(), spec.boot.ssh_user.clone()); + provider_metadata.insert("offer_id".into(), instance.offer_id.to_string()); + provider_metadata.insert("gpu_name".into(), instance.gpu_name); + provider_metadata.insert("dph_total".into(), instance.dph_total.to_string()); + if let Some(host_id) = instance.host_id { + provider_metadata.insert("host_id".into(), host_id.to_string()); + } + if let Some(gpu_ram) = instance.gpu_ram { + provider_metadata.insert("gpu_ram".into(), gpu_ram.to_string()); + } + + LeaseFacts { + provider: ProviderKind::VastAi, + lease_id: lease_id.clone(), + provider_contract_id: contract_id.clone(), + offer_id: provider_metadata.get("offer_id").cloned(), + destroy_handle: DestroyHandle { + provider: ProviderKind::VastAi, + lease_id, + provider_contract_id: contract_id, + }, + provider_metadata, + } + } + + fn contract_id(lease: &LeaseFacts) -> Result { + if lease.provider != ProviderKind::VastAi { + return Err(ProviderError::new( + "vastai provider received non-vastai lease", + )); + } + lease + .provider_contract_id + .parse::() + .map_err(|e| ProviderError::new(format!("invalid vastai contract id: {e}"))) + } +} + +impl ProviderPlugin for VastAiProviderPlugin +where + C: VastAiLeaseClient, +{ + fn create_lease( + &mut self, + request: CreateLeaseRequest, + ) -> Result { + if request.spec.provider != ProviderKind::VastAi { + return Err(ProviderError::new( + "vastai provider received non-vastai node spec", + )); + } + let label = self.label_for(&request.spec); + let provision_request = self.build_request(&request.spec, label.clone()); + let instance = self + .client + .provision_one(provision_request) + .map_err(ProviderError::new)?; + let lease = Self::lease_from_instance(&request.spec, &label, instance); + Ok(CreateLeaseResult { + lease, + endpoint: None, + }) + } + + fn lookup_endpoint( + &mut self, + lease: &LeaseFacts, + ) -> Result, ProviderError> { + let contract_id = Self::contract_id(lease)?; + let label = lease + .provider_metadata + .get("label") + .ok_or_else(|| ProviderError::new("vastai lease missing label"))?; + let ssh_user = lease + .provider_metadata + .get("ssh_user") + .map(String::as_str) + .unwrap_or(&self.config.ssh_user); + let endpoint = self + .client + .ssh_endpoint(contract_id, label, &self.config.lifecycle, ssh_user) + .map_err(ProviderError::new)?; + Ok(Some(SshEndpoint { + host: endpoint.host, + port: endpoint.port, + user: endpoint.user, + auth_ref: format!("vastai:{contract_id}:ssh"), + })) + } + + fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError> { + if handle.provider != ProviderKind::VastAi { + return Err(ProviderError::new( + "vastai provider received non-vastai destroy handle", + )); + } + let contract_id = handle + .provider_contract_id + .parse::() + .map_err(|e| ProviderError::new(format!("invalid vastai contract id: {e}")))?; + self.client + .destroy_contract(contract_id) + .map_err(ProviderError::new) + } +} + +pub trait VastAiBootstrapLauncher: Send { + type Handle: Send; + + fn start_bootstrap( + &mut self, + spec: NodeProvisionSpec, + endpoint: VastAiSshEndpoint, + sink: PluginSink, + producer: Option, + ) -> Result; + + fn stop_bootstrap(&mut self, handle: &mut Self::Handle); +} + +#[derive(Default, Debug, Clone, Copy)] +pub struct SshCommandBootstrapLauncher; + +pub struct SshCommandBootstrapHandle { + child: Arc>, +} + +impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher { + type Handle = SshCommandBootstrapHandle; + + fn start_bootstrap( + &mut self, + spec: NodeProvisionSpec, + endpoint: VastAiSshEndpoint, + sink: PluginSink, + producer: Option, + ) -> Result { + if spec.args.is_empty() { + return Err(format!( + "VastAI node {} SSH bootstrap command missing", + spec.node_id + )); + } + + let mut command = Command::new("ssh"); + command + .arg("-p") + .arg(endpoint.port.to_string()) + .arg("-o") + .arg("BatchMode=yes") + .arg(format!("{}@{}", endpoint.user, endpoint.host)) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + command.arg(spec.args.join(" ")); + + let mut child = command + .spawn() + .map_err(|e| format!("spawn VastAI SSH bootstrap {}: {e}", spec.node_id))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| format!("VastAI node {} SSH stdout missing", spec.node_id))?; + let stderr = child + .stderr + .take() + .ok_or_else(|| format!("VastAI node {} SSH stderr missing", spec.node_id))?; + + let bridge = BootstrapDatastreamBridge::new(spec, sink, producer); + bridge.spawn_stdout_reader(stdout); + bridge.spawn_stderr_reader(stderr); + + Ok(SshCommandBootstrapHandle { + child: Arc::new(Mutex::new(child)), + }) + } + + fn stop_bootstrap(&mut self, handle: &mut Self::Handle) { + let mut child = handle.child.lock(); + let _ = child.kill(); + let _ = child.wait(); + } +} + +pub struct VastAiProvisioningPlugin +where + C: VastAiLeaseClient, + B: VastAiBootstrapLauncher, +{ + client: C, + bootstrap: B, + config: VastAiProvisioningConfig, + bootstrap_producer: Option, + next_handle_id: u64, + nodes: BTreeMap>, +} + +struct VastAiNode { + contract_id: u64, + bootstrap: Option, +} + +impl VastAiProvisioningPlugin +where + C: VastAiLeaseClient, + B: VastAiBootstrapLauncher, +{ + pub fn new(client: C, bootstrap: B, config: VastAiProvisioningConfig) -> Self { + Self { + client, + bootstrap, + config, + bootstrap_producer: None, + next_handle_id: 1, + nodes: BTreeMap::new(), + } + } + + pub fn with_bootstrap_producer(mut self, producer: DatastreamProducer) -> Self { + self.bootstrap_producer = Some(producer); + self + } + + pub fn client(&self) -> &C { + &self.client + } + + pub fn client_mut(&mut self) -> &mut C { + &mut self.client + } + + pub fn bootstrap(&self) -> &B { + &self.bootstrap + } + + pub fn bootstrap_mut(&mut self) -> &mut B { + &mut self.bootstrap + } + + pub fn config(&self) -> &VastAiProvisioningConfig { + &self.config + } + + pub fn active_contract_count(&self) -> usize { + self.nodes.len() + } + + fn label_for(&self, spec: &NodeProvisionSpec) -> String { + format!( + "{}-{}-{}", + self.config.label_prefix, spec.run_id, spec.node_id + ) + } + + fn build_request(&self, spec: &NodeProvisionSpec, label: String) -> ProvisionRequest { + let env = spec.env.iter().cloned().collect::>(); + ProvisionRequest { + count: 1, + image: spec.image.clone(), + label: Some(label), + disk_gb: self.config.disk_gb, + env, + per_instance_env: vec![BTreeMap::new()], + onstart: self + .config + .onstart + .clone() + .or_else(|| (!spec.args.is_empty()).then(|| spec.args.join(" "))), + selection: self.config.selection.clone(), + lifecycle: self.config.lifecycle.clone(), + confirm_lease: self.config.confirm_lease, + } + } + + fn cleanup_contract_after_start_error(&mut self, contract_id: u64, reason: String) -> String { + match self.client.destroy_contract(contract_id) { + Ok(()) => reason, + Err(cleanup) => format!("{reason}; cleanup destroy {contract_id} failed: {cleanup}"), + } + } +} + +impl ProvisionPlugin for VastAiProvisioningPlugin +where + C: VastAiLeaseClient + 'static, + B: VastAiBootstrapLauncher + 'static, +{ + fn start_node( + &mut self, + spec: NodeProvisionSpec, + sink: PluginSink, + ) -> Result { + let stream_id = node_stream_id(spec.run_id, spec.node_id); + let label = self.label_for(&spec); + sink.observe(PluginObservation::ProviderLine { + run_id: spec.run_id, + node_id: spec.node_id, + line: format!("vastai provisioning label={label} stream={stream_id}"), + }); + + let request = self.build_request(&spec, label.clone()); + let instance = self + .client + .provision_one(request) + .map_err(|e| format!("vastai provision node {}: {e}", spec.node_id))?; + sink.observe(PluginObservation::ProviderLine { + run_id: spec.run_id, + node_id: spec.node_id, + line: format!( + "vastai contract {} ready for SSH lookup", + instance.contract_id + ), + }); + + let endpoint = match self.client.ssh_endpoint( + instance.contract_id, + &label, + &self.config.lifecycle, + &self.config.ssh_user, + ) { + Ok(endpoint) => endpoint, + Err(error) => { + return Err(self.cleanup_contract_after_start_error( + instance.contract_id, + format!("vastai SSH endpoint node {}: {error}", spec.node_id), + )); + } + }; + + let bootstrap = match self.bootstrap.start_bootstrap( + spec.clone(), + endpoint, + sink, + self.bootstrap_producer.clone(), + ) { + Ok(handle) => handle, + Err(error) => { + return Err(self.cleanup_contract_after_start_error( + instance.contract_id, + format!("vastai bootstrap node {}: {error}", spec.node_id), + )); + } + }; + + let handle = PluginNodeHandle { + id: self.next_handle_id, + provider_process_id: None, + }; + self.next_handle_id = self.next_handle_id.wrapping_add(1).max(1); + self.nodes.insert( + handle.id, + VastAiNode { + contract_id: instance.contract_id, + bootstrap: Some(bootstrap), + }, + ); + Ok(handle) + } + + fn stop_node(&mut self, handle: &PluginNodeHandle) -> Result<(), String> { + let Some(mut node) = self.nodes.remove(&handle.id) else { + return Ok(()); + }; + if let Some(mut bootstrap) = node.bootstrap.take() { + self.bootstrap.stop_bootstrap(&mut bootstrap); + } + self.client.destroy_contract(node.contract_id) + } +} diff --git a/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py b/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py index 7b4a00a..ca3ecdf 100755 --- a/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py +++ b/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py @@ -94,6 +94,22 @@ def configure_role(cmd: dict[str, Any]) -> None: control(type="RoleConfigured", role_id=role["role_id"]) +def load_weights(cmd: dict[str, Any]) -> None: + if not role: + fatal("RoleNotConfigured") + role["model_id"] = cmd["model_id"] + role["gguf_source"] = cmd["gguf_source"] + role["tokenizer"] = cmd["tokenizer"] + role["weight_layer_start"] = int(cmd["layer_start"]) + role["weight_layer_end_exclusive"] = int(cmd["layer_end_exclusive"]) + control( + type="WeightsLoaded", + model_id=role["model_id"], + layer_start=role["weight_layer_start"], + layer_end_exclusive=role["weight_layer_end_exclusive"], + ) + + def parse_record(ring: dict[str, Any]) -> tuple[int, int, int, bytes]: view = require_arena() base = ring["data_offset"] @@ -236,6 +252,7 @@ HANDLERS = { "InitializeWorker": initialize, "InstallRing": install_ring, "ConfigureRole": configure_role, + "LoadWeights": load_weights, "RingReadable": ring_readable, "ExecuteStep": execute_step, "ReleaseDeviceObject": release_device_object, diff --git a/crates/mvp-system/tests/one_node_chat_e2e.rs b/crates/mvp-system/tests/one_node_chat_e2e.rs new file mode 100644 index 0000000..378b217 --- /dev/null +++ b/crates/mvp-system/tests/one_node_chat_e2e.rs @@ -0,0 +1,362 @@ +use std::io::{Read, Write}; +use std::net::TcpStream; +use std::process::{Child, Command, Stdio}; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::{Duration, Instant}; + +use serde_json::Value; + +#[cfg(target_os = "linux")] +use std::os::unix::process::CommandExt; + +const TEST_TIMEOUT: Duration = Duration::from_secs(1_800); +const PROMPT_TIMEOUT: Duration = Duration::from_secs(600); +const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(60); +const DASHBOARD_ADDR: &str = "127.0.0.1:9090"; +const DEFAULT_CONTAINER: &str = "mvp-orch-one-node-1-1"; + +#[test] +fn one_node_chat_docker_cuda_e2e() { + let root = workspace_root(); + require_docker(&root); + + let mut command = Command::new("cargo"); + command + .current_dir(&root) + .args(["mvp-chat"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + #[cfg(target_os = "linux")] + unsafe { + command.pre_exec(|| { + if libc::setpgid(0, 0) == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } + }); + } + + let mut child = command.spawn().expect("spawn cargo mvp-chat"); + let mut stdin = child.stdin.take().expect("cargo mvp-chat stdin"); + let stdout = Arc::new(Mutex::new(String::new())); + let stderr = Arc::new(Mutex::new(String::new())); + let stdout_reader = spawn_capture(child.stdout.take().expect("stdout"), Arc::clone(&stdout)); + let stderr_reader = spawn_capture(child.stderr.take().expect("stderr"), Arc::clone(&stderr)); + + let result = run_full_flow(&mut child, &mut stdin, &stdout, &stderr); + if result.is_err() { + request_child_interrupt(&child); + let _ = wait_child(&mut child, SHUTDOWN_TIMEOUT); + let _ = child.kill(); + let _ = child.wait(); + } + let _ = stdout_reader.join(); + let _ = stderr_reader.join(); + assert_container_removed(&root, DEFAULT_CONTAINER); + + if let Err(error) = result { + panic!( + "{error}\nstdout:\n{}\nstderr:\n{}\ndashboard frames:\n{}", + snapshot(&stdout), + snapshot(&stderr), + dashboard_snapshot().unwrap_or_else(|err| format!("")) + ); + } +} + +fn run_full_flow( + child: &mut Child, + stdin: &mut impl Write, + stdout: &Arc>, + stderr: &Arc>, +) -> Result<(), String> { + wait_for_child_or(TEST_TIMEOUT, child, || { + stderr_contains(stderr, "dashboard_ready") + }) + .map_err(|e| format!("dashboard_ready not observed: {e}"))?; + wait_for_child_or(TEST_TIMEOUT, child, dashboard_responding) + .map_err(|e| format!("dashboard API not live: {e}"))?; + wait_for_child_or(TEST_TIMEOUT, child, || { + dashboard_has_channel_or_payload("mvp.provisioning.logs", "mvp-entrypoint") + }) + .map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?; + wait_for_child_or(TEST_TIMEOUT, child, || { + dashboard_has_frame("mvp.worker.weights", "GgufDownloadProgress") + }) + .map_err(|e| format!("GGUF download progress not visible in dashboard: {e}"))?; + wait_for_child_or(TEST_TIMEOUT, child, || { + dashboard_has_frame("mvp.worker.weights", "WeightsLoaded") + }) + .map_err(|e| format!("WeightsLoaded not visible in dashboard: {e}"))?; + wait_for_child_or(TEST_TIMEOUT, child, || { + stderr_contains(stderr, "prompt_loop_ready") + }) + .map_err(|e| format!("prompt_loop_ready not observed: {e}"))?; + + writeln!(stdin, "hello from full cargo mvp-chat e2e") + .map_err(|e| format!("write prompt: {e}"))?; + stdin.flush().map_err(|e| format!("flush prompt: {e}"))?; + wait_for_child_or(PROMPT_TIMEOUT, child, || stdout_contains(stdout, ">")) + .map_err(|e| format!("prompt marker not visible: {e}"))?; + wait_for_child_or(PROMPT_TIMEOUT, child, || response_text_visible(stdout)) + .map_err(|e| format!("decoded response text not visible: {e}"))?; + wait_for_child_or(PROMPT_TIMEOUT, child, || { + stderr_contains(stderr, "done request=") + }) + .map_err(|e| format!("prompt did not complete: {e}"))?; + wait_for_child_or(PROMPT_TIMEOUT, child, || { + dashboard_has_frame("mvp.worker.prompt", "PromptCompleted") + }) + .map_err(|e| format!("prompt result not visible in dashboard: {e}"))?; + + request_child_interrupt(child); + let status = wait_child(child, SHUTDOWN_TIMEOUT) + .ok_or_else(|| "cargo mvp-chat did not exit after Ctrl-C".to_owned())?; + if status.success() || status.code() == Some(130) || status.signal_name() == Some("SIGINT") { + Ok(()) + } else { + Err(format!("cargo mvp-chat exited with {status}")) + } +} + +fn wait_for_child_or( + timeout: Duration, + child: &mut Child, + mut predicate: impl FnMut() -> bool, +) -> Result<(), String> { + let start = Instant::now(); + while start.elapsed() < timeout { + if predicate() { + return Ok(()); + } + if let Some(status) = child.try_wait().map_err(|e| format!("poll child: {e}"))? { + return Err(format!("child exited early with {status}")); + } + thread::sleep(Duration::from_millis(250)); + } + Err("timed out".to_owned()) +} + +fn spawn_capture( + mut reader: impl Read + Send + 'static, + out: Arc>, +) -> thread::JoinHandle<()> { + thread::spawn(move || { + let mut buf = [0_u8; 8192]; + loop { + match reader.read(&mut buf) { + Ok(0) => break, + Ok(n) => out + .lock() + .expect("capture mutex") + .push_str(&String::from_utf8_lossy(&buf[..n])), + Err(_) => break, + } + } + }) +} + +fn dashboard_responding() -> bool { + dashboard_frames().is_ok() +} + +fn dashboard_has_channel_or_payload(channel_substr: &str, payload_substr: &str) -> bool { + dashboard_frames() + .map(|frames| { + frames.iter().any(|frame| { + frame.channel.contains(channel_substr) || frame.payload.contains(payload_substr) + }) + }) + .unwrap_or(false) +} + +fn dashboard_has_frame(channel_substr: &str, payload_substr: &str) -> bool { + dashboard_frames() + .map(|frames| { + frames.iter().any(|frame| { + frame.channel.contains(channel_substr) && frame.payload.contains(payload_substr) + }) + }) + .unwrap_or(false) +} + +#[derive(Debug)] +struct SeenFrame { + channel: String, + payload: String, +} + +fn dashboard_frames() -> Result, String> { + let response = http_get("/api/frames")?; + let (_, body) = response + .split_once("\r\n\r\n") + .ok_or_else(|| "HTTP response missing body".to_owned())?; + let values = serde_json::from_str::>(body) + .map_err(|e| format!("parse dashboard frames JSON: {e}; body={body:?}"))?; + Ok(values + .into_iter() + .map(|value| SeenFrame { + channel: value + .get("channel") + .and_then(Value::as_str) + .unwrap_or_default() + .to_owned(), + payload: decode_payload(value.get("payload")).unwrap_or_default(), + }) + .collect()) +} + +fn dashboard_snapshot() -> Result { + let mut frames = dashboard_frames()?; + let keep = frames.len().saturating_sub(40); + frames.drain(0..keep); + Ok(frames + .into_iter() + .map(|frame| format!("{} {}", frame.channel, frame.payload)) + .collect::>() + .join("\n")) +} + +fn decode_payload(value: Option<&Value>) -> Option { + let bytes = value? + .as_array()? + .iter() + .map(|byte| byte.as_u64().map(|n| n as u8)) + .collect::>>()?; + Some(String::from_utf8_lossy(&bytes).to_string()) +} + +fn http_get(path: &str) -> Result { + let mut stream = TcpStream::connect(DASHBOARD_ADDR) + .map_err(|e| format!("connect dashboard {DASHBOARD_ADDR}: {e}"))?; + stream + .set_read_timeout(Some(Duration::from_secs(2))) + .map_err(|e| format!("set read timeout: {e}"))?; + write!( + stream, + "GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" + ) + .map_err(|e| format!("write HTTP request: {e}"))?; + let mut response = String::new(); + stream + .read_to_string(&mut response) + .map_err(|e| format!("read HTTP response: {e}"))?; + if response.starts_with("HTTP/1.1 200") { + Ok(response) + } else { + Err(format!("non-200 dashboard response: {response:?}")) + } +} + +fn stdout_contains(stdout: &Arc>, needle: &str) -> bool { + snapshot(stdout).contains(needle) +} + +fn response_text_visible(stdout: &Arc>) -> bool { + let text = snapshot(stdout); + text.lines() + .any(|line| line.trim_start_matches('>').trim().len() > 8) +} + +fn stderr_contains(stderr: &Arc>, needle: &str) -> bool { + snapshot(stderr).contains(needle) +} + +fn snapshot(buf: &Arc>) -> String { + buf.lock().expect("capture mutex").clone() +} + +fn wait_child(child: &mut Child, timeout: Duration) -> Option { + let start = Instant::now(); + while start.elapsed() < timeout { + if let Some(status) = child.try_wait().expect("poll child") { + return Some(status); + } + thread::sleep(Duration::from_millis(100)); + } + None +} + +fn request_child_interrupt(child: &Child) { + #[cfg(target_os = "linux")] + unsafe { + let _ = libc::kill(-(child.id() as libc::pid_t), libc::SIGINT); + } + #[cfg(not(target_os = "linux"))] + { + let _ = child; + } +} + +trait ExitStatusSignalName { + fn signal_name(&self) -> Option<&'static str>; +} + +impl ExitStatusSignalName for std::process::ExitStatus { + fn signal_name(&self) -> Option<&'static str> { + #[cfg(target_os = "linux")] + { + use std::os::unix::process::ExitStatusExt; + match self.signal() { + Some(libc::SIGINT) => Some("SIGINT"), + Some(libc::SIGTERM) => Some("SIGTERM"), + _ => None, + } + } + #[cfg(not(target_os = "linux"))] + { + let _ = self; + None + } + } +} + +fn require_docker(root: &std::path::Path) { + let version = Command::new("docker") + .current_dir(root) + .arg("version") + .output() + .expect("run docker version"); + assert!( + version.status.success(), + "docker is not available\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&version.stdout), + String::from_utf8_lossy(&version.stderr) + ); +} + +fn assert_container_removed(root: &std::path::Path, container: &str) { + let output = Command::new("docker") + .current_dir(root) + .args([ + "ps", + "-a", + "--filter", + &format!("name=^{container}$"), + "--format", + "{{.Names}}", + ]) + .output() + .expect("run docker ps"); + assert!( + output.status.success(), + "docker ps failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert!( + String::from_utf8_lossy(&output.stdout).trim().is_empty(), + "container {container} still exists" + ); +} + +fn workspace_root() -> std::path::PathBuf { + std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .and_then(std::path::Path::parent) + .expect("workspace root") + .to_path_buf() +} diff --git a/mvp-node-image-gguf-shape.md b/mvp-node-image-gguf-shape.md new file mode 100644 index 0000000..688fe3f --- /dev/null +++ b/mvp-node-image-gguf-shape.md @@ -0,0 +1,167 @@ +# MVP node image + GGUF deployment shape + +## Goal + +Prepare the first deployment-test node image for one rented GPU node. + +The image should be minimal: swactor MVP node binary plus tinygrad with CUDA support. The node should stay at the swactor level; Python is only the compute helper, not the top-level application runtime. + +The first GGUF target is a small ~1B parameter model fetched whole from HuggingFace and cached on the node. + +## Existing reference + +The closest existing image is in `apps/old-pipeline-parallel-inference`: + +- `Dockerfile.base` + - CUDA runtime/base image layering; + - Python, tinygrad, numpy; + - NVRTC runtime; + - CUDA headers copied from a builder stage; + - sshd; + - PID-1 entrypoint. + +- `Dockerfile` + - thin code layer over the base image; + - copies Rust binaries and worker script. + +- `pp_entrypoint.sh` + - starts sshd deterministically; + - accepts VastAI-injected public keys; + - runs the worker as a child; + - keeps the container alive after worker exit for postmortem. + +- `pp_tinygrad_worker.py` + - contains practical GGUF/tinygrad loading lessons; + - fetches model artifacts; + - uses tinygrad tokenizer/model code; + - old runtime enters the Python worker protocol directly. + +The new image should reuse the packaging/runtime lessons, not the old pipeline-parallel application shape. + +## Runtime architecture + +```text +VastAI or Docker provisioner + -> container entrypoint starts sshd + -> entrypoint starts mvp node binary + -> mvp node joins orchestrator over iroh/swactor + -> mvp node registers node actor(s) + -> mvp node starts tinygrad helper as a child process + -> helper fetches/caches GGUF + -> helper loads model on CUDA + -> orchestrator drives prompt/inference via swactor messages +``` + +## Binary shape + +Add or identify a dedicated deployment node binary, for example: + +```text +crates/mvp-system/src/bin/mvp-node.rs +``` + +Responsibilities: + +- read deployment env/config; +- start iroh/swactor runtime; +- join the orchestrator; +- register the MVP node actor path; +- start and supervise the tinygrad helper process; +- emit a ready JSON line on stdout once the swactor node is reachable; +- continue pumping runtime until shutdown or fault. + +This should be the deployment form of the existing local node role, not coupled to local-only TCP/test assumptions. + +## Image shape + +Layering should follow the old image pattern: + +```text +mvp-node-base + CUDA runtime + NVRTC runtime library + CUDA headers required by tinygrad runtime JIT + python3 + tinygrad + numpy + openssh-server + mvp entrypoint + +mvp-node + /usr/local/bin/mvp-node + /usr/local/share/mvp/tinygrad_worker.py +``` + +NVRTC is needed because tinygrad's CUDA backend JIT-compiles kernels at runtime. The old image installs the runtime library in the final image and copies CUDA headers from a builder stage; keep that lesson unless a better tinygrad-compatible base is chosen. + +## Entrypoint shape + +The entrypoint should: + +- install/accept SSH public keys from VastAI-compatible env; +- start `sshd` before the node binary; +- print bootstrap progress to stdout/stderr for SSH datastream capture; +- run `/usr/local/bin/mvp-node` as a child; +- tee node output to a log file readable over SSH; +- keep PID 1 alive after node crash for postmortem. + +This mirrors the old `pp_entrypoint.sh` behavior, with MVP names and paths. + +## GGUF shape + +First deployment target: + +- whole-file HuggingFace GGUF fetch; +- small ~1B parameter model; +- node-local cache path; +- progress/fault observations surfaced through the node/datastream path; +- tokenizer from GGUF or adjacent tokenizer source, depending on what the chosen model path supports. + +Existing MVP planning code already has concepts for this: + +- `run_plan.rs::GgufSource::{LocalPath, HuggingFaceGguf}`; +- `TokenizerSource::{EmbeddedGguf, LocalPath}`. + +Known gap: + +- actor/runtime weight commands currently collapse to test artifacts in the node-agent/stage-controller path. +- real `GgufSource` must be carried through to the node/worker load command before this deployment path is real. + +## Tinygrad helper shape + +Python helper remains subordinate to swactor. + +Responsibilities: + +- import tinygrad lazily; +- verify CUDA with a small tensor realization; +- fetch/cache the requested GGUF; +- load the model/tokenizer; +- execute inference operations requested by the swactor node; +- emit structured lifecycle/progress/fault events. + +Do not make Python the deployment app entrypoint. The app entrypoint is the swactor node binary. + +## Local verification target + +This workstation has a GPU and Docker CUDA should work. The deployment image should be verified locally before VastAI. + +Target E2E: + +1. Build base image. +2. Build thin MVP node image. +3. Run with `docker run --gpus all`. +4. Confirm bootstrap stdout/stderr are capturable. +5. Confirm `mvp-node` starts and joins the orchestrator path. +6. Confirm native datastream frame collection from the container. +7. Confirm tinygrad imports and realizes a CUDA tensor. +8. Confirm the chosen GGUF is fetched/cached from HuggingFace. +9. Confirm one prompt/inference request completes. + +The Docker E2E should use the same image shape expected by VastAI; only provisioning differs. + +## Non-goals for this doc + +- designing the full prompt loop schema; +- full VastAI provisioning policy; +- sharded/range GGUF fetching; +- multi-node pipeline-parallel layout. diff --git a/vastai-datastream-plugin-shape.md b/vastai-datastream-plugin-shape.md new file mode 100644 index 0000000..4797d77 --- /dev/null +++ b/vastai-datastream-plugin-shape.md @@ -0,0 +1,110 @@ +# VastAI datastream + provisioning plugin shape + +## Goal + +First deployment preflight needs observability from the moment a VastAI node is provisioned. + +The node should produce datastream frames in two phases: + +1. **Bootstrap phase:** orchestrator reaches the node over SSH and captures remote stdout/stderr as datastream frames. +2. **Native phase:** once the remote swactor node is live, the node sends datastream frames back directly over the normal remote transport path. + +The key invariant: the logical node stream identity stays stable across both phases. The source changes from SSH bridge to native remote datastream; the node stream does not. + +## Motivation + +If the node fails before swactor starts, native actor/datastream paths are unavailable. We still need early boot logs, image startup errors, dependency failures, worker launch errors, and ready-line parsing in the same observation surface used after handoff. + +This avoids a blind gap between `vast.ai contract created` and `remote swactor joined`. + +## Existing pieces to reuse + +- `crates/mvp-system/src/provisioning.rs` + - `ProvisionPlugin` + - `PluginSink` + - `PluginObservation::{ProviderLine, StdoutLine, StderrLine, RuntimeReady, Failed, Exited}` + - Docker already parses stdout ready JSON. + +- `crates/mvp-system/src/actors/provisioner.rs` + - already turns plugin observations into provisioning reports and datastream records when given a `DatastreamProducer`. + +- `crates/mvp-system/src/telemetry.rs` + - provisioning events/log channels already exist. + +- `crates/datastream` + - inspect and reuse existing remote/transport APIs before adding anything new. // USER: Don't add anything new, you should not need to. + - `StreamId` already includes node id + `Lifetime`; do not invent a second lifetime concept. + +- `tools/vastai` + - existing VastAI client/provisioning utility. + - plugin design should wrap this, not duplicate provider API logic. + +## Proposed code shape + +### `crates/mvp-system/src/vastai_provisioning.rs` + +Add an MVP VastAI provisioning plugin around `tools/vastai`. + +Responsibilities: + +- build a one-node VastAI provision request from MVP config/spec inputs; +- create/track the contract handle; +- obtain SSH endpoint details; +- start the SSH bootstrap/datastream bridge; +- emit provider/stdout/stderr/ready observations through `PluginSink`; +- destroy the known contract on stop. + +Keep this focused. Full plugin details need their own design pass: offer policy, spend guard, retries, replacement, recovery, labels, and held-cluster behavior. + +### `crates/mvp-system/src/bootstrap_datastream.rs` + +Small bridge for pre-swactor visibility. + +Responsibilities: + +- read remote stdout/stderr lines from an SSH session or equivalent stream; +- submit those lines as datastream frames/records for the assigned node stream; +- also forward lines to `PluginSink` so existing provisioner reports/dashboard behavior still works; +- parse the same ready JSON shape Docker uses and emit `RuntimeReady`. + +This module should not be VastAI-specific. + +### Remote/native datastream hookup + +Before implementing new APIs, inspect `crates/datastream` for existing remote transport support. + +Desired behavior: + +- provisioner assigns the node datastream identity before lease/bootstrap; +- bootstrap env passes that identity to the remote node; +- remote node starts native datastream emission once swactor/iroh is live; +- remote node emits a native-ready marker; +- orchestrator overlaps SSH capture briefly, then closes the SSH tail. + +If an mvp-system adapter is needed, keep it thin and local to remote datastream receiver/handoff glue. + +## Handoff model + +States: + +- `BootstrapSsh`: SSH bridge is authoritative. +- `Overlap`: first valid native frame or native-ready marker observed; keep SSH briefly. +- `NativeIroh`: native remote datastream is authoritative; SSH tail is closed. + +Do not require native datastream to be available before bootstrap logs start. + +## Test ladder + +1. Fake VastAI provision returns a contract + SSH endpoint. +2. Fake SSH stdout/stderr line becomes a datastream observation. +3. Ready JSON on stdout emits `RuntimeReady` through the existing plugin/provisioner path. +4. Native-ready handoff moves from SSH bridge to native and closes SSH after overlap. +5. Stop destroys the known VastAI contract exactly once. +6. Dockerized E2E submits remote datastream frames over the remote transport and the orchestrator collects them. +7. VastAI plugin uses the same remote-frame receiver path; only lease/SSH acquisition differs. + +## Non-goals for this doc + +- exact VastAI offer-selection policy; +- exact spend/confirmation UX; +- full recovery of abandoned contracts;