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 <zacheryasc@gmail.com>
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-01 12:44:25 +04:00
parent 8363be74cd
commit 6608824cb0
36 changed files with 5445 additions and 128 deletions

View file

@ -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"

View file

@ -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

2
Cargo.lock generated
View file

@ -2477,10 +2477,12 @@ dependencies = [
"iroh",
"iroh-driver",
"libc",
"parking_lot",
"serde",
"serde_json",
"swactor",
"swactor-transport",
"swactor-vastai",
"tokio",
]

15
apps/mvp-node/Dockerfile Normal file
View file

@ -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

View file

@ -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"]

View file

@ -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

34
apps/mvp-node/mvp_entrypoint.sh Executable file
View file

@ -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

299
apps/mvp-node/tinygrad_worker.py Executable file
View file

@ -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)

View file

@ -1,6 +1,6 @@
mod live_explorer;
mod server;
mod store;
mod live_explorer;
pub mod swactor;
pub mod view;

View file

@ -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"

View file

@ -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);
}

View file

@ -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<StageCommandWire>,
events: Vec<StageLifecycleWire>,
@ -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,
},

View file

@ -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::<Vec<_>>();
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<edge_actor::ObjectIdAllocator> = 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<IrohDriver, String> {
@ -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<ChildStdout>,
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<NodeStdoutLine>,
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<NodeStdoutLine>) -> usize {
@ -1044,6 +1100,26 @@ fn drain_node_stdout(rx: &Receiver<NodeStdoutLine>) -> 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()

View file

@ -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<Duration, String> {
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<Duration, String> {
}
fn main() -> ExitCode {
install_signal_handlers();
let args = std::env::args().collect::<Vec<_>>();
let result = if args.iter().any(|arg| arg == "--role=node") {
run_node_role(&args)
@ -243,6 +259,26 @@ struct ProvisionedDockerNode {
provider_process_id: Option<u32>,
provisioner: LocalDockerNodeProvisioner,
events: Receiver<LocalDockerNodeEvent>,
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::<u64, Vec<u8>>::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<IrohDriver, String> {
@ -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<LoadedObject, String> {
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(

View file

@ -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::<PromptEvent>(&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<Self, String> {
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<u64, String> {
match std::env::var(name).ok().filter(|value| !value.is_empty()) {
Some(value) => value
.parse::<u64>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
match std::env::var(name).ok().filter(|value| !value.is_empty()) {
Some(value) => value
.parse::<u32>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn parse_next<T>(args: &mut impl Iterator<Item = String>, name: &str) -> Result<T, String>
where
T: std::str::FromStr,
T::Err: std::fmt::Display,
{
let value = args
.next()
.ok_or_else(|| format!("missing value after {name}"))?;
value
.parse::<T>()
.map_err(|e| format!("invalid {name}={value:?}: {e}"))
}

View file

@ -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::<NodeAgentReport>()
.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<Box<dyn FrameSink>> = 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<dyn FrameSink> = 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<Box<dyn FrameSink>>,
}
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<Self> {
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<EndpointAddr>,
orchestrator_actor: Option<ActorAddress>,
datastream_sink_actor: Option<ActorAddress>,
datastream_frame_log: Option<String>,
relay_mode: iroh::RelayMode,
worker_script: String,
device: String,
model_id: String,
gguf_source: GgufSource,
tokenizer: TokenizerSource,
self_test_prompt: Option<String>,
self_test_layer_end: u32,
self_test_max_tokens: u32,
max_runtime: Duration,
}
impl DeploymentConfig {
fn from_env() -> Result<Self, String> {
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<ChildStdout>,
}
impl TinygradWorker {
fn spawn(config: &DeploymentConfig) -> Result<Self, String> {
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<Value, String> {
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<Option<std::process::ExitStatus>, 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<Value, String> {
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<Value, String> {
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<String> {
std::env::var(name)
.ok()
.map(|value| value.trim().to_owned())
.filter(|value| !value.is_empty())
}
fn env_u64(name: &str, default: u64) -> Result<u64, String> {
match env_optional(name) {
Some(value) => value
.parse::<u64>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
match env_optional(name) {
Some(value) => value
.parse::<u32>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_json<T>(name: &str) -> Result<Option<T>, 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<iroh::RelayMode, String> {
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)
}

View file

@ -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<String>,
rpc_addr: String,
node_image: String,
max_tokens: u32,
timeout_ms: u64,
dashboard: bool,
build_image: bool,
}
impl Config {
fn from_args() -> Result<Self, String> {
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<ChildStdin>,
ready_rx: mpsc::Receiver<String>,
cleaned: bool,
}
impl OrchChild {
fn spawn(config: &Config) -> Result<Self, String> {
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::<Value>(&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<String, String> {
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::<String>();
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::<PromptEvent>(&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<PathBuf, String> {
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<PathBuf, String> {
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<bool, String> {
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<u64, String> {
match std::env::var(name).ok().filter(|value| !value.is_empty()) {
Some(value) => value
.parse::<u64>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
match std::env::var(name).ok().filter(|value| !value.is_empty()) {
Some(value) => value
.parse::<u32>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn next_arg(args: &mut impl Iterator<Item = String>, name: &str) -> Result<String, String> {
args.next()
.ok_or_else(|| format!("missing value after {name}"))
}
fn parse_next<T>(args: &mut impl Iterator<Item = String>, name: &str) -> Result<T, String>
where
T: std::str::FromStr,
T::Err: std::fmt::Display,
{
let value = next_arg(args, name)?;
value
.parse::<T>()
.map_err(|e| format!("invalid {name}={value:?}: {e}"))
}

View file

@ -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::<PromptEvent>()
.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::<PromptWork>();
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::<PluginObservation>();
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<Self, String> {
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<NodeProvisionSpec, String> {
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<PromptEvent>,
}
struct ActivePrompt {
request: SubmitPrompt,
events: mpsc::Sender<PromptEvent>,
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<u8>,
) {
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<Option<Self>, 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::<u16>()
.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<Option<Self>, 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<mpsc::Sender<PluginObservation>>,
}
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<PromptWork>,
default_max_tokens: u32,
default_timeout_ms: u64,
) -> Result<SocketAddr, String> {
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<PromptWork>,
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<PluginObservation>,
frame_rx: &mpsc::Receiver<(StreamId, Frame)>,
dashboard: Option<&DashboardSupport>,
orch_datastream: &mut OrchDatastream,
) -> Result<RuntimeReady, String> {
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<PluginObservation>,
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<PluginObservation>,
frame_rx: &mpsc::Receiver<(StreamId, Frame)>,
work_rx: &mpsc::Receiver<PromptWork>,
prompt_events: &swactor::runtime::Inbox<PromptEvent>,
stop_rx: &mpsc::Receiver<()>,
dashboard: Option<&DashboardSupport>,
orch_datastream: &mut OrchDatastream,
node_actor: ActorAddress,
reply_to: ActorAddress,
) -> Result<(), String> {
let mut active: Option<ActivePrompt> = 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<PluginObservation>,
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::<Value>(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<String> {
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<bool, String> {
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<u64, String> {
match env_optional(name) {
Some(value) => value
.parse::<u64>()
.map_err(|e| format!("invalid {name}={value:?}: {e}")),
None => Ok(default),
}
}
fn env_u32(name: &str, default: u32) -> Result<u32, String> {
match env_optional(name) {
Some(value) => value
.parse::<u32>()
.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<iroh::RelayMode, String> {
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<Item = String>, name: &str) -> Result<String, String> {
args.next()
.ok_or_else(|| format!("missing value after {name}"))
}
fn parse_next<T>(args: &mut impl Iterator<Item = String>, name: &str) -> Result<T, String>
where
T: std::str::FromStr,
T::Err: std::fmt::Display,
{
let value = next_arg(args, name)?;
value
.parse::<T>()
.map_err(|e| format!("invalid {name}={value:?}: {e}"))
}

View file

@ -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<DatastreamProducer>,
}
impl BootstrapDatastreamBridge {
pub fn new(
spec: NodeProvisionSpec,
sink: PluginSink,
producer: Option<DatastreamProducer>,
) -> 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<String>) {
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<String>) {
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<String>) {
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<R>(&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<R>(&self, stderr: R) -> JoinHandle<()>
where
R: Read + Send + 'static,
{
let bridge = self.clone();
thread::spawn(move || bridge.read_stderr(stderr))
}
fn read_stdout<R>(&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<R>(&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<PluginObservation> {
let ready = serde_json::from_str::<RuntimeReadyLine>(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)
}

View file

@ -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;

View file

@ -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::<PromptEvent, _>(JsonCodec::<PromptEvent>::default());
}
pub fn write_json_line<T: Serialize>(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<Option<SubmitPrompt>, 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::<SubmitPrompt>(&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()
);
}
}

View file

@ -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::<RuntimeReadyLine>(&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);
}

View file

@ -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 {

View file

@ -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<String>,
gguf_source: GgufSource,
tokenizer: TokenizerSource,
) -> Self {
Self {
model_id: model_id.into(),
gguf_source,
tokenizer,
}
}
pub fn embedded_gguf(model_id: impl Into<String>, path: impl Into<String>) -> Self {
Self::new(
model_id,
GgufSource::LocalPath(path.into()),
TokenizerSource::EmbeddedGguf,
)
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]

View file

@ -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<Vec<PluginObservation>>,
}
impl RecordingSink {
fn observations(&self) -> Vec<PluginObservation> {
self.observations.lock().clone()
}
}
impl PluginObservationSink for RecordingSink {
fn observe(&self, observation: PluginObservation) {
self.observations.lock().push(observation);
}
}
fn recording_sink() -> (Arc<RecordingSink>, 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
)));
}

View file

@ -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,
),
}
}

View file

@ -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;

View file

@ -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();

View file

@ -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);

View file

@ -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"),
}
}

View file

@ -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<Vec<PluginObservation>>,
}
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<ProvisionRequest>,
endpoint_lookups: Vec<(u64, String, String)>,
endpoint_results: VecDeque<Result<VastAiSshEndpoint, String>>,
destroyed: Vec<u64>,
destroy_result: Option<Result<(), String>>,
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<ProvisionedInstance, String> {
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<VastAiSshEndpoint, String> {
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<usize>,
fail: Option<String>,
next_handle: usize,
}
impl VastAiBootstrapLauncher for FakeBootstrap {
type Handle = usize;
fn start_bootstrap(
&mut self,
spec: NodeProvisionSpec,
endpoint: VastAiSshEndpoint,
_sink: PluginSink,
_producer: Option<datastream::DatastreamProducer>,
) -> Result<Self::Handle, String> {
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]);
}

View file

@ -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<String>,
}
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<ProvisionedInstance, String>;
fn ssh_endpoint(
&mut self,
contract_id: u64,
label: &str,
lifecycle: &LifecyclePolicy,
ssh_user: &str,
) -> Result<VastAiSshEndpoint, String>;
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<Self, String> {
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<String>) -> Result<Self, String> {
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<ProvisionedInstance, String> {
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<VastAiSshEndpoint, String> {
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<VastAiSshEndpoint, String> {
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<C>
where
C: VastAiLeaseClient,
{
client: C,
config: VastAiProvisioningConfig,
}
impl<C> VastAiProviderPlugin<C>
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<u64, ProviderError> {
if lease.provider != ProviderKind::VastAi {
return Err(ProviderError::new(
"vastai provider received non-vastai lease",
));
}
lease
.provider_contract_id
.parse::<u64>()
.map_err(|e| ProviderError::new(format!("invalid vastai contract id: {e}")))
}
}
impl<C> ProviderPlugin for VastAiProviderPlugin<C>
where
C: VastAiLeaseClient,
{
fn create_lease(
&mut self,
request: CreateLeaseRequest,
) -> Result<CreateLeaseResult, ProviderError> {
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<Option<SshEndpoint>, 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::<u64>()
.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<DatastreamProducer>,
) -> Result<Self::Handle, String>;
fn stop_bootstrap(&mut self, handle: &mut Self::Handle);
}
#[derive(Default, Debug, Clone, Copy)]
pub struct SshCommandBootstrapLauncher;
pub struct SshCommandBootstrapHandle {
child: Arc<Mutex<Child>>,
}
impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
type Handle = SshCommandBootstrapHandle;
fn start_bootstrap(
&mut self,
spec: NodeProvisionSpec,
endpoint: VastAiSshEndpoint,
sink: PluginSink,
producer: Option<DatastreamProducer>,
) -> Result<Self::Handle, String> {
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<C, B>
where
C: VastAiLeaseClient,
B: VastAiBootstrapLauncher,
{
client: C,
bootstrap: B,
config: VastAiProvisioningConfig,
bootstrap_producer: Option<DatastreamProducer>,
next_handle_id: u64,
nodes: BTreeMap<u64, VastAiNode<B::Handle>>,
}
struct VastAiNode<H> {
contract_id: u64,
bootstrap: Option<H>,
}
impl<C, B> VastAiProvisioningPlugin<C, B>
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::<BTreeMap<_, _>>();
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<C, B> ProvisionPlugin for VastAiProvisioningPlugin<C, B>
where
C: VastAiLeaseClient + 'static,
B: VastAiBootstrapLauncher + 'static,
{
fn start_node(
&mut self,
spec: NodeProvisionSpec,
sink: PluginSink,
) -> Result<PluginNodeHandle, String> {
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)
}
}

View file

@ -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,

View file

@ -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!("<dashboard unavailable: {err}>"))
);
}
}
fn run_full_flow(
child: &mut Child,
stdin: &mut impl Write,
stdout: &Arc<Mutex<String>>,
stderr: &Arc<Mutex<String>>,
) -> 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<Mutex<String>>,
) -> 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<Vec<SeenFrame>, 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::<Vec<Value>>(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<String, String> {
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::<Vec<_>>()
.join("\n"))
}
fn decode_payload(value: Option<&Value>) -> Option<String> {
let bytes = value?
.as_array()?
.iter()
.map(|byte| byte.as_u64().map(|n| n as u8))
.collect::<Option<Vec<_>>>()?;
Some(String::from_utf8_lossy(&bytes).to_string())
}
fn http_get(path: &str) -> Result<String, String> {
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<Mutex<String>>, needle: &str) -> bool {
snapshot(stdout).contains(needle)
}
fn response_text_visible(stdout: &Arc<Mutex<String>>) -> bool {
let text = snapshot(stdout);
text.lines()
.any(|line| line.trim_start_matches('>').trim().len() > 8)
}
fn stderr_contains(stderr: &Arc<Mutex<String>>, needle: &str) -> bool {
snapshot(stderr).contains(needle)
}
fn snapshot(buf: &Arc<Mutex<String>>) -> String {
buf.lock().expect("capture mutex").clone()
}
fn wait_child(child: &mut Child, timeout: Duration) -> Option<std::process::ExitStatus> {
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()
}

View file

@ -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.

View file

@ -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;