feat(mvp-system): expand one-node mvp binaries and provisioning
Flesh out mvp_node, mvp_one_node_chat, and mvp_orch_one_node binaries. Add node_image and relay_provisioning; grow provisioning and vastai. Tune datastream mux/timing/emit. Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
parent
6608824cb0
commit
a97864f7ec
32 changed files with 5794 additions and 487 deletions
|
|
@ -1,4 +1,4 @@
|
|||
[alias]
|
||||
xtask = "run --package xtask --"
|
||||
mvp-chat = "run -p mvp-system --features local-e2e --bin mvp-one-node-chat --"
|
||||
mvp-chat = "run --package xtask -- mvp-chat"
|
||||
mvp-chat-test = "test -p mvp-system --features local-e2e --test one_node_chat_e2e -- --nocapture"
|
||||
|
|
|
|||
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -6,6 +6,8 @@ __pycache__
|
|||
fuzz/artifacts/**
|
||||
corpus
|
||||
.loop/
|
||||
.model-cache/
|
||||
|
||||
|
||||
# Analysis artifacts (depgraph + spectral)
|
||||
**/deps.dot
|
||||
|
|
|
|||
1
Cargo.lock
generated
1
Cargo.lock
generated
|
|
@ -2471,6 +2471,7 @@ dependencies = [
|
|||
name = "mvp-system"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"blake3",
|
||||
"dashboard",
|
||||
"datastream",
|
||||
"distribution",
|
||||
|
|
|
|||
|
|
@ -1,10 +1,14 @@
|
|||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
import hashlib
|
||||
import json
|
||||
import linecache
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
import urllib.parse
|
||||
|
|
@ -18,6 +22,218 @@ model: Any = None
|
|||
tokenizer: Any = None
|
||||
role: dict[str, Any] = {}
|
||||
loaded: dict[str, Any] = {}
|
||||
_nvidia_smi_available: bool | None = None
|
||||
|
||||
|
||||
class CpuLineSampler:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
phase: str,
|
||||
request_id: int | None,
|
||||
model_id: str | None,
|
||||
interval_secs: float,
|
||||
) -> None:
|
||||
self.phase = phase
|
||||
self.request_id = request_id
|
||||
self.model_id = model_id
|
||||
self.interval_secs = interval_secs
|
||||
self.target_thread_id = threading.get_ident()
|
||||
self.samples: dict[tuple[str, int, str], int] = {}
|
||||
self.wall_start = time.perf_counter()
|
||||
self.process_cpu_start = time.process_time()
|
||||
self._running = True
|
||||
self._thread = threading.Thread(target=self._run, name="cpu-line-sampler", daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def _run(self) -> None:
|
||||
while self._running:
|
||||
frame = sys._current_frames().get(self.target_thread_id)
|
||||
if frame is not None:
|
||||
code = frame.f_code
|
||||
key = (code.co_filename, frame.f_lineno, code.co_name)
|
||||
self.samples[key] = self.samples.get(key, 0) + 1
|
||||
time.sleep(self.interval_secs)
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
self._thread.join(timeout=max(0.25, self.interval_secs * 4.0))
|
||||
wall_elapsed_ms = (time.perf_counter() - self.wall_start) * 1000.0
|
||||
process_cpu_elapsed_ms = (time.process_time() - self.process_cpu_start) * 1000.0
|
||||
total_samples = sum(self.samples.values())
|
||||
top = []
|
||||
for (filename, line, function), count in sorted(
|
||||
self.samples.items(), key=lambda item: item[1], reverse=True
|
||||
)[:32]:
|
||||
top.append(
|
||||
{
|
||||
"file": filename,
|
||||
"line": line,
|
||||
"function": function,
|
||||
"source": linecache.getline(filename, line).strip(),
|
||||
"samples": count,
|
||||
"percent": round((count * 100.0 / total_samples), 2) if total_samples else 0.0,
|
||||
}
|
||||
)
|
||||
control(
|
||||
type="CpuLineProfileSummary",
|
||||
phase=self.phase,
|
||||
request_id=self.request_id,
|
||||
model_id=self.model_id,
|
||||
interval_ms=round(self.interval_secs * 1000.0, 3),
|
||||
wall_elapsed_ms=round(wall_elapsed_ms, 3),
|
||||
process_cpu_elapsed_ms=round(process_cpu_elapsed_ms, 3),
|
||||
process_cpu_over_wall=round(process_cpu_elapsed_ms / wall_elapsed_ms, 4)
|
||||
if wall_elapsed_ms > 0.0
|
||||
else 0.0,
|
||||
total_samples=total_samples,
|
||||
top=top,
|
||||
)
|
||||
|
||||
|
||||
def start_cpu_line_sampler(
|
||||
*,
|
||||
phase: str,
|
||||
request_id: int | None,
|
||||
model_id: str | None,
|
||||
) -> CpuLineSampler | None:
|
||||
raw = os.environ.get("MVP_CPU_LINE_PROFILE")
|
||||
if not env_flag("MVP_CPU_LINE_PROFILE", False):
|
||||
control(
|
||||
type="CpuLineProfileSkipped",
|
||||
phase=phase,
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
env_value=raw,
|
||||
)
|
||||
return None
|
||||
interval_ms = float(os.environ.get("MVP_CPU_LINE_PROFILE_INTERVAL_MS", "2"))
|
||||
interval_secs = max(0.0005, interval_ms / 1000.0)
|
||||
control(
|
||||
type="CpuLineProfileStarted",
|
||||
phase=phase,
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
interval_ms=round(interval_secs * 1000.0, 3),
|
||||
)
|
||||
return CpuLineSampler(
|
||||
phase=phase,
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
interval_secs=interval_secs,
|
||||
)
|
||||
|
||||
|
||||
def stop_cpu_line_sampler(sampler: CpuLineSampler | None) -> None:
|
||||
if sampler is not None:
|
||||
sampler.stop()
|
||||
|
||||
|
||||
def env_flag(name: str, default: bool = True) -> bool:
|
||||
raw = os.environ.get(name)
|
||||
if raw is None:
|
||||
return default
|
||||
return raw.strip().lower() not in {"0", "false", "no", "off"}
|
||||
|
||||
|
||||
def parse_int(value: str) -> int | None:
|
||||
stripped = value.strip()
|
||||
if not stripped or stripped == "[Not Supported]":
|
||||
return None
|
||||
try:
|
||||
return int(float(stripped))
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def run_nvidia_smi(args: list[str]) -> tuple[bool, str, str]:
|
||||
global _nvidia_smi_available
|
||||
if _nvidia_smi_available is False:
|
||||
return False, "", "nvidia-smi unavailable"
|
||||
try:
|
||||
result = subprocess.run(
|
||||
["nvidia-smi", *args],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=float(os.environ.get("MVP_GPU_SAMPLE_TIMEOUT_SECS", "2")),
|
||||
)
|
||||
except FileNotFoundError:
|
||||
_nvidia_smi_available = False
|
||||
return False, "", "nvidia-smi not found"
|
||||
except Exception as exc:
|
||||
return False, "", str(exc)
|
||||
_nvidia_smi_available = True
|
||||
return result.returncode == 0, result.stdout, result.stderr.strip()
|
||||
|
||||
|
||||
def parse_gpu_rows(raw: str) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for row in csv.reader(raw.splitlines()):
|
||||
if len(row) < 5:
|
||||
continue
|
||||
memory_total = parse_int(row[2])
|
||||
memory_used = parse_int(row[3])
|
||||
utilization = parse_int(row[4])
|
||||
rows.append(
|
||||
{
|
||||
"index": parse_int(row[0]),
|
||||
"name": row[1].strip(),
|
||||
"memory_total_mib": memory_total,
|
||||
"memory_used_mib": memory_used,
|
||||
"utilization_gpu_percent": utilization,
|
||||
}
|
||||
)
|
||||
return rows
|
||||
|
||||
|
||||
def parse_process_rows(raw: str) -> list[dict[str, Any]]:
|
||||
rows: list[dict[str, Any]] = []
|
||||
for row in csv.reader(raw.splitlines()):
|
||||
if len(row) < 2:
|
||||
continue
|
||||
pid = parse_int(row[0])
|
||||
memory_used = parse_int(row[1])
|
||||
if pid is None:
|
||||
continue
|
||||
rows.append({"pid": pid, "used_memory_mib": memory_used})
|
||||
return rows
|
||||
|
||||
|
||||
def gpu_sample(label: str, **fields: Any) -> None:
|
||||
if not env_flag("MVP_GPU_SAMPLE", True):
|
||||
control(type="GpuSample", label=label, pid=os.getpid(), enabled=False, **fields)
|
||||
return
|
||||
gpu_ok, gpu_stdout, gpu_error = run_nvidia_smi(
|
||||
[
|
||||
"--query-gpu=index,name,memory.total,memory.used,utilization.gpu",
|
||||
"--format=csv,noheader,nounits",
|
||||
]
|
||||
)
|
||||
proc_ok, proc_stdout, proc_error = run_nvidia_smi(
|
||||
[
|
||||
"--query-compute-apps=pid,used_memory",
|
||||
"--format=csv,noheader,nounits",
|
||||
]
|
||||
)
|
||||
pid = os.getpid()
|
||||
processes = parse_process_rows(proc_stdout) if proc_ok else []
|
||||
worker_processes = [process for process in processes if process.get("pid") == pid]
|
||||
control(
|
||||
type="GpuSample",
|
||||
label=label,
|
||||
pid=pid,
|
||||
enabled=True,
|
||||
nvidia_smi_available=(_nvidia_smi_available is True),
|
||||
gpu_query_ok=gpu_ok,
|
||||
process_query_ok=proc_ok,
|
||||
gpu_query_error=None if gpu_ok else gpu_error,
|
||||
process_query_error=None if proc_ok else proc_error,
|
||||
gpus=parse_gpu_rows(gpu_stdout) if gpu_ok else [],
|
||||
processes=processes,
|
||||
worker_processes=worker_processes,
|
||||
**fields,
|
||||
)
|
||||
|
||||
|
||||
def control(**event: Any) -> None:
|
||||
|
|
@ -45,16 +261,20 @@ def initialize(cmd: dict[str, Any]) -> None:
|
|||
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)
|
||||
control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV"))
|
||||
from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes
|
||||
|
||||
control(type="TinygradImportReady", requested_device=device, env_DEV=os.environ.get("DEV"))
|
||||
Tensor = TinyTensor
|
||||
dtypes = tiny_dtypes
|
||||
control(type="TinygradDeviceProbeStarted", requested_device=device)
|
||||
value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist()
|
||||
control(type="TinygradDeviceProbeReady", requested_device=device, probe_result=value)
|
||||
gpu_sample("after_worker_probe", requested_device=device)
|
||||
control(
|
||||
type="WorkerReady",
|
||||
pid=os.getpid(),
|
||||
backend={"device": device},
|
||||
backend={"requested_device": device, "env_DEV": os.environ.get("DEV"), "tinygrad_device": device},
|
||||
cuda_probe=value,
|
||||
elapsed_ms=int((time.monotonic() - started) * 1000),
|
||||
)
|
||||
|
|
@ -98,6 +318,14 @@ def source_path(source: dict[str, Any]) -> Path | None:
|
|||
return Path(str(source["LocalPath"])).expanduser()
|
||||
|
||||
|
||||
def source_kind(source: dict[str, Any]) -> str:
|
||||
if "LocalPath" in source:
|
||||
return "LocalPath"
|
||||
if "HuggingFaceGguf" in source:
|
||||
return "HuggingFaceGguf"
|
||||
return "Unknown"
|
||||
|
||||
|
||||
def cache_path_for(url: str) -> Path:
|
||||
parsed = urllib.parse.urlparse(url)
|
||||
basename = Path(parsed.path).name or "model.gguf"
|
||||
|
|
@ -116,9 +344,11 @@ def request_headers() -> dict[str, str]:
|
|||
def fetch_whole(source: dict[str, Any]) -> Path:
|
||||
local = source_path(source)
|
||||
if local is not None:
|
||||
control(type="GgufLocalPathStatStarted", path=str(local))
|
||||
if not local.is_file():
|
||||
fatal("GgufLocalPathMissing", path=str(local))
|
||||
control(type="GgufCacheReady", path=str(local), cache_hit=True, source="local")
|
||||
stat = local.stat()
|
||||
control(type="GgufCacheReady", path=str(local), bytes=stat.st_size, cache_hit=True, source="local")
|
||||
return local
|
||||
|
||||
url = source_url(source)
|
||||
|
|
@ -182,6 +412,14 @@ def load_weights(cmd: dict[str, Any]) -> None:
|
|||
TensorCls = require_tinygrad()
|
||||
started = time.monotonic()
|
||||
model_id = str(cmd["model_id"])
|
||||
source = cmd["gguf_source"]
|
||||
control(
|
||||
type="LoadWeightsStarted",
|
||||
model_id=model_id,
|
||||
source_kind=source_kind(source),
|
||||
layer_start=int(cmd.get("layer_start", 0)),
|
||||
layer_end_exclusive=int(cmd.get("layer_end_exclusive", 0)),
|
||||
)
|
||||
if test_mode():
|
||||
model = {"test_mode": True}
|
||||
tokenizer = {"test_mode": True}
|
||||
|
|
@ -200,17 +438,43 @@ def load_weights(cmd: dict[str, Any]) -> None:
|
|||
elapsed_ms=int((time.monotonic() - started) * 1000),
|
||||
)
|
||||
return
|
||||
source = cmd["gguf_source"]
|
||||
control(type="GgufResolveStarted", model_id=model_id, source_kind=source_kind(source))
|
||||
path = fetch_whole(source)
|
||||
model_bytes = path.stat().st_size
|
||||
control(type="GgufResolveReady", model_id=model_id, path=str(path), bytes=model_bytes)
|
||||
try:
|
||||
control(type="TinygradLlmImportStarted", model_id=model_id)
|
||||
from tinygrad.apps.llm import SimpleTokenizer, Transformer
|
||||
|
||||
control(type="TinygradLlmImportReady", model_id=model_id)
|
||||
max_context_raw = os.environ.get("MVP_MAX_CONTEXT", "512")
|
||||
max_context = int(max_context_raw) if max_context_raw else 512
|
||||
gpu_sample("before_model_load", model_id=model_id, path=str(path), bytes=model_bytes)
|
||||
control(
|
||||
type="TransformerFromGgufStarted",
|
||||
model_id=model_id,
|
||||
path=str(path),
|
||||
bytes=model_bytes,
|
||||
max_context=max_context,
|
||||
realize=True,
|
||||
requested_device=os.environ.get("DEV"),
|
||||
)
|
||||
model, kv = Transformer.from_gguf(TensorCls(path), max_context=max_context, realize=True)
|
||||
control(
|
||||
type="TransformerFromGgufReady",
|
||||
model_id=model_id,
|
||||
path=str(path),
|
||||
bytes=model_bytes,
|
||||
max_context=max_context,
|
||||
realize=True,
|
||||
requested_device=os.environ.get("DEV"),
|
||||
)
|
||||
gpu_sample("after_model_load", model_id=model_id, path=str(path), bytes=model_bytes)
|
||||
tok_src = cmd.get("tokenizer", {"EmbeddedGguf": None})
|
||||
if "EmbeddedGguf" in tok_src:
|
||||
control(type="TokenizerBuildStarted", model_id=model_id, source="EmbeddedGguf")
|
||||
tokenizer = SimpleTokenizer.from_gguf_kv(kv)
|
||||
control(type="TokenizerBuildReady", model_id=model_id, source="EmbeddedGguf")
|
||||
else:
|
||||
fatal("UnsupportedTokenizerSource", tokenizer=tok_src)
|
||||
except SystemExit:
|
||||
|
|
@ -234,16 +498,106 @@ def load_weights(cmd: dict[str, Any]) -> None:
|
|||
)
|
||||
|
||||
|
||||
|
||||
def decode_greedy_device_resident(
|
||||
prompt_tokens: list[int],
|
||||
max_tokens: int,
|
||||
*,
|
||||
request_id: int | None,
|
||||
model_id: str | None,
|
||||
progress_every: int,
|
||||
) -> list[int]:
|
||||
if max_tokens <= 0:
|
||||
return []
|
||||
max_context = int(getattr(model, "max_context", len(prompt_tokens) + max_tokens))
|
||||
generation_limit = min(max_tokens, max(0, max_context - len(prompt_tokens)))
|
||||
if generation_limit <= 0:
|
||||
control(
|
||||
type="DecodeContextFull",
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
max_context=max_context,
|
||||
)
|
||||
return []
|
||||
if generation_limit < max_tokens:
|
||||
control(
|
||||
type="DecodeLimitedByContext",
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
requested_tokens=max_tokens,
|
||||
generation_limit=generation_limit,
|
||||
max_context=max_context,
|
||||
)
|
||||
|
||||
TensorCls = require_tinygrad()
|
||||
from tinygrad.uop.ops import UOp
|
||||
|
||||
if hasattr(model, "forward_jit"):
|
||||
model.forward_jit.reset()
|
||||
use_symbolic_pos = os.environ.get("SYM", "1").strip().lower() not in {"0", "false", "no", "off"}
|
||||
pos_upper_bound = max(1, max_context - 1)
|
||||
symbolic_start_pos = UOp.variable("start_pos", 1, pos_upper_bound)
|
||||
next_token = model(TensorCls([prompt_tokens], dtype="int32"), 0).realize()
|
||||
generated_tensors = []
|
||||
|
||||
for token_index in range(generation_limit):
|
||||
generated_tensors.append(next_token.clone().realize())
|
||||
tokens_generated = token_index + 1
|
||||
if tokens_generated == 1:
|
||||
control(
|
||||
type="FirstTokenReady",
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
token_index=1,
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
)
|
||||
gpu_sample("after_first_token", request_id=request_id, model_id=model_id)
|
||||
elif progress_every > 0 and tokens_generated % progress_every == 0:
|
||||
control(
|
||||
type="TokenProgress",
|
||||
request_id=request_id,
|
||||
model_id=model_id,
|
||||
tokens_generated=tokens_generated,
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
)
|
||||
if tokens_generated >= generation_limit:
|
||||
break
|
||||
start_pos = len(prompt_tokens) + token_index
|
||||
pos = symbolic_start_pos.bind(start_pos) if use_symbolic_pos else start_pos
|
||||
next_token = model(next_token, pos).realize()
|
||||
generated_tensor = (
|
||||
generated_tensors[0]
|
||||
if len(generated_tensors) == 1
|
||||
else generated_tensors[0].cat(*generated_tensors[1:], dim=1)
|
||||
)
|
||||
generated_array = generated_tensor.numpy().reshape(-1).tolist()
|
||||
return [int(token) for token in generated_array]
|
||||
|
||||
|
||||
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))
|
||||
request_id_raw = cmd.get("request_id")
|
||||
request_id = int(request_id_raw) if request_id_raw is not None else None
|
||||
started = time.monotonic()
|
||||
control(
|
||||
type="PromptStarted",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_bytes=len(prompt.encode("utf-8")),
|
||||
prompt_chars=len(prompt),
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
gpu_sample("before_prompt", request_id=request_id, model_id=loaded.get("model_id"))
|
||||
if test_mode():
|
||||
text = f"mvp-test response: {prompt}"
|
||||
control(
|
||||
type="PromptCompleted",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_tokens=[],
|
||||
generated_tokens=list(range(min(max_tokens, 3))),
|
||||
|
|
@ -252,15 +606,59 @@ def infer_prompt(cmd: dict[str, Any]) -> None:
|
|||
elapsed_ms=int((time.monotonic() - started) * 1000),
|
||||
)
|
||||
return
|
||||
control(type="PromptEncodeStarted", request_id=request_id, model_id=loaded.get("model_id"))
|
||||
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
|
||||
control(
|
||||
type="PromptEncodeReady",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_bytes=len(prompt.encode("utf-8")),
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
)
|
||||
progress_every = int(os.environ.get("MVP_TOKEN_PROGRESS_EVERY", "16") or "16")
|
||||
control(
|
||||
type="DecodeStarted",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
max_tokens=max_tokens,
|
||||
decode_impl="device_resident_greedy",
|
||||
)
|
||||
cpu_sampler = start_cpu_line_sampler(
|
||||
phase="decode",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
)
|
||||
try:
|
||||
generated = decode_greedy_device_resident(
|
||||
prompt_tokens,
|
||||
max_tokens,
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
progress_every=progress_every,
|
||||
)
|
||||
finally:
|
||||
stop_cpu_line_sampler(cpu_sampler)
|
||||
control(
|
||||
type="DecodeReady",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_tokens=len(prompt_tokens),
|
||||
tokens_generated=len(generated),
|
||||
)
|
||||
gpu_sample("after_decode", request_id=request_id, model_id=loaded.get("model_id"))
|
||||
control(type="TextDecodeStarted", request_id=request_id, model_id=loaded.get("model_id"), tokens_generated=len(generated))
|
||||
text = tokenizer.decode(generated) if generated else ""
|
||||
control(
|
||||
type="TextDecodeReady",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
tokens_generated=len(generated),
|
||||
text_bytes=len(text.encode("utf-8")),
|
||||
)
|
||||
control(
|
||||
type="PromptCompleted",
|
||||
request_id=request_id,
|
||||
model_id=loaded.get("model_id"),
|
||||
prompt_tokens=prompt_tokens,
|
||||
generated_tokens=generated,
|
||||
|
|
|
|||
|
|
@ -85,6 +85,17 @@ impl DatastreamEmitter {
|
|||
self.mux.dropped()
|
||||
}
|
||||
|
||||
/// Enable or disable optional sidecar timing samples for newly submitted
|
||||
/// frames from this emitter and its cloned submit handles.
|
||||
pub fn set_frame_timing_enabled(&self, enabled: bool) {
|
||||
self.mux.set_frame_timing_enabled(enabled);
|
||||
}
|
||||
|
||||
/// Whether this emitter's mux currently emits sidecar frame timing samples.
|
||||
pub fn frame_timing_enabled(&self) -> bool {
|
||||
self.mux.frame_timing_enabled()
|
||||
}
|
||||
|
||||
/// Submit a typed record defined by the caller's crate.
|
||||
pub fn submit_record<R: Record>(&self, record: &R) -> Position {
|
||||
self.mux.submit(R::channel(), record.encode())
|
||||
|
|
|
|||
|
|
@ -258,6 +258,17 @@ impl DatastreamEndpoint {
|
|||
&self.mux
|
||||
}
|
||||
|
||||
/// Enable or disable optional sidecar timing samples for newly submitted
|
||||
/// frames from this endpoint's producers.
|
||||
pub fn set_frame_timing_enabled(&self, enabled: bool) {
|
||||
self.mux.set_frame_timing_enabled(enabled);
|
||||
}
|
||||
|
||||
/// Whether this endpoint currently emits sidecar frame timing samples.
|
||||
pub fn frame_timing_enabled(&self) -> bool {
|
||||
self.mux.frame_timing_enabled()
|
||||
}
|
||||
|
||||
pub fn producer(&self) -> DatastreamProducer {
|
||||
DatastreamProducer {
|
||||
mux: Arc::clone(&self.mux),
|
||||
|
|
@ -343,6 +354,16 @@ impl DatastreamProducer {
|
|||
self.mux.submit(channel, bytes)
|
||||
}
|
||||
|
||||
/// Enable or disable optional sidecar timing samples for this producer's mux.
|
||||
pub fn set_frame_timing_enabled(&self, enabled: bool) {
|
||||
self.mux.set_frame_timing_enabled(enabled);
|
||||
}
|
||||
|
||||
/// Whether this producer's mux currently emits sidecar frame timing samples.
|
||||
pub fn frame_timing_enabled(&self) -> bool {
|
||||
self.mux.frame_timing_enabled()
|
||||
}
|
||||
|
||||
pub fn process_observer_with<F>(&self, channel_for: F) -> Arc<dyn ProcessOutputObserver>
|
||||
where
|
||||
F: Fn(&str, bool) -> ChannelId + Send + Sync + 'static,
|
||||
|
|
|
|||
|
|
@ -158,9 +158,9 @@ impl fmt::Display for StreamId {
|
|||
/// interpret it (spec §4.1). A typed event and a log line are the same
|
||||
/// kind of thing here: bytes on a channel.
|
||||
///
|
||||
/// Per the `// USER:` annotation on spec §4.1/§5.2 there is no per-frame
|
||||
/// wall-clock timestamp: frames are ordered and correlated by position
|
||||
/// alone.
|
||||
/// Frames do not carry wall-clock timestamps. Optional frame-construction
|
||||
/// timing rides as sidecar records keyed by stream-local position, preserving
|
||||
/// the core frame and wire shape.
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub struct Frame {
|
||||
/// The lane these bytes belong to.
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ pub mod mux;
|
|||
pub mod record;
|
||||
pub mod sink_actor;
|
||||
pub mod store;
|
||||
pub mod timing;
|
||||
pub mod transport;
|
||||
pub mod views;
|
||||
pub mod wire;
|
||||
|
|
@ -53,5 +54,6 @@ pub use mux::Mux;
|
|||
pub use record::{ChannelKind, ChannelRegistry, Record};
|
||||
pub use sink_actor::{DATASTREAM_SINK_NAME, DatastreamSink};
|
||||
pub use store::{GapSpan, Store, StoredStream};
|
||||
pub use timing::{FRAME_TIME_CHANNEL, FrameTimeSample};
|
||||
pub use transport::{Delivery, Reorder, ScriptedTransport, StreamScript};
|
||||
pub use views::{Body, LogEntry, MergedFrame};
|
||||
|
|
|
|||
|
|
@ -24,9 +24,12 @@
|
|||
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Mutex;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use super::frame::{ChannelId, Frame, Position, StreamId};
|
||||
use super::record::Record;
|
||||
use super::timing::{FRAME_TIME_CHANNEL, FrameTimeSample};
|
||||
|
||||
/// A node's single position authority and outgoing telemetry buffer.
|
||||
///
|
||||
|
|
@ -37,6 +40,7 @@ pub struct Mux {
|
|||
stream: StreamId,
|
||||
next: AtomicU64,
|
||||
dropped: AtomicU64,
|
||||
frame_timing_enabled: AtomicBool,
|
||||
capacity: usize,
|
||||
buffer: Mutex<VecDeque<Frame>>,
|
||||
}
|
||||
|
|
@ -50,6 +54,7 @@ impl Mux {
|
|||
stream,
|
||||
next: AtomicU64::new(0),
|
||||
dropped: AtomicU64::new(0),
|
||||
frame_timing_enabled: AtomicBool::new(true),
|
||||
capacity,
|
||||
buffer: Mutex::new(VecDeque::new()),
|
||||
}
|
||||
|
|
@ -76,8 +81,15 @@ impl Mux {
|
|||
// Assign first, unconditionally: numbering is independent of
|
||||
// whether the frame survives the buffer (spec §5.2).
|
||||
let position = Position(self.next.fetch_add(1, Ordering::Relaxed));
|
||||
let channel = channel.into();
|
||||
let timing_sample = self
|
||||
.frame_timing_enabled
|
||||
.load(Ordering::Relaxed)
|
||||
.then(|| now_unix_ns())
|
||||
.filter(|_| channel.as_str() != FRAME_TIME_CHANNEL)
|
||||
.map(|created_at_unix_ns| FrameTimeSample::new(position, created_at_unix_ns));
|
||||
let frame = Frame {
|
||||
channel: channel.into(),
|
||||
channel,
|
||||
position,
|
||||
payload,
|
||||
};
|
||||
|
|
@ -85,6 +97,9 @@ impl Mux {
|
|||
let mut buffer = self.buffer.lock().expect("mux buffer poisoned");
|
||||
if buffer.len() < self.capacity {
|
||||
buffer.push_back(frame);
|
||||
if let Some(sample) = timing_sample {
|
||||
self.push_timing_sample_if_room(&mut buffer, sample);
|
||||
}
|
||||
} else {
|
||||
// Overflow: drop the frame that does not fit. Its position is
|
||||
// already spent, so it will read as a gap, not a renumber.
|
||||
|
|
@ -101,8 +116,32 @@ impl Mux {
|
|||
buffer.drain(..).collect()
|
||||
}
|
||||
|
||||
/// Enable or disable optional sidecar timing samples for newly submitted
|
||||
/// frames. The core frame shape and wire envelope remain unchanged.
|
||||
pub fn set_frame_timing_enabled(&self, enabled: bool) {
|
||||
self.frame_timing_enabled.store(enabled, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Whether this mux currently emits sidecar frame timing samples.
|
||||
pub fn frame_timing_enabled(&self) -> bool {
|
||||
self.frame_timing_enabled.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
fn push_timing_sample_if_room(&self, buffer: &mut VecDeque<Frame>, sample: FrameTimeSample) {
|
||||
if buffer.len() >= self.capacity {
|
||||
return;
|
||||
}
|
||||
|
||||
let position = Position(self.next.fetch_add(1, Ordering::Relaxed));
|
||||
buffer.push_back(Frame {
|
||||
channel: FRAME_TIME_CHANNEL.into(),
|
||||
position,
|
||||
payload: sample.encode(),
|
||||
});
|
||||
}
|
||||
|
||||
/// How many positions have been assigned — the gap-free high-water mark
|
||||
/// (spec §5.2). Equal to the number of `submit` calls.
|
||||
/// (spec §5.2). Sidecar timing samples, when enabled, are frames too.
|
||||
pub fn assigned(&self) -> u64 {
|
||||
self.next.load(Ordering::Relaxed)
|
||||
}
|
||||
|
|
@ -113,3 +152,10 @@ impl Mux {
|
|||
self.dropped.load(Ordering::Relaxed)
|
||||
}
|
||||
}
|
||||
|
||||
fn now_unix_ns() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| u64::try_from(duration.as_nanos()).unwrap_or(u64::MAX))
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
|
|
|||
29
crates/datastream/src/timing.rs
Normal file
29
crates/datastream/src/timing.rs
Normal file
|
|
@ -0,0 +1,29 @@
|
|||
//! Optional frame-construction timing sidecar records.
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::frame::Position;
|
||||
use crate::record::Record;
|
||||
|
||||
/// Reserved channel carrying optional timing samples for frames in the same stream.
|
||||
pub const FRAME_TIME_CHANNEL: &str = "datastream.frame_time";
|
||||
|
||||
/// Sidecar timing sample keyed by the target frame's stream-local position.
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct FrameTimeSample {
|
||||
pub target_position: u64,
|
||||
pub created_at_unix_ns: u64,
|
||||
}
|
||||
|
||||
impl FrameTimeSample {
|
||||
pub fn new(target_position: Position, created_at_unix_ns: u64) -> Self {
|
||||
Self {
|
||||
target_position: target_position.0,
|
||||
created_at_unix_ns,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Record for FrameTimeSample {
|
||||
const CHANNEL: &'static str = FRAME_TIME_CHANNEL;
|
||||
}
|
||||
|
|
@ -20,14 +20,14 @@ use datastream::store::{GapSpan, StoredStream};
|
|||
use datastream::transport::{Delivery, Reorder, ScriptedTransport, StreamScript};
|
||||
use datastream::views::{self, Body, LogEntry};
|
||||
use datastream::wire::{WireError, decode_delivery, encode_delivery};
|
||||
use datastream::{ChannelKind, ChannelRegistry, Record};
|
||||
use datastream::{ChannelKind, ChannelRegistry, FRAME_TIME_CHANNEL, FrameTimeSample, Record};
|
||||
|
||||
use support::reference::TimelineItem;
|
||||
use support::schema::{
|
||||
self as catalog, ActorRuntimeDetail, DatastreamHealth, DistributionState, IdentityRecord,
|
||||
ProcStream, ResourceSample,
|
||||
};
|
||||
use support::{Node, payloads, reference};
|
||||
use support::{payloads, reference};
|
||||
|
||||
/// The frames of `sent` that survive dropping `dropped`, in send order —
|
||||
/// the scenario's delivered set, derived without running the pipe.
|
||||
|
|
@ -92,6 +92,49 @@ fn test_stream() -> StreamId {
|
|||
StreamId::new(NodeId::new("node-alpha"), Lifetime(1))
|
||||
}
|
||||
|
||||
fn mux_without_frame_timing(stream: StreamId) -> Mux {
|
||||
let mux = Mux::unbounded(stream);
|
||||
mux.set_frame_timing_enabled(false);
|
||||
mux
|
||||
}
|
||||
|
||||
fn bounded_mux_without_frame_timing(stream: StreamId, capacity: usize) -> Mux {
|
||||
let mux = Mux::new(stream, capacity);
|
||||
mux.set_frame_timing_enabled(false);
|
||||
mux
|
||||
}
|
||||
|
||||
struct Node {
|
||||
mux: Mux,
|
||||
}
|
||||
|
||||
impl Node {
|
||||
fn new(stream: StreamId) -> Self {
|
||||
Node {
|
||||
mux: mux_without_frame_timing(stream),
|
||||
}
|
||||
}
|
||||
|
||||
fn emit<R: Record>(&self, record: &R) -> Position {
|
||||
self.mux.submit(R::channel(), record.encode())
|
||||
}
|
||||
|
||||
fn emit_text(&self, label: &str, stream: ProcStream, line: &str) -> Position {
|
||||
self.mux.submit(
|
||||
catalog::process_output(label, stream),
|
||||
line.as_bytes().to_vec(),
|
||||
)
|
||||
}
|
||||
|
||||
fn emit_opaque(&self, channel: &str, bytes: &[u8]) -> Position {
|
||||
self.mux.submit(ChannelId::new(channel), bytes.to_vec())
|
||||
}
|
||||
|
||||
fn sent(&self) -> Vec<Frame> {
|
||||
self.mux.drain()
|
||||
}
|
||||
}
|
||||
|
||||
fn test_registry() -> ChannelRegistry {
|
||||
ChannelRegistry::new()
|
||||
.with_record::<IdentityRecord>()
|
||||
|
|
@ -148,7 +191,7 @@ impl Record for ExternalPluginRecord {
|
|||
#[test]
|
||||
fn external_record_full_pipe_round_trips_without_datastream_catalog() {
|
||||
let stream = test_stream();
|
||||
let mux = Mux::unbounded(stream.clone());
|
||||
let mux = mux_without_frame_timing(stream.clone());
|
||||
let record = ExternalPluginRecord {
|
||||
value: 42,
|
||||
label: "owned outside datastream".into(),
|
||||
|
|
@ -380,7 +423,7 @@ fn records_name_their_own_typed_channel() {
|
|||
/// exactly, and `submit` returns each position in order.
|
||||
#[test]
|
||||
fn mux_numbers_monotonic_and_gap_free() {
|
||||
let mux = Mux::unbounded(test_stream());
|
||||
let mux = mux_without_frame_timing(test_stream());
|
||||
let k = 64u64;
|
||||
for i in 0..k {
|
||||
let pos = mux.submit(catalog::HOST_RESOURCE, payloads::resource(i).encode());
|
||||
|
|
@ -410,7 +453,7 @@ fn mux_numbers_monotonic_and_gap_free() {
|
|||
/// A typed event and a log line share the single timeline (spec §5.1).
|
||||
#[test]
|
||||
fn mux_seam_preserves_every_submission_byte_identical() {
|
||||
let mux = Mux::unbounded(test_stream());
|
||||
let mux = mux_without_frame_timing(test_stream());
|
||||
|
||||
// A realistic interleaving of typed records and raw process output —
|
||||
// the same kind of thing on one stream (spec §4.2).
|
||||
|
|
@ -465,13 +508,128 @@ fn mux_seam_preserves_every_submission_byte_identical() {
|
|||
}
|
||||
}
|
||||
|
||||
/// Timing sidecars are enabled by default, and callers can opt out to preserve
|
||||
/// the existing one-submission-to-one-frame stream shape.
|
||||
#[test]
|
||||
fn mux_timing_is_enabled_by_default_and_can_be_disabled() {
|
||||
let mux = Mux::unbounded(test_stream());
|
||||
|
||||
assert!(mux.frame_timing_enabled());
|
||||
let timed_position = mux.submit(catalog::HOST_RESOURCE, payloads::resource(0).encode());
|
||||
let timed_frames = mux.drain();
|
||||
|
||||
assert_eq!(timed_position, Position(0));
|
||||
assert_eq!(mux.assigned(), 2);
|
||||
assert_eq!(timed_frames.len(), 2);
|
||||
assert_eq!(timed_frames[0].position, Position(0));
|
||||
assert_eq!(
|
||||
timed_frames[0].channel,
|
||||
ChannelId::new(catalog::HOST_RESOURCE)
|
||||
);
|
||||
assert_eq!(timed_frames[0].payload, payloads::resource(0).encode());
|
||||
assert_eq!(timed_frames[1].position, Position(1));
|
||||
assert_eq!(timed_frames[1].channel, ChannelId::new(FRAME_TIME_CHANNEL));
|
||||
let sample = FrameTimeSample::decode(&timed_frames[1].payload).expect("timing sidecar decodes");
|
||||
assert_eq!(sample.target_position, timed_position.0);
|
||||
|
||||
mux.set_frame_timing_enabled(false);
|
||||
assert!(!mux.frame_timing_enabled());
|
||||
let untimed_position = mux.submit(catalog::HOST_RESOURCE, payloads::resource(1).encode());
|
||||
let untimed_frames = mux.drain();
|
||||
|
||||
assert_eq!(untimed_position, Position(2));
|
||||
assert_eq!(mux.assigned(), 3);
|
||||
assert_eq!(untimed_frames.len(), 1);
|
||||
assert_eq!(untimed_frames[0].position, Position(2));
|
||||
assert_eq!(
|
||||
untimed_frames[0].channel,
|
||||
ChannelId::new(catalog::HOST_RESOURCE)
|
||||
);
|
||||
assert_eq!(untimed_frames[0].payload, payloads::resource(1).encode());
|
||||
}
|
||||
|
||||
/// With timing enabled by default, a submitted data frame is immediately
|
||||
/// followed by a timing sidecar whose payload keys the sample to the data
|
||||
/// frame's position, while the reserved timing channel itself does not recurse.
|
||||
#[test]
|
||||
fn mux_timing_sidecar_decodes_and_timing_channel_does_not_recurse() {
|
||||
let mux = Mux::unbounded(test_stream());
|
||||
assert!(mux.frame_timing_enabled());
|
||||
let submitted_timing = FrameTimeSample::new(Position(42), 123);
|
||||
|
||||
let data_position = mux.submit(catalog::HOST_RESOURCE, payloads::resource(7).encode());
|
||||
let timing_position = mux.submit(FRAME_TIME_CHANNEL, submitted_timing.encode());
|
||||
let frames = mux.drain();
|
||||
|
||||
assert_eq!(data_position, Position(0));
|
||||
assert_eq!(timing_position, Position(2));
|
||||
assert_eq!(mux.assigned(), 3);
|
||||
assert_eq!(frames.len(), 3);
|
||||
assert_eq!(frames[0].position, Position(0));
|
||||
assert_eq!(frames[0].channel, ChannelId::new(catalog::HOST_RESOURCE));
|
||||
assert_eq!(frames[0].payload, payloads::resource(7).encode());
|
||||
assert_eq!(frames[1].position, Position(1));
|
||||
assert_eq!(frames[1].channel, ChannelId::new(FRAME_TIME_CHANNEL));
|
||||
|
||||
let sample = FrameTimeSample::decode(&frames[1].payload).expect("timing sidecar decodes");
|
||||
assert_eq!(sample.target_position, data_position.0);
|
||||
assert!(
|
||||
sample.created_at_unix_ns > 0,
|
||||
"sidecar records a concrete creation timestamp"
|
||||
);
|
||||
|
||||
assert_eq!(frames[2].position, Position(2));
|
||||
assert_eq!(frames[2].channel, ChannelId::new(FRAME_TIME_CHANNEL));
|
||||
let decoded =
|
||||
FrameTimeSample::decode(&frames[2].payload).expect("submitted timing frame decodes");
|
||||
assert_eq!(decoded, submitted_timing);
|
||||
}
|
||||
|
||||
/// When timing sidecars fill the buffer with a data frame, a later overflowing
|
||||
/// data submission must not leave behind a sidecar for the dropped position.
|
||||
#[test]
|
||||
fn mux_timing_overflow_does_not_sample_dropped_data_frame() {
|
||||
let mux = Mux::new(test_stream(), 2);
|
||||
assert!(mux.frame_timing_enabled());
|
||||
|
||||
mux.submit(catalog::HOST_RESOURCE, payloads::resource(0).encode());
|
||||
let dropped_position = mux.submit(catalog::HOST_RESOURCE, payloads::resource(1).encode());
|
||||
let frames = mux.drain();
|
||||
|
||||
assert_eq!(dropped_position, Position(2));
|
||||
assert_eq!(mux.assigned(), 3);
|
||||
assert_eq!(mux.dropped(), 1);
|
||||
assert_eq!(frames.len(), 2);
|
||||
assert_eq!(
|
||||
frames
|
||||
.iter()
|
||||
.map(|frame| frame.position.0)
|
||||
.collect::<Vec<_>>(),
|
||||
vec![0, 1]
|
||||
);
|
||||
assert_eq!(frames[0].channel, ChannelId::new(catalog::HOST_RESOURCE));
|
||||
assert_eq!(frames[1].channel, ChannelId::new(FRAME_TIME_CHANNEL));
|
||||
|
||||
let timing_targets: Vec<u64> = frames
|
||||
.iter()
|
||||
.filter(|frame| frame.channel.as_str() == FRAME_TIME_CHANNEL)
|
||||
.map(|frame| {
|
||||
FrameTimeSample::decode(&frame.payload)
|
||||
.expect("timing sidecar decodes")
|
||||
.target_position
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(timing_targets, vec![0]);
|
||||
assert!(!timing_targets.contains(&dropped_position.0));
|
||||
}
|
||||
|
||||
/// Spec §5.3 — on overflow the mux drops the frame, but its position is
|
||||
/// already spent, so the loss surfaces as a missing position (a detectable
|
||||
/// gap), never a silent renumber. The reference-model gap oracle confirms
|
||||
/// the interior gap.
|
||||
#[test]
|
||||
fn mux_overflow_drops_surface_as_a_gap_not_a_renumber() {
|
||||
let mux = Mux::new(test_stream(), 2); // tiny buffer
|
||||
let mux = bounded_mux_without_frame_timing(test_stream(), 2); // tiny buffer
|
||||
|
||||
mux.submit(catalog::HOST_RESOURCE, payloads::resource(0).encode()); // pos 0 -> buffered
|
||||
mux.submit(catalog::HOST_RESOURCE, payloads::resource(1).encode()); // pos 1 -> buffered
|
||||
|
|
@ -514,7 +672,7 @@ fn mux_overflow_drops_surface_as_a_gap_not_a_renumber() {
|
|||
/// duplicate and no gap. This is the single-ordering-authority guarantee.
|
||||
#[test]
|
||||
fn mux_serializes_concurrent_producers_without_collision() {
|
||||
let mux = Arc::new(Mux::unbounded(test_stream()));
|
||||
let mux = Arc::new(mux_without_frame_timing(test_stream()));
|
||||
let threads = 8u64;
|
||||
let per_thread = 500u64;
|
||||
|
||||
|
|
@ -567,7 +725,7 @@ fn mux_serializes_concurrent_producers_without_collision() {
|
|||
/// real mux with a mix of typed records, raw process output, and a channel
|
||||
/// the consumer does not know, then take its output.
|
||||
fn realistic_stream(stream: &StreamId) -> Vec<Frame> {
|
||||
let mux = Mux::unbounded(stream.clone());
|
||||
let mux = mux_without_frame_timing(stream.clone());
|
||||
mux.submit(
|
||||
catalog::IDENTITY,
|
||||
payloads::identity(stream.node.as_str(), stream.life.0).encode(),
|
||||
|
|
@ -819,7 +977,7 @@ fn view_decodes_each_channel_and_degrades_gracefully() {
|
|||
#[test]
|
||||
fn metric_projection_decodes_one_typed_channel_into_a_series() {
|
||||
let stream = test_stream();
|
||||
let mux = Mux::unbounded(stream.clone());
|
||||
let mux = mux_without_frame_timing(stream.clone());
|
||||
mux.submit(
|
||||
catalog::IDENTITY,
|
||||
payloads::identity("node-alpha", 1).encode(),
|
||||
|
|
@ -853,7 +1011,7 @@ fn metric_projection_decodes_one_typed_channel_into_a_series() {
|
|||
#[test]
|
||||
fn metric_projection_recovers_each_consolidated_record() {
|
||||
let stream = test_stream();
|
||||
let mux = Mux::unbounded(stream.clone());
|
||||
let mux = mux_without_frame_timing(stream.clone());
|
||||
|
||||
let mut dist: Vec<(Position, DistributionState)> = Vec::new();
|
||||
let mut actors: Vec<(Position, ActorRuntimeDetail)> = Vec::new();
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use datastream::{
|
||||
ChannelId, DatastreamEndpoint, DeliveryFanout, Frame, Lifetime, NodeId, Position, StreamId,
|
||||
ChannelId, DatastreamEndpoint, DeliveryFanout, FRAME_TIME_CHANNEL, Frame, FrameTimeSample,
|
||||
Lifetime, NodeId, Position, Record, StreamId,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use swactor::actor::ActorAddress;
|
||||
|
|
@ -15,6 +16,7 @@ fn stream() -> StreamId {
|
|||
fn endpoint_without_subscribers_drains_to_bitbucket() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 4);
|
||||
let producer = endpoint.producer();
|
||||
producer.set_frame_timing_enabled(false);
|
||||
|
||||
producer.submit_text("runtime.log", "before");
|
||||
let tick = endpoint.tick();
|
||||
|
|
@ -31,6 +33,7 @@ fn endpoint_without_subscribers_drains_to_bitbucket() {
|
|||
fn subscription_receives_only_future_frames() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 4);
|
||||
let producer = endpoint.producer();
|
||||
producer.set_frame_timing_enabled(false);
|
||||
|
||||
producer.submit_text("runtime.log", "pre-subscription");
|
||||
endpoint.tick();
|
||||
|
|
@ -51,6 +54,7 @@ fn subscription_receives_only_future_frames() {
|
|||
fn endpoint_fans_out_ordered_frames_to_multiple_subscribers() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 4);
|
||||
let producer = endpoint.producer();
|
||||
producer.set_frame_timing_enabled(false);
|
||||
let left = endpoint.subscribe_all("left");
|
||||
let right = endpoint.subscribe_all("right");
|
||||
|
||||
|
|
@ -70,6 +74,7 @@ fn endpoint_fans_out_ordered_frames_to_multiple_subscribers() {
|
|||
fn slow_subscriber_drops_without_blocking_fast_subscriber() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 8);
|
||||
let producer = endpoint.producer();
|
||||
producer.set_frame_timing_enabled(false);
|
||||
let slow = endpoint.subscribe_all_with_capacity("slow", 1);
|
||||
let fast = endpoint.subscribe_all_with_capacity("fast", 8);
|
||||
|
||||
|
|
@ -111,10 +116,49 @@ fn delivery_fanout_can_publish_collector_deliveries_without_a_mux() {
|
|||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn endpoint_producer_timing_sidecars_are_fanned_out_when_enabled() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 4);
|
||||
let producer = endpoint.producer();
|
||||
let sub = endpoint.subscribe_all("test");
|
||||
|
||||
producer.set_frame_timing_enabled(true);
|
||||
let data_position = producer.submit_text("runtime.log", "visible");
|
||||
let tick = endpoint.tick();
|
||||
|
||||
assert!(endpoint.frame_timing_enabled());
|
||||
assert!(producer.frame_timing_enabled());
|
||||
assert_eq!(data_position, Position(0));
|
||||
assert_eq!(tick.drained, 2);
|
||||
assert_eq!(tick.delivered, 2);
|
||||
|
||||
let deliveries = sub.drain_available();
|
||||
assert_eq!(deliveries.len(), 2);
|
||||
assert_eq!(deliveries[0].stream, stream());
|
||||
assert_eq!(deliveries[0].frame.position, Position(0));
|
||||
assert_eq!(deliveries[0].frame.channel, ChannelId::new("runtime.log"));
|
||||
assert_eq!(deliveries[0].frame.payload, b"visible");
|
||||
assert_eq!(deliveries[1].stream, stream());
|
||||
assert_eq!(deliveries[1].frame.position, Position(1));
|
||||
assert_eq!(
|
||||
deliveries[1].frame.channel,
|
||||
ChannelId::new(FRAME_TIME_CHANNEL)
|
||||
);
|
||||
|
||||
let sample =
|
||||
FrameTimeSample::decode(&deliveries[1].frame.payload).expect("timing sidecar decodes");
|
||||
assert_eq!(sample.target_position, data_position.0);
|
||||
assert!(
|
||||
sample.created_at_unix_ns > 0,
|
||||
"sidecar records a concrete creation timestamp"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn process_observer_adapter_submits_configured_channels() {
|
||||
let endpoint = DatastreamEndpoint::with_capacity(stream(), 8, 4);
|
||||
let producer = endpoint.producer();
|
||||
producer.set_frame_timing_enabled(false);
|
||||
let observer = producer.process_observer_with(|label, is_stderr| {
|
||||
ChannelId::new(format!(
|
||||
"proc.{label}.{}",
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ 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"
|
||||
blake3 = "1"
|
||||
|
||||
[target.'cfg(target_os = "linux")'.dependencies]
|
||||
libc = "0.2"
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use datastream::DatastreamProducer;
|
||||
use datastream::{ChannelId, DatastreamProducer};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use swactor::actor::{ActorAddress, ActorInterface};
|
||||
use swactor::runtime::{Ctx, ExternalSender};
|
||||
|
|
@ -100,6 +100,7 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
kind: ProvisionEventKind::ProvisionStart,
|
||||
provider: None,
|
||||
message: None,
|
||||
});
|
||||
let sink = PluginSink::new(Arc::new(ActorPluginSink {
|
||||
|
|
@ -148,6 +149,7 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
run_id,
|
||||
node_id,
|
||||
kind: ProvisionEventKind::NodeStopped,
|
||||
provider: None,
|
||||
message,
|
||||
});
|
||||
}
|
||||
|
|
@ -196,6 +198,9 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
line,
|
||||
},
|
||||
),
|
||||
PluginObservation::DatastreamFrame {
|
||||
channel, payload, ..
|
||||
} => self.emit_datastream_frame(channel, payload),
|
||||
PluginObservation::RuntimeReady {
|
||||
run_id,
|
||||
node_id,
|
||||
|
|
@ -209,6 +214,7 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
run_id,
|
||||
node_id,
|
||||
kind: ProvisionEventKind::NodeLive,
|
||||
provider: None,
|
||||
message: None,
|
||||
});
|
||||
let _ = ctx.send(
|
||||
|
|
@ -314,6 +320,7 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
run_id,
|
||||
node_id,
|
||||
kind: ProvisionEventKind::ProvisionFailed,
|
||||
provider: None,
|
||||
message: Some(reason.to_owned()),
|
||||
});
|
||||
}
|
||||
|
|
@ -352,6 +359,12 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
producer.submit_bytes(channel, payload);
|
||||
}
|
||||
}
|
||||
|
||||
fn emit_datastream_frame(&self, channel: String, payload: String) {
|
||||
if let Some(producer) = &self.telemetry {
|
||||
producer.submit_bytes(ChannelId::new(channel), payload.into_bytes());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<P: ProvisionPlugin + 'static> ActorInterface for ProvisionerActor<P> {
|
||||
|
|
|
|||
|
|
@ -39,6 +39,9 @@ use mvp_system::node_provisioning as node_provision;
|
|||
use mvp_system::observability_surface as obs;
|
||||
use mvp_system::orchestrator_run_fsm as fsm;
|
||||
use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream};
|
||||
use mvp_system::relay_provisioning::{
|
||||
LocalShimRelayProvider, RelayProvider, RelayProvisionRequest, RelayPurpose,
|
||||
};
|
||||
use mvp_system::run_plan as plan;
|
||||
use mvp_system::stage_controller as stage;
|
||||
use mvp_system::tx_rx_edge_actor as edge_actor;
|
||||
|
|
@ -1224,11 +1227,17 @@ fn run_node_role(args: &[String]) -> Result<(), String> {
|
|||
}
|
||||
|
||||
fn new_driver(handle: tokio::runtime::Handle) -> Result<IrohDriver, String> {
|
||||
let mut relay_provider = LocalShimRelayProvider;
|
||||
let relay = relay_provider.provision_relay(RelayProvisionRequest {
|
||||
run_id: RUN_ID,
|
||||
purpose: RelayPurpose::Combined,
|
||||
})?;
|
||||
let relay_mode = relay_provider.relay_mode(&relay)?;
|
||||
IrohDriver::with_handle(
|
||||
handle,
|
||||
IrohDriverConfig {
|
||||
secret_key: None,
|
||||
relay_mode: iroh::RelayMode::Disabled,
|
||||
relay_mode,
|
||||
node: DistributedNodeConfig::default(),
|
||||
peer_auth: None,
|
||||
additional_alpns: vec![EDGE_ALPN.to_vec()],
|
||||
|
|
@ -2285,6 +2294,7 @@ fn local_docker_spec(
|
|||
"--orchestrator-actor".to_owned(),
|
||||
orchestrator_actor_json.to_owned(),
|
||||
],
|
||||
mounts: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -3,7 +3,8 @@ use std::thread::{self, JoinHandle};
|
|||
|
||||
use datastream::{ChannelId, DatastreamProducer, Lifetime, NodeId, StreamId};
|
||||
use iroh::EndpointAddr;
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use swactor::actor::ActorAddress;
|
||||
|
||||
use crate::provisioning::{
|
||||
|
|
@ -49,6 +50,10 @@ impl BootstrapDatastreamBridge {
|
|||
|
||||
pub fn observe_stdout_line(&self, line: impl Into<String>) {
|
||||
let line = line.into();
|
||||
if let Some(frame) = parse_stdio_datastream_frame(&self.spec, &line) {
|
||||
self.sink.observe(frame);
|
||||
return;
|
||||
}
|
||||
self.submit_log(ProvisionLogStream::Stdout, &line);
|
||||
self.sink.observe(PluginObservation::StdoutLine {
|
||||
run_id: self.spec.run_id,
|
||||
|
|
@ -154,6 +159,30 @@ impl BootstrapDatastreamBridge {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Serialize)]
|
||||
struct StdioDatastreamFrame {
|
||||
mvp_stdio_event: u32,
|
||||
kind: String,
|
||||
channel: String,
|
||||
payload: Value,
|
||||
}
|
||||
|
||||
pub fn parse_stdio_datastream_frame(
|
||||
spec: &NodeProvisionSpec,
|
||||
line: &str,
|
||||
) -> Option<PluginObservation> {
|
||||
let frame = serde_json::from_str::<StdioDatastreamFrame>(line).ok()?;
|
||||
if frame.mvp_stdio_event != 1 || frame.kind != "datastream_frame" {
|
||||
return None;
|
||||
}
|
||||
Some(PluginObservation::DatastreamFrame {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
channel: frame.channel,
|
||||
payload: frame.payload.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RuntimeReadyLine {
|
||||
#[serde(rename = "type")]
|
||||
|
|
|
|||
|
|
@ -18,12 +18,14 @@ pub mod gpu_worker_ingress_parser;
|
|||
pub mod gpu_worker_process_adapter;
|
||||
pub mod membership_pool_readiness;
|
||||
pub mod node_boot_lifecycle;
|
||||
pub mod node_image;
|
||||
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 relay_provisioning;
|
||||
pub mod resource_inventory;
|
||||
pub mod run_plan;
|
||||
pub mod shard_fetch;
|
||||
|
|
|
|||
635
crates/mvp-system/src/node_image.rs
Normal file
635
crates/mvp-system/src/node_image.rs
Normal file
|
|
@ -0,0 +1,635 @@
|
|||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::fs::{self, File};
|
||||
use std::io::Read;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, Stdio};
|
||||
|
||||
const IMAGE_SOURCE_INPUTS: &[&str] = &[
|
||||
"Cargo.lock",
|
||||
"Cargo.toml",
|
||||
"src",
|
||||
"crates/datastream/Cargo.toml",
|
||||
"crates/datastream/src",
|
||||
"crates/distribution/Cargo.toml",
|
||||
"crates/distribution/src",
|
||||
"crates/iroh-driver/Cargo.toml",
|
||||
"crates/iroh-driver/src",
|
||||
"crates/mvp-system/Cargo.toml",
|
||||
"crates/mvp-system/src",
|
||||
"crates/transport/Cargo.toml",
|
||||
"crates/transport/src",
|
||||
"tools/vastai/Cargo.toml",
|
||||
"tools/vastai/src",
|
||||
"apps/mvp-node/Dockerfile",
|
||||
"apps/mvp-node/Dockerfile.base",
|
||||
"apps/mvp-node/mvp_entrypoint.sh",
|
||||
"apps/mvp-node/tinygrad_worker.py",
|
||||
];
|
||||
|
||||
const BASE_IMAGE_SOURCE_INPUTS: &[&str] = &[
|
||||
"apps/mvp-node/Dockerfile.base",
|
||||
"apps/mvp-node/mvp_entrypoint.sh",
|
||||
];
|
||||
const NODE_IMAGE_TAG_LABEL: &str = "org.swactor.mvp.node-image-tag";
|
||||
const NODE_IMAGE_SOURCE_HASH_LABEL: &str = "org.swactor.mvp.node.source-hash";
|
||||
const NODE_IMAGE_WORKER_HASH_LABEL: &str = "org.swactor.mvp.node.worker-hash";
|
||||
const NODE_IMAGE_BASE_HASH_LABEL: &str = "org.swactor.mvp.node.base-hash";
|
||||
const BASE_IMAGE_SOURCE_HASH_LABEL: &str = "org.swactor.mvp.base.source-hash";
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum NodeImageProvider {
|
||||
Docker,
|
||||
VastAi,
|
||||
}
|
||||
|
||||
impl NodeImageProvider {
|
||||
fn requires_remote_image(self) -> bool {
|
||||
matches!(self, Self::VastAi)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct NodeImageRequest {
|
||||
pub requested_image: String,
|
||||
pub base_image: String,
|
||||
pub node_bin: PathBuf,
|
||||
pub provider: NodeImageProvider,
|
||||
pub extra_tag: Option<String>,
|
||||
pub push: bool,
|
||||
pub force_refresh: bool,
|
||||
pub enabled: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct PreparedNodeImage {
|
||||
pub image_ref: String,
|
||||
pub tag: String,
|
||||
pub already_available: bool,
|
||||
pub built: bool,
|
||||
pub pushed: bool,
|
||||
}
|
||||
|
||||
pub fn prepare_node_image(request: NodeImageRequest) -> Result<PreparedNodeImage, String> {
|
||||
if !request.enabled {
|
||||
return Ok(PreparedNodeImage {
|
||||
image_ref: request.requested_image,
|
||||
tag: String::new(),
|
||||
already_available: false,
|
||||
built: false,
|
||||
pushed: false,
|
||||
});
|
||||
}
|
||||
|
||||
let root = workspace_root()?;
|
||||
let image = ImageName::parse(&request.requested_image)?;
|
||||
if request.provider.requires_remote_image() && !looks_registry_reachable(&image.repository) {
|
||||
return Err(format!(
|
||||
"VastAI node image {:?} must include a registry namespace",
|
||||
image.repository
|
||||
));
|
||||
}
|
||||
|
||||
let tag = image_version_tag(&root)?;
|
||||
let image_ref = image.ref_for_tag(&tag);
|
||||
let source_hash = source_content_hash(&root)?;
|
||||
let worker_hash = file_content_hash(&root, Path::new("apps/mvp-node/tinygrad_worker.py"))?;
|
||||
let base_hash = content_hash_for_inputs(&root, BASE_IMAGE_SOURCE_INPUTS)?;
|
||||
let expected_node_labels = node_image_labels(&tag, &source_hash, &worker_hash, &base_hash);
|
||||
let expected_base_labels = base_image_labels(&base_hash);
|
||||
let alias_tags = alias_tags(&image, request.extra_tag.as_deref(), &tag)?;
|
||||
let remote_required = request.provider.requires_remote_image() || request.push;
|
||||
|
||||
let local_image_matches = docker_image_labels_match(&root, &image_ref, &expected_node_labels)?;
|
||||
let remote_available = remote_required && docker_manifest_exists(&root, &image_ref);
|
||||
if !request.force_refresh && remote_required && remote_available {
|
||||
let pushed = ensure_aliases_for_remote(&root, &image_ref, &image, &alias_tags)?;
|
||||
return Ok(PreparedNodeImage {
|
||||
image_ref,
|
||||
tag,
|
||||
already_available: true,
|
||||
built: false,
|
||||
pushed,
|
||||
});
|
||||
}
|
||||
if !request.force_refresh && remote_required && local_image_matches {
|
||||
ensure_aliases_local(&root, &image_ref, &image, &alias_tags)?;
|
||||
push_image(&root, &image_ref)?;
|
||||
for alias in alias_refs(&image, &alias_tags) {
|
||||
push_image(&root, &alias)?;
|
||||
}
|
||||
return Ok(PreparedNodeImage {
|
||||
image_ref,
|
||||
tag,
|
||||
already_available: true,
|
||||
built: false,
|
||||
pushed: true,
|
||||
});
|
||||
}
|
||||
if !request.force_refresh && !remote_required && local_image_matches {
|
||||
ensure_aliases_local(&root, &image_ref, &image, &alias_tags)?;
|
||||
return Ok(PreparedNodeImage {
|
||||
image_ref,
|
||||
tag,
|
||||
already_available: true,
|
||||
built: false,
|
||||
pushed: false,
|
||||
});
|
||||
}
|
||||
|
||||
run_status(
|
||||
&root,
|
||||
"cargo",
|
||||
&["build", "--quiet", "-p", "mvp-system", "--bin", "mvp-node"],
|
||||
"build mvp-node",
|
||||
)?;
|
||||
let base_image_matches =
|
||||
docker_image_labels_match(&root, &request.base_image, &expected_base_labels)?;
|
||||
if !base_image_matches {
|
||||
run_status_vec(
|
||||
&root,
|
||||
"docker",
|
||||
vec![
|
||||
"build".to_owned(),
|
||||
"-f".to_owned(),
|
||||
"apps/mvp-node/Dockerfile.base".to_owned(),
|
||||
"--label".to_owned(),
|
||||
format!("{BASE_IMAGE_SOURCE_HASH_LABEL}={base_hash}"),
|
||||
"-t".to_owned(),
|
||||
request.base_image.clone(),
|
||||
".".to_owned(),
|
||||
],
|
||||
"build mvp node base image",
|
||||
)?;
|
||||
}
|
||||
|
||||
let node_bin = request.node_bin.to_string_lossy().to_string();
|
||||
let mut build_args = vec![
|
||||
"build".to_owned(),
|
||||
"-f".to_owned(),
|
||||
"apps/mvp-node/Dockerfile".to_owned(),
|
||||
"--build-arg".to_owned(),
|
||||
format!("BASE_IMAGE={}", request.base_image),
|
||||
"--build-arg".to_owned(),
|
||||
format!("MVP_NODE_BIN={node_bin}"),
|
||||
];
|
||||
for (key, value) in &expected_node_labels {
|
||||
build_args.push("--label".to_owned());
|
||||
build_args.push(format!("{key}={value}"));
|
||||
}
|
||||
build_args.extend(["-t".to_owned(), image_ref.clone(), ".".to_owned()]);
|
||||
run_status_vec(&root, "docker", build_args, "build mvp node image")?;
|
||||
ensure_aliases_local(&root, &image_ref, &image, &alias_tags)?;
|
||||
|
||||
let mut pushed = false;
|
||||
if remote_required {
|
||||
push_image(&root, &image_ref)?;
|
||||
pushed = true;
|
||||
for alias in alias_refs(&image, &alias_tags) {
|
||||
push_image(&root, &alias)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(PreparedNodeImage {
|
||||
image_ref,
|
||||
tag,
|
||||
already_available: false,
|
||||
built: true,
|
||||
pushed,
|
||||
})
|
||||
}
|
||||
|
||||
fn workspace_root() -> Result<PathBuf, String> {
|
||||
let output = Command::new("git")
|
||||
.args(["rev-parse", "--show-toplevel"])
|
||||
.stdin(Stdio::null())
|
||||
.output()
|
||||
.map_err(|e| format!("locate repository root with git: {e}"))?;
|
||||
if !output.status.success() {
|
||||
return Err(format!(
|
||||
"git rev-parse --show-toplevel failed with {}: {}",
|
||||
output.status,
|
||||
String::from_utf8_lossy(&output.stderr).trim()
|
||||
));
|
||||
}
|
||||
Ok(PathBuf::from(
|
||||
String::from_utf8_lossy(&output.stdout).trim(),
|
||||
))
|
||||
}
|
||||
|
||||
fn image_version_tag(root: &Path) -> Result<String, String> {
|
||||
if git_worktree_clean(root)? {
|
||||
let sha = git_capture(root, &["rev-parse", "--short=12", "HEAD"])?;
|
||||
Ok(format!("git-{}", sha.trim()))
|
||||
} else {
|
||||
Ok(format!("dirty-{}", dirty_content_hash(root)?))
|
||||
}
|
||||
}
|
||||
|
||||
fn git_worktree_clean(root: &Path) -> Result<bool, String> {
|
||||
Ok(git_capture(root, &["status", "--porcelain"])?
|
||||
.trim()
|
||||
.is_empty())
|
||||
}
|
||||
|
||||
fn git_capture(root: &Path, args: &[&str]) -> Result<String, String> {
|
||||
let output = Command::new("git")
|
||||
.current_dir(root)
|
||||
.args(args)
|
||||
.stdin(Stdio::null())
|
||||
.output()
|
||||
.map_err(|e| format!("run git {}: {e}", args.join(" ")))?;
|
||||
if output.status.success() {
|
||||
Ok(String::from_utf8_lossy(&output.stdout).to_string())
|
||||
} else {
|
||||
Err(format!(
|
||||
"git {} failed with {}: {}",
|
||||
args.join(" "),
|
||||
output.status,
|
||||
String::from_utf8_lossy(&output.stderr).trim()
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn dirty_content_hash(root: &Path) -> Result<String, String> {
|
||||
source_content_hash(root)
|
||||
}
|
||||
|
||||
fn source_content_hash(root: &Path) -> Result<String, String> {
|
||||
content_hash_for_inputs(root, IMAGE_SOURCE_INPUTS)
|
||||
}
|
||||
|
||||
fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result<String, String> {
|
||||
let mut files = Vec::new();
|
||||
for input in inputs {
|
||||
let path = root.join(input);
|
||||
collect_hash_inputs(root, &path, &mut files)?;
|
||||
}
|
||||
files.sort();
|
||||
files.dedup();
|
||||
hash_relative_files(root, files)
|
||||
}
|
||||
|
||||
fn file_content_hash(root: &Path, path: &Path) -> Result<String, String> {
|
||||
hash_relative_files(root, vec![relative_path(root, &root.join(path))?])
|
||||
}
|
||||
|
||||
fn hash_relative_files(root: &Path, files: Vec<PathBuf>) -> Result<String, String> {
|
||||
let mut hasher = blake3::Hasher::new();
|
||||
for relative in files {
|
||||
let full = root.join(&relative);
|
||||
hasher.update(relative.to_string_lossy().as_bytes());
|
||||
hasher.update(b"\0");
|
||||
hash_file_content(root, &full, &mut hasher)?;
|
||||
hasher.update(b"\0");
|
||||
}
|
||||
let hash = hasher.finalize().to_hex().to_string();
|
||||
Ok(hash[..16].to_owned())
|
||||
}
|
||||
|
||||
fn hash_file_content(root: &Path, path: &Path, hasher: &mut blake3::Hasher) -> Result<(), String> {
|
||||
let display = display_workspace_path(root, path);
|
||||
let mut file = File::open(path).map_err(|e| format!("open {display}: {e}"))?;
|
||||
let mut buf = [0_u8; 64 * 1024];
|
||||
loop {
|
||||
let n = file
|
||||
.read(&mut buf)
|
||||
.map_err(|e| format!("read {display}: {e}"))?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
hasher.update(&buf[..n]);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn collect_hash_inputs(root: &Path, path: &Path, out: &mut Vec<PathBuf>) -> Result<(), String> {
|
||||
if !path.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
let display = display_workspace_path(root, path);
|
||||
let metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?;
|
||||
if metadata.is_file() {
|
||||
if !skip_file(path) {
|
||||
out.push(relative_path(root, path)?);
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
if !metadata.is_dir() || skip_dir(path) {
|
||||
return Ok(());
|
||||
}
|
||||
let entries = fs::read_dir(path).map_err(|e| format!("read dir {display}: {e}"))?;
|
||||
for entry in entries {
|
||||
let entry = entry.map_err(|e| format!("read dir entry {display}: {e}"))?;
|
||||
collect_hash_inputs(root, &entry.path(), out)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn relative_path(root: &Path, path: &Path) -> Result<PathBuf, String> {
|
||||
path.strip_prefix(root).map(Path::to_path_buf).map_err(|e| {
|
||||
format!(
|
||||
"make {} relative to {}: {e}",
|
||||
display_workspace_path(root, path),
|
||||
"."
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn display_workspace_path(root: &Path, path: &Path) -> String {
|
||||
match path.strip_prefix(root) {
|
||||
Ok(relative) if relative.as_os_str().is_empty() => ".".to_owned(),
|
||||
Ok(relative) => format!("./{}", relative.display()),
|
||||
Err(_) => path.display().to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn skip_dir(path: &Path) -> bool {
|
||||
matches!(
|
||||
path.file_name().and_then(|name| name.to_str()),
|
||||
Some(".git" | "target" | "__pycache__")
|
||||
)
|
||||
}
|
||||
|
||||
fn skip_file(path: &Path) -> bool {
|
||||
matches!(path.extension().and_then(|ext| ext.to_str()), Some("pyc"))
|
||||
}
|
||||
|
||||
fn alias_tags(
|
||||
image: &ImageName,
|
||||
extra_tag: Option<&str>,
|
||||
version_tag: &str,
|
||||
) -> Result<BTreeSet<String>, String> {
|
||||
let mut tags = BTreeSet::new();
|
||||
if let Some(tag) = image.requested_tag.as_deref() {
|
||||
insert_alias_tag(&mut tags, tag, version_tag)?;
|
||||
}
|
||||
if let Some(tag) = extra_tag {
|
||||
insert_alias_tag(&mut tags, tag, version_tag)?;
|
||||
}
|
||||
Ok(tags)
|
||||
}
|
||||
|
||||
fn insert_alias_tag(
|
||||
tags: &mut BTreeSet<String>,
|
||||
tag: &str,
|
||||
version_tag: &str,
|
||||
) -> Result<(), String> {
|
||||
let tag = tag.trim();
|
||||
if tag.is_empty() {
|
||||
return Err("node image tag must not be empty".to_owned());
|
||||
}
|
||||
if tag != version_tag {
|
||||
tags.insert(tag.to_owned());
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn node_image_labels<'a>(
|
||||
tag: &'a str,
|
||||
source_hash: &'a str,
|
||||
worker_hash: &'a str,
|
||||
base_hash: &'a str,
|
||||
) -> Vec<(&'static str, &'a str)> {
|
||||
vec![
|
||||
(NODE_IMAGE_TAG_LABEL, tag),
|
||||
(NODE_IMAGE_SOURCE_HASH_LABEL, source_hash),
|
||||
(NODE_IMAGE_WORKER_HASH_LABEL, worker_hash),
|
||||
(NODE_IMAGE_BASE_HASH_LABEL, base_hash),
|
||||
]
|
||||
}
|
||||
|
||||
fn base_image_labels<'a>(base_hash: &'a str) -> Vec<(&'static str, &'a str)> {
|
||||
vec![(BASE_IMAGE_SOURCE_HASH_LABEL, base_hash)]
|
||||
}
|
||||
|
||||
fn ensure_aliases_local(
|
||||
root: &Path,
|
||||
source_ref: &str,
|
||||
image: &ImageName,
|
||||
alias_tags: &BTreeSet<String>,
|
||||
) -> Result<(), String> {
|
||||
for alias in alias_refs(image, alias_tags) {
|
||||
if alias != source_ref {
|
||||
run_status(
|
||||
root,
|
||||
"docker",
|
||||
&["tag", source_ref, &alias],
|
||||
"tag mvp node image",
|
||||
)?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_aliases_for_remote(
|
||||
root: &Path,
|
||||
source_ref: &str,
|
||||
image: &ImageName,
|
||||
alias_tags: &BTreeSet<String>,
|
||||
) -> Result<bool, String> {
|
||||
if alias_tags.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
if !docker_image_exists(root, source_ref) {
|
||||
run_status(root, "docker", &["pull", source_ref], "pull mvp node image")?;
|
||||
}
|
||||
ensure_aliases_local(root, source_ref, image, alias_tags)?;
|
||||
for alias in alias_refs(image, alias_tags) {
|
||||
push_image(root, &alias)?;
|
||||
}
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
fn alias_refs(image: &ImageName, alias_tags: &BTreeSet<String>) -> Vec<String> {
|
||||
alias_tags
|
||||
.iter()
|
||||
.map(|tag| image.ref_for_tag(tag))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn push_image(root: &Path, image_ref: &str) -> Result<(), String> {
|
||||
run_status(root, "docker", &["push", image_ref], "push mvp node image")
|
||||
}
|
||||
|
||||
fn docker_image_exists(root: &Path, image_ref: &str) -> bool {
|
||||
Command::new("docker")
|
||||
.current_dir(root)
|
||||
.args(["image", "inspect", image_ref])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.map(|status| status.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn docker_image_labels_match(
|
||||
root: &Path,
|
||||
image_ref: &str,
|
||||
expected: &[(&str, &str)],
|
||||
) -> Result<bool, String> {
|
||||
let output = Command::new("docker")
|
||||
.current_dir(root)
|
||||
.args([
|
||||
"image",
|
||||
"inspect",
|
||||
"--format",
|
||||
"{{ json .Config.Labels }}",
|
||||
image_ref,
|
||||
])
|
||||
.stdin(Stdio::null())
|
||||
.output()
|
||||
.map_err(|e| format!("inspect docker image {image_ref}: {e}"))?;
|
||||
if !output.status.success() {
|
||||
return Ok(false);
|
||||
}
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let labels: Option<BTreeMap<String, String>> = serde_json::from_str(stdout.trim())
|
||||
.map_err(|e| format!("parse docker labels for {image_ref}: {e}"))?;
|
||||
let labels = labels.unwrap_or_default();
|
||||
Ok(expected
|
||||
.iter()
|
||||
.all(|(key, value)| labels.get(*key).map(String::as_str) == Some(*value)))
|
||||
}
|
||||
|
||||
fn docker_manifest_exists(root: &Path, image_ref: &str) -> bool {
|
||||
Command::new("docker")
|
||||
.current_dir(root)
|
||||
.args(["manifest", "inspect", image_ref])
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.map(|status| status.success())
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn run_status(root: &Path, program: &str, args: &[&str], label: &str) -> Result<(), String> {
|
||||
eprintln!("mvp-node-image: {label}");
|
||||
let status = Command::new(program)
|
||||
.current_dir(root)
|
||||
.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 run_status_vec(
|
||||
root: &Path,
|
||||
program: &str,
|
||||
args: Vec<String>,
|
||||
label: &str,
|
||||
) -> Result<(), String> {
|
||||
eprintln!("mvp-node-image: {label}");
|
||||
let status = Command::new(program)
|
||||
.current_dir(root)
|
||||
.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 looks_registry_reachable(repository: &str) -> bool {
|
||||
let first = repository.split('/').next().unwrap_or(repository);
|
||||
repository.contains('/') || first.contains('.') || first.contains(':') || first == "localhost"
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct ImageName {
|
||||
repository: String,
|
||||
requested_tag: Option<String>,
|
||||
}
|
||||
|
||||
impl ImageName {
|
||||
fn parse(raw: &str) -> Result<Self, String> {
|
||||
let raw = raw.trim();
|
||||
if raw.is_empty() {
|
||||
return Err("node image must not be empty".to_owned());
|
||||
}
|
||||
if raw.contains('@') {
|
||||
return Err(format!(
|
||||
"node image {raw:?} uses a digest; use a repository/tag base for image preparation"
|
||||
));
|
||||
}
|
||||
let last_slash = raw.rfind('/');
|
||||
let last_colon = raw.rfind(':');
|
||||
let has_tag = match (last_slash, last_colon) {
|
||||
(_, None) => false,
|
||||
(None, Some(_)) => true,
|
||||
(Some(slash), Some(colon)) => colon > slash,
|
||||
};
|
||||
let (repository, requested_tag) = if has_tag {
|
||||
let colon = last_colon.expect("has tag colon");
|
||||
let repository = raw[..colon].to_owned();
|
||||
let tag = raw[colon + 1..].to_owned();
|
||||
if tag.is_empty() {
|
||||
return Err(format!("node image {raw:?} has an empty tag"));
|
||||
}
|
||||
(repository, Some(tag))
|
||||
} else {
|
||||
(raw.to_owned(), None)
|
||||
};
|
||||
if repository.is_empty() {
|
||||
return Err(format!("node image {raw:?} has an empty repository"));
|
||||
}
|
||||
Ok(Self {
|
||||
repository,
|
||||
requested_tag,
|
||||
})
|
||||
}
|
||||
|
||||
fn ref_for_tag(&self, tag: &str) -> String {
|
||||
format!("{}:{tag}", self.repository)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn image_name_splits_tag_after_last_slash() {
|
||||
let image = ImageName::parse("localhost:5000/team/mvp-node:trial").unwrap();
|
||||
|
||||
assert_eq!(image.repository, "localhost:5000/team/mvp-node");
|
||||
assert_eq!(image.requested_tag.as_deref(), Some("trial"));
|
||||
assert_eq!(
|
||||
image.ref_for_tag("git-abcdef"),
|
||||
"localhost:5000/team/mvp-node:git-abcdef"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_name_keeps_registry_port_without_tag() {
|
||||
let image = ImageName::parse("localhost:5000/team/mvp-node").unwrap();
|
||||
|
||||
assert_eq!(image.repository, "localhost:5000/team/mvp-node");
|
||||
assert_eq!(image.requested_tag, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn alias_tags_include_requested_and_extra_without_version_duplicate() {
|
||||
let image = ImageName::parse("ghcr.io/team/mvp-node:latest").unwrap();
|
||||
|
||||
let aliases = alias_tags(&image, Some("smoke"), "dirty-1234").unwrap();
|
||||
|
||||
assert_eq!(
|
||||
aliases.into_iter().collect::<Vec<_>>(),
|
||||
vec!["latest".to_owned(), "smoke".to_owned()]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -42,6 +42,26 @@ pub enum ProviderKind {
|
|||
VastAi,
|
||||
}
|
||||
|
||||
impl ProviderKind {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Mock => "mock",
|
||||
Self::Docker => "docker",
|
||||
Self::VastAi => "vastai",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn parse_deploy(value: &str) -> Result<Self, String> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"docker" | "local_docker" | "local-docker" => Ok(Self::Docker),
|
||||
"vastai" | "vast_ai" | "vast-ai" => Ok(Self::VastAi),
|
||||
other => Err(format!(
|
||||
"unsupported provider {other:?}; use docker or vastai"
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
|
||||
pub struct DesiredNodeShape {
|
||||
pub image: String,
|
||||
|
|
|
|||
|
|
@ -4,10 +4,13 @@
|
|||
//! report observations back to the provisioner actor through [`PluginSink`].
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::fs;
|
||||
use std::io::Write;
|
||||
use std::path::Path;
|
||||
use std::process::{ChildStdin, Command, Stdio};
|
||||
use std::sync::Arc;
|
||||
use std::thread;
|
||||
use std::time::UNIX_EPOCH;
|
||||
|
||||
use iroh::EndpointAddr;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
|
@ -23,6 +26,15 @@ pub struct NodeProvisionSpec {
|
|||
pub image: String,
|
||||
pub env: Vec<(String, String)>,
|
||||
pub args: Vec<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub mounts: Vec<ProviderMount>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct ProviderMount {
|
||||
pub host_path: String,
|
||||
pub container_path: String,
|
||||
pub readonly: bool,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
|
|
@ -30,6 +42,8 @@ pub struct ProvisionEvent {
|
|||
pub run_id: u64,
|
||||
pub node_id: u64,
|
||||
pub kind: ProvisionEventKind,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider: Option<String>,
|
||||
pub message: Option<String>,
|
||||
}
|
||||
|
||||
|
|
@ -68,6 +82,12 @@ pub enum PluginObservation {
|
|||
node_id: u64,
|
||||
line: String,
|
||||
},
|
||||
DatastreamFrame {
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
channel: String,
|
||||
payload: String,
|
||||
},
|
||||
ProviderLine {
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
|
|
@ -148,6 +168,164 @@ impl LocalDockerPlugin {
|
|||
}
|
||||
}
|
||||
|
||||
fn docker_mount_arg(mount: &ProviderMount) -> String {
|
||||
let mut arg = format!(
|
||||
"type=bind,src={},dst={}",
|
||||
mount.host_path, mount.container_path
|
||||
);
|
||||
if mount.readonly {
|
||||
arg.push_str(",readonly");
|
||||
}
|
||||
arg
|
||||
}
|
||||
|
||||
fn docker_runtime_mount_arg(image: &str, mount: &ProviderMount) -> Result<String, String> {
|
||||
let host_path = Path::new(&mount.host_path);
|
||||
if host_path.is_file() {
|
||||
return prepare_docker_file_volume(image, mount, host_path);
|
||||
}
|
||||
Ok(docker_mount_arg(mount))
|
||||
}
|
||||
|
||||
fn prepare_docker_file_volume(
|
||||
image: &str,
|
||||
mount: &ProviderMount,
|
||||
host_path: &Path,
|
||||
) -> Result<String, String> {
|
||||
let container_path = Path::new(&mount.container_path);
|
||||
let container_dir = container_path
|
||||
.parent()
|
||||
.and_then(|path| path.to_str())
|
||||
.filter(|path| !path.is_empty())
|
||||
.ok_or_else(|| {
|
||||
format!(
|
||||
"cached model container path has no parent: {}",
|
||||
mount.container_path
|
||||
)
|
||||
})?;
|
||||
let file_name = container_path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.filter(|name| !name.is_empty())
|
||||
.ok_or_else(|| {
|
||||
format!(
|
||||
"cached model container path has no file name: {}",
|
||||
mount.container_path
|
||||
)
|
||||
})?;
|
||||
let volume = docker_file_cache_volume_name(host_path)?;
|
||||
docker_status(
|
||||
&["volume", "create", &volume],
|
||||
"create cached model docker volume",
|
||||
)?;
|
||||
let loader_name = format!("mvp-cache-load-{}-{volume}", std::process::id());
|
||||
let _ = Command::new("docker")
|
||||
.args(["rm", "-f", &loader_name])
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status();
|
||||
docker_status_vec(
|
||||
vec![
|
||||
"create".to_owned(),
|
||||
"--name".to_owned(),
|
||||
loader_name.clone(),
|
||||
"--mount".to_owned(),
|
||||
format!("type=volume,src={volume},dst={container_dir}"),
|
||||
"--entrypoint".to_owned(),
|
||||
"/bin/sh".to_owned(),
|
||||
image.to_owned(),
|
||||
"-c".to_owned(),
|
||||
"true".to_owned(),
|
||||
],
|
||||
"create cached model loader container",
|
||||
)?;
|
||||
let copy_result = docker_status_vec(
|
||||
vec![
|
||||
"cp".to_owned(),
|
||||
host_path.to_string_lossy().to_string(),
|
||||
format!("{loader_name}:{container_dir}/{file_name}"),
|
||||
],
|
||||
"copy cached model into docker volume",
|
||||
);
|
||||
let _ = Command::new("docker")
|
||||
.args(["rm", &loader_name])
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status();
|
||||
copy_result?;
|
||||
|
||||
// tinygrad opens GGUF files read-write even when it does not intend to mutate
|
||||
// the original artifact. The Docker volume is a copied cache, so keeping the
|
||||
// runtime mount writable preserves the host cache while satisfying the loader.
|
||||
Ok(format!("type=volume,src={volume},dst={container_dir}"))
|
||||
}
|
||||
|
||||
fn docker_file_cache_volume_name(path: &Path) -> Result<String, String> {
|
||||
let metadata = fs::metadata(path).map_err(|e| format!("stat {}: {e}", path.display()))?;
|
||||
let modified = metadata
|
||||
.modified()
|
||||
.ok()
|
||||
.and_then(|time| time.duration_since(UNIX_EPOCH).ok())
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or(0);
|
||||
let file = path
|
||||
.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.map(safe_docker_volume_component)
|
||||
.unwrap_or_else(|| "model".to_owned());
|
||||
Ok(format!("mvp-cache-{file}-{}-{modified}", metadata.len()))
|
||||
}
|
||||
|
||||
fn safe_docker_volume_component(value: &str) -> String {
|
||||
let mut out = String::with_capacity(value.len().min(40));
|
||||
for ch in value.chars() {
|
||||
if out.len() >= 40 {
|
||||
break;
|
||||
}
|
||||
if ch.is_ascii_alphanumeric() {
|
||||
out.push(ch.to_ascii_lowercase());
|
||||
} else if !out.ends_with('-') {
|
||||
out.push('-');
|
||||
}
|
||||
}
|
||||
let trimmed = out.trim_matches('-');
|
||||
if trimmed.is_empty() {
|
||||
"model".to_owned()
|
||||
} else {
|
||||
trimmed.to_owned()
|
||||
}
|
||||
}
|
||||
|
||||
fn docker_status(args: &[&str], label: &str) -> Result<(), String> {
|
||||
let status = Command::new("docker")
|
||||
.args(args)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::inherit())
|
||||
.status()
|
||||
.map_err(|e| format!("{label}: {e}"))?;
|
||||
if status.success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("{label} failed with {status}"))
|
||||
}
|
||||
}
|
||||
|
||||
fn docker_status_vec(args: Vec<String>, label: &str) -> Result<(), String> {
|
||||
let status = Command::new("docker")
|
||||
.args(&args)
|
||||
.stdin(Stdio::null())
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::inherit())
|
||||
.status()
|
||||
.map_err(|e| format!("{label}: {e}"))?;
|
||||
if status.success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!("{label} failed with {status}"))
|
||||
}
|
||||
}
|
||||
|
||||
impl ProvisionPlugin for LocalDockerPlugin {
|
||||
fn start_node(
|
||||
&mut self,
|
||||
|
|
@ -167,17 +345,33 @@ impl ProvisionPlugin for LocalDockerPlugin {
|
|||
.arg("--name")
|
||||
.arg(&container_name)
|
||||
.arg("-i");
|
||||
let spec_gpus = spec
|
||||
let docker_gpus = spec
|
||||
.env
|
||||
.iter()
|
||||
.find(|(key, _)| key == "MVP_DOCKER_GPUS")
|
||||
.map(|(_, value)| value.clone());
|
||||
if let Some(gpus) = spec_gpus
|
||||
.map(|(_, value)| value.clone())
|
||||
.or_else(|| std::env::var("MVP_DOCKER_GPUS").ok())
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
{
|
||||
.filter(|value| !value.trim().is_empty());
|
||||
sink.observe(PluginObservation::DatastreamFrame {
|
||||
run_id: spec.run_id,
|
||||
node_id: spec.node_id,
|
||||
channel: "mvp.node.bootstrap".to_owned(),
|
||||
payload: serde_json::json!({
|
||||
"type":"DockerGpuConfigResolved",
|
||||
"provider":"Docker",
|
||||
"gpus_arg":docker_gpus.as_deref(),
|
||||
"has_gpus_arg":docker_gpus.is_some(),
|
||||
})
|
||||
.to_string(),
|
||||
});
|
||||
if let Some(gpus) = docker_gpus.as_deref() {
|
||||
command.arg("--gpus").arg(gpus);
|
||||
}
|
||||
for mount in &spec.mounts {
|
||||
command
|
||||
.arg("--mount")
|
||||
.arg(docker_runtime_mount_arg(&spec.image, mount)?);
|
||||
}
|
||||
for (key, value) in &spec.env {
|
||||
command.arg("-e").arg(format!("{key}={value}"));
|
||||
}
|
||||
|
|
@ -278,3 +472,78 @@ fn spawn_stderr_reader(
|
|||
) {
|
||||
BootstrapDatastreamBridge::new(spec, sink, None).spawn_stderr_reader(stderr);
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec_with_mounts(mounts: Vec<ProviderMount>) -> NodeProvisionSpec {
|
||||
NodeProvisionSpec {
|
||||
run_id: 17,
|
||||
node_id: 23,
|
||||
stage_index: Some(2),
|
||||
image: "swactor-mvp-node:test".to_owned(),
|
||||
env: vec![("MVP_RUN_ID".to_owned(), "17".to_owned())],
|
||||
args: vec!["--serve".to_owned()],
|
||||
mounts,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn node_provision_spec_serde_preserves_readonly_mounts() {
|
||||
let spec = spec_with_mounts(vec![ProviderMount {
|
||||
host_path: "/cache/models/model.gguf".to_owned(),
|
||||
container_path: "/models/cached/model.gguf".to_owned(),
|
||||
readonly: true,
|
||||
}]);
|
||||
|
||||
let json = serde_json::to_string(&spec).expect("serialize mounted node spec");
|
||||
let decoded: NodeProvisionSpec =
|
||||
serde_json::from_str(&json).expect("deserialize mounted node spec");
|
||||
|
||||
assert_eq!(decoded, spec);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn node_provision_spec_serde_omits_empty_mounts_and_accepts_missing_mounts() {
|
||||
let spec = spec_with_mounts(Vec::new());
|
||||
|
||||
let json = serde_json::to_value(&spec).expect("serialize unmounted node spec");
|
||||
assert_eq!(json.get("mounts"), None);
|
||||
|
||||
let decoded: NodeProvisionSpec = serde_json::from_value(serde_json::json!({
|
||||
"run_id": 17,
|
||||
"node_id": 23,
|
||||
"stage_index": 2,
|
||||
"image": "swactor-mvp-node:test",
|
||||
"env": [["MVP_RUN_ID", "17"]],
|
||||
"args": ["--serve"]
|
||||
}))
|
||||
.expect("deserialize node spec written before mounts existed");
|
||||
|
||||
assert!(decoded.mounts.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_mount_arg_uses_bind_src_dst_and_readonly_flag() {
|
||||
let mount = ProviderMount {
|
||||
host_path: "/cache/models/model.gguf".to_owned(),
|
||||
container_path: "/models/cached/model.gguf".to_owned(),
|
||||
readonly: true,
|
||||
};
|
||||
|
||||
assert_eq!(
|
||||
docker_mount_arg(&mount),
|
||||
"type=bind,src=/cache/models/model.gguf,dst=/models/cached/model.gguf,readonly"
|
||||
);
|
||||
|
||||
let writable_mount = ProviderMount {
|
||||
readonly: false,
|
||||
..mount
|
||||
};
|
||||
assert_eq!(
|
||||
docker_mount_arg(&writable_mount),
|
||||
"type=bind,src=/cache/models/model.gguf,dst=/models/cached/model.gguf"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
191
crates/mvp-system/src/relay_provisioning.rs
Normal file
191
crates/mvp-system/src/relay_provisioning.rs
Normal file
|
|
@ -0,0 +1,191 @@
|
|||
//! Relay provisioning shims for MVP runtimes.
|
||||
//!
|
||||
//! The current implementations are deliberately small: local tests get the same
|
||||
//! abstraction without an external relay, and deploy runs can point every node at
|
||||
//! one operator-managed relay URL. A future provider can replace the static shim
|
||||
//! with real relay leases without changing node provisioning.
|
||||
|
||||
use iroh::{RelayMode, RelayUrl};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
pub const MVP_IROH_RELAY_MODE_ENV: &str = "MVP_IROH_RELAY_MODE";
|
||||
pub const MVP_IROH_RELAY_URL_ENV: &str = "MVP_IROH_RELAY_URL";
|
||||
pub const SWACTOR_IROH_RELAY_URL_ENV: &str = "SWACTOR_IROH_RELAY_URL";
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum RelayPurpose {
|
||||
Combined,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RelayProvisionRequest {
|
||||
pub run_id: u64,
|
||||
pub purpose: RelayPurpose,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RelayLeaseId(pub String);
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub enum RelayProviderKind {
|
||||
LocalShim,
|
||||
Static,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RelayEndpoint {
|
||||
pub url: String,
|
||||
pub provider: RelayProviderKind,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
|
||||
pub struct RelayLease {
|
||||
pub id: RelayLeaseId,
|
||||
pub endpoints: Vec<RelayEndpoint>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RelayRuntimeConfig {
|
||||
pub mode: RelayMode,
|
||||
pub url: Option<String>,
|
||||
}
|
||||
|
||||
pub trait RelayProvider: Send {
|
||||
fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result<RelayLease, String>;
|
||||
|
||||
fn relay_mode(&self, lease: &RelayLease) -> Result<RelayMode, String>;
|
||||
|
||||
fn release_relay(&mut self, _lease: RelayLease) -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct LocalShimRelayProvider;
|
||||
|
||||
impl RelayProvider for LocalShimRelayProvider {
|
||||
fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result<RelayLease, String> {
|
||||
Ok(RelayLease {
|
||||
id: RelayLeaseId(format!("local-shim:{}", request.run_id)),
|
||||
endpoints: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
fn relay_mode(&self, _lease: &RelayLease) -> Result<RelayMode, String> {
|
||||
Ok(RelayMode::Disabled)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct StaticRelayProvider {
|
||||
url: RelayUrl,
|
||||
}
|
||||
|
||||
impl StaticRelayProvider {
|
||||
pub fn new(url: RelayUrl) -> Self {
|
||||
Self { url }
|
||||
}
|
||||
|
||||
pub fn from_url_str(raw: &str) -> Result<Self, String> {
|
||||
parse_relay_url(raw).map(Self::new)
|
||||
}
|
||||
|
||||
pub fn from_env() -> Result<Option<Self>, String> {
|
||||
selected_relay_url_from_env()
|
||||
.map(|url| Self::from_url_str(&url).map(Some))
|
||||
.unwrap_or(Ok(None))
|
||||
}
|
||||
|
||||
pub fn url(&self) -> String {
|
||||
self.url.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
impl RelayProvider for StaticRelayProvider {
|
||||
fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result<RelayLease, String> {
|
||||
Ok(RelayLease {
|
||||
id: RelayLeaseId(format!("static-relay:{}:{}", request.run_id, self.url)),
|
||||
endpoints: vec![RelayEndpoint {
|
||||
url: self.url.to_string(),
|
||||
provider: RelayProviderKind::Static,
|
||||
}],
|
||||
})
|
||||
}
|
||||
|
||||
fn relay_mode(&self, lease: &RelayLease) -> Result<RelayMode, String> {
|
||||
let urls = lease
|
||||
.endpoints
|
||||
.iter()
|
||||
.map(|endpoint| parse_relay_url(&endpoint.url))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
if urls.is_empty() {
|
||||
return Err("static relay lease has no endpoints".to_owned());
|
||||
}
|
||||
Ok(RelayMode::custom(urls))
|
||||
}
|
||||
}
|
||||
|
||||
pub fn relay_runtime_config_from_env(run_id: u64) -> Result<RelayRuntimeConfig, String> {
|
||||
match relay_mode_setting_from_env().as_deref() {
|
||||
Some("disabled") => Ok(RelayRuntimeConfig {
|
||||
mode: RelayMode::Disabled,
|
||||
url: None,
|
||||
}),
|
||||
None | Some("default") => relay_runtime_config_from_optional_static_provider(run_id),
|
||||
Some(other) => Err(format!(
|
||||
"unsupported {MVP_IROH_RELAY_MODE_ENV}={other:?}; use disabled or default"
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn relay_mode_env_value(mode: &RelayMode) -> &'static str {
|
||||
match mode {
|
||||
RelayMode::Disabled => "disabled",
|
||||
_ => "default",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn selected_relay_url_from_env() -> Option<String> {
|
||||
env_optional(MVP_IROH_RELAY_URL_ENV).or_else(|| env_optional(SWACTOR_IROH_RELAY_URL_ENV))
|
||||
}
|
||||
|
||||
fn relay_runtime_config_from_optional_static_provider(
|
||||
run_id: u64,
|
||||
) -> Result<RelayRuntimeConfig, String> {
|
||||
let Some(mut provider) = StaticRelayProvider::from_env()? else {
|
||||
return Ok(RelayRuntimeConfig {
|
||||
mode: RelayMode::Default,
|
||||
url: None,
|
||||
});
|
||||
};
|
||||
let lease = provider.provision_relay(RelayProvisionRequest {
|
||||
run_id,
|
||||
purpose: RelayPurpose::Combined,
|
||||
})?;
|
||||
let mode = provider.relay_mode(&lease)?;
|
||||
Ok(RelayRuntimeConfig {
|
||||
mode,
|
||||
url: Some(provider.url()),
|
||||
})
|
||||
}
|
||||
|
||||
fn relay_mode_setting_from_env() -> Option<String> {
|
||||
env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase())
|
||||
}
|
||||
|
||||
fn env_optional(name: &str) -> Option<String> {
|
||||
std::env::var(name)
|
||||
.ok()
|
||||
.map(|value| value.trim().to_owned())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn parse_relay_url(raw: &str) -> Result<RelayUrl, String> {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() {
|
||||
return Err("relay URL cannot be empty".to_owned());
|
||||
}
|
||||
trimmed
|
||||
.parse::<RelayUrl>()
|
||||
.map_err(|e| format!("invalid relay URL {trimmed:?}: {e}"))
|
||||
}
|
||||
|
|
@ -41,6 +41,7 @@ fn spec() -> NodeProvisionSpec {
|
|||
image: "worker:latest".to_owned(),
|
||||
env: Vec::new(),
|
||||
args: Vec::new(),
|
||||
mounts: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ mod node_provisioning_guarantees;
|
|||
mod observability_surface_guarantees;
|
||||
mod orchestrator_run_fsm_guarantees;
|
||||
mod orchestrator_token_endpoint_guarantees;
|
||||
mod relay_provisioning_guarantees;
|
||||
mod resource_inventory_guarantees;
|
||||
mod run_plan_guarantees;
|
||||
mod shared_ring_helper_abi_guarantees;
|
||||
|
|
|
|||
171
crates/mvp-system/src/tests/relay_provisioning_guarantees.rs
Normal file
171
crates/mvp-system/src/tests/relay_provisioning_guarantees.rs
Normal file
|
|
@ -0,0 +1,171 @@
|
|||
use std::ffi::OsString;
|
||||
use std::sync::Mutex;
|
||||
|
||||
use iroh::{RelayMode, RelayUrl};
|
||||
use mvp_system::relay_provisioning::{
|
||||
LocalShimRelayProvider, MVP_IROH_RELAY_MODE_ENV, MVP_IROH_RELAY_URL_ENV, RelayProvider,
|
||||
RelayProviderKind, RelayProvisionRequest, RelayPurpose, SWACTOR_IROH_RELAY_URL_ENV,
|
||||
StaticRelayProvider, relay_runtime_config_from_env,
|
||||
};
|
||||
|
||||
static ENV_LOCK: Mutex<()> = Mutex::new(());
|
||||
|
||||
const RELAY_ENV_KEYS: &[&str] = &[
|
||||
MVP_IROH_RELAY_MODE_ENV,
|
||||
MVP_IROH_RELAY_URL_ENV,
|
||||
SWACTOR_IROH_RELAY_URL_ENV,
|
||||
];
|
||||
|
||||
struct RestoreEnv {
|
||||
saved: Vec<(&'static str, Option<OsString>)>,
|
||||
}
|
||||
|
||||
impl Drop for RestoreEnv {
|
||||
fn drop(&mut self) {
|
||||
for (key, value) in &self.saved {
|
||||
match value {
|
||||
Some(value) => unsafe { std::env::set_var(key, value) },
|
||||
None => unsafe { std::env::remove_var(key) },
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn with_relay_env<T>(settings: &[(&'static str, &'static str)], test: impl FnOnce() -> T) -> T {
|
||||
let _lock = ENV_LOCK
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
let saved = RELAY_ENV_KEYS
|
||||
.iter()
|
||||
.map(|&key| (key, std::env::var_os(key)))
|
||||
.collect::<Vec<_>>();
|
||||
for key in RELAY_ENV_KEYS {
|
||||
unsafe { std::env::remove_var(key) };
|
||||
}
|
||||
for (key, value) in settings {
|
||||
assert!(
|
||||
RELAY_ENV_KEYS.contains(key),
|
||||
"test env key {key} must be restored"
|
||||
);
|
||||
unsafe { std::env::set_var(key, value) };
|
||||
}
|
||||
let _restore = RestoreEnv { saved };
|
||||
test()
|
||||
}
|
||||
|
||||
fn provision_request(run_id: u64) -> RelayProvisionRequest {
|
||||
RelayProvisionRequest {
|
||||
run_id,
|
||||
purpose: RelayPurpose::Combined,
|
||||
}
|
||||
}
|
||||
|
||||
fn canonical_relay_url(raw: &str) -> String {
|
||||
raw.parse::<RelayUrl>()
|
||||
.expect("fixture relay URL parses")
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn assert_custom_relay_mode(mode: RelayMode, expected_url: &str) {
|
||||
let expected_url = expected_url
|
||||
.parse::<RelayUrl>()
|
||||
.expect("fixture relay URL parses");
|
||||
match mode {
|
||||
RelayMode::Custom(relay_map) => {
|
||||
assert_eq!(relay_map.len(), 1, "custom relay map must contain one URL");
|
||||
assert!(
|
||||
relay_map.contains(&expected_url),
|
||||
"custom relay map must contain {expected_url}"
|
||||
);
|
||||
}
|
||||
other => panic!("expected custom relay mode for {expected_url}, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_shim_provisions_disabled_lease_without_endpoints() {
|
||||
let mut provider = LocalShimRelayProvider;
|
||||
|
||||
let lease = provider
|
||||
.provision_relay(provision_request(77))
|
||||
.expect("local shim provisioning succeeds");
|
||||
|
||||
assert_eq!(lease.endpoints, Vec::new());
|
||||
assert!(matches!(
|
||||
provider
|
||||
.relay_mode(&lease)
|
||||
.expect("local shim relay mode resolves"),
|
||||
RelayMode::Disabled
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn static_provider_provisions_one_endpoint_and_custom_relay_mode() {
|
||||
const RELAY_URL: &str = "https://relay-static.example.com";
|
||||
let mut provider =
|
||||
StaticRelayProvider::from_url_str(RELAY_URL).expect("static relay URL parses");
|
||||
let expected_url = canonical_relay_url(RELAY_URL);
|
||||
|
||||
let lease = provider
|
||||
.provision_relay(provision_request(88))
|
||||
.expect("static relay provisioning succeeds");
|
||||
|
||||
assert_eq!(lease.endpoints.len(), 1);
|
||||
assert_eq!(lease.endpoints[0].url, expected_url);
|
||||
assert_eq!(lease.endpoints[0].provider, RelayProviderKind::Static);
|
||||
assert_custom_relay_mode(
|
||||
provider
|
||||
.relay_mode(&lease)
|
||||
.expect("static relay mode resolves"),
|
||||
&expected_url,
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn default_relay_mode_uses_configured_mvp_or_swactor_relay_url() {
|
||||
const MVP_URL: &str = "https://relay-mvp.example.com";
|
||||
const SWACTOR_URL: &str = "https://relay-swactor.example.com";
|
||||
|
||||
for (name, settings, expected_url) in [
|
||||
(
|
||||
"mvp relay URL",
|
||||
[
|
||||
(MVP_IROH_RELAY_MODE_ENV, "default"),
|
||||
(MVP_IROH_RELAY_URL_ENV, MVP_URL),
|
||||
],
|
||||
MVP_URL,
|
||||
),
|
||||
(
|
||||
"swactor relay URL fallback",
|
||||
[
|
||||
(MVP_IROH_RELAY_MODE_ENV, "default"),
|
||||
(SWACTOR_IROH_RELAY_URL_ENV, SWACTOR_URL),
|
||||
],
|
||||
SWACTOR_URL,
|
||||
),
|
||||
] {
|
||||
with_relay_env(&settings, || {
|
||||
let config = relay_runtime_config_from_env(901).unwrap_or_else(|error| {
|
||||
panic!("{name} should resolve custom relay config: {error}")
|
||||
});
|
||||
let expected_url = canonical_relay_url(expected_url);
|
||||
assert_eq!(config.url.as_deref(), Some(expected_url.as_str()));
|
||||
assert_custom_relay_mode(config.mode, &expected_url);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn disabled_relay_mode_ignores_configured_url() {
|
||||
with_relay_env(
|
||||
&[
|
||||
(MVP_IROH_RELAY_MODE_ENV, "disabled"),
|
||||
(MVP_IROH_RELAY_URL_ENV, "https://ignored-relay.example.com"),
|
||||
],
|
||||
|| {
|
||||
let config = relay_runtime_config_from_env(902).expect("disabled relay mode resolves");
|
||||
assert!(matches!(config.mode, RelayMode::Disabled));
|
||||
assert_eq!(config.url, None);
|
||||
},
|
||||
);
|
||||
}
|
||||
|
|
@ -42,8 +42,16 @@ fn provisioning_records_round_trip_on_owned_datastream_channels() {
|
|||
run_id: 77,
|
||||
node_id: 11,
|
||||
kind: ProvisionEventKind::NodeLive,
|
||||
provider: Some("docker".to_owned()),
|
||||
message: None,
|
||||
});
|
||||
let event_without_provider = MvpProvisionEventRecord::new(ProvisionEvent {
|
||||
run_id: 77,
|
||||
node_id: 12,
|
||||
kind: ProvisionEventKind::ProvisionStart,
|
||||
provider: None,
|
||||
message: Some("queued".to_owned()),
|
||||
});
|
||||
let log = MvpProvisionLogRecord::new(ProvisionLogLine {
|
||||
run_id: 77,
|
||||
node_id: 11,
|
||||
|
|
@ -64,6 +72,10 @@ fn provisioning_records_round_trip_on_owned_datastream_channels() {
|
|||
event
|
||||
);
|
||||
assert_eq!(MvpProvisionLogRecord::decode(&log.encode()).unwrap(), log);
|
||||
assert_eq!(
|
||||
MvpProvisionEventRecord::decode(&event_without_provider.encode()).unwrap(),
|
||||
event_without_provider
|
||||
);
|
||||
assert_eq!(
|
||||
telemetry::mvp_provision_log_channel(11, ProvisionLogStream::Stdout).as_str(),
|
||||
"mvp.provisioning.logs.node.11.stdout"
|
||||
|
|
|
|||
|
|
@ -124,6 +124,7 @@ fn spec() -> NodeProvisionSpec {
|
|||
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()],
|
||||
mounts: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,7 +1,11 @@
|
|||
use parking_lot::Mutex;
|
||||
use std::collections::BTreeMap;
|
||||
use std::process::{Child, Command, Stdio};
|
||||
use std::sync::Arc;
|
||||
use std::sync::{
|
||||
Arc,
|
||||
atomic::{AtomicBool, Ordering},
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
use datastream::DatastreamProducer;
|
||||
use swactor_vastai::{LifecyclePolicy, ProvisionRequest, ProvisionedInstance, SelectionPolicy};
|
||||
|
|
@ -374,6 +378,7 @@ pub struct SshCommandBootstrapLauncher;
|
|||
|
||||
pub struct SshCommandBootstrapHandle {
|
||||
child: Arc<Mutex<Child>>,
|
||||
stopping: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
|
||||
|
|
@ -416,22 +421,72 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher {
|
|||
.take()
|
||||
.ok_or_else(|| format!("VastAI node {} SSH stderr missing", spec.node_id))?;
|
||||
|
||||
let run_id = spec.run_id;
|
||||
let node_id = spec.node_id;
|
||||
let exit_sink = sink.clone();
|
||||
let child = Arc::new(Mutex::new(child));
|
||||
let stopping = Arc::new(AtomicBool::new(false));
|
||||
let bridge = BootstrapDatastreamBridge::new(spec, sink, producer);
|
||||
bridge.spawn_stdout_reader(stdout);
|
||||
bridge.spawn_stderr_reader(stderr);
|
||||
spawn_ssh_exit_watcher(run_id, node_id, child.clone(), stopping.clone(), exit_sink);
|
||||
|
||||
Ok(SshCommandBootstrapHandle {
|
||||
child: Arc::new(Mutex::new(child)),
|
||||
})
|
||||
Ok(SshCommandBootstrapHandle { child, stopping })
|
||||
}
|
||||
|
||||
fn stop_bootstrap(&mut self, handle: &mut Self::Handle) {
|
||||
handle.stopping.store(true, Ordering::SeqCst);
|
||||
let mut child = handle.child.lock();
|
||||
let _ = child.kill();
|
||||
let _ = child.wait();
|
||||
}
|
||||
}
|
||||
|
||||
fn spawn_ssh_exit_watcher(
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
child: Arc<Mutex<Child>>,
|
||||
stopping: Arc<AtomicBool>,
|
||||
sink: PluginSink,
|
||||
) {
|
||||
std::thread::spawn(move || {
|
||||
loop {
|
||||
match child.lock().try_wait() {
|
||||
Ok(Some(status)) => {
|
||||
if stopping.load(Ordering::SeqCst) {
|
||||
return;
|
||||
}
|
||||
if status.success() {
|
||||
sink.observe(PluginObservation::Exited {
|
||||
run_id,
|
||||
node_id,
|
||||
status: status.code(),
|
||||
});
|
||||
} else {
|
||||
sink.observe(PluginObservation::Failed {
|
||||
run_id,
|
||||
node_id,
|
||||
reason: format!("VastAI SSH bootstrap exited: {status}"),
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
Ok(None) => std::thread::sleep(Duration::from_millis(100)),
|
||||
Err(error) => {
|
||||
if !stopping.load(Ordering::SeqCst) {
|
||||
sink.observe(PluginObservation::Failed {
|
||||
run_id,
|
||||
node_id,
|
||||
reason: format!("wait VastAI SSH bootstrap: {error}"),
|
||||
});
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub struct VastAiProvisioningPlugin<C, B>
|
||||
where
|
||||
C: VastAiLeaseClient,
|
||||
|
|
@ -540,6 +595,9 @@ where
|
|||
spec: NodeProvisionSpec,
|
||||
sink: PluginSink,
|
||||
) -> Result<PluginNodeHandle, String> {
|
||||
if !spec.mounts.is_empty() {
|
||||
return Err("vastai provider does not support host file mounts".to_owned());
|
||||
}
|
||||
let stream_id = node_stream_id(spec.run_id, spec.node_id);
|
||||
let label = self.label_for(&spec);
|
||||
sink.observe(PluginObservation::ProviderLine {
|
||||
|
|
|
|||
|
|
@ -25,6 +25,8 @@ fn one_node_chat_docker_cuda_e2e() {
|
|||
command
|
||||
.current_dir(&root)
|
||||
.args(["mvp-chat"])
|
||||
.env("MVP_RUNTIME_CONFIG", "local")
|
||||
.env("MVP_IROH_RELAY_MODE", "disabled")
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
|
|
@ -46,7 +48,7 @@ fn one_node_chat_docker_cuda_e2e() {
|
|||
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);
|
||||
let mut result = run_full_flow(&mut child, &mut stdin, &stdout);
|
||||
if result.is_err() {
|
||||
request_child_interrupt(&child);
|
||||
let _ = wait_child(&mut child, SHUTDOWN_TIMEOUT);
|
||||
|
|
@ -55,6 +57,9 @@ fn one_node_chat_docker_cuda_e2e() {
|
|||
}
|
||||
let _ = stdout_reader.join();
|
||||
let _ = stderr_reader.join();
|
||||
if result.is_ok() {
|
||||
result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr);
|
||||
}
|
||||
assert_container_removed(&root, DEFAULT_CONTAINER);
|
||||
|
||||
if let Err(error) = result {
|
||||
|
|
@ -71,12 +76,7 @@ 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, || {
|
||||
|
|
@ -91,27 +91,26 @@ fn run_full_flow(
|
|||
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}"))?;
|
||||
wait_for_child_or(TEST_TIMEOUT, child, || prompt_visible(stdout))
|
||||
.map_err(|e| format!("chat prompt not visible: {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, || {
|
||||
stdout_contains(stdout, "decoding...")
|
||||
})
|
||||
.map_err(|e| format!("prompt was not submitted to chat loop: {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}"))?;
|
||||
|
||||
wait_for_child_or(PROMPT_TIMEOUT, child, dashboard_has_orch_prompt_lifecycle)
|
||||
.map_err(|e| format!("orchestrator prompt lifecycle not visible in dashboard: {e}"))?;
|
||||
wait_for_child_or(PROMPT_TIMEOUT, child, || prompt_count(stdout) >= 2)
|
||||
.map_err(|e| format!("chat prompt did not return after response: {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())?;
|
||||
|
|
@ -183,6 +182,79 @@ fn dashboard_has_frame(channel_substr: &str, payload_substr: &str) -> bool {
|
|||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn dashboard_has_orch_prompt_lifecycle() -> bool {
|
||||
dashboard_frames()
|
||||
.map(|frames| {
|
||||
let events = frames
|
||||
.iter()
|
||||
.filter_map(orch_prompt_observation)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
events.iter().any(|prompt_work| {
|
||||
prompt_work.phase == "prompt_work"
|
||||
&& prompt_work.status == "observed"
|
||||
&& events.iter().any(|event| {
|
||||
same_prompt(prompt_work, event)
|
||||
&& event.phase == "node_prompt_send"
|
||||
&& event.status == "ready"
|
||||
})
|
||||
&& events.iter().any(|event| {
|
||||
same_prompt(prompt_work, event)
|
||||
&& event.phase == "node_prompt_event"
|
||||
&& event.status == "observed"
|
||||
&& event.detail_event.as_deref() == Some("Done")
|
||||
})
|
||||
&& events.iter().any(|event| {
|
||||
same_prompt(prompt_work, event)
|
||||
&& event.phase == "prompt_complete"
|
||||
&& event.status == "ready"
|
||||
&& event.detail_event.as_deref() == Some("Done")
|
||||
})
|
||||
})
|
||||
})
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn orch_prompt_observation(frame: &SeenFrame) -> Option<OrchPromptObservation> {
|
||||
if frame.channel != "mvp.orch.prompt" {
|
||||
return None;
|
||||
}
|
||||
|
||||
let value = serde_json::from_str::<Value>(&frame.payload).ok()?;
|
||||
if value.get("type").and_then(Value::as_str)? != "OrchPromptEvent" {
|
||||
return None;
|
||||
}
|
||||
let detail = value.get("detail")?;
|
||||
|
||||
Some(OrchPromptObservation {
|
||||
phase: value.get("phase").and_then(Value::as_str)?.to_owned(),
|
||||
status: value.get("status").and_then(Value::as_str)?.to_owned(),
|
||||
run_id: value.get("run_id").and_then(Value::as_u64)?,
|
||||
node_id: value.get("node_id").and_then(Value::as_u64)?,
|
||||
request_id: value.get("request_id").and_then(Value::as_u64)?,
|
||||
detail_event: detail
|
||||
.get("event")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned),
|
||||
})
|
||||
}
|
||||
|
||||
fn same_prompt(left: &OrchPromptObservation, right: &OrchPromptObservation) -> bool {
|
||||
left.run_id == right.run_id
|
||||
&& left.node_id == right.node_id
|
||||
&& left.request_id == right.request_id
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct OrchPromptObservation {
|
||||
phase: String,
|
||||
status: String,
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
request_id: u64,
|
||||
detail_event: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct SeenFrame {
|
||||
channel: String,
|
||||
|
|
@ -255,16 +327,53 @@ fn stdout_contains(stdout: &Arc<Mutex<String>>, needle: &str) -> bool {
|
|||
snapshot(stdout).contains(needle)
|
||||
}
|
||||
|
||||
fn prompt_visible(stdout: &Arc<Mutex<String>>) -> bool {
|
||||
snapshot(stdout).contains("prompt:> ")
|
||||
}
|
||||
|
||||
fn prompt_count(stdout: &Arc<Mutex<String>>) -> usize {
|
||||
snapshot(stdout).matches("prompt:> ").count()
|
||||
}
|
||||
|
||||
fn response_text_visible(stdout: &Arc<Mutex<String>>) -> bool {
|
||||
let text = snapshot(stdout);
|
||||
text.lines()
|
||||
.any(|line| line.trim_start_matches('>').trim().len() > 8)
|
||||
snapshot(stdout)
|
||||
.split("Response: ")
|
||||
.skip(1)
|
||||
.any(|text| !text.lines().next().unwrap_or_default().trim().is_empty())
|
||||
}
|
||||
|
||||
fn stderr_contains(stderr: &Arc<Mutex<String>>, needle: &str) -> bool {
|
||||
snapshot(stderr).contains(needle)
|
||||
fn assert_no_lower_layer_terminal_leaks(
|
||||
stdout: &Arc<Mutex<String>>,
|
||||
stderr: &Arc<Mutex<String>>,
|
||||
) -> Result<(), String> {
|
||||
let leaks = lower_layer_leak_lines("stdout", &snapshot(stdout))
|
||||
.into_iter()
|
||||
.chain(lower_layer_leak_lines("stderr", &snapshot(stderr)))
|
||||
.collect::<Vec<_>>();
|
||||
if leaks.is_empty() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(format!(
|
||||
"lower-layer runtime output leaked to terminal:\n{}",
|
||||
leaks.join("\n")
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
fn lower_layer_leak_lines(stream: &str, output: &str) -> Vec<String> {
|
||||
output
|
||||
.lines()
|
||||
.filter_map(|line| {
|
||||
let trimmed = line.trim_start();
|
||||
let leaked = trimmed.contains("prompt_loop_ready")
|
||||
|| trimmed.contains("dashboard_ready")
|
||||
|| trimmed.starts_with("mvp-orch-one-node:")
|
||||
|| trimmed.starts_with("mvp-node:")
|
||||
|| trimmed.starts_with("mvp_tinygrad_worker:");
|
||||
leaked.then(|| format!("{stream}: {line}"))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
fn snapshot(buf: &Arc<Mutex<String>>) -> String {
|
||||
buf.lock().expect("capture mutex").clone()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,36 @@
|
|||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, ExitCode};
|
||||
use std::time::Instant;
|
||||
use std::time::{Instant, SystemTime};
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
use std::os::unix::process::CommandExt;
|
||||
|
||||
struct TestStep {
|
||||
label: &'static str,
|
||||
args: &'static [&'static str],
|
||||
}
|
||||
|
||||
const MVP_CHAT_REBUILD_INPUTS: &[&str] = &[
|
||||
"Cargo.lock",
|
||||
"Cargo.toml",
|
||||
"src",
|
||||
"crates/datastream/Cargo.toml",
|
||||
"crates/datastream/src",
|
||||
"crates/dashboard/Cargo.toml",
|
||||
"crates/dashboard/src",
|
||||
"crates/distribution/Cargo.toml",
|
||||
"crates/distribution/src",
|
||||
"crates/iroh-driver/Cargo.toml",
|
||||
"crates/iroh-driver/src",
|
||||
"crates/mvp-system/Cargo.toml",
|
||||
"crates/mvp-system/src",
|
||||
"crates/transport/Cargo.toml",
|
||||
"crates/transport/src",
|
||||
"tools/vastai/Cargo.toml",
|
||||
"tools/vastai/src",
|
||||
];
|
||||
|
||||
const BASIC_TESTS: &[TestStep] = &[
|
||||
TestStep {
|
||||
label: "root crate",
|
||||
|
|
@ -58,10 +83,13 @@ fn cargo_bin() -> String {
|
|||
fn print_usage() {
|
||||
println!(
|
||||
"\
|
||||
USAGE: cargo xtask test
|
||||
USAGE: cargo xtask <command>
|
||||
|
||||
COMMANDS:
|
||||
test Run all basic non-binding tests. This includes the root crate with\n `cargo test` plus each non-binding workspace package with `cargo test -p`.\n Feature-gated E2E/bin tests are intentionally excluded."
|
||||
mvp-chat [args...] Run mvp-one-node-chat, building it only when tracked inputs are newer.
|
||||
test Run all basic non-binding tests. This includes the root crate with
|
||||
`cargo test` plus each non-binding repository package with `cargo test -p`.
|
||||
Feature-gated E2E/bin tests are intentionally excluded."
|
||||
);
|
||||
}
|
||||
|
||||
|
|
@ -100,11 +128,135 @@ fn run_tests() -> ExitCode {
|
|||
ExitCode::SUCCESS
|
||||
}
|
||||
|
||||
fn run_mvp_chat(args: Vec<String>) -> ExitCode {
|
||||
match ensure_mvp_chat_binary() {
|
||||
Ok(bin) => exec_mvp_chat(&bin, args),
|
||||
Err(error) => {
|
||||
eprintln!("cargo mvp-chat: {error}");
|
||||
ExitCode::from(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_mvp_chat_binary() -> Result<PathBuf, String> {
|
||||
let root = workspace_root();
|
||||
let bin = root
|
||||
.join("target")
|
||||
.join("debug")
|
||||
.join(format!("mvp-one-node-chat{}", std::env::consts::EXE_SUFFIX));
|
||||
if rebuild_needed(&bin, &root, MVP_CHAT_REBUILD_INPUTS)? {
|
||||
eprintln!("cargo mvp-chat: building mvp-one-node-chat");
|
||||
let status = Command::new(cargo_bin())
|
||||
.current_dir(&root)
|
||||
.args([
|
||||
"build",
|
||||
"--quiet",
|
||||
"-p",
|
||||
"mvp-system",
|
||||
"--features",
|
||||
"local-e2e",
|
||||
"--bin",
|
||||
"mvp-one-node-chat",
|
||||
])
|
||||
.status()
|
||||
.map_err(|e| format!("run cargo build for mvp-one-node-chat: {e}"))?;
|
||||
if !status.success() {
|
||||
return Err(format!("build mvp-one-node-chat failed with {status}"));
|
||||
}
|
||||
} else {
|
||||
eprintln!("cargo mvp-chat: mvp-one-node-chat is up to date; skipping cargo build");
|
||||
}
|
||||
Ok(bin)
|
||||
}
|
||||
|
||||
fn exec_mvp_chat(bin: &Path, args: Vec<String>) -> ExitCode {
|
||||
let mut command = Command::new(bin);
|
||||
command.args(args);
|
||||
#[cfg(target_os = "linux")]
|
||||
{
|
||||
let error = command.exec();
|
||||
eprintln!("exec {}: {error}", bin.display());
|
||||
ExitCode::from(1)
|
||||
}
|
||||
#[cfg(not(target_os = "linux"))]
|
||||
{
|
||||
match command.status() {
|
||||
Ok(status) if status.success() => ExitCode::SUCCESS,
|
||||
Ok(status) => ExitCode::from(status.code().unwrap_or(1) as u8),
|
||||
Err(error) => {
|
||||
eprintln!("run {}: {error}", bin.display());
|
||||
ExitCode::from(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn rebuild_needed(bin: &Path, root: &Path, inputs: &[&str]) -> Result<bool, String> {
|
||||
if !bin.is_file() {
|
||||
return Ok(true);
|
||||
}
|
||||
let bin_mtime = modified_time(root, bin)?;
|
||||
for input in inputs {
|
||||
let path = root.join(input);
|
||||
if latest_mtime(root, &path)? > bin_mtime {
|
||||
return Ok(true);
|
||||
}
|
||||
}
|
||||
Ok(false)
|
||||
}
|
||||
|
||||
fn latest_mtime(root: &Path, path: &Path) -> Result<SystemTime, String> {
|
||||
let display = display_workspace_path(root, path);
|
||||
let metadata = fs::metadata(path).map_err(|e| format!("stat {display}: {e}"))?;
|
||||
let mut latest = metadata
|
||||
.modified()
|
||||
.map_err(|e| format!("modified time {display}: {e}"))?;
|
||||
if metadata.is_dir() {
|
||||
for entry in fs::read_dir(path).map_err(|e| format!("read dir {display}: {e}"))? {
|
||||
let entry = entry.map_err(|e| format!("read dir entry {display}: {e}"))?;
|
||||
let entry_mtime = latest_mtime(root, &entry.path())?;
|
||||
if entry_mtime > latest {
|
||||
latest = entry_mtime;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(latest)
|
||||
}
|
||||
|
||||
fn modified_time(root: &Path, path: &Path) -> Result<SystemTime, String> {
|
||||
let display = display_workspace_path(root, path);
|
||||
fs::metadata(path)
|
||||
.map_err(|e| format!("stat {display}: {e}"))?
|
||||
.modified()
|
||||
.map_err(|e| format!("modified time {display}: {e}"))
|
||||
}
|
||||
|
||||
fn display_workspace_path(root: &Path, path: &Path) -> String {
|
||||
match path.strip_prefix(root) {
|
||||
Ok(relative) if relative.as_os_str().is_empty() => ".".to_owned(),
|
||||
Ok(relative) => format!("./{}", relative.display()),
|
||||
Err(_) => path.display().to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn workspace_root() -> PathBuf {
|
||||
let git_root = Command::new("git")
|
||||
.args(["rev-parse", "--show-toplevel"])
|
||||
.output();
|
||||
if let Ok(output) = git_root {
|
||||
if output.status.success() {
|
||||
return PathBuf::from(String::from_utf8_lossy(&output.stdout).trim());
|
||||
}
|
||||
}
|
||||
std::env::current_dir().expect("current directory is available")
|
||||
}
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let mut args = std::env::args().skip(1);
|
||||
match (args.next().as_deref(), args.next()) {
|
||||
(Some("test"), None) => run_tests(),
|
||||
(Some("help" | "--help" | "-h"), None) | (None, None) => {
|
||||
match args.next().as_deref() {
|
||||
Some("test") if args.next().is_none() => run_tests(),
|
||||
Some("mvp-chat") => run_mvp_chat(args.collect()),
|
||||
Some("help" | "--help" | "-h") | None => {
|
||||
print_usage();
|
||||
ExitCode::SUCCESS
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue