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:
parent
8363be74cd
commit
6608824cb0
36 changed files with 5445 additions and 128 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
2
Cargo.lock
generated
|
|
@ -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
15
apps/mvp-node/Dockerfile
Normal 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
|
||||
68
apps/mvp-node/Dockerfile.base
Normal file
68
apps/mvp-node/Dockerfile.base
Normal 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"]
|
||||
54
apps/mvp-node/local_docker_e2e.sh
Executable file
54
apps/mvp-node/local_docker_e2e.sh
Executable 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
34
apps/mvp-node/mvp_entrypoint.sh
Executable 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
299
apps/mvp-node/tinygrad_worker.py
Executable 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)
|
||||
|
|
@ -1,6 +1,6 @@
|
|||
mod live_explorer;
|
||||
mod server;
|
||||
mod store;
|
||||
mod live_explorer;
|
||||
pub mod swactor;
|
||||
pub mod view;
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
174
crates/mvp-system/src/bin/mvp_chat.rs
Normal file
174
crates/mvp-system/src/bin/mvp_chat.rs
Normal 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}"))
|
||||
}
|
||||
794
crates/mvp-system/src/bin/mvp_node.rs
Normal file
794
crates/mvp-system/src/bin/mvp_node.rs
Normal 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)
|
||||
}
|
||||
512
crates/mvp-system/src/bin/mvp_one_node_chat.rs
Normal file
512
crates/mvp-system/src/bin/mvp_one_node_chat.rs
Normal 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}"))
|
||||
}
|
||||
983
crates/mvp-system/src/bin/mvp_orch_one_node.rs
Normal file
983
crates/mvp-system/src/bin/mvp_orch_one_node.rs
Normal 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}"))
|
||||
}
|
||||
183
crates/mvp-system/src/bootstrap_datastream.rs
Normal file
183
crates/mvp-system/src/bootstrap_datastream.rs
Normal 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)
|
||||
}
|
||||
|
|
@ -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;
|
||||
|
||||
|
|
|
|||
136
crates/mvp-system/src/prompt_rpc.rs
Normal file
136
crates/mvp-system/src/prompt_rpc.rs
Normal 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()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
115
crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs
Normal file
115
crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs
Normal 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
|
||||
)));
|
||||
}
|
||||
|
|
@ -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,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
351
crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs
Normal file
351
crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs
Normal 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]);
|
||||
}
|
||||
619
crates/mvp-system/src/vastai_provisioning.rs
Normal file
619
crates/mvp-system/src/vastai_provisioning.rs
Normal 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)
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
362
crates/mvp-system/tests/one_node_chat_e2e.rs
Normal file
362
crates/mvp-system/tests/one_node_chat_e2e.rs
Normal 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()
|
||||
}
|
||||
167
mvp-node-image-gguf-shape.md
Normal file
167
mvp-node-image-gguf-shape.md
Normal 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.
|
||||
110
vastai-datastream-plugin-shape.md
Normal file
110
vastai-datastream-plugin-shape.md
Normal 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;
|
||||
Loading…
Reference in a new issue