diff --git a/.cargo/config.toml b/.cargo/config.toml index 65e7c3b..22edee3 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -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" diff --git a/.gitignore b/.gitignore index 9016cb4..6364d67 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,8 @@ __pycache__ fuzz/artifacts/** corpus .loop/ +.model-cache/ + # Analysis artifacts (depgraph + spectral) **/deps.dot diff --git a/Cargo.lock b/Cargo.lock index 9e9f761..81ed545 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2471,6 +2471,7 @@ dependencies = [ name = "mvp-system" version = "0.1.0" dependencies = [ + "blake3", "dashboard", "datastream", "distribution", diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py index e685cbc..038aff3 100755 --- a/apps/mvp-node/tinygrad_worker.py +++ b/apps/mvp-node/tinygrad_worker.py @@ -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, diff --git a/crates/datastream/src/emit.rs b/crates/datastream/src/emit.rs index 5dd2e5a..e982e53 100644 --- a/crates/datastream/src/emit.rs +++ b/crates/datastream/src/emit.rs @@ -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(&self, record: &R) -> Position { self.mux.submit(R::channel(), record.encode()) diff --git a/crates/datastream/src/endpoint.rs b/crates/datastream/src/endpoint.rs index 9e95fba..af38747 100644 --- a/crates/datastream/src/endpoint.rs +++ b/crates/datastream/src/endpoint.rs @@ -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(&self, channel_for: F) -> Arc where F: Fn(&str, bool) -> ChannelId + Send + Sync + 'static, diff --git a/crates/datastream/src/frame.rs b/crates/datastream/src/frame.rs index 1e3af22..ca014a7 100644 --- a/crates/datastream/src/frame.rs +++ b/crates/datastream/src/frame.rs @@ -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. diff --git a/crates/datastream/src/lib.rs b/crates/datastream/src/lib.rs index f98bde1..485685a 100644 --- a/crates/datastream/src/lib.rs +++ b/crates/datastream/src/lib.rs @@ -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}; diff --git a/crates/datastream/src/mux.rs b/crates/datastream/src/mux.rs index 9eec08c..2b28847 100644 --- a/crates/datastream/src/mux.rs +++ b/crates/datastream/src/mux.rs @@ -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>, } @@ -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, 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) +} diff --git a/crates/datastream/src/timing.rs b/crates/datastream/src/timing.rs new file mode 100644 index 0000000..c90f952 --- /dev/null +++ b/crates/datastream/src/timing.rs @@ -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; +} diff --git a/crates/datastream/tests/t_datastream.rs b/crates/datastream/tests/t_datastream.rs index 70f7d7b..18b44a9 100644 --- a/crates/datastream/tests/t_datastream.rs +++ b/crates/datastream/tests/t_datastream.rs @@ -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(&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 { + self.mux.drain() + } +} + fn test_registry() -> ChannelRegistry { ChannelRegistry::new() .with_record::() @@ -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![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 = 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 { - 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(); diff --git a/crates/datastream/tests/t_datastream_endpoint.rs b/crates/datastream/tests/t_datastream_endpoint.rs index 5e3516c..ed71d94 100644 --- a/crates/datastream/tests/t_datastream_endpoint.rs +++ b/crates/datastream/tests/t_datastream_endpoint.rs @@ -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}.{}", diff --git a/crates/mvp-system/Cargo.toml b/crates/mvp-system/Cargo.toml index 727900e..057c609 100644 --- a/crates/mvp-system/Cargo.toml +++ b/crates/mvp-system/Cargo.toml @@ -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" diff --git a/crates/mvp-system/src/actors/provisioner.rs b/crates/mvp-system/src/actors/provisioner.rs index dc9e8eb..1ea94ea 100644 --- a/crates/mvp-system/src/actors/provisioner.rs +++ b/crates/mvp-system/src/actors/provisioner.rs @@ -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 ProvisionerActor

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

{ run_id, node_id, kind: ProvisionEventKind::NodeStopped, + provider: None, message, }); } @@ -196,6 +198,9 @@ impl ProvisionerActor

{ line, }, ), + PluginObservation::DatastreamFrame { + channel, payload, .. + } => self.emit_datastream_frame(channel, payload), PluginObservation::RuntimeReady { run_id, node_id, @@ -209,6 +214,7 @@ impl ProvisionerActor

{ run_id, node_id, kind: ProvisionEventKind::NodeLive, + provider: None, message: None, }); let _ = ctx.send( @@ -314,6 +320,7 @@ impl ProvisionerActor

{ run_id, node_id, kind: ProvisionEventKind::ProvisionFailed, + provider: None, message: Some(reason.to_owned()), }); } @@ -352,6 +359,12 @@ impl ProvisionerActor

{ 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 ActorInterface for ProvisionerActor

{ diff --git a/crates/mvp-system/src/bin/local_e2e_cluster.rs b/crates/mvp-system/src/bin/local_e2e_cluster.rs index fb3dbcb..3b296f5 100644 --- a/crates/mvp-system/src/bin/local_e2e_cluster.rs +++ b/crates/mvp-system/src/bin/local_e2e_cluster.rs @@ -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 { + 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(), } } diff --git a/crates/mvp-system/src/bin/mvp_node.rs b/crates/mvp-system/src/bin/mvp_node.rs index 58910a2..1843a55 100644 --- a/crates/mvp-system/src/bin/mvp_node.rs +++ b/crates/mvp-system/src/bin/mvp_node.rs @@ -8,8 +8,8 @@ use std::sync::{ use std::thread; use std::time::{Duration, Instant}; +use datastream::ChannelId; use datastream::emit::{ClusterFrameSink, DatastreamEmitter, EmitterConfig, FrameSink, NoopSink}; -use datastream::{ChannelId, DATASTREAM_SINK_NAME}; use distribution::node::DistributedNodeConfig; use iroh::EndpointAddr; @@ -20,6 +20,7 @@ use mvp_system::actors::node_agent::{ use mvp_system::actors::register_mvp_actor_codecs; use mvp_system::distribution_stack::DistributionRuntimeStack; use mvp_system::prompt_rpc::PromptEvent; +use mvp_system::relay_provisioning::relay_runtime_config_from_env; use mvp_system::run_plan::{GgufSource, TokenizerSource}; use mvp_system::stage_controller as stage; use serde_json::{Value, json}; @@ -31,6 +32,65 @@ const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF"; const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; const PUMP_INTERVAL: Duration = Duration::from_millis(10); +const NODE_BOOTSTRAP_CHANNEL: &str = "mvp.node.bootstrap"; +const NODE_RUNTIME_CHANNEL: &str = "mvp.node.runtime"; +const NODE_STAGE_CHANNEL: &str = "mvp.node.stage"; +const NODE_WORKER_CHANNEL: &str = "mvp.node.worker"; +const NODE_PROMPT_CHANNEL: &str = "mvp.node.prompt"; +const NODE_SHUTDOWN_CHANNEL: &str = "mvp.node.shutdown"; + +fn node_event_payload( + config: &DeploymentConfig, + phase: &str, + status: &str, + detail: Value, +) -> Value { + json!({ + "type":"NodeEvent", + "phase":phase, + "status":status, + "run_id":config.run_id, + "node_id":config.logical_node_id, + "stage_index":config.stage_index, + "detail":detail, + }) +} + +fn emit_stdio_node_event( + config: &DeploymentConfig, + channel: &str, + phase: &str, + status: &str, + detail: Value, +) -> Result<(), String> { + println!( + "{}", + json!({ + "mvp_stdio_event":1, + "kind":"datastream_frame", + "channel":channel, + "payload":node_event_payload(config, phase, status, detail), + }) + ); + std::io::stdout() + .flush() + .map_err(|e| format!("flush stdio node event: {e}")) +} + +fn emit_node_event( + datastream: &mut DatastreamEmitter, + config: &DeploymentConfig, + channel: &str, + phase: &str, + status: &str, + detail: Value, +) { + datastream.submit_text( + ChannelId::new(channel), + node_event_payload(config, phase, status, detail).to_string(), + ); + datastream.tick(); +} fn main() -> ExitCode { match run() { @@ -44,17 +104,53 @@ fn main() -> ExitCode { fn run() -> Result<(), String> { let config = DeploymentConfig::from_env()?; - eprintln!( - "mvp-node: boot run={} logical_node={} stage={} worker={} device={}", - config.run_id, - config.logical_node_id, - config.stage_index, - config.worker_script, - config.device - ); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "config", + "ready", + json!({ + "worker_script":&config.worker_script, + "device":&config.device, + "model_id":&config.model_id, + "has_coordinator_endpoint":config.coordinator_endpoint.is_some(), + "has_orchestrator_actor":config.orchestrator_actor.is_some(), + "has_datastream_sink_actor":config.datastream_sink_actor.is_some(), + "self_test_enabled":config.self_test_prompt.is_some(), + "max_runtime_secs":config.max_runtime.as_secs(), + }), + )?; + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "process", + "started", + json!({"binary":"mvp-node","pid":std::process::id()}), + )?; - let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; - let mut driver = IrohDriver::with_handle( + let tokio = match tokio::runtime::Runtime::new() { + Ok(runtime) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "tokio_runtime", + "ready", + json!({"runtime":"tokio"}), + )?; + runtime + } + Err(error) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "tokio_runtime", + "failed", + json!({"error":error.to_string()}), + )?; + return Err(format!("tokio runtime: {error}")); + } + }; + let mut driver = match IrohDriver::with_handle( tokio.handle().clone(), IrohDriverConfig { secret_key: None, @@ -63,13 +159,45 @@ fn run() -> Result<(), String> { peer_auth: None, additional_alpns: vec![], }, - ) - .map_err(|e| format!("create iroh driver: {e}"))?; + ) { + Ok(driver) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "iroh_driver", + "ready", + json!({"endpoint":driver.endpoint_addr(),"relay_mode":format!("{:?}", config.relay_mode)}), + )?; + driver + } + Err(error) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "iroh_driver", + "failed", + json!({"error":error.to_string()}), + )?; + return Err(format!("create iroh driver: {error}")); + } + }; if let Some(coordinator) = &config.coordinator_endpoint { - eprintln!("mvp-node: joining coordinator {coordinator:?}"); driver.join(std::slice::from_ref(coordinator)); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "coordinator_join", + "started", + json!({"endpoint":coordinator}), + )?; } else { - eprintln!("mvp-node: no MVP_COORDINATOR_ENDPOINT set; running standalone until joined"); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "coordinator_join", + "skipped", + json!({"reason":"MVP_COORDINATOR_ENDPOINT not set","mode":"standalone"}), + )?; } let stack = DistributionRuntimeStack::new_with_codecs( @@ -80,6 +208,20 @@ fn run() -> Result<(), String> { datastream::wire::register_datastream_codec(registry); }, ); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "distribution_stack", + "ready", + json!({"actors":"initialized","route_view":"initialized","swim":"initialized","outbox":"initialized"}), + )?; + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "codecs", + "ready", + json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), + )?; driver.enable_actor_bridge( stack.runtime.clone(), stack.codec.clone(), @@ -88,33 +230,167 @@ fn run() -> Result<(), String> { stack.relay_mirror.clone(), stack.route_view.clone(), ); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "actor_bridge", + "ready", + json!({"transport":"iroh","routes":"attached"}), + )?; let mut datastream = node_datastream(&config, &stack); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "datastream_emitter", + "ready", + config.datastream_sink_detail(), + )?; - let reports = stack - .runtime - .new_inbox::() - .map_err(|e| format!("node report inbox: {e}"))?; - let orchestrator = config - .orchestrator_actor - .unwrap_or_else(ActorAddress::new_random); - let node_actor = stack - .runtime - .spawn(NodeAgentActor::new( - stage::NodeId(config.logical_node_id), - orchestrator, - Some(*reports.addr()), - )) - .map_err(|e| format!("spawn node agent: {e}"))?; + let reports = match stack.runtime.new_inbox::() { + Ok(inbox) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "node_report_inbox", + "ready", + json!({"actor":inbox.addr()}), + )?; + inbox + } + Err(error) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "node_report_inbox", + "failed", + json!({"error":error.to_string()}), + )?; + return Err(format!("node report inbox: {error}")); + } + }; + let (orchestrator, orchestrator_source) = match config.orchestrator_actor { + Some(actor) => (actor, "env"), + None => (ActorAddress::new_random(), "generated_fallback"), + }; + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "orchestrator_actor", + "ready", + json!({"actor":orchestrator,"source":orchestrator_source}), + )?; + let node_actor = match stack.runtime.spawn(NodeAgentActor::new( + stage::NodeId(config.logical_node_id), + orchestrator, + Some(*reports.addr()), + )) { + Ok(actor) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "node_agent", + "ready", + json!({"node_actor":actor}), + )?; + actor + } + Err(error) => { + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "node_agent", + "failed", + json!({"error":error.to_string()}), + )?; + return Err(format!("spawn node agent: {error}")); + } + }; stack.register_local_actor(driver.register_actor(node_actor, 1)); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "node_actor_registration", + "ready", + json!({"node_actor":node_actor,"network_reachable":true}), + )?; - let mut worker = TinygradWorker::spawn(&config)?; + emit_stdio_node_event( + &config, + NODE_WORKER_CHANNEL, + "worker_process", + "started", + json!({ + "program":"python3", + "script":&config.worker_script, + "device":&config.device, + "stdin":"piped", + "stdout":"piped", + "stderr":"piped", + }), + )?; + let mut worker = match TinygradWorker::spawn(&config) { + Ok(worker) => worker, + Err(error) => { + emit_stdio_node_event( + &config, + NODE_WORKER_CHANNEL, + "worker_process", + "failed", + json!({"error":error}), + )?; + return Err(error); + } + }; + emit_stdio_node_event( + &config, + NODE_WORKER_CHANNEL, + "worker_initialize", + "started", + json!({"command":"InitializeWorker","helper_abi_version":1,"device":&config.device}), + )?; let mut initial_pump = || {}; - worker.initialize(&config.device, &mut datastream, &mut initial_pump)?; - stack + match worker.initialize(&config.device, &config, &mut datastream, &mut initial_pump) { + Ok(()) => emit_stdio_node_event( + &config, + NODE_WORKER_CHANNEL, + "worker_initialize", + "ready", + json!({"worker_event_type":"WorkerReady"}), + )?, + Err(error) => { + emit_stdio_node_event( + &config, + NODE_WORKER_CHANNEL, + "worker_initialize", + "failed", + json!({"error":error}), + )?; + return Err(error); + } + } + match stack .runtime .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) - .map_err(|e| format!("mark initialized worker ready: {e}"))?; + { + Ok(()) => emit_stdio_node_event( + &config, + NODE_RUNTIME_CHANNEL, + "mark_worker_ready", + "ready", + json!({"sent":"NodeAgentMsg::MarkWorkerReady","node_actor":node_actor}), + )?, + Err(error) => { + emit_stdio_node_event( + &config, + NODE_RUNTIME_CHANNEL, + "mark_worker_ready", + "failed", + json!({"error":error.to_string()}), + )?; + return Err(format!("mark initialized worker ready: {error}")); + } + } let ready = json!({ "type":"ready", @@ -124,8 +400,28 @@ fn run() -> Result<(), String> { "logical_node_id": config.logical_node_id, "stage_index": config.stage_index, }); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "runtime_ready", + "ready", + json!({ + "endpoint":driver.endpoint_addr(), + "node_actor":node_actor, + "logical_node_id":config.logical_node_id, + "stage_index":config.stage_index, + "stdio_ready_line_emitted":true, + }), + )?; println!("{ready}"); datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string()); + emit_stdio_node_event( + &config, + NODE_BOOTSTRAP_CHANNEL, + "datastream_handoff", + "ready", + json!({"from":"stdio_envelope","to":"cluster_datastream","channel":NODE_BOOTSTRAP_CHANNEL}), + )?; std::io::stdout() .flush() .map_err(|e| format!("flush ready line: {e}"))?; @@ -135,13 +431,34 @@ fn run() -> Result<(), String> { } let shutdown_rx = spawn_stdin_shutdown_listener(); + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "stdin_shutdown_listener", + "ready", + json!({"command":"shutdown"}), + ); + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "main_loop", + "started", + json!({ + "poll_interval_ms":PUMP_INTERVAL.as_millis(), + "checks":["network","datastream","node_reports","stdin_shutdown","worker_health","max_runtime"], + }), + ); let started = Instant::now(); loop { pump_network(&mut driver, &stack); datastream.tick(); + worker.drain_stderr(&config, &mut datastream); while let Some(report) = reports.try_recv() { handle_node_report( report, + &config, &stack, &mut driver, node_actor, @@ -150,21 +467,78 @@ fn run() -> Result<(), String> { )?; } if shutdown_rx.try_recv().is_ok() { - eprintln!("mvp-node: shutdown requested on stdin"); + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "shutdown", + "started", + json!({"source":"stdin","command":"shutdown"}), + ); let mut pump = || pump_network(&mut driver, &stack); - let _ = worker.shutdown(&mut datastream, &mut pump); + match worker.shutdown(&config, &mut datastream, &mut pump) { + Ok(()) => { + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "worker_shutdown", + "ready", + json!({"worker_event_type":"WorkerStopped"}), + ); + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "node_exit", + "ready", + json!({"result":"ok"}), + ); + } + Err(error) => emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "worker_shutdown", + "failed", + json!({"error":error}), + ), + } return Ok(()); } if let Some(status) = worker.try_wait()? { + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "worker_process", + "failed", + json!({"exit_status":status.to_string()}), + ); let _ = stack .runtime .send_to(node_actor, NodeAgentMsg::WorkerCrashed); return Err(format!("tinygrad helper exited with {status}")); } if started.elapsed() > config.max_runtime && config.max_runtime != Duration::ZERO { - eprintln!("mvp-node: max runtime elapsed; shutting down cleanly"); + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "max_runtime", + "started", + json!({"max_runtime_secs":config.max_runtime.as_secs()}), + ); let mut pump = || pump_network(&mut driver, &stack); - let _ = worker.shutdown(&mut datastream, &mut pump); + let _ = worker.shutdown(&config, &mut datastream, &mut pump); + emit_node_event( + &mut datastream, + &config, + NODE_SHUTDOWN_CHANNEL, + "node_exit", + "ready", + json!({"result":"max_runtime_elapsed"}), + ); return Ok(()); } thread::sleep(PUMP_INTERVAL); @@ -186,7 +560,6 @@ fn node_datastream( if let Some(actor) = config.datastream_sink_actor { let sink_addr = Arc::new(OnceLock::new()); let _ = sink_addr.set(actor); - eprintln!("mvp-node: native datastream targeting {DATASTREAM_SINK_NAME} actor {actor:?}"); sinks.push(Box::new(ClusterFrameSink::new( stack.runtime.clone(), sink_addr, @@ -195,16 +568,12 @@ fn node_datastream( if let Some(path) = &config.datastream_frame_log { match JsonlFrameSink::open(path) { Ok(sink) => { - eprintln!("mvp-node: native datastream frame log {path}"); sinks.push(Box::new(sink)); } - Err(error) => { - eprintln!("mvp-node: failed to open MVP_DATASTREAM_FRAME_LOG={path:?}: {error}"); - } + Err(_error) => {} } } if sinks.is_empty() { - eprintln!("mvp-node: no datastream sink configured; native datastream drops locally"); sinks.push(Box::new(NoopSink)); } let sink: Box = if sinks.len() == 1 { @@ -262,23 +631,46 @@ impl FrameSink for JsonlFrameSink { fn handle_node_report( report: NodeAgentReport, + config: &DeploymentConfig, stack: &DistributionRuntimeStack, driver: &mut IrohDriver, node_actor: ActorAddress, worker: &mut TinygradWorker, datastream: &mut DatastreamEmitter, ) -> Result<(), String> { + let kind = match &report { + NodeAgentReport::Command(_) => "Command", + NodeAgentReport::Lifecycle(_) => "Lifecycle", + NodeAgentReport::PromptRequested { .. } => "PromptRequested", + NodeAgentReport::Snapshot { .. } => "Snapshot", + }; + emit_node_event( + datastream, + config, + NODE_RUNTIME_CHANNEL, + "node_report", + "observed", + json!({"kind":kind}), + ); match report { - NodeAgentReport::Command(command) => { - handle_stage_command(command, stack, driver, node_actor, worker, datastream) - } + NodeAgentReport::Command(command) => handle_stage_command( + command, config, stack, driver, node_actor, worker, datastream, + ), NodeAgentReport::Lifecycle(event) => { - let record = json!({"type":"node_lifecycle","event":format!("{event:?}")}); - println!("{record}"); - datastream.submit_text(ChannelId::new("mvp.node.lifecycle"), record.to_string()); - std::io::stdout() - .flush() - .map_err(|e| format!("flush lifecycle: {e}")) + let event = format!("{event:?}"); + datastream.submit_text( + ChannelId::new("mvp.node.lifecycle"), + json!({"type":"node_lifecycle","event":event}).to_string(), + ); + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "lifecycle", + "observed", + json!({"event":event}), + ); + Ok(()) } NodeAgentReport::PromptRequested { request_id, @@ -286,7 +678,7 @@ fn handle_node_report( max_tokens, reply_to, } => handle_prompt_request( - request_id, prompt, max_tokens, reply_to, stack, driver, worker, datastream, + request_id, prompt, max_tokens, reply_to, config, stack, driver, worker, datastream, ), NodeAgentReport::Snapshot { .. } => Ok(()), } @@ -297,61 +689,169 @@ fn handle_prompt_request( prompt: String, max_tokens: u32, reply_to: ActorAddress, + config: &DeploymentConfig, stack: &DistributionRuntimeStack, driver: &mut IrohDriver, worker: &mut TinygradWorker, datastream: &mut DatastreamEmitter, ) -> Result<(), String> { let started = Instant::now(); + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_requested", + "started", + json!({"request_id":request_id,"max_tokens":max_tokens,"reply_to":reply_to,"prompt_bytes":prompt.len()}), + ); + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "infer_prompt", + "started", + json!({"request_id":request_id,"command":"InferPrompt","max_tokens":max_tokens}), + ); let mut pump = || pump_network(driver, stack); - match worker.infer_prompt(&prompt, max_tokens, datastream, &mut pump) { + match worker.infer_prompt( + request_id, &prompt, max_tokens, config, datastream, &mut pump, + ) { Ok(result) => { let text = result .get("text") .and_then(Value::as_str) .unwrap_or_default() .to_owned(); + let text_bytes = text.len(); + let worker_result_payload_bytes = result.to_string().len(); + let prompt_tokens = result + .get("prompt_tokens") + .and_then(Value::as_array) + .map_or(0, |tokens| tokens.len() as u32); let tokens_generated = result .get("generated_tokens") .and_then(Value::as_array) .map_or(0, |tokens| tokens.len() as u32); + let elapsed_ms = result + .get("elapsed_ms") + .and_then(Value::as_u64) + .unwrap_or_else(|| started.elapsed().as_millis() as u64); + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "infer_prompt", + "ready", + json!({"request_id":request_id,"worker_event_type":"PromptCompleted","prompt_tokens":prompt_tokens,"tokens_generated":tokens_generated,"elapsed_ms":elapsed_ms,"text_bytes":text_bytes,"worker_result_payload_bytes":worker_result_payload_bytes}), + ); if !text.is_empty() { - stack - .runtime - .send_to( - reply_to, - PromptEvent::TextDelta { - request_id, - text: text.clone(), - }, - ) - .map_err(|e| format!("send prompt text delta: {e}"))?; - } - stack - .runtime - .send_to( + match stack.runtime.send_to( reply_to, - PromptEvent::Done { + PromptEvent::TextDelta { request_id, - final_text: text, - tokens_generated, - elapsed_ms: result - .get("elapsed_ms") - .and_then(Value::as_u64) - .unwrap_or_else(|| started.elapsed().as_millis() as u64), + text: text.clone(), }, - ) - .map_err(|e| format!("send prompt done: {e}")) + ) { + Ok(()) => emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "ready", + json!({"request_id":request_id,"event":"TextDelta","bytes":text_bytes,"reply_to":reply_to}), + ), + Err(error) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "failed", + json!({"request_id":request_id,"event":"TextDelta","error":error.to_string()}), + ); + return Err(format!("send prompt text delta: {error}")); + } + } + } + match stack.runtime.send_to( + reply_to, + PromptEvent::Done { + request_id, + final_text: text, + tokens_generated, + elapsed_ms, + }, + ) { + Ok(()) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "ready", + json!({"request_id":request_id,"event":"Done","tokens_generated":tokens_generated,"elapsed_ms":elapsed_ms,"final_text_bytes":text_bytes,"reply_to":reply_to}), + ); + Ok(()) + } + Err(error) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "failed", + json!({"request_id":request_id,"event":"Done","error":error.to_string()}), + ); + Err(format!("send prompt done: {error}")) + } + } + } + Err(error) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "infer_prompt", + "failed", + json!({"request_id":request_id,"error":error}), + ); + match stack.runtime.send_to( + reply_to, + PromptEvent::Fault { + request_id, + error: error.clone(), + }, + ) { + Ok(()) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "ready", + json!({"request_id":request_id,"event":"Fault","reply_to":reply_to}), + ); + Ok(()) + } + Err(send_error) => { + emit_node_event( + datastream, + config, + NODE_PROMPT_CHANNEL, + "prompt_response", + "failed", + json!({"request_id":request_id,"event":"Fault","error":send_error.to_string()}), + ); + Err(format!("send prompt fault: {send_error}")) + } + } } - Err(error) => stack - .runtime - .send_to(reply_to, PromptEvent::Fault { request_id, error }) - .map_err(|e| format!("send prompt fault: {e}")), } } fn handle_stage_command( command: StageCommandWire, + config: &DeploymentConfig, stack: &DistributionRuntimeStack, driver: &mut IrohDriver, node_actor: ActorAddress, @@ -365,15 +865,44 @@ fn handle_stage_command( layer_start, layer_end_exclusive, } => { + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "configure_worker_role", + "started", + json!({"run_id":run_id,"stage_index":stage_index,"layer_range":{"start":layer_start,"end_exclusive":layer_end_exclusive}}), + ); let mut pump = || pump_network(driver, stack); - worker.configure_role( + match worker.configure_role( run_id, stage_index, layer_start, layer_end_exclusive, + config, datastream, &mut pump, - )?; + ) { + Ok(()) => emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "configure_worker_role", + "ready", + json!({"worker_event_type":"RoleConfigured"}), + ), + Err(error) => { + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "configure_worker_role", + "failed", + json!({"error":error}), + ); + return Err(error); + } + } stack .runtime .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) @@ -386,16 +915,53 @@ fn handle_stage_command( layer_start, layer_end_exclusive, } => { + let gguf_source_kind = match &gguf_source { + GgufSource::LocalPath(_) => "local_path", + GgufSource::HuggingFaceGguf { .. } => "huggingface", + }; + let tokenizer_kind = match &tokenizer { + TokenizerSource::EmbeddedGguf => "gguf", + TokenizerSource::LocalPath(_) => "local_path", + }; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "load_weights", + "started", + json!({"model_id":&model_id,"gguf_source":gguf_source_kind,"tokenizer":tokenizer_kind,"layer_range":{"start":layer_start,"end_exclusive":layer_end_exclusive}}), + ); let mut pump = || pump_network(driver, stack); - worker.load_weights( - model_id, + match worker.load_weights( + model_id.clone(), gguf_source, tokenizer, layer_start, layer_end_exclusive, + config, datastream, &mut pump, - )?; + ) { + Ok(()) => emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "load_weights", + "ready", + json!({"worker_event_type":"WeightsLoaded","model_id":model_id}), + ), + Err(error) => { + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "load_weights", + "failed", + json!({"error":error}), + ); + return Err(error); + } + } stack .runtime .send_to(node_actor, NodeAgentMsg::MarkWeightsReady) @@ -409,7 +975,16 @@ fn handle_stage_command( stack .runtime .send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id }) - .map_err(|e| format!("mark worker rings quiesced: {e}")) + .map_err(|e| format!("mark worker rings quiesced: {e}"))?; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "local_edges", + "ready", + json!({"run_id":run_id,"sent":["LocalEdgesStopped","WorkerRingsQuiesced"]}), + ); + Ok(()) } StageCommandWire::ReleaseRunDeviceObjects { run_id } => { stack @@ -419,20 +994,81 @@ fn handle_stage_command( stack .runtime .send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id }) - .map_err(|e| format!("mark worker role reset: {e}")) + .map_err(|e| format!("mark worker role reset: {e}"))?; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "device_objects", + "ready", + json!({"run_id":run_id,"sent":["DeviceObjectsReleased","WorkerRoleReset"]}), + ); + Ok(()) + } + StageCommandWire::EstablishInboundEdge { edge_id } => { + stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id }) + .map_err(|e| format!("mark inbound edge ready: {e}"))?; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "inbound_edge", + "ready", + json!({"edge_id":format!("{edge_id:?}")}), + ); + Ok(()) + } + StageCommandWire::EstablishOutboundEdge { edge_id } => { + stack + .runtime + .send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id }) + .map_err(|e| format!("mark outbound edge ready: {e}"))?; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "outbound_edge", + "ready", + json!({"edge_id":format!("{edge_id:?}")}), + ); + Ok(()) + } + StageCommandWire::RewireEdge { .. } => { + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "rewire_edge", + "skipped", + json!({"reason":"not implemented in mvp-node image path"}), + ); + Ok(()) + } + StageCommandWire::ReleaseInputHandle { .. } => { + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "release_input_handle", + "skipped", + json!({"reason":"not implemented in mvp-node image path"}), + ); + Ok(()) + } + StageCommandWire::ExecuteStep { .. } => { + let error = "mvp-node image path does not carry ring payload commands yet"; + emit_node_event( + datastream, + config, + NODE_STAGE_CHANNEL, + "execute_step", + "failed", + json!({"error":error}), + ); + Err(format!("{error}: {command:?}")) } - StageCommandWire::EstablishInboundEdge { edge_id } => stack - .runtime - .send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id }) - .map_err(|e| format!("mark inbound edge ready: {e}")), - StageCommandWire::EstablishOutboundEdge { edge_id } => stack - .runtime - .send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id }) - .map_err(|e| format!("mark outbound edge ready: {e}")), - StageCommandWire::RewireEdge { .. } | StageCommandWire::ReleaseInputHandle { .. } => Ok(()), - StageCommandWire::ExecuteStep { .. } => Err(format!( - "mvp-node image path does not carry ring payload commands yet: {command:?}" - )), } } @@ -442,13 +1078,21 @@ fn run_self_test( prompt: &str, datastream: &mut DatastreamEmitter, ) -> Result<(), String> { - eprintln!("mvp-node: running self-test prompt"); + emit_node_event( + datastream, + config, + NODE_RUNTIME_CHANNEL, + "self_test", + "started", + json!({"prompt_bytes":prompt.len()}), + ); let mut pump = || {}; worker.configure_role( config.run_id, config.stage_index, 0, config.self_test_layer_end, + config, datastream, &mut pump, )?; @@ -458,20 +1102,29 @@ fn run_self_test( config.tokenizer.clone(), 0, config.self_test_layer_end, + config, datastream, &mut pump, )?; - let result = worker.infer_prompt(prompt, config.self_test_max_tokens, datastream, &mut pump)?; - let record = json!({ - "type":"self_test_completed", - "prompt":prompt, - "result":result, - }); - println!("{record}"); + let result = worker.infer_prompt( + 0, + prompt, + config.self_test_max_tokens, + config, + datastream, + &mut pump, + )?; + let record = json!({"type":"self_test_completed","prompt_bytes":prompt.len(),"result":result}); datastream.submit_text(ChannelId::new("mvp.node.self_test"), record.to_string()); - std::io::stdout() - .flush() - .map_err(|e| format!("flush self-test: {e}")) + emit_node_event( + datastream, + config, + NODE_RUNTIME_CHANNEL, + "self_test", + "ready", + json!({"prompt_bytes":prompt.len()}), + ); + Ok(()) } #[derive(Clone)] @@ -497,15 +1150,17 @@ struct DeploymentConfig { impl DeploymentConfig { fn from_env() -> Result { + let run_id = env_u64("MVP_RUN_ID", 1)?; + let relay = relay_runtime_config_from_env(run_id)?; Ok(Self { - run_id: env_u64("MVP_RUN_ID", 1)?, + run_id, logical_node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, stage_index: env_u32("MVP_STAGE_INDEX", 0)?, coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?, datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), - relay_mode: relay_mode_from_env()?, + relay_mode: relay.mode, worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), device: env_string("DEV", DEFAULT_DEVICE), model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), @@ -517,12 +1172,30 @@ impl DeploymentConfig { max_runtime: Duration::from_secs(env_u64("MVP_NODE_MAX_RUNTIME_SECS", 0)?), }) } + + fn datastream_sink_detail(&self) -> Value { + let sink = match ( + self.datastream_sink_actor.is_some(), + self.datastream_frame_log.is_some(), + ) { + (true, true) => "tee", + (true, false) => "cluster_actor", + (false, true) => "frame_log", + (false, false) => "noop", + }; + json!({ + "sink":sink, + "cluster_actor":self.datastream_sink_actor, + "frame_log":self.datastream_frame_log, + }) + } } struct TinygradWorker { child: Child, stdin: ChildStdin, stdout: BufReader, + stderr_rx: Receiver, } impl TinygradWorker { @@ -532,7 +1205,7 @@ impl TinygradWorker { .env("DEV", &config.device) .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::inherit()) + .stderr(Stdio::piped()) .spawn() .map_err(|e| format!("spawn tinygrad helper {}: {e}", config.worker_script))?; let stdin = child @@ -543,22 +1216,37 @@ impl TinygradWorker { .stdout .take() .ok_or_else(|| "tinygrad helper stdout missing".to_owned())?; + let stderr = child + .stderr + .take() + .ok_or_else(|| "tinygrad helper stderr missing".to_owned())?; + let (stderr_tx, stderr_rx) = mpsc::channel(); + thread::spawn(move || { + for line in BufReader::new(stderr).lines().map_while(Result::ok) { + if stderr_tx.send(line).is_err() { + break; + } + } + }); Ok(Self { child, stdin, stdout: BufReader::new(stdout), + stderr_rx, }) } fn initialize( &mut self, device: &str, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, pump: &mut dyn FnMut(), ) -> Result<(), String> { self.command( json!({"type":"InitializeWorker","helper_abi_version":1,"backend":{"device":device}}), "WorkerReady", + config, datastream, ChannelId::new("mvp.worker.initialize"), pump, @@ -572,6 +1260,7 @@ impl TinygradWorker { stage_index: u32, layer_start: u32, layer_end_exclusive: u32, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, pump: &mut dyn FnMut(), ) -> Result<(), String> { @@ -587,6 +1276,7 @@ impl TinygradWorker { } }), "RoleConfigured", + config, datastream, ChannelId::new("mvp.worker.role"), pump, @@ -601,6 +1291,7 @@ impl TinygradWorker { tokenizer: TokenizerSource, layer_start: u32, layer_end_exclusive: u32, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, pump: &mut dyn FnMut(), ) -> Result<(), String> { @@ -614,6 +1305,7 @@ impl TinygradWorker { "layer_end_exclusive":layer_end_exclusive, }), "WeightsLoaded", + config, datastream, ChannelId::new("mvp.worker.weights"), pump, @@ -623,14 +1315,17 @@ impl TinygradWorker { fn infer_prompt( &mut self, + request_id: u64, prompt: &str, max_tokens: u32, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, pump: &mut dyn FnMut(), ) -> Result { self.command( - json!({"type":"InferPrompt","prompt":prompt,"max_tokens":max_tokens}), + json!({"type":"InferPrompt","request_id":request_id,"prompt":prompt,"max_tokens":max_tokens}), "PromptCompleted", + config, datastream, ChannelId::new("mvp.worker.prompt"), pump, @@ -639,12 +1334,14 @@ impl TinygradWorker { fn shutdown( &mut self, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, pump: &mut dyn FnMut(), ) -> Result<(), String> { self.command( json!({"type":"ShutdownWorker"}), "WorkerStopped", + config, datastream, ChannelId::new("mvp.worker.shutdown"), pump, @@ -658,51 +1355,178 @@ impl TinygradWorker { .map_err(|e| format!("poll tinygrad helper: {e}")) } + fn drain_stderr(&mut self, config: &DeploymentConfig, datastream: &mut DatastreamEmitter) { + let mut emitted = false; + while let Ok(line) = self.stderr_rx.try_recv() { + let payload = + node_event_payload(config, "worker_stderr", "observed", json!({"line":line})); + datastream.submit_text(ChannelId::new("mvp.worker.stderr"), payload.to_string()); + emitted = true; + } + if emitted { + datastream.tick(); + } + } + fn command( &mut self, command: Value, expected: &str, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, channel: ChannelId, pump: &mut dyn FnMut(), ) -> Result { - writeln!(self.stdin, "{command}").map_err(|e| format!("write helper command: {e}"))?; - self.stdin - .flush() - .map_err(|e| format!("flush helper command: {e}"))?; - self.expect_event(expected, datastream, channel, pump) + let command_type = command + .get("type") + .and_then(Value::as_str) + .unwrap_or("unknown") + .to_owned(); + let command_text = command.to_string(); + let command_bytes = command_text.len() + 1; + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_command_write", + "started", + json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes}), + ); + if let Err(error) = writeln!(self.stdin, "{command_text}") { + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_command_write", + "failed", + json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes,"error":error.to_string()}), + ); + return Err(format!("write helper command: {error}")); + } + if let Err(error) = self.stdin.flush() { + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_command_write", + "failed", + json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes,"error":error.to_string()}), + ); + return Err(format!("flush helper command: {error}")); + } + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_command_write", + "ready", + json!({"command_type":command_type.as_str(),"expected_event_type":expected,"command_bytes":command_bytes}), + ); + self.expect_event(expected, config, datastream, channel, pump) } fn expect_event( &mut self, expected: &str, + config: &DeploymentConfig, datastream: &mut DatastreamEmitter, channel: ChannelId, pump: &mut dyn FnMut(), ) -> Result { loop { let mut line = String::new(); - let n = self - .stdout - .read_line(&mut line) - .map_err(|e| format!("read helper stdout: {e}"))?; + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_read", + "started", + json!({"expected_event_type":expected,"channel":channel.as_str()}), + ); + let n = match self.stdout.read_line(&mut line) { + Ok(n) => n, + Err(error) => { + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_read", + "failed", + json!({"expected_event_type":expected,"channel":channel.as_str(),"error":error.to_string()}), + ); + return Err(format!("read helper stdout: {error}")); + } + }; + self.drain_stderr(config, datastream); if n == 0 { + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_read", + "failed", + json!({"expected_event_type":expected,"channel":channel.as_str(),"line_bytes":0,"error":"stdout closed"}), + ); return Err(format!( "tinygrad helper stdout closed while waiting for {expected}" )); } - let value: Value = serde_json::from_str(&line) - .map_err(|e| format!("parse helper stdout {line:?}: {e}"))?; + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_read", + "ready", + json!({"expected_event_type":expected,"channel":channel.as_str(),"line_bytes":n}), + ); + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_parse", + "started", + json!({"expected_event_type":expected,"channel":channel.as_str(),"line_bytes":n}), + ); + let value: Value = match serde_json::from_str(&line) { + Ok(value) => value, + Err(error) => { + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_parse", + "failed", + json!({"expected_event_type":expected,"channel":channel.as_str(),"line_bytes":n,"error":error.to_string()}), + ); + return Err(format!("parse helper stdout {line:?}: {error}")); + } + }; + let worker_event_type = value + .get("type") + .and_then(Value::as_str) + .unwrap_or("unknown"); + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_stdout_parse", + "ready", + json!({"expected_event_type":expected,"channel":channel.as_str(),"line_bytes":n,"worker_event_type":worker_event_type}), + ); datastream.submit_text(channel.clone(), value.to_string()); datastream.tick(); pump(); if value.get("type").and_then(Value::as_str) == Some(expected) { return Ok(value); } - println!("{}", json!({"type":"tinygrad_event","event":value})); - std::io::stdout() - .flush() - .map_err(|e| format!("flush helper passthrough: {e}"))?; + emit_node_event( + datastream, + config, + NODE_WORKER_CHANNEL, + "worker_event", + "observed", + json!({"command_waiting_for":expected,"worker_event_type":worker_event_type,"event":value}), + ); } } } @@ -766,16 +1590,6 @@ where .transpose() } -fn relay_mode_from_env() -> Result { - match env_string("MVP_IROH_RELAY_MODE", "disabled").as_str() { - "disabled" => Ok(iroh::RelayMode::Disabled), - "default" => Ok(iroh::RelayMode::Default), - other => Err(format!( - "unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default" - )), - } -} - fn gguf_source_from_env() -> GgufSource { if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { return GgufSource::LocalPath(path); diff --git a/crates/mvp-system/src/bin/mvp_one_node_chat.rs b/crates/mvp-system/src/bin/mvp_one_node_chat.rs index 11848fe..8f588b6 100644 --- a/crates/mvp-system/src/bin/mvp_one_node_chat.rs +++ b/crates/mvp-system/src/bin/mvp_one_node_chat.rs @@ -1,12 +1,16 @@ -use std::io::{self, BufRead, BufReader, Write}; -use std::net::TcpStream; -use std::path::PathBuf; +use std::collections::HashSet; +use std::fs::{self, File}; +use std::io::{self, BufRead, BufReader, Read, Seek, SeekFrom, Write}; +use std::net::{Shutdown, TcpStream}; +use std::path::{Path, PathBuf}; use std::process::{Child, ChildStdin, Command, ExitCode, Stdio}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::mpsc; use std::thread; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime}; +use mvp_system::node_image::{NodeImageProvider, NodeImageRequest, prepare_node_image}; +use mvp_system::node_provisioning::ProviderKind; use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, write_json_line}; use serde_json::Value; @@ -16,13 +20,39 @@ use std::os::unix::process::CommandExt; const DEFAULT_RPC_ADDR: &str = "127.0.0.1:19777"; const DEFAULT_NODE_IMAGE: &str = "swactor-mvp-node:latest"; const BASE_NODE_IMAGE: &str = "swactor-mvp-node-base:cuda12.6"; +const MVP_RUNTIME_CONFIG_ENV: &str = "MVP_RUNTIME_CONFIG"; +const CACHED_MODEL_HOST_ENV: &str = "MVP_CACHED_MODEL_HOST_PATH"; +const DEFAULT_CACHED_MODEL_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; +const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; const DEFAULT_MAX_TOKENS: u32 = 64; const DEFAULT_TIMEOUT_MS: u64 = 120_000; +const ORCH_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 ORCH_READY_TIMEOUT: Duration = Duration::from_secs(1_200); const ORCH_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30); +const INTERRUPT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +const FORCE_KILL_DELAY: Duration = Duration::from_millis(500); const CHAT_READ_TIMEOUT: Duration = Duration::from_millis(100); static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); +static STOP_ACKNOWLEDGED: AtomicBool = AtomicBool::new(false); fn main() -> ExitCode { install_signal_handlers(); @@ -36,13 +66,32 @@ fn main() -> ExitCode { } fn run() -> Result<(), String> { - let config = Config::from_args()?; - prepare_runtime(&config)?; - let mut orch = OrchChild::spawn(&config)?; - let rpc_addr = orch.wait_ready(config.rpc_addr.clone())?; - eprintln!("mvp-one-node-chat: prompt loop ready at {rpc_addr}"); + let mut config = Config::from_args()?; + let image_ref = prepare_runtime(&config)?; + let frame_log = configure_progress_frame_log(&mut config)?; + let mut progress = StartupProgress::new(&frame_log); + let mut orch = OrchChild::spawn(&config, &image_ref)?; + let rpc_addr = match orch.wait_ready(config.rpc_addr.clone(), &mut progress) { + Ok(addr) => addr, + Err(_) if STOP_REQUESTED.load(Ordering::SeqCst) => { + acknowledge_stop(); + let forced = orch.shutdown(true); + report_interrupt_shutdown(forced); + return Ok(()); + } + Err(error) => { + progress.poll(); + return Err(progress.failure_summary().unwrap_or(error)); + } + }; + progress.poll(); + println!("model successfully loaded."); let result = run_chat_loop(&rpc_addr, config.max_tokens, config.timeout_ms); - orch.shutdown(); + let interrupted = STOP_REQUESTED.load(Ordering::SeqCst); + let forced = orch.shutdown(interrupted); + if interrupted { + report_interrupt_shutdown(forced); + } result } @@ -51,14 +100,22 @@ struct Config { orch_args: Vec, rpc_addr: String, node_image: String, + config_profile: RuntimeConfigProfile, + provider: ProviderKind, + relay_mode: iroh::RelayMode, max_tokens: u32, timeout_ms: u64, dashboard: bool, build_image: bool, + image_tag: Option, + push_image: bool, + force_image_refresh: bool, + cached_model: Option, + datastream_frame_log: Option, } - impl Config { fn from_args() -> Result { + let config_profile = RuntimeConfigProfile::from_env()?; let mut config = Self { orch_bin: default_orch_bin()?, orch_args: Vec::new(), @@ -67,10 +124,20 @@ impl Config { .unwrap_or_else(|_| DEFAULT_RPC_ADDR.to_owned()), node_image: std::env::var("MVP_NODE_IMAGE") .unwrap_or_else(|_| DEFAULT_NODE_IMAGE.to_owned()), + config_profile, + provider: provider_from_env(config_profile)?, + relay_mode: relay_mode_from_env()?, max_tokens: env_u32("MVP_PROMPT_MAX_TOKENS", DEFAULT_MAX_TOKENS)?, timeout_ms: env_u64("MVP_PROMPT_TIMEOUT_MS", DEFAULT_TIMEOUT_MS)?, dashboard: env_bool("MVP_DASHBOARD", true)?, build_image: env_bool("MVP_BUILD_NODE_IMAGE", true)?, + image_tag: std::env::var("MVP_NODE_IMAGE_TAG") + .ok() + .filter(|value| !value.trim().is_empty()), + push_image: env_bool("MVP_PUSH_NODE_IMAGE", false)?, + force_image_refresh: env_bool("MVP_FORCE_NODE_IMAGE_REFRESH", false)?, + cached_model: None, + datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG").map(PathBuf::from), }; let mut args = std::env::args().skip(1).peekable(); @@ -79,16 +146,27 @@ impl Config { "--orch-bin" => config.orch_bin = PathBuf::from(next_arg(&mut args, "--orch-bin")?), "--addr" => config.rpc_addr = next_arg(&mut args, "--addr")?, "--image" => { - let image = next_arg(&mut args, "--image")?; - config.node_image = image.clone(); - config.orch_args.push("--image".to_owned()); - config.orch_args.push(image); + config.node_image = next_arg(&mut args, "--image")?; } "--max-tokens" => config.max_tokens = parse_next(&mut args, "--max-tokens")?, "--timeout-ms" => config.timeout_ms = parse_next(&mut args, "--timeout-ms")?, + "--datastream-frame-log" => { + config.datastream_frame_log = Some(PathBuf::from(next_arg( + &mut args, + "--datastream-frame-log", + )?)); + } "--dashboard" => config.dashboard = true, "--no-dashboard" => config.dashboard = false, "--no-build-image" => config.build_image = false, + "--image-tag" => config.image_tag = Some(next_arg(&mut args, "--image-tag")?), + "--push-image" => config.push_image = true, + "--no-push-image" => config.push_image = false, + "--force-image-refresh" => config.force_image_refresh = true, + "--cached-model" => { + let path = args.next_if(|value| !value.starts_with("--")); + config.cached_model = Some(CachedModelConfig::from_arg(path)?); + } "--" => { config.orch_args.extend(args); break; @@ -100,27 +178,266 @@ impl Config { } } +struct FrameLogConfig { + path: PathBuf, + start_offset: u64, + remove_on_drop: bool, +} + +fn configure_progress_frame_log(config: &mut Config) -> Result { + if let Some(path) = &config.datastream_frame_log { + let start_offset = fs::metadata(path) + .map(|metadata| metadata.len()) + .unwrap_or(0); + return Ok(FrameLogConfig { + path: path.clone(), + start_offset, + remove_on_drop: false, + }); + } + + let path = PathBuf::from("target") + .join("mvp-chat") + .join(format!("startup-{}.frames.jsonl", std::process::id())); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent) + .map_err(|e| format!("create startup frame log dir {}: {e}", parent.display()))?; + } + File::create(&path).map_err(|e| format!("create startup frame log {}: {e}", path.display()))?; + config.datastream_frame_log = Some(path.clone()); + Ok(FrameLogConfig { + path, + start_offset: 0, + remove_on_drop: true, + }) +} + +impl Drop for FrameLogConfig { + fn drop(&mut self) { + if self.remove_on_drop { + let _ = fs::remove_file(&self.path); + } + } +} + +struct StartupProgress { + path: PathBuf, + offset: u64, + partial: String, + printed: HashSet, + last_error_line: Option, + last_failure: Option, + last_download_bucket: Option, +} + +impl StartupProgress { + fn new(frame_log: &FrameLogConfig) -> Self { + Self { + path: frame_log.path.clone(), + offset: frame_log.start_offset, + partial: String::new(), + printed: HashSet::new(), + last_error_line: None, + last_failure: None, + last_download_bucket: None, + } + } + + fn poll(&mut self) { + let mut file = match File::open(&self.path) { + Ok(file) => file, + Err(_) => return, + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut chunk = String::new(); + if file.read_to_string(&mut chunk).is_err() || chunk.is_empty() { + return; + } + self.offset += chunk.as_bytes().len() as u64; + self.partial.push_str(&chunk); + while let Some(newline) = self.partial.find('\n') { + let line: String = self.partial.drain(..=newline).collect(); + let line = line.trim(); + if !line.is_empty() { + self.observe_archive_line(line); + } + } + } + + fn failure_summary(&self) -> Option { + if let Some(line) = &self.last_error_line { + return Some(format!("node provisioning failed: {line}")); + } + self.last_failure.clone() + } + + fn observe_archive_line(&mut self, line: &str) { + let Ok(record) = serde_json::from_str::(line) else { + return; + }; + let Some(channel) = record.get("channel").and_then(Value::as_str) else { + return; + }; + let Some(payload_text) = record + .get("payload") + .and_then(|payload| payload.get("value")) + .and_then(Value::as_str) + else { + return; + }; + let Ok(payload) = serde_json::from_str::(payload_text) else { + return; + }; + + if channel == "mvp.orch.bootstrap" { + self.observe_bootstrap(&payload); + } else if channel == "mvp.provisioning.events" { + self.observe_provision_event(&payload); + } else if channel == "mvp.worker.weights" { + self.observe_worker_weights(&payload); + } else if channel.starts_with("mvp.provisioning.logs.") { + self.observe_provision_log(&payload); + } + } + + fn observe_bootstrap(&mut self, payload: &Value) { + let phase = payload.get("phase").and_then(Value::as_str).unwrap_or(""); + let status = payload.get("status").and_then(Value::as_str).unwrap_or(""); + if status == "failed" { + let error = payload + .get("detail") + .and_then(|detail| detail.get("error")) + .and_then(Value::as_str) + .unwrap_or("unknown error"); + self.last_failure = Some(format!("{phase} failed: {error}")); + return; + } + match (phase, status) { + ("provider_start", "started") => { + self.print_once("starting_docker_node", "starting docker node") + } + ("node_runtime_ready", "started") => { + self.print_once("node_runtime_ready_started", "waiting for node runtime") + } + ("node_runtime_ready", "ready") => { + self.print_once("node_runtime_ready", "node runtime ready") + } + ("stage_provision", "started") => { + self.print_once("stage_provision", "configuring model stage") + } + ("weights_loaded", "started") => { + self.print_once("weights_loaded_started", "loading model") + } + ("prompt_rpc", "ready") | ("prompt_loop", "ready") => { + self.print_once("prompt_ready", "prompt RPC ready") + } + _ => {} + } + } + + fn observe_provision_event(&mut self, payload: &Value) { + let Some(event) = payload.get("event") else { + return; + }; + match event.get("kind").and_then(Value::as_str).unwrap_or("") { + "ProvisionStart" => self.print_once("starting_docker_node", "starting docker node"), + "NodeLive" => self.print_once("node_runtime_ready", "node runtime ready"), + "ProvisionFailed" => { + let message = event + .get("message") + .and_then(Value::as_str) + .unwrap_or("provisioning failed"); + self.last_failure = Some(format!("node provisioning failed: {message}")); + } + _ => {} + } + } + + fn observe_worker_weights(&mut self, payload: &Value) { + match payload.get("type").and_then(Value::as_str).unwrap_or("") { + "GgufDownloadStarted" => { + self.print_once("download_started", "downloading model weights") + } + "GgufDownloadProgress" => { + let done = payload + .get("bytes_done") + .and_then(Value::as_u64) + .unwrap_or(0); + let total = payload + .get("bytes_total") + .and_then(Value::as_u64) + .unwrap_or(0); + if total == 0 { + return; + } + let pct = done.saturating_mul(100).saturating_div(total).min(100); + let bucket = pct / 10; + if self.last_download_bucket != Some(bucket) { + self.last_download_bucket = Some(bucket); + eprintln!("mvp-one-node-chat: downloading model weights {pct}%"); + } + } + _ => {} + } + } + + fn observe_provision_log(&mut self, payload: &Value) { + let Some(line) = payload + .get("line") + .and_then(|line| line.get("line")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|line| !line.is_empty()) + else { + return; + }; + if line.starts_with("docker:") || line.contains("Error response") || line.contains("error") + { + self.last_error_line = Some(line.to_owned()); + } + } + + fn print_once(&mut self, key: &str, message: &str) { + if self.printed.insert(key.to_owned()) { + eprintln!("mvp-one-node-chat: {message}"); + } + } +} + struct OrchChild { child: Child, stdin: Option, - ready_rx: mpsc::Receiver, cleaned: bool, } impl OrchChild { - fn spawn(config: &Config) -> Result { + fn spawn(config: &Config, image_ref: &str) -> Result { let mut command = Command::new(&config.orch_bin); command .args(&config.orch_args) - .env("MVP_NODE_IMAGE", &config.node_image) + .env("MVP_NODE_PROVIDER", config.provider.as_str()) + .env(MVP_RUNTIME_CONFIG_ENV, config.config_profile.as_str()) + .env("MVP_NODE_IMAGE", image_ref) + .env( + "MVP_IROH_RELAY_MODE", + relay_mode_env_value(&config.relay_mode), + ) .env("MVP_PROMPT_RPC_BIND", &config.rpc_addr) .env("MVP_PROMPT_RPC_ADDR", &config.rpc_addr) .env("MVP_PROMPT_MAX_TOKENS", config.max_tokens.to_string()) .env("MVP_PROMPT_TIMEOUT_MS", config.timeout_ms.to_string()) .env("MVP_DASHBOARD", if config.dashboard { "1" } else { "0" }) .stdin(Stdio::piped()) - .stdout(Stdio::piped()) + .stdout(Stdio::null()) .stderr(Stdio::inherit()); + if let Some(cached_model) = &config.cached_model { + command.env(CACHED_MODEL_HOST_ENV, &cached_model.host_path); + } + if let Some(path) = &config.datastream_frame_log { + command.env("MVP_DATASTREAM_FRAME_LOG", path); + } #[cfg(target_os = "linux")] unsafe { command.pre_exec(|| { @@ -131,142 +448,131 @@ impl OrchChild { } }); } + let display_orch = display_user_path(&config.orch_bin); let mut child = command .spawn() - .map_err(|e| format!("spawn {}: {e}", config.orch_bin.display()))?; + .map_err(|e| format!("spawn {display_orch}: {e}"))?; let stdin = child.stdin.take(); - let stdout = child - .stdout - .take() - .ok_or_else(|| "orchestrator stdout missing".to_owned())?; - let (ready_tx, ready_rx) = mpsc::channel(); - thread::spawn(move || { - for line in BufReader::new(stdout).lines().map_while(Result::ok) { - eprintln!("mvp-orch-one-node: {line}"); - if let Ok(value) = serde_json::from_str::(&line) - && value.get("type").and_then(Value::as_str) == Some("prompt_loop_ready") - && let Some(addr) = value.get("addr").and_then(Value::as_str) - { - let _ = ready_tx.send(addr.to_owned()); - } - } - }); Ok(Self { child, stdin, - ready_rx, cleaned: false, }) } - fn wait_ready(&mut self, fallback_addr: String) -> Result { + fn wait_ready( + &mut self, + rpc_addr: String, + progress: &mut StartupProgress, + ) -> Result { let start = Instant::now(); loop { + progress.poll(); if STOP_REQUESTED.load(Ordering::SeqCst) { + acknowledge_stop(); return Err("interrupted before orchestrator became ready".to_owned()); } - if let Ok(addr) = self.ready_rx.recv_timeout(Duration::from_millis(100)) { - return Ok(addr); + match TcpStream::connect(&rpc_addr) { + Ok(stream) => { + let _ = stream.shutdown(Shutdown::Both); + return Ok(rpc_addr); + } + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::ConnectionRefused + | std::io::ErrorKind::TimedOut + | std::io::ErrorKind::AddrNotAvailable + ) => {} + Err(error) => return Err(format!("connect prompt RPC {rpc_addr}: {error}")), } if let Some(status) = self .child .try_wait() .map_err(|e| format!("poll orchestrator: {e}"))? { + progress.poll(); + if let Some(failure) = progress.failure_summary() { + return Err(failure); + } return Err(format!( - "orchestrator exited before prompt loop ready: {status}" + "orchestrator exited before prompt RPC ready: {status}" )); } if start.elapsed() > ORCH_READY_TIMEOUT { + progress.poll(); + if let Some(failure) = progress.failure_summary() { + return Err(failure); + } return Err(format!( - "timed out waiting for orchestrator prompt loop; try connecting to {fallback_addr} if it is still booting" + "timed out waiting for orchestrator prompt RPC at {rpc_addr}" )); } + thread::sleep(Duration::from_millis(100)); } } - fn shutdown(&mut self) { + fn shutdown(&mut self, interrupt: bool) -> bool { if self.cleaned { - return; + return false; } self.cleaned = true; - if let Some(stdin) = self.stdin.as_mut() { + if let Some(mut stdin) = self.stdin.take() { let _ = writeln!(stdin, "shutdown"); let _ = stdin.flush(); } + let timeout = if interrupt { + INTERRUPT_SHUTDOWN_TIMEOUT + } else { + ORCH_SHUTDOWN_TIMEOUT + }; let start = Instant::now(); - while start.elapsed() < ORCH_SHUTDOWN_TIMEOUT { + while start.elapsed() < timeout { match self.child.try_wait() { - Ok(Some(_)) => return, + Ok(Some(_)) => return false, Ok(None) => thread::sleep(Duration::from_millis(100)), Err(_) => break, } } - terminate_process_group(self.child.id()); + terminate_process_group(self.child.id(), FORCE_KILL_DELAY); let _ = self.child.wait(); + true } } impl Drop for OrchChild { fn drop(&mut self) { - self.shutdown(); + let _ = self.shutdown(false); } } -fn prepare_runtime(config: &Config) -> Result<(), String> { - run_status( - cargo_command(), - &[ - "build", - "-p", - "mvp-system", - "--features", - "local-e2e", - "--bin", - "mvp-orch-one-node", - ], - "build mvp-orch-one-node", - )?; +fn prepare_runtime(config: &Config) -> Result { + ensure_orch_binary(config)?; if !config.build_image { - eprintln!("mvp-one-node-chat: skipping node image rebuild (--no-build-image)"); - return Ok(()); + eprintln!("mvp-one-node-chat: skipping node image preparation (--no-build-image)"); + return Ok(config.node_image.clone()); } - run_status( - cargo_command(), - &["build", "-p", "mvp-system", "--bin", "mvp-node"], - "build mvp-node", - )?; - if !docker_image_exists(BASE_NODE_IMAGE) { - run_status( - "docker", - &[ - "build", - "-f", - "apps/mvp-node/Dockerfile.base", - "-t", - BASE_NODE_IMAGE, - ".", - ], - "build mvp node base image", - )?; + if let Some(cached_model) = &config.cached_model { + eprintln!( + "mvp-one-node-chat: using cached model {}", + cached_model.display_path.display() + ); } - let node_bin = node_bin_for_current_profile()?; - run_status( - "docker", - &[ - "build", - "-f", - "apps/mvp-node/Dockerfile", - "--build-arg", - &format!("BASE_IMAGE={BASE_NODE_IMAGE}"), - "--build-arg", - &format!("MVP_NODE_BIN={}", node_bin.display()), - "-t", - &config.node_image, - ".", - ], - "build mvp node image", - ) + let prepared = prepare_node_image(NodeImageRequest { + requested_image: config.node_image.clone(), + base_image: BASE_NODE_IMAGE.to_owned(), + node_bin: node_bin_for_current_profile()?, + provider: node_image_provider(config.provider)?, + extra_tag: config.image_tag.clone(), + push: config.push_image, + force_refresh: config.force_image_refresh, + enabled: true, + })?; + eprintln!( + "mvp-one-node-chat: using node image {} ({})", + prepared.image_ref, prepared.tag + ); + Ok(prepared.image_ref) } fn run_chat_loop(addr: &str, max_tokens: u32, timeout_ms: u64) -> Result<(), String> { @@ -291,19 +597,21 @@ fn run_chat_loop(addr: &str, max_tokens: u32, timeout_ms: u64) -> Result<(), Str }); let mut next_request_id = 1_u64; eprintln!( - "mvp-one-node-chat: Ctrl-C cleans up the orchestrator and Docker node; /exit exits cleanly" + "mvp-one-node-chat: Ctrl-C cleans up the orchestrator and provider node; /exit exits cleanly" ); loop { if STOP_REQUESTED.load(Ordering::SeqCst) { + acknowledge_stop(); return Ok(()); } - print!("> "); + print!("prompt:> "); io::stdout() .flush() .map_err(|e| format!("flush prompt: {e}"))?; let prompt = loop { if STOP_REQUESTED.load(Ordering::SeqCst) { + acknowledge_stop(); return Ok(()); } match input_rx.recv_timeout(Duration::from_millis(100)) { @@ -330,9 +638,12 @@ fn run_chat_loop(addr: &str, max_tokens: u32, timeout_ms: u64) -> Result<(), Str timeout_ms, }, )?; + println!("decoding..."); + let mut response_started = false; loop { if STOP_REQUESTED.load(Ordering::SeqCst) { + acknowledge_stop(); return Ok(()); } let mut line = String::new(); @@ -356,29 +667,30 @@ fn run_chat_loop(addr: &str, max_tokens: u32, timeout_ms: u64) -> Result<(), Str request_id: seen, text, } if seen == request_id => { + if !response_started { + print!("Response: "); + response_started = true; + } print!("{text}"); io::stdout() .flush() .map_err(|e| format!("flush response text: {e}"))?; } PromptEvent::Done { - request_id: seen, - tokens_generated, - elapsed_ms, - .. + request_id: seen, .. } if seen == request_id => { - println!(); - eprintln!( - "mvp-one-node-chat: done request={} tokens={} elapsed_ms={}", - seen, tokens_generated, elapsed_ms - ); + if response_started { + println!(); + } else { + println!("Response: "); + } break; } PromptEvent::Fault { request_id: seen, error, } if seen == request_id => { - eprintln!("mvp-one-node-chat: fault request={seen}: {error}"); + eprintln!("error: {error}"); break; } _ => {} @@ -411,6 +723,142 @@ fn cargo_command() -> &'static str { "cargo" } +fn ensure_orch_binary(config: &Config) -> Result<(), String> { + let default_orch = default_orch_bin()?; + if config.orch_bin != default_orch { + if config.orch_bin.is_file() { + let display_orch = display_workspace_path(&workspace_root(), &config.orch_bin); + eprintln!( + "mvp-one-node-chat: using custom orchestrator binary {display_orch}; skipping cargo build" + ); + return Ok(()); + } + let display_orch = display_workspace_path(&workspace_root(), &config.orch_bin); + return Err(format!( + "custom orchestrator binary {display_orch} does not exist" + )); + } + + let root = workspace_root(); + let rebuild_needed = orch_rebuild_needed(&config.orch_bin, &root, ORCH_REBUILD_INPUTS)?; + let dashboard_feature_stale = + config.dashboard && orch_local_e2e_marker_stale(&config.orch_bin, &root)?; + if !rebuild_needed && !dashboard_feature_stale { + eprintln!("mvp-one-node-chat: mvp-orch-one-node is up to date; skipping cargo build"); + return Ok(()); + } + + run_status( + cargo_command(), + &[ + "build", + "--quiet", + "-p", + "mvp-system", + "--features", + "local-e2e", + "--bin", + "mvp-orch-one-node", + ], + "build mvp-orch-one-node", + )?; + write_orch_local_e2e_marker(&config.orch_bin, &root) +} + +fn orch_rebuild_needed(bin: &Path, root: &Path, inputs: &[&str]) -> Result { + 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 orch_local_e2e_marker(bin: &Path) -> PathBuf { + let mut marker = bin.to_path_buf(); + let file_name = bin + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("mvp-orch-one-node"); + marker.set_file_name(format!("{file_name}.local-e2e")); + marker +} + +fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result { + if !bin.is_file() { + return Ok(true); + } + let marker = orch_local_e2e_marker(bin); + if !marker.is_file() { + return Ok(true); + } + Ok(modified_time(root, &marker)? < modified_time(root, bin)?) +} + +fn write_orch_local_e2e_marker(bin: &Path, root: &Path) -> Result<(), String> { + let marker = orch_local_e2e_marker(bin); + let display = display_workspace_path(root, &marker); + fs::write(&marker, b"local-e2e\n").map_err(|e| format!("write {display}: {e}")) +} + +fn latest_mtime(root: &Path, path: &Path) -> Result { + 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 { + 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 display_user_path(path: &Path) -> String { + let root = workspace_root(); + if let Ok(relative) = path.strip_prefix(&root) { + if relative.as_os_str().is_empty() { + return ".".to_owned(); + } + return format!("./{}", relative.display()); + } + if let Ok(cwd) = std::env::current_dir() { + if let Ok(relative) = path.strip_prefix(&cwd) { + if relative.as_os_str().is_empty() { + return ".".to_owned(); + } + return format!("./{}", relative.display()); + } + } + path.display().to_string() +} + fn run_status(program: &str, args: &[&str], label: &str) -> Result<(), String> { eprintln!("mvp-one-node-chat: {label}"); let status = Command::new(program) @@ -427,17 +875,6 @@ fn run_status(program: &str, args: &[&str], label: &str) -> Result<(), String> { } } -fn docker_image_exists(image: &str) -> bool { - Command::new("docker") - .args(["image", "inspect", image]) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map(|status| status.success()) - .unwrap_or(false) -} - extern "C" fn request_stop(_: libc::c_int) { STOP_REQUESTED.store(true, Ordering::SeqCst); } @@ -450,17 +887,153 @@ fn install_signal_handlers() { } } -fn terminate_process_group(pid: u32) { +fn acknowledge_stop() { + if !STOP_ACKNOWLEDGED.swap(true, Ordering::SeqCst) { + eprintln!("mvp-one-node-chat: Ctrl-C received; stopping runtime..."); + } +} + +fn report_interrupt_shutdown(forced: bool) { + if forced { + eprintln!("mvp-one-node-chat: runtime did not stop in 5s; force-killed"); + } else { + eprintln!("mvp-one-node-chat: runtime stopped"); + } +} + +fn terminate_process_group(pid: u32, kill_after: Duration) { #[cfg(target_os = "linux")] unsafe { let pgid = -(pid as libc::pid_t); let _ = libc::kill(pgid, libc::SIGTERM); - thread::sleep(Duration::from_secs(2)); + thread::sleep(kill_after); let _ = libc::kill(pgid, libc::SIGKILL); } #[cfg(not(target_os = "linux"))] { - let _ = pid; + let _ = (pid, kill_after); + } +} + +#[derive(Clone, Debug)] +struct CachedModelConfig { + host_path: PathBuf, + display_path: PathBuf, +} + +impl CachedModelConfig { + fn from_arg(path: Option) -> Result { + let requested = path + .map(PathBuf::from) + .unwrap_or_else(default_cached_model_path); + let host_path = requested + .canonicalize() + .map_err(|e| format!("resolve --cached-model path {}: {e}", requested.display()))?; + let metadata = fs::metadata(&host_path) + .map_err(|e| format!("stat cached model {}: {e}", requested.display()))?; + if !metadata.is_file() { + return Err(format!( + "--cached-model must point at a file: {}", + requested.display() + )); + } + Ok(Self { + host_path, + display_path: requested, + }) + } +} + +fn default_cached_model_path() -> PathBuf { + PathBuf::from(".") + .join(REPO_MODEL_CACHE_DIR) + .join(DEFAULT_CACHED_MODEL_FILE) +} + +fn workspace_root() -> PathBuf { + let git_root = Command::new("git") + .args(["rev-parse", "--show-toplevel"]) + .stdin(Stdio::null()) + .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") +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RuntimeConfigProfile { + Local, + Deploy, +} + +impl RuntimeConfigProfile { + fn from_env() -> Result { + match env_optional(MVP_RUNTIME_CONFIG_ENV).as_deref() { + None | Some("local") => Ok(Self::Local), + Some("deploy") => Ok(Self::Deploy), + Some(other) => Err(format!( + "unsupported {MVP_RUNTIME_CONFIG_ENV}={other:?}; use local or deploy" + )), + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Local => "local", + Self::Deploy => "deploy", + } + } + + fn default_provider(self) -> ProviderKind { + match self { + Self::Local => ProviderKind::Docker, + Self::Deploy => ProviderKind::VastAi, + } + } +} + +fn provider_from_env(config_profile: RuntimeConfigProfile) -> Result { + match env_optional("MVP_NODE_PROVIDER").or_else(|| env_optional("MVP_PROVIDER")) { + Some(value) => ProviderKind::parse_deploy(&value), + None => Ok(config_profile.default_provider()), + } +} + +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn relay_mode_from_env() -> Result { + match env_optional("MVP_IROH_RELAY_MODE") + .as_deref() + .unwrap_or("default") + { + "disabled" => Ok(iroh::RelayMode::Disabled), + "default" => Ok(iroh::RelayMode::Default), + other => Err(format!( + "unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default" + )), + } +} + +fn relay_mode_env_value(mode: &iroh::RelayMode) -> &'static str { + match mode { + iroh::RelayMode::Disabled => "disabled", + _ => "default", + } +} + +fn node_image_provider(provider: ProviderKind) -> Result { + match provider { + ProviderKind::Docker => Ok(NodeImageProvider::Docker), + ProviderKind::VastAi => Ok(NodeImageProvider::VastAi), + ProviderKind::Mock => Err("mvp-one-node-chat does not support mock provider".to_owned()), } } @@ -510,3 +1083,120 @@ where .parse::() .map_err(|e| format!("invalid {name}={value:?}: {e}")) } + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering}; + + static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0); + + struct TempWorkspace { + root: PathBuf, + } + + impl TempWorkspace { + fn new(name: &str) -> Self { + let counter = TEMP_COUNTER.fetch_add(1, AtomicOrdering::Relaxed); + let root = std::env::temp_dir().join(format!( + "mvp-one-node-chat-{name}-{}-{counter}", + std::process::id() + )); + let _ = fs::remove_dir_all(&root); + fs::create_dir_all(&root).expect("create temp workspace"); + Self { root } + } + + fn path(&self, relative: &str) -> PathBuf { + self.root.join(relative) + } + + fn write(&self, relative: &str, contents: &[u8]) -> PathBuf { + let path = self.path(relative); + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).expect("create temp parent directory"); + } + fs::write(&path, contents).expect("write temp file"); + path + } + } + + impl Drop for TempWorkspace { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.root); + } + } + + fn write_file_newer_than(path: &Path, contents: &[u8], older_than: SystemTime) { + if let Some(parent) = path.parent() { + fs::create_dir_all(parent).expect("create temp parent directory"); + } + + let started = Instant::now(); + loop { + fs::write(path, contents).expect("write temp file"); + let mtime = fs::metadata(path) + .expect("stat temp file") + .modified() + .expect("read temp file mtime"); + if mtime > older_than { + return; + } + assert!( + started.elapsed() <= Duration::from_secs(3), + "filesystem did not record a newer mtime for {}", + path.display() + ); + thread::sleep(Duration::from_millis(20)); + } + } + + #[test] + fn orch_rebuild_missing_binary_requires_rebuild() { + let workspace = TempWorkspace::new("missing-binary"); + workspace.write("src/main.rs", b"fn main() {}\n"); + let bin = workspace.path("target/debug/mvp-orch-one-node"); + + let needed = orch_rebuild_needed(&bin, &workspace.root, &["src/main.rs"]) + .expect("missing binary check succeeds"); + + assert!(needed, "missing orchestrator binary must trigger rebuild"); + } + + #[test] + fn orch_rebuild_binary_newer_than_input_skips_rebuild() { + let workspace = TempWorkspace::new("fresh-binary"); + let input = workspace.write("src/main.rs", b"fn main() {}\n"); + let input_mtime = modified_time(&workspace.root, &input).expect("read input mtime"); + let bin = workspace.path("target/debug/mvp-orch-one-node"); + write_file_newer_than(&bin, b"orchestrator binary\n", input_mtime); + + let needed = orch_rebuild_needed(&bin, &workspace.root, &["src/main.rs"]) + .expect("fresh binary check succeeds"); + + assert!( + !needed, + "binary newer than every tracked input must skip rebuild" + ); + } + + #[test] + fn orch_rebuild_nested_directory_input_newer_than_binary_requires_rebuild() { + let workspace = TempWorkspace::new("nested-newer-input"); + let nested_input = workspace.write("src/nested/orchestrator.rs", b"old source\n"); + let src_mtime = + latest_mtime(&workspace.root, &workspace.path("src")).expect("read source tree mtime"); + let bin = workspace.path("target/debug/mvp-orch-one-node"); + write_file_newer_than(&bin, b"orchestrator binary\n", src_mtime); + let bin_mtime = modified_time(&workspace.root, &bin).expect("read binary mtime"); + write_file_newer_than(&nested_input, b"new source\n", bin_mtime); + + let needed = orch_rebuild_needed(&bin, &workspace.root, &["src"]) + .expect("stale binary check succeeds"); + + assert!( + needed, + "newer file inside a tracked directory must trigger rebuild" + ); + } +} diff --git a/crates/mvp-system/src/bin/mvp_orch_one_node.rs b/crates/mvp-system/src/bin/mvp_orch_one_node.rs index f145b2a..9b1bf47 100644 --- a/crates/mvp-system/src/bin/mvp_orch_one_node.rs +++ b/crates/mvp-system/src/bin/mvp_orch_one_node.rs @@ -1,11 +1,13 @@ -use std::io::{BufRead, BufReader}; +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Write}; use std::net::{SocketAddr, TcpListener, TcpStream}; +use std::path::{Path, PathBuf}; use std::process::ExitCode; use std::sync::{Arc, mpsc}; use std::thread; use std::time::{Duration, Instant}; -use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, NodeId, Position, StreamId}; +use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, Mux, NodeId, StreamId}; use distribution::node::DistributedNodeConfig; use iroh::EndpointAddr; use iroh_driver::{IrohDriver, IrohDriverConfig}; @@ -14,21 +16,35 @@ use mvp_system::actors::register_mvp_actor_codecs; #[cfg(feature = "local-e2e")] use mvp_system::dashboard_view::MvpClusterDashboardView; use mvp_system::distribution_stack::DistributionRuntimeStack; +use mvp_system::node_provisioning::ProviderKind; use mvp_system::prompt_rpc::{PromptEvent, SubmitPrompt, read_submit_prompt, write_json_line}; use mvp_system::provisioning::{ LocalDockerPlugin, NodeProvisionSpec, PluginObservation, PluginObservationSink, PluginSink, - ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, ProvisionPlugin, + ProviderMount, ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, + ProvisionPlugin, +}; +#[cfg(test)] +use mvp_system::relay_provisioning::SWACTOR_IROH_RELAY_URL_ENV; +use mvp_system::relay_provisioning::{ + MVP_IROH_RELAY_URL_ENV, RelayRuntimeConfig, relay_mode_env_value, relay_runtime_config_from_env, }; use mvp_system::run_plan::{GgufSource, TokenizerSource}; use mvp_system::telemetry::{ MVP_PROVISIONING_EVENTS, MvpProvisionEventRecord, MvpProvisionLogRecord, mvp_provision_log_channel, }; +use mvp_system::vastai_provisioning::{ + SshCommandBootstrapLauncher, ToolsVastAiLeaseClient, VastAiProvisioningConfig, + VastAiProvisioningPlugin, +}; use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::ActorAddress; const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest"; +const MVP_RUNTIME_CONFIG_ENV: &str = "MVP_RUNTIME_CONFIG"; +const CACHED_MODEL_HOST_ENV: &str = "MVP_CACHED_MODEL_HOST_PATH"; +const CACHED_MODEL_CONTAINER_DIR: &str = "/models/cached"; const DEFAULT_RPC_BIND: &str = "127.0.0.1:19777"; const DEFAULT_HF_REPO: &str = "bartowski/Llama-3.2-1B-Instruct-GGUF"; const DEFAULT_HF_FILE: &str = "Llama-3.2-1B-Instruct-Q4_K_M.gguf"; @@ -39,6 +55,9 @@ const PUMP_INTERVAL: Duration = Duration::from_millis(10); const BOOT_TIMEOUT: Duration = Duration::from_secs(180); const ROUTE_TIMEOUT: Duration = Duration::from_secs(30); const WEIGHT_TIMEOUT: Duration = Duration::from_secs(900); +const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; +const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; +const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; fn main() -> ExitCode { match run() { @@ -52,18 +71,84 @@ fn main() -> ExitCode { fn run() -> Result<(), String> { let config = Config::from_env_and_args()?; - let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; - let mut driver = IrohDriver::with_handle( + let mut orch_datastream = + OrchDatastream::new(config.run_id, config.datastream_frame_log.as_deref())?; + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "config", + "ready", + json!({ + "config_profile":config.config_profile.as_str(), + "image":&config.image, + "provider":config.provider.as_str(), + "rpc_bind":config.rpc_bind.to_string(), + "model_id":&config.model_id, + "stage_index":config.stage_index, + "layer_end_exclusive":config.layer_end_exclusive, + "relay_mode":format!("{:?}", config.relay.mode), + "provider_config":config.provider_datastream_detail(), + }), + ); + + let tokio = match tokio::runtime::Runtime::new() { + Ok(runtime) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "tokio_runtime", + "ready", + json!({"runtime":"tokio"}), + ); + runtime + } + Err(error) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "tokio_runtime", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("tokio runtime: {error}")); + } + }; + let mut driver = match IrohDriver::with_handle( tokio.handle().clone(), IrohDriverConfig { secret_key: None, - relay_mode: config.relay_mode.clone(), + relay_mode: config.relay.mode.clone(), node: DistributedNodeConfig::default(), peer_auth: None, additional_alpns: vec![], }, - ) - .map_err(|e| format!("create iroh driver: {e}"))?; + ) { + Ok(driver) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "iroh_driver", + "ready", + json!({"relay_mode":format!("{:?}", config.relay.mode)}), + ); + driver + } + Err(error) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "iroh_driver", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("create iroh driver: {error}")); + } + }; let stack = DistributionRuntimeStack::new_with_codecs( driver.node_id(), DistributedNodeConfig::default(), @@ -72,6 +157,22 @@ fn run() -> Result<(), String> { datastream::wire::register_datastream_codec(registry); }, ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "distribution_stack", + "ready", + json!({"actors":"initialized","route_view":"initialized","swim":"initialized"}), + ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "codecs", + "ready", + json!({"registered":["node_agent","orchestrator","provisioner","prompt_rpc","datastream"]}), + ); driver.enable_actor_bridge( stack.runtime.clone(), stack.codec.clone(), @@ -80,82 +181,344 @@ fn run() -> Result<(), String> { stack.relay_mirror.clone(), stack.route_view.clone(), ); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "actor_bridge", + "ready", + json!({"transport":"iroh","routes":"attached"}), + ); let (frame_tx, frame_rx) = mpsc::channel::<(StreamId, Frame)>(); - let datastream_sink = stack + let datastream_sink = match stack .runtime .spawn(DatastreamSink::new(move |stream, frame| { let _ = frame_tx.send((stream, frame)); - })) - .map_err(|e| format!("spawn datastream sink: {e}"))?; + })) { + Ok(actor) => actor, + Err(error) => { + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "datastream_sink", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("spawn datastream sink: {error}")); + } + }; stack.register_local_actor(driver.register_actor(datastream_sink, 1)); + orch_datastream.emit_bootstrap( + None, + config.run_id, + config.node_id, + "datastream_sink", + "ready", + json!({"actor":datastream_sink,"channel":"local_mpsc"}), + ); let dashboard = DashboardSupport::start_from_env()?; - let mut orch_datastream = OrchDatastream::new(config.run_id); - - let prompt_events = stack - .runtime - .new_inbox::() - .map_err(|e| format!("prompt event inbox: {e}"))?; - let prompt_reply_actor = *prompt_events.addr(); - stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); - - let (work_tx, work_rx) = mpsc::channel::(); - let rpc_addr = spawn_prompt_rpc( - config.rpc_bind, - work_tx, - config.default_max_tokens, - config.default_timeout_ms, - )?; - println!( - "{}", - json!({"type":"prompt_rpc_ready","addr":rpc_addr.to_string()}) + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "dashboard", + "ready", + json!({"enabled":dashboard.is_some()}), ); - let mut docker = LocalDockerPlugin::new("mvp-orch-one-node"); + let prompt_events = match stack.runtime.new_inbox::() { + Ok(inbox) => inbox, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_reply_actor", + "failed", + json!({"error":error.to_string()}), + ); + return Err(format!("prompt event inbox: {error}")); + } + }; + let prompt_reply_actor = *prompt_events.addr(); + stack.register_local_actor(driver.register_actor(prompt_reply_actor, 1)); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_reply_actor", + "ready", + json!({"actor":prompt_reply_actor}), + ); + + let (work_tx, work_rx) = mpsc::channel::(); + let stop_rx = spawn_stop_listener(); + + let mut provisioner = config.build_provisioner()?; + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_provisioner", + "ready", + json!({ + "provider":config.provider.as_str(), + "owner":"mvp-orch-one-node", + "config":config.provider_datastream_detail(), + }), + ); let (obs_tx, obs_rx) = mpsc::channel::(); let sink = PluginSink::new(Arc::new(ChannelObservationSink { tx: Mutex::new(obs_tx), })); + let node_spec = config.node_spec(driver.endpoint_addr(), datastream_sink)?; orch_datastream.emit_event( dashboard.as_ref(), ProvisionEvent { run_id: config.run_id, node_id: config.node_id, kind: ProvisionEventKind::ProvisionStart, - message: Some(format!("starting Docker image {}", config.image)), + provider: Some(config.provider.as_str().to_owned()), + message: Some(format!( + "starting {} image {}", + config.provider.as_str(), + config.image + )), }, ); - let handle = docker.start_node( - config.node_spec(driver.endpoint_addr(), datastream_sink)?, - sink, - )?; + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_spec", + "ready", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "relay_mode":relay_mode_env_value(&config.relay.mode), + "docker_gpus":if config.provider == ProviderKind::Docker { Some(config.docker_gpus.as_str()) } else { None }, + "provider_config":config.provider_datastream_detail(), + "env_keys":config.node_spec_env_keys(), + }), + ); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_start", + "started", + json!({ + "provider":config.provider.as_str(), + "image":&config.image, + "node_id":config.node_id, + "stage_index":config.stage_index, + }), + ); + let handle = match provisioner.start_node(node_spec, sink) { + Ok(handle) => handle, + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_start", + "failed", + json!({"provider":config.provider.as_str(),"error":error}), + ); + return Err(error); + } + }; - let ready = wait_for_runtime_ready( + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_runtime_ready", + "started", + json!({"timeout_ms":BOOT_TIMEOUT.as_millis()}), + ); + let ready = match wait_for_runtime_ready( &mut driver, &stack, &obs_rx, &frame_rx, dashboard.as_ref(), &mut orch_datastream, - )?; + config.provider, + ) { + Ok(ready) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_runtime_ready", + "ready", + json!({"endpoint":&ready.endpoint,"node_actor":ready.node_actor}), + ); + ready + } + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_runtime_ready", + "failed", + json!({"error":error}), + ); + return Err(error); + } + }; driver.join(std::slice::from_ref(&ready.endpoint)); - wait_for_route(&mut driver, &stack, ready.node_actor)?; - provision_stage(&stack, ready.node_actor, &config)?; - wait_for_weights_loaded( + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_join", + "ready", + json!({"endpoint":&ready.endpoint}), + ); + match wait_for_route(&mut driver, &stack, ready.node_actor) { + Ok(()) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_route", + "ready", + json!({"node_actor":ready.node_actor,"timeout_ms":ROUTE_TIMEOUT.as_millis()}), + ), + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "node_route", + "failed", + json!({"node_actor":ready.node_actor,"timeout_ms":ROUTE_TIMEOUT.as_millis(),"error":error}), + ); + return Err(error); + } + } + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "stage_provision", + "started", + json!({ + "run_id":config.run_id, + "node_id":config.node_id, + "stage_index":config.stage_index, + "stage_count":1, + "layer_range":{"start":0,"end_exclusive":config.layer_end_exclusive}, + "model_id":&config.model_id, + }), + ); + match provision_stage(&stack, ready.node_actor, &config) { + Ok(()) => {} + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "stage_provision", + "failed", + json!({"error":error}), + ); + return Err(error); + } + } + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "weights_loaded", + "started", + json!({"model_id":&config.model_id,"timeout_ms":WEIGHT_TIMEOUT.as_millis()}), + ); + match wait_for_weights_loaded( &mut driver, &stack, &obs_rx, &frame_rx, dashboard.as_ref(), &mut orch_datastream, - )?; + config.provider, + ) { + Ok(()) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "weights_loaded", + "ready", + json!({"source_channel":"mvp.worker.weights","model_id":&config.model_id}), + ), + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "weights_loaded", + "failed", + json!({"error":error}), + ); + return Err(error); + } + } - println!( - "{}", - json!({"type":"prompt_loop_ready","addr":rpc_addr.to_string(),"node_actor":ready.node_actor}) + let rpc_addr = match spawn_prompt_rpc( + config.rpc_bind, + work_tx, + config.default_max_tokens, + config.default_timeout_ms, + ) { + Ok(addr) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_rpc", + "ready", + json!({ + "addr":addr.to_string(), + "default_max_tokens":config.default_max_tokens, + "default_timeout_ms":config.default_timeout_ms, + }), + ); + addr + } + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_rpc", + "failed", + json!({"error":error}), + ); + return Err(error); + } + }; + + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "prompt_loop", + "ready", + json!({"addr":rpc_addr.to_string(),"node_actor":ready.node_actor}), ); - let stop_rx = spawn_stop_listener(); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "serve_prompts", + "started", + json!({"mode":"single_active_prompt","poll_interval_ms":PUMP_INTERVAL.as_millis()}), + ); let result = serve_prompts( &mut driver, &stack, @@ -166,17 +529,225 @@ fn run() -> Result<(), String> { &stop_rx, dashboard.as_ref(), &mut orch_datastream, + config.run_id, + config.node_id, ready.node_actor, prompt_reply_actor, + config.provider, ); - let stop_result = docker.stop_node(&handle); + if let Err(error) = &result { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "serve_prompts", + "failed", + json!({"error":error}), + ); + } + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "started", + json!({"provider":config.provider.as_str(),"node_id":config.node_id}), + ); + let stop_result = provisioner.stop_node(&handle); + match &stop_result { + Ok(()) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "ready", + json!({"provider":config.provider.as_str(),"node_id":config.node_id}), + ), + Err(error) => orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "provider_stop", + "failed", + json!({"provider":config.provider.as_str(),"node_id":config.node_id,"error":error}), + ), + } + if result.is_ok() && stop_result.is_ok() { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "orch_exit", + "ready", + json!({"result":"ok"}), + ); + } result.and(stop_result) } +#[derive(Clone)] +struct VastAiRuntimeConfig { + api_key: Option, + provisioning: VastAiProvisioningConfig, + bootstrap_command: Option, +} + +impl VastAiRuntimeConfig { + fn from_env() -> Result { + let mut provisioning = VastAiProvisioningConfig::default(); + provisioning.disk_gb = env_u32("MVP_VASTAI_DISK_GB", provisioning.disk_gb)?; + provisioning.ssh_user = env_string("MVP_VASTAI_SSH_USER", &provisioning.ssh_user); + provisioning.confirm_lease = + env_bool("MVP_VASTAI_CONFIRM_LEASE", provisioning.confirm_lease)?; + provisioning.onstart = env_optional("MVP_VASTAI_ONSTART"); + provisioning.selection.gpu_name = env_optional("MVP_VASTAI_GPU_NAME"); + if let Some(min_gpu_ram_mb) = env_optional_u64("MVP_VASTAI_MIN_GPU_RAM_MB")? { + provisioning.selection.min_gpu_ram_mb = Some(min_gpu_ram_mb); + } + if let Some(min_down_mbps) = env_optional_f64("MVP_VASTAI_MIN_DOWN_MBPS")? { + provisioning.selection.min_down_mbps = min_down_mbps; + } + if let Some(min_up_mbps) = env_optional_f64("MVP_VASTAI_MIN_UP_MBPS")? { + provisioning.selection.min_up_mbps = Some(min_up_mbps); + } + if let Some(min_reliability) = env_optional_f64("MVP_VASTAI_MIN_RELIABILITY")? { + provisioning.selection.min_reliability = min_reliability; + } + provisioning.selection.require_verified = env_bool( + "MVP_VASTAI_REQUIRE_VERIFIED", + provisioning.selection.require_verified, + )?; + if let Some(poll_interval_secs) = env_optional_u64("MVP_VASTAI_POLL_INTERVAL_SECS")? { + provisioning.lifecycle.poll_interval = Duration::from_secs(poll_interval_secs); + } + if let Some(max_polls) = env_optional_u32("MVP_VASTAI_MAX_POLLS")? { + provisioning.lifecycle.max_polls = max_polls; + } + if let Some(max_create_attempts) = env_optional_u32("MVP_VASTAI_MAX_CREATE_ATTEMPTS")? { + provisioning.lifecycle.max_create_attempts = max_create_attempts; + } + Ok(Self { + api_key: env_optional("MVP_VASTAI_API_KEY").or_else(|| env_optional("VASTAI_API_KEY")), + provisioning, + bootstrap_command: env_optional("MVP_VASTAI_BOOTSTRAP_COMMAND"), + }) + } + + fn datastream_detail(&self) -> Value { + json!({ + "disk_gb": self.provisioning.disk_gb, + "ssh_user": &self.provisioning.ssh_user, + "gpu_name": &self.provisioning.selection.gpu_name, + "min_gpu_ram_mb": self.provisioning.selection.min_gpu_ram_mb, + "min_down_mbps": self.provisioning.selection.min_down_mbps, + "min_up_mbps": self.provisioning.selection.min_up_mbps, + "min_reliability": self.provisioning.selection.min_reliability, + "require_verified": self.provisioning.selection.require_verified, + "confirm_lease": self.provisioning.confirm_lease, + "has_api_key": self.api_key.is_some(), + "has_onstart": self.provisioning.onstart.is_some(), + "has_bootstrap_command": self.bootstrap_command.is_some(), + }) + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RuntimeConfigProfile { + Local, + Deploy, +} + +impl RuntimeConfigProfile { + fn from_env() -> Result { + match env_optional(MVP_RUNTIME_CONFIG_ENV).as_deref() { + None | Some("local") => Ok(Self::Local), + Some("deploy") => Ok(Self::Deploy), + Some(other) => Err(format!( + "unsupported {MVP_RUNTIME_CONFIG_ENV}={other:?}; use local or deploy" + )), + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Local => "local", + Self::Deploy => "deploy", + } + } + + fn default_provider(self) -> ProviderKind { + match self { + Self::Local => ProviderKind::Docker, + Self::Deploy => ProviderKind::VastAi, + } + } +} + +#[derive(Clone)] +struct CachedModelConfig { + host_path: PathBuf, + container_path: String, +} + +impl CachedModelConfig { + fn from_env(provider: ProviderKind) -> Result, String> { + let Some(raw_host_path) = env_optional(CACHED_MODEL_HOST_ENV) else { + return Ok(None); + }; + if provider != ProviderKind::Docker { + return Err(format!( + "{CACHED_MODEL_HOST_ENV} is a host-local cache path and requires provider=docker" + )); + } + let requested = PathBuf::from(&raw_host_path); + let host_path = requested.canonicalize().map_err(|e| { + format!( + "resolve {CACHED_MODEL_HOST_ENV} path {}: {e}", + requested.display() + ) + })?; + if !host_path.is_file() { + return Err(format!( + "{CACHED_MODEL_HOST_ENV} must point at a file: {}", + host_path.display() + )); + } + let container_path = cached_model_container_path(&host_path)?; + Ok(Some(Self { + host_path, + container_path, + })) + } + + fn datastream_detail(&self) -> Value { + json!({ + "host_path_present": true, + "file": self.host_path.file_name().and_then(|name| name.to_str()), + "container_path": &self.container_path, + }) + } +} + +fn cached_model_container_path(host_path: &Path) -> Result { + let file_name = host_path + .file_name() + .and_then(|name| name.to_str()) + .filter(|name| !name.is_empty()) + .ok_or_else(|| { + format!( + "cached model path has no file name: {}", + host_path.display() + ) + })?; + Ok(format!("{CACHED_MODEL_CONTAINER_DIR}/{file_name}")) +} + #[derive(Clone)] struct Config { + config_profile: RuntimeConfigProfile, image: String, docker_gpus: String, + provider: ProviderKind, rpc_bind: SocketAddr, run_id: u64, node_id: u64, @@ -187,28 +758,54 @@ struct Config { tokenizer: TokenizerSource, default_max_tokens: u32, default_timeout_ms: u64, - relay_mode: iroh::RelayMode, + relay: RelayRuntimeConfig, + vastai: Option, + cached_model: Option, + datastream_frame_log: Option, } impl Config { fn from_env_and_args() -> Result { - let mut args = std::env::args().skip(1); + Self::from_env_and_args_iter(std::env::args().skip(1)) + } + + fn from_env_and_args_iter(args: impl IntoIterator) -> Result { + let mut args = args.into_iter(); + let config_profile = RuntimeConfigProfile::from_env()?; + let provider = provider_from_env(config_profile)?; + let cached_model = CachedModelConfig::from_env(provider)?; + let vastai = if provider == ProviderKind::VastAi { + Some(VastAiRuntimeConfig::from_env()?) + } else { + None + }; + let gguf_source = cached_model + .as_ref() + .map(|cached_model| GgufSource::LocalPath(cached_model.container_path.clone())) + .unwrap_or_else(gguf_source_from_env); + let run_id = env_u64("MVP_RUN_ID", 1)?; + let relay = relay_runtime_config_from_env(run_id)?; let mut config = Self { + config_profile, image: env_string("MVP_NODE_IMAGE", DEFAULT_IMAGE), docker_gpus: env_string("MVP_DOCKER_GPUS", "all"), + provider, rpc_bind: env_string("MVP_PROMPT_RPC_BIND", DEFAULT_RPC_BIND) .parse() .map_err(|e| format!("invalid MVP_PROMPT_RPC_BIND: {e}"))?, - run_id: env_u64("MVP_RUN_ID", 1)?, + run_id, node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, stage_index: env_u32("MVP_STAGE_INDEX", 0)?, layer_end_exclusive: env_u32("MVP_LAYER_END_EXCLUSIVE", 16)?, model_id: env_string("MVP_MODEL_ID", DEFAULT_MODEL_ID), - gguf_source: gguf_source_from_env(), + cached_model, + gguf_source, tokenizer: tokenizer_from_env(), default_max_tokens: env_u32("MVP_PROMPT_MAX_TOKENS", DEFAULT_MAX_TOKENS)?, default_timeout_ms: env_u64("MVP_PROMPT_TIMEOUT_MS", DEFAULT_TIMEOUT_MS)?, - relay_mode: relay_mode_from_env()?, + relay, + vastai, + datastream_frame_log: env_optional(DATASTREAM_FRAME_LOG_ENV).map(PathBuf::from), }; while let Some(arg) = args.next() { @@ -228,6 +825,12 @@ impl Config { "--timeout-ms" => { config.default_timeout_ms = parse_next(&mut args, "--timeout-ms")? } + "--datastream-frame-log" => { + config.datastream_frame_log = Some(PathBuf::from(next_arg( + &mut args, + "--datastream-frame-log", + )?)); + } "--model-id" => config.model_id = next_arg(&mut args, "--model-id")?, "--gguf-local-path" => { config.gguf_source = @@ -273,6 +876,108 @@ impl Config { Ok(config) } + fn provider_datastream_detail(&self) -> Value { + match self.provider { + ProviderKind::Docker => json!({ + "docker_gpus": &self.docker_gpus, + "cached_model": self.cached_model.as_ref().map(CachedModelConfig::datastream_detail), + }), + ProviderKind::VastAi => self + .vastai + .as_ref() + .map_or_else(|| json!({}), VastAiRuntimeConfig::datastream_detail), + ProviderKind::Mock => json!({}), + } + } + + fn build_provisioner(&self) -> Result, String> { + match self.provider { + ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new("mvp-orch-one-node"))), + ProviderKind::VastAi => { + let vastai = self.vastai.as_ref().ok_or_else(|| { + "VastAI config was not resolved for provider vastai".to_owned() + })?; + if vastai.bootstrap_command.is_none() { + return Err( + "MVP_VASTAI_BOOTSTRAP_COMMAND is required when MVP_NODE_PROVIDER=vastai" + .to_owned(), + ); + } + let api_key = vastai.api_key.clone().ok_or_else(|| { + "MVP_VASTAI_API_KEY or VASTAI_API_KEY is required when MVP_NODE_PROVIDER=vastai" + .to_owned() + })?; + let client = ToolsVastAiLeaseClient::from_api_key(api_key)?; + Ok(Box::new(VastAiProvisioningPlugin::new( + client, + SshCommandBootstrapLauncher, + vastai.provisioning.clone(), + ))) + } + ProviderKind::Mock => { + Err("mvp-orch-one-node does not support mock provider".to_owned()) + } + } + } + + fn node_spec_env_keys(&self) -> Vec<&'static str> { + let mut keys = vec![ + "MVP_RUN_ID", + "MVP_LOGICAL_NODE_ID", + "MVP_NODE_PROVIDER", + "MVP_STAGE_INDEX", + "MVP_COORDINATOR_ENDPOINT", + "MVP_DATASTREAM_SINK_ACTOR", + "MVP_MODEL_ID", + "MVP_NODE_MAX_RUNTIME_SECS", + "MVP_IROH_RELAY_MODE", + ]; + if self.relay.url.is_some() { + keys.push(MVP_IROH_RELAY_URL_ENV); + } + if self.provider == ProviderKind::Docker { + keys.push("MVP_DOCKER_GPUS"); + } + if std::env::var_os("MVP_TINYGRAD_TEST_MODE").is_some() { + keys.push("MVP_TINYGRAD_TEST_MODE"); + } + if std::env::var_os("MVP_CPU_LINE_PROFILE").is_some() { + keys.push("MVP_CPU_LINE_PROFILE"); + } + if std::env::var_os("MVP_CPU_LINE_PROFILE_INTERVAL_MS").is_some() { + keys.push("MVP_CPU_LINE_PROFILE_INTERVAL_MS"); + } + if std::env::var_os("MVP_GPU_SAMPLE").is_some() { + keys.push("MVP_GPU_SAMPLE"); + } + if std::env::var_os("MVP_TOKEN_PROGRESS_EVERY").is_some() { + keys.push("MVP_TOKEN_PROGRESS_EVERY"); + } + if std::env::var_os("CUDA_DEVICE_SCHEDULE").is_some() { + keys.push("CUDA_DEVICE_SCHEDULE"); + } + if std::env::var_os("MVP_MODEL_CACHE_DIR").is_some() { + keys.push("MVP_MODEL_CACHE_DIR"); + } + if std::env::var_os("HF_TOKEN").is_some() { + keys.push("HF_TOKEN"); + } + match &self.gguf_source { + GgufSource::LocalPath(_) => keys.push("MVP_GGUF_LOCAL_PATH"), + GgufSource::HuggingFaceGguf { revision, .. } => { + keys.push("MVP_GGUF_REPO"); + keys.push("MVP_GGUF_FILE"); + if revision.is_some() { + keys.push("MVP_GGUF_REVISION"); + } + } + } + if matches!(self.tokenizer, TokenizerSource::LocalPath(_)) { + keys.push("MVP_TOKENIZER_LOCAL_PATH"); + } + keys + } + fn node_spec( &self, coordinator: EndpointAddr, @@ -282,6 +987,10 @@ impl Config { ("MVP_RUN_ID".to_owned(), self.run_id.to_string()), ("MVP_LOGICAL_NODE_ID".to_owned(), self.node_id.to_string()), ("MVP_STAGE_INDEX".to_owned(), self.stage_index.to_string()), + ( + "MVP_NODE_PROVIDER".to_owned(), + self.provider.as_str().to_owned(), + ), ( "MVP_COORDINATOR_ENDPOINT".to_owned(), serde_json::to_string(&coordinator) @@ -294,9 +1003,23 @@ impl Config { ), ("MVP_MODEL_ID".to_owned(), self.model_id.clone()), ("MVP_NODE_MAX_RUNTIME_SECS".to_owned(), "0".to_owned()), - ("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone()), + ( + "MVP_IROH_RELAY_MODE".to_owned(), + relay_mode_env_value(&self.relay.mode).to_owned(), + ), ]; + if let Some(url) = &self.relay.url { + env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), url.clone())); + } + if self.provider == ProviderKind::Docker { + env.push(("MVP_DOCKER_GPUS".to_owned(), self.docker_gpus.clone())); + } env.extend(optional_env("MVP_TINYGRAD_TEST_MODE")); + env.extend(optional_env("MVP_CPU_LINE_PROFILE")); + env.extend(optional_env("MVP_CPU_LINE_PROFILE_INTERVAL_MS")); + env.extend(optional_env("MVP_GPU_SAMPLE")); + env.extend(optional_env("MVP_TOKEN_PROGRESS_EVERY")); + env.extend(optional_env("CUDA_DEVICE_SCHEDULE")); env.extend(optional_env("MVP_MODEL_CACHE_DIR")); env.extend(optional_env("HF_TOKEN")); match &self.gguf_source { @@ -318,13 +1041,35 @@ impl Config { if let TokenizerSource::LocalPath(path) = &self.tokenizer { env.push(("MVP_TOKENIZER_LOCAL_PATH".to_owned(), path.clone())); } + let args = match self.provider { + ProviderKind::VastAi => self + .vastai + .as_ref() + .and_then(|vastai| vastai.bootstrap_command.clone()) + .into_iter() + .collect(), + ProviderKind::Docker => Vec::new(), + ProviderKind::Mock => { + return Err("mvp-orch-one-node does not support mock provider".to_owned()); + } + }; + let mounts = if let Some(cached_model) = &self.cached_model { + vec![ProviderMount { + host_path: cached_model.host_path.to_string_lossy().to_string(), + container_path: cached_model.container_path.clone(), + readonly: false, + }] + } else { + Vec::new() + }; Ok(NodeProvisionSpec { run_id: self.run_id, node_id: self.node_id, stage_index: Some(self.stage_index), image: self.image.clone(), env, - args: Vec::new(), + args, + mounts, }) } } @@ -346,17 +1091,62 @@ struct ActivePrompt { deadline: Instant, } +struct FrameArchive { + file: File, + next_seq: u64, +} + +impl FrameArchive { + fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() + && !parent.as_os_str().is_empty() + { + std::fs::create_dir_all(parent).map_err(|e| { + format!("create datastream frame log dir {}: {e}", parent.display()) + })?; + } + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .map_err(|e| format!("open datastream frame log {}: {e}", path.display()))?; + Ok(Self { file, next_seq: 0 }) + } + + fn record(&mut self, source: &str, stream: &StreamId, frame: &Frame) { + let payload = match std::str::from_utf8(&frame.payload) { + Ok(text) => json!({"encoding":"utf8","value":text}), + Err(_) => json!({"encoding":"bytes","value":frame.payload}), + }; + let record = json!({ + "arrival_seq":self.next_seq, + "source":source, + "stream":stream.to_string(), + "channel":frame.channel.as_str(), + "position":frame.position.0, + "payload":payload, + }); + self.next_seq += 1; + let _ = serde_json::to_writer(&mut self.file, &record); + let _ = writeln!(self.file); + let _ = self.file.flush(); + } +} + struct OrchDatastream { stream: StreamId, - next_position: u64, + mux: Mux, + archive: Option, } impl OrchDatastream { - fn new(run_id: u64) -> Self { - Self { - stream: StreamId::new(NodeId::new("mvp-orchestrator"), Lifetime(run_id)), - next_position: 0, - } + fn new(run_id: u64, frame_log: Option<&Path>) -> Result { + let stream = StreamId::new(NodeId::new("mvp-orchestrator"), Lifetime(run_id)); + Ok(Self { + stream: stream.clone(), + mux: Mux::unbounded(stream), + archive: frame_log.map(FrameArchive::open).transpose()?, + }) } fn emit_event(&mut self, dashboard: Option<&DashboardSupport>, event: ProvisionEvent) { @@ -372,20 +1162,81 @@ impl OrchDatastream { self.emit_bytes(dashboard, channel, payload); } + fn emit_bootstrap( + &mut self, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, + phase: &str, + status: &str, + detail: Value, + ) { + let payload = serde_json::to_vec(&json!({ + "type":"OrchBootstrap", + "phase":phase, + "status":status, + "run_id":run_id, + "node_id":node_id, + "detail":detail, + })) + .expect("serialize orch bootstrap event"); + self.emit_bytes(dashboard, ChannelId::new(MVP_ORCH_BOOTSTRAP), payload); + } + + fn emit_prompt( + &mut self, + dashboard: Option<&DashboardSupport>, + run_id: u64, + node_id: u64, + request_id: u64, + phase: &str, + status: &str, + detail: Value, + ) { + let payload = serde_json::to_vec(&json!({ + "type":"OrchPromptEvent", + "phase":phase, + "status":status, + "run_id":run_id, + "node_id":node_id, + "request_id":request_id, + "detail":detail, + })) + .expect("serialize orch prompt event"); + self.emit_bytes(dashboard, ChannelId::new(MVP_ORCH_PROMPT), payload); + } + fn emit_bytes( &mut self, dashboard: Option<&DashboardSupport>, channel: ChannelId, payload: Vec, ) { - let frame = Frame::new(channel, Position(self.next_position), payload); - self.next_position += 1; - ingest_dashboard_frame(dashboard, &self.stream, &frame); - eprintln!( - "mvp-orch-one-node: datastream {} {}", - frame.channel.as_str(), - String::from_utf8_lossy(&frame.payload) - ); + self.emit_bytes_from(dashboard, channel, payload, "orchestrator"); + } + + fn emit_bytes_from( + &mut self, + dashboard: Option<&DashboardSupport>, + channel: ChannelId, + payload: Vec, + source: &str, + ) { + self.mux.submit(channel, payload); + self.flush(dashboard, source); + } + + fn flush(&mut self, dashboard: Option<&DashboardSupport>, source: &str) { + for frame in self.mux.drain() { + ingest_dashboard_frame(dashboard, &self.stream, &frame); + self.archive_frame(source, &self.stream.clone(), &frame); + } + } + + fn archive_frame(&mut self, source: &str, stream: &StreamId, frame: &Frame) { + if let Some(archive) = &mut self.archive { + archive.record(source, stream, frame); + } } } @@ -406,11 +1257,9 @@ impl DashboardSupport { .parse::() .map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?; } - let url = format!("http://127.0.0.1:{}/view/datastream/live", config.port); let handle = dashboard::start_dashboard(config); handle.register_view(Arc::new(MvpClusterDashboardView::new())); handle.start_http_standalone(); - println!("{}", json!({"type":"dashboard_ready","url":url})); Ok(Some(Self { handle })) } @@ -462,17 +1311,15 @@ fn spawn_prompt_rpc( Ok(stream) => { let tx = work_tx.clone(); thread::spawn(move || { - if let Err(error) = handle_prompt_connection( + let _ = handle_prompt_connection( stream, tx, default_max_tokens, default_timeout_ms, - ) { - eprintln!("mvp-orch-one-node: prompt connection closed: {error}"); - } + ); }); } - Err(error) => eprintln!("mvp-orch-one-node: accept prompt RPC: {error}"), + Err(_) => break, } } }); @@ -524,13 +1371,14 @@ fn wait_for_runtime_ready( frame_rx: &mpsc::Receiver<(StreamId, Frame)>, dashboard: Option<&DashboardSupport>, orch_datastream: &mut OrchDatastream, + provider: ProviderKind, ) -> Result { let start = Instant::now(); loop { pump(driver, stack); - drain_frames(frame_rx, dashboard); + drain_frames(frame_rx, dashboard, orch_datastream); while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, &observation); + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); match observation { PluginObservation::RuntimeReady { endpoint, @@ -542,11 +1390,10 @@ fn wait_for_runtime_ready( node_actor, }); } - PluginObservation::ProviderLine { line, .. } - | PluginObservation::StdoutLine { line, .. } - | PluginObservation::StderrLine { line, .. } => { - eprintln!("mvp-orch-one-node: node: {line}"); - } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} PluginObservation::Failed { reason, .. } => return Err(reason), PluginObservation::Exited { status, .. } => { return Err(format!("node exited before ready: {status:?}")); @@ -568,12 +1415,12 @@ fn wait_for_route( let start = Instant::now(); while start.elapsed() <= ROUTE_TIMEOUT { pump(driver, stack); - if stack + let ready = stack .route_view .read() - .expect("route view poisoned") - .contains_key(&actor) - { + .map(|view| view.contains_key(&actor)) + .unwrap_or(false); + if ready { return Ok(()); } thread::sleep(PUMP_INTERVAL); @@ -617,33 +1464,29 @@ fn wait_for_weights_loaded( frame_rx: &mpsc::Receiver<(StreamId, Frame)>, dashboard: Option<&DashboardSupport>, orch_datastream: &mut OrchDatastream, + provider: ProviderKind, ) -> Result<(), String> { let start = Instant::now(); loop { pump(driver, stack); while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, &observation); + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); match observation { PluginObservation::Failed { reason, .. } => return Err(reason), PluginObservation::Exited { status, .. } => { return Err(format!("node exited while loading weights: {status:?}")); } - PluginObservation::ProviderLine { line, .. } - | PluginObservation::StdoutLine { line, .. } - | PluginObservation::StderrLine { line, .. } => { - eprintln!("mvp-orch-one-node: node: {line}"); - } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} PluginObservation::RuntimeReady { .. } => {} } } while let Ok((stream, frame)) = frame_rx.try_recv() { ingest_dashboard_frame(dashboard, &stream, &frame); + orch_datastream.archive_frame("node_cluster", &stream, &frame); let payload = String::from_utf8_lossy(&frame.payload); - eprintln!( - "mvp-orch-one-node: datastream {} {}", - frame.channel.as_str(), - payload - ); if frame.channel == ChannelId::new("mvp.worker.weights") && json_type_is(&payload, "WeightsLoaded") { @@ -670,68 +1513,233 @@ fn serve_prompts( stop_rx: &mpsc::Receiver<()>, dashboard: Option<&DashboardSupport>, orch_datastream: &mut OrchDatastream, + run_id: u64, + node_id: u64, node_actor: ActorAddress, reply_to: ActorAddress, + provider: ProviderKind, ) -> Result<(), String> { let mut active: Option = None; loop { pump(driver, stack); - drain_observations(obs_rx, dashboard, orch_datastream)?; - drain_frames(frame_rx, dashboard); + drain_observations(obs_rx, dashboard, orch_datastream, provider)?; + drain_frames(frame_rx, dashboard, orch_datastream); if stop_rx.try_recv().is_ok() { - eprintln!("mvp-orch-one-node: stop requested"); + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "shutdown", + "started", + json!({"source":"stdin"}), + ); return Ok(()); } - if active.is_none() { - if let Ok(work) = work_rx.try_recv() { - let request = work.request; - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::InferPrompt { - request_id: request.request_id, - prompt: request.prompt_text.clone(), - max_tokens: request.max_tokens, - reply_to, - }, - ) - .map_err(|e| format!("send prompt request: {e}"))?; - active = Some(ActivePrompt { - deadline: Instant::now() + Duration::from_millis(request.timeout_ms), - request, - events: work.events, - }); + if active.is_none() + && let Ok(work) = work_rx.try_recv() + { + let request = work.request; + let request_id = request.request_id; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "prompt_work", + "observed", + json!({ + "prompt_bytes":request.prompt_text.len(), + "max_tokens":request.max_tokens, + "timeout_ms":request.timeout_ms, + }), + ); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "started", + json!({"node_actor":node_actor,"reply_to":reply_to}), + ); + match stack.runtime.send_to( + node_actor, + NodeAgentMsg::InferPrompt { + request_id, + prompt: request.prompt_text.clone(), + max_tokens: request.max_tokens, + reply_to, + }, + ) { + Ok(()) => { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "ready", + json!({"node_actor":node_actor,"reply_to":reply_to}), + ); + active = Some(ActivePrompt { + deadline: Instant::now() + Duration::from_millis(request.timeout_ms), + request, + events: work.events, + }); + } + Err(error) => { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_send", + "failed", + json!({"node_actor":node_actor,"reply_to":reply_to,"error":error.to_string()}), + ); + return Err(format!("send prompt request: {error}")); + } } } while let Some(event) = prompt_events.try_recv() { - if let Some(current) = active.as_ref() { - if event.request_id() == current.request.request_id { - let terminal = event.is_terminal(); - let _ = current.events.send(event); - if terminal { - active = None; - } - } + let request_id = event.request_id(); + let Some(current) = active.as_ref() else { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "dropped", + json!({"reason":"no_active_prompt","event":prompt_event_name(&event)}), + ); + continue; + }; + if request_id != current.request.request_id { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "dropped", + json!({ + "reason":"request_mismatch", + "event":prompt_event_name(&event), + "active_request_id":current.request.request_id, + }), + ); + continue; + } + + let terminal = event.is_terminal(); + let completion_status = prompt_completion_status(&event); + let completion_detail = prompt_completion_detail(&event); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "node_prompt_event", + "observed", + prompt_event_detail(&event), + ); + let _ = current.events.send(event); + if terminal { + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + request_id, + "prompt_complete", + completion_status, + completion_detail, + ); + active = None; } } if let Some(current) = active.as_ref() && Instant::now() >= current.deadline { + let current = active.take().expect("active prompt checked above"); + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + current.request.request_id, + "prompt_timeout", + "failed", + json!({"error":"prompt timed out","timeout_ms":current.request.timeout_ms}), + ); let _ = current.events.send(PromptEvent::Fault { request_id: current.request.request_id, error: "prompt timed out".to_owned(), }); - active = None; + orch_datastream.emit_prompt( + dashboard, + run_id, + node_id, + current.request.request_id, + "prompt_complete", + "failed", + json!({"event":"Timeout","error":"prompt timed out"}), + ); } thread::sleep(PUMP_INTERVAL); } } +fn prompt_event_name(event: &PromptEvent) -> &'static str { + match event { + PromptEvent::TextDelta { .. } => "TextDelta", + PromptEvent::Done { .. } => "Done", + PromptEvent::Fault { .. } => "Fault", + } +} + +fn prompt_event_detail(event: &PromptEvent) -> Value { + match event { + PromptEvent::TextDelta { text, .. } => { + json!({"event":"TextDelta","terminal":false,"text_bytes":text.len()}) + } + PromptEvent::Done { + final_text, + tokens_generated, + elapsed_ms, + .. + } => json!({ + "event":"Done", + "terminal":true, + "tokens_generated":tokens_generated, + "elapsed_ms":elapsed_ms, + "final_text_bytes":final_text.len(), + }), + PromptEvent::Fault { error, .. } => { + json!({"event":"Fault","terminal":true,"error":error}) + } + } +} + +fn prompt_completion_status(event: &PromptEvent) -> &'static str { + match event { + PromptEvent::Done { .. } => "ready", + PromptEvent::Fault { .. } => "failed", + PromptEvent::TextDelta { .. } => "observed", + } +} + +fn prompt_completion_detail(event: &PromptEvent) -> Value { + match event { + PromptEvent::Done { .. } => json!({"event":"Done"}), + PromptEvent::Fault { error, .. } => json!({"event":"Fault","error":error}), + PromptEvent::TextDelta { .. } => json!({"event":"TextDelta"}), + } +} + fn spawn_stop_listener() -> mpsc::Receiver<()> { let (tx, rx) = mpsc::channel(); thread::spawn(move || { @@ -754,19 +1762,19 @@ fn drain_observations( obs_rx: &mpsc::Receiver, dashboard: Option<&DashboardSupport>, orch_datastream: &mut OrchDatastream, + provider: ProviderKind, ) -> Result<(), String> { while let Ok(observation) = obs_rx.try_recv() { - emit_plugin_observation(orch_datastream, dashboard, &observation); + emit_plugin_observation(orch_datastream, dashboard, provider, &observation); match observation { PluginObservation::Failed { reason, .. } => return Err(reason), PluginObservation::Exited { status, .. } => { return Err(format!("node exited: {status:?}")); } - PluginObservation::ProviderLine { line, .. } - | PluginObservation::StdoutLine { line, .. } - | PluginObservation::StderrLine { line, .. } => { - eprintln!("mvp-orch-one-node: node: {line}"); - } + PluginObservation::DatastreamFrame { .. } => {} + PluginObservation::ProviderLine { .. } + | PluginObservation::StdoutLine { .. } + | PluginObservation::StderrLine { .. } => {} PluginObservation::RuntimeReady { .. } => {} } } @@ -776,6 +1784,7 @@ fn drain_observations( fn emit_plugin_observation( orch_datastream: &mut OrchDatastream, dashboard: Option<&DashboardSupport>, + provider: ProviderKind, observation: &PluginObservation, ) { match observation { @@ -818,6 +1827,14 @@ fn emit_plugin_observation( line: line.clone(), }, ), + PluginObservation::DatastreamFrame { + channel, payload, .. + } => orch_datastream.emit_bytes_from( + dashboard, + ChannelId::new(channel), + payload.as_bytes().to_vec(), + "node_bootstrap_stdio", + ), PluginObservation::RuntimeReady { run_id, node_id, .. } => orch_datastream.emit_event( @@ -826,6 +1843,7 @@ fn emit_plugin_observation( run_id: *run_id, node_id: *node_id, kind: ProvisionEventKind::NodeLive, + provider: Some(provider.as_str().to_owned()), message: None, }, ), @@ -839,6 +1857,7 @@ fn emit_plugin_observation( run_id: *run_id, node_id: *node_id, kind: ProvisionEventKind::NodeStopped, + provider: Some(provider.as_str().to_owned()), message: Some(format!("node process exited with {status:?}")), }, ), @@ -852,29 +1871,27 @@ fn emit_plugin_observation( run_id: *run_id, node_id: *node_id, kind: ProvisionEventKind::ProvisionFailed, + provider: Some(provider.as_str().to_owned()), message: Some(reason.clone()), }, ), } } -fn ingest_dashboard_frame(dashboard: Option<&DashboardSupport>, stream: &StreamId, frame: &Frame) { - if let Some(dashboard) = dashboard { - dashboard.ingest(stream, frame); - } -} - fn drain_frames( frame_rx: &mpsc::Receiver<(StreamId, Frame)>, dashboard: Option<&DashboardSupport>, + orch_datastream: &mut OrchDatastream, ) { while let Ok((stream, frame)) = frame_rx.try_recv() { ingest_dashboard_frame(dashboard, &stream, &frame); - eprintln!( - "mvp-orch-one-node: datastream {} {}", - frame.channel.as_str(), - String::from_utf8_lossy(&frame.payload) - ); + orch_datastream.archive_frame("node_cluster", &stream, &frame); + } +} + +fn ingest_dashboard_frame(dashboard: Option<&DashboardSupport>, stream: &StreamId, frame: &Frame) { + if let Some(dashboard) = dashboard { + dashboard.ingest(stream, frame); } } @@ -908,6 +1925,13 @@ fn env_string(name: &str, default: &str) -> String { env_optional(name).unwrap_or_else(|| default.to_owned()) } +fn provider_from_env(config_profile: RuntimeConfigProfile) -> Result { + match env_optional("MVP_NODE_PROVIDER").or_else(|| env_optional("MVP_PROVIDER")) { + Some(value) => ProviderKind::parse_deploy(&value), + None => Ok(config_profile.default_provider()), + } +} + fn env_bool(name: &str, default: bool) -> Result { match env_optional(name) { None => Ok(default), @@ -938,6 +1962,35 @@ fn env_u32(name: &str, default: u32) -> Result { None => Ok(default), } } +fn env_optional_u32(name: &str) -> Result, String> { + env_optional(name) + .map(|value| { + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) + }) + .transpose() +} + +fn env_optional_u64(name: &str) -> Result, String> { + env_optional(name) + .map(|value| { + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) + }) + .transpose() +} + +fn env_optional_f64(name: &str) -> Result, String> { + env_optional(name) + .map(|value| { + value + .parse::() + .map_err(|e| format!("invalid {name}={value:?}: {e}")) + }) + .transpose() +} fn gguf_source_from_env() -> GgufSource { if let Some(path) = env_optional("MVP_GGUF_LOCAL_PATH") { @@ -956,14 +2009,9 @@ fn tokenizer_from_env() -> TokenizerSource { .unwrap_or(TokenizerSource::EmbeddedGguf) } +#[cfg(test)] fn relay_mode_from_env() -> Result { - match env_string("MVP_IROH_RELAY_MODE", "disabled").as_str() { - "disabled" => Ok(iroh::RelayMode::Disabled), - "default" => Ok(iroh::RelayMode::Default), - other => Err(format!( - "unsupported MVP_IROH_RELAY_MODE={other:?}; use disabled or default" - )), - } + relay_runtime_config_from_env(env_u64("MVP_RUN_ID", 1)?).map(|relay| relay.mode) } fn next_arg(args: &mut impl Iterator, name: &str) -> Result { @@ -981,3 +2029,371 @@ where .parse::() .map_err(|e| format!("invalid {name}={value:?}: {e}")) } + +#[cfg(test)] +mod tests { + use super::*; + use std::{ffi::OsString, path::PathBuf}; + + static ENV_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + + const ENV_KEYS: &[&str] = &[ + CACHED_MODEL_HOST_ENV, + "HF_TOKEN", + "MVP_CPU_LINE_PROFILE", + "MVP_CPU_LINE_PROFILE_INTERVAL_MS", + "CUDA_DEVICE_SCHEDULE", + "MVP_DOCKER_GPUS", + "MVP_GGUF_FILE", + "MVP_GPU_SAMPLE", + "MVP_GGUF_LOCAL_PATH", + "MVP_GGUF_REPO", + "MVP_GGUF_REVISION", + "MVP_IROH_RELAY_MODE", + MVP_IROH_RELAY_URL_ENV, + "MVP_LAYER_END_EXCLUSIVE", + "MVP_LOGICAL_NODE_ID", + "MVP_MODEL_CACHE_DIR", + "MVP_MODEL_ID", + "MVP_NODE_IMAGE", + "MVP_NODE_PROVIDER", + "MVP_PROVIDER", + "MVP_PROMPT_MAX_TOKENS", + "MVP_PROMPT_RPC_BIND", + "MVP_PROMPT_TIMEOUT_MS", + "MVP_RUN_ID", + "MVP_RUNTIME_CONFIG", + "MVP_STAGE_INDEX", + "MVP_TOKEN_PROGRESS_EVERY", + "MVP_TINYGRAD_TEST_MODE", + "MVP_TOKENIZER_LOCAL_PATH", + "MVP_VASTAI_API_KEY", + "MVP_VASTAI_BOOTSTRAP_COMMAND", + "MVP_VASTAI_CONFIRM_LEASE", + "MVP_VASTAI_DISK_GB", + "MVP_VASTAI_GPU_NAME", + "MVP_VASTAI_MAX_CREATE_ATTEMPTS", + "MVP_VASTAI_MAX_POLLS", + "MVP_VASTAI_MIN_DOWN_MBPS", + "MVP_VASTAI_MIN_GPU_RAM_MB", + "MVP_VASTAI_MIN_RELIABILITY", + "MVP_VASTAI_MIN_UP_MBPS", + "MVP_VASTAI_ONSTART", + "MVP_VASTAI_POLL_INTERVAL_SECS", + "MVP_VASTAI_REQUIRE_VERIFIED", + "MVP_VASTAI_SSH_USER", + "VASTAI_API_KEY", + SWACTOR_IROH_RELAY_URL_ENV, + ]; + + struct RestoreEnv { + saved: Vec<(&'static str, Option)>, + } + + 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_clean_env(settings: &[(&'static str, &'static str)], test: impl FnOnce() -> T) -> T { + let settings = settings + .iter() + .map(|(key, value)| (*key, OsString::from(value))) + .collect::>(); + with_clean_env_os(&settings, test) + } + + fn with_clean_env_os(settings: &[(&'static str, OsString)], test: impl FnOnce() -> T) -> T { + let _lock = ENV_LOCK + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let saved = ENV_KEYS + .iter() + .map(|&key| (key, std::env::var_os(key))) + .collect::>(); + for key in ENV_KEYS { + unsafe { std::env::remove_var(key) }; + } + for (key, value) in settings { + assert!( + ENV_KEYS.contains(key), + "test env key {key} must be restored" + ); + unsafe { std::env::set_var(key, value) }; + } + let _restore = RestoreEnv { saved }; + test() + } + + fn selected_provider(settings: &[(&'static str, &'static str)]) -> ProviderKind { + with_clean_env(settings, || { + let profile = RuntimeConfigProfile::from_env().expect("runtime config parses"); + provider_from_env(profile).expect("provider parses") + }) + } + + fn canonical_relay_url(raw: &str) -> String { + raw.parse::() + .expect("fixture relay URL parses") + .to_string() + } + + fn node_spec_env(settings: &[(&'static str, &'static str)]) -> Vec<(String, String)> { + with_clean_env(settings, || { + let config = Config::from_env_and_args_iter(std::iter::empty::()) + .expect("config parses"); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[9; 32]).public()); + let datastream_sink = ActorAddress([11; 32]); + config + .node_spec(coordinator, datastream_sink) + .expect("node spec builds") + .env + }) + } + + fn env_value<'a>(env: &'a [(String, String)], key: &str) -> Option<&'a str> { + env.iter() + .find(|(env_key, _)| env_key == key) + .map(|(_, value)| value.as_str()) + } + + struct TempModelFile { + root: PathBuf, + raw_path: PathBuf, + canonical_path: PathBuf, + } + + impl TempModelFile { + fn new(file_name: &str) -> Self { + let root = std::env::temp_dir().join(format!( + "mvp-cached-model-test-{}-{}", + std::process::id(), + std::thread::current().name().unwrap_or("unnamed") + )); + let _ = std::fs::remove_dir_all(&root); + std::fs::create_dir_all(root.join("nested")).expect("create temp model dir"); + let canonical_path = root.join(file_name); + std::fs::write(&canonical_path, b"fake gguf bytes").expect("write temp model file"); + let raw_path = root.join("nested").join("..").join(file_name); + Self { + root, + raw_path, + canonical_path: canonical_path + .canonicalize() + .expect("canonicalize temp model file"), + } + } + } + + impl Drop for TempModelFile { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.root); + } + } + + #[test] + fn frame_archive_writes_jsonl_records_for_text_and_binary_payloads() { + static NEXT_TEMP_FILE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + + let suffix = NEXT_TEMP_FILE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "mvp-frame-archive-test-{}-{suffix}.jsonl", + std::process::id() + )); + let _ = std::fs::remove_file(&path); + + let stream = StreamId::new("test-node", Lifetime(42)); + let mut archive = FrameArchive::open(&path).expect("frame archive opens"); + archive.record( + "orchestrator", + &stream, + &Frame::new( + "stdout", + datastream::Position(7), + b"hello \xce\xbb".to_vec(), + ), + ); + archive.record( + "orchestrator", + &stream, + &Frame::new("stderr", datastream::Position(8), vec![0xff, 0x00, b'A']), + ); + drop(archive); + + let contents = std::fs::read_to_string(&path).expect("read frame archive jsonl"); + let records = contents + .lines() + .map(|line| serde_json::from_str::(line).expect("archive line is json")) + .collect::>(); + let _ = std::fs::remove_file(&path); + + assert_eq!( + records, + vec![ + json!({ + "arrival_seq":0, + "source":"orchestrator", + "stream":"test-node#42", + "channel":"stdout", + "position":7, + "payload":{"encoding":"utf8","value":"hello λ"}, + }), + json!({ + "arrival_seq":1, + "source":"orchestrator", + "stream":"test-node#42", + "channel":"stderr", + "position":8, + "payload":{"encoding":"bytes","value":[255,0,65]}, + }), + ] + ); + } + + #[test] + fn runtime_profile_selects_provider_and_node_provider_takes_precedence() { + assert_eq!( + selected_provider(&[("MVP_RUNTIME_CONFIG", "local")]), + ProviderKind::Docker + ); + assert_eq!( + selected_provider(&[("MVP_RUNTIME_CONFIG", "deploy")]), + ProviderKind::VastAi + ); + assert_eq!( + selected_provider(&[ + ("MVP_RUNTIME_CONFIG", "deploy"), + ("MVP_NODE_PROVIDER", "docker"), + ]), + ProviderKind::Docker + ); + } + + #[test] + fn relay_mode_env_uses_default_relay_and_accepts_disabled() { + with_clean_env(&[], || { + assert!(matches!( + relay_mode_from_env().expect("unset relay mode parses"), + iroh::RelayMode::Default + )); + }); + with_clean_env(&[("MVP_IROH_RELAY_MODE", "disabled")], || { + assert!(matches!( + relay_mode_from_env().expect("disabled relay mode parses"), + iroh::RelayMode::Disabled + )); + }); + } + + #[test] + fn node_spec_env_propagates_relay_url_only_for_custom_relay_config() { + const RELAY_URL: &str = "https://relay-node-spec.example.com"; + + let custom_env = node_spec_env(&[(MVP_IROH_RELAY_URL_ENV, RELAY_URL)]); + let expected_url = canonical_relay_url(RELAY_URL); + assert_eq!( + env_value(&custom_env, MVP_IROH_RELAY_URL_ENV), + Some(expected_url.as_str()) + ); + + let disabled_env = node_spec_env(&[ + ("MVP_IROH_RELAY_MODE", "disabled"), + (MVP_IROH_RELAY_URL_ENV, RELAY_URL), + ]); + assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None); + } + + #[test] + fn docker_config_construction_ignores_malformed_vastai_environment() { + let config = with_clean_env( + &[ + ("MVP_RUNTIME_CONFIG", "local"), + ("MVP_NODE_PROVIDER", "docker"), + ("MVP_VASTAI_CONFIRM_LEASE", "definitely-not-a-bool"), + ("MVP_VASTAI_DISK_GB", "not-a-u32"), + ("MVP_VASTAI_MIN_DOWN_MBPS", "not-a-float"), + ], + || { + Config::from_env_and_args_iter(std::iter::empty::()) + .expect("docker config ignores VastAI-only env") + }, + ); + + assert_eq!(config.provider, ProviderKind::Docker); + assert!(config.vastai.is_none()); + } + + #[test] + fn docker_cached_model_builds_local_gguf_env_and_writable_file_mount_from_canonical_host_path() + { + let model = TempModelFile::new("weights-q4.gguf"); + let config = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("local")), + ("MVP_NODE_PROVIDER", OsString::from("docker")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ], + || { + Config::from_env_and_args_iter(std::iter::empty::()) + .expect("docker cached model config parses") + }, + ); + let coordinator = EndpointAddr::new(iroh::SecretKey::from_bytes(&[7; 32]).public()); + let datastream_sink = ActorAddress([13; 32]); + let spec = config + .node_spec(coordinator, datastream_sink) + .expect("cached model node spec builds"); + + assert_eq!( + env_value(&spec.env, "MVP_GGUF_LOCAL_PATH"), + Some("/models/cached/weights-q4.gguf") + ); + assert_eq!(env_value(&spec.env, "MVP_GGUF_REPO"), None); + assert_eq!(env_value(&spec.env, "MVP_GGUF_FILE"), None); + assert_eq!( + spec.mounts, + vec![ProviderMount { + host_path: model.canonical_path.to_string_lossy().to_string(), + container_path: "/models/cached/weights-q4.gguf".to_owned(), + readonly: false, + }] + ); + } + + #[test] + fn cached_model_with_deploy_provider_is_rejected_before_vastai_env_is_parsed() { + let model = TempModelFile::new("deploy-rejected.gguf"); + let error = with_clean_env_os( + &[ + ("MVP_RUNTIME_CONFIG", OsString::from("deploy")), + (CACHED_MODEL_HOST_ENV, model.raw_path.as_os_str().to_owned()), + ( + "MVP_VASTAI_CONFIRM_LEASE", + OsString::from("definitely-not-a-bool"), + ), + ("MVP_VASTAI_DISK_GB", OsString::from("not-a-u32")), + ], + || match Config::from_env_and_args_iter(std::iter::empty::()) { + Ok(_) => panic!("deploy cached model must be rejected"), + Err(error) => error, + }, + ); + + assert!( + error.contains( + "MVP_CACHED_MODEL_HOST_PATH is a host-local cache path and requires provider=docker" + ), + "unexpected error: {error}" + ); + assert!( + !error.contains("MVP_VASTAI_CONFIRM_LEASE") && !error.contains("MVP_VASTAI_DISK_GB"), + "cached-model rejection should not require valid VastAI env, got: {error}" + ); + } +} diff --git a/crates/mvp-system/src/bootstrap_datastream.rs b/crates/mvp-system/src/bootstrap_datastream.rs index 4f8114b..eb0357c 100644 --- a/crates/mvp-system/src/bootstrap_datastream.rs +++ b/crates/mvp-system/src/bootstrap_datastream.rs @@ -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) { 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 { + let frame = serde_json::from_str::(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")] diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 4600e8e..aa65b71 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -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; diff --git a/crates/mvp-system/src/node_image.rs b/crates/mvp-system/src/node_image.rs new file mode 100644 index 0000000..bf84de6 --- /dev/null +++ b/crates/mvp-system/src/node_image.rs @@ -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, + 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 { + 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 { + 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 { + 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 { + Ok(git_capture(root, &["status", "--porcelain"])? + .trim() + .is_empty()) +} + +fn git_capture(root: &Path, args: &[&str]) -> Result { + 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 { + source_content_hash(root) +} + +fn source_content_hash(root: &Path) -> Result { + content_hash_for_inputs(root, IMAGE_SOURCE_INPUTS) +} + +fn content_hash_for_inputs(root: &Path, inputs: &[&str]) -> Result { + 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 { + hash_relative_files(root, vec![relative_path(root, &root.join(path))?]) +} + +fn hash_relative_files(root: &Path, files: Vec) -> Result { + 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) -> 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 { + 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, 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, + 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, +) -> 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, +) -> Result { + 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) -> Vec { + 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 { + 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> = 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, + 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, +} + +impl ImageName { + fn parse(raw: &str) -> Result { + 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!["latest".to_owned(), "smoke".to_owned()] + ); + } +} diff --git a/crates/mvp-system/src/node_provisioning.rs b/crates/mvp-system/src/node_provisioning.rs index cb20809..a0123a8 100644 --- a/crates/mvp-system/src/node_provisioning.rs +++ b/crates/mvp-system/src/node_provisioning.rs @@ -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 { + 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, diff --git a/crates/mvp-system/src/provisioning.rs b/crates/mvp-system/src/provisioning.rs index 82f8262..0b497d0 100644 --- a/crates/mvp-system/src/provisioning.rs +++ b/crates/mvp-system/src/provisioning.rs @@ -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, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub mounts: Vec, +} + +#[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, pub message: Option, } @@ -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 { + 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 { + 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 { + 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, 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) -> 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" + ); + } +} diff --git a/crates/mvp-system/src/relay_provisioning.rs b/crates/mvp-system/src/relay_provisioning.rs new file mode 100644 index 0000000..1a3842b --- /dev/null +++ b/crates/mvp-system/src/relay_provisioning.rs @@ -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, +} + +#[derive(Clone, Debug)] +pub struct RelayRuntimeConfig { + pub mode: RelayMode, + pub url: Option, +} + +pub trait RelayProvider: Send { + fn provision_relay(&mut self, request: RelayProvisionRequest) -> Result; + + fn relay_mode(&self, lease: &RelayLease) -> Result; + + 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 { + Ok(RelayLease { + id: RelayLeaseId(format!("local-shim:{}", request.run_id)), + endpoints: Vec::new(), + }) + } + + fn relay_mode(&self, _lease: &RelayLease) -> Result { + 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 { + parse_relay_url(raw).map(Self::new) + } + + pub fn from_env() -> Result, 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 { + 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 { + let urls = lease + .endpoints + .iter() + .map(|endpoint| parse_relay_url(&endpoint.url)) + .collect::, _>>()?; + 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 { + 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 { + 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 { + 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 { + env_optional(MVP_IROH_RELAY_MODE_ENV).map(|value| value.to_ascii_lowercase()) +} + +fn env_optional(name: &str) -> Option { + std::env::var(name) + .ok() + .map(|value| value.trim().to_owned()) + .filter(|value| !value.is_empty()) +} + +fn parse_relay_url(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err("relay URL cannot be empty".to_owned()); + } + trimmed + .parse::() + .map_err(|e| format!("invalid relay URL {trimmed:?}: {e}")) +} diff --git a/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs index 17c949d..f0268c2 100644 --- a/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs +++ b/crates/mvp-system/src/tests/bootstrap_datastream_guarantees.rs @@ -41,6 +41,7 @@ fn spec() -> NodeProvisionSpec { image: "worker:latest".to_owned(), env: Vec::new(), args: Vec::new(), + mounts: Vec::new(), } } diff --git a/crates/mvp-system/src/tests/mod.rs b/crates/mvp-system/src/tests/mod.rs index a006607..04a24de 100644 --- a/crates/mvp-system/src/tests/mod.rs +++ b/crates/mvp-system/src/tests/mod.rs @@ -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; diff --git a/crates/mvp-system/src/tests/relay_provisioning_guarantees.rs b/crates/mvp-system/src/tests/relay_provisioning_guarantees.rs new file mode 100644 index 0000000..def07fd --- /dev/null +++ b/crates/mvp-system/src/tests/relay_provisioning_guarantees.rs @@ -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)>, +} + +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(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::>(); + 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::() + .expect("fixture relay URL parses") + .to_string() +} + +fn assert_custom_relay_mode(mode: RelayMode, expected_url: &str) { + let expected_url = expected_url + .parse::() + .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); + }, + ); +} diff --git a/crates/mvp-system/src/tests/telemetry_guarantees.rs b/crates/mvp-system/src/tests/telemetry_guarantees.rs index a610929..be59ca2 100644 --- a/crates/mvp-system/src/tests/telemetry_guarantees.rs +++ b/crates/mvp-system/src/tests/telemetry_guarantees.rs @@ -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" diff --git a/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs index 4c2e07d..c12dcef 100644 --- a/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs +++ b/crates/mvp-system/src/tests/vastai_provisioning_guarantees.rs @@ -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(), } } diff --git a/crates/mvp-system/src/vastai_provisioning.rs b/crates/mvp-system/src/vastai_provisioning.rs index dfc2676..9f221d2 100644 --- a/crates/mvp-system/src/vastai_provisioning.rs +++ b/crates/mvp-system/src/vastai_provisioning.rs @@ -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>, + stopping: Arc, } 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>, + stopping: Arc, + 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 where C: VastAiLeaseClient, @@ -540,6 +595,9 @@ where spec: NodeProvisionSpec, sink: PluginSink, ) -> Result { + 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 { diff --git a/crates/mvp-system/tests/one_node_chat_e2e.rs b/crates/mvp-system/tests/one_node_chat_e2e.rs index 378b217..964070d 100644 --- a/crates/mvp-system/tests/one_node_chat_e2e.rs +++ b/crates/mvp-system/tests/one_node_chat_e2e.rs @@ -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>, - stderr: &Arc>, ) -> Result<(), String> { - wait_for_child_or(TEST_TIMEOUT, child, || { - stderr_contains(stderr, "dashboard_ready") - }) - .map_err(|e| format!("dashboard_ready not observed: {e}"))?; wait_for_child_or(TEST_TIMEOUT, child, dashboard_responding) .map_err(|e| format!("dashboard API not live: {e}"))?; wait_for_child_or(TEST_TIMEOUT, child, || { @@ -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::>(); + + 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 { + if frame.channel != "mvp.orch.prompt" { + return None; + } + + let value = serde_json::from_str::(&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, +} + #[derive(Debug)] struct SeenFrame { channel: String, @@ -255,16 +327,53 @@ fn stdout_contains(stdout: &Arc>, needle: &str) -> bool { snapshot(stdout).contains(needle) } +fn prompt_visible(stdout: &Arc>) -> bool { + snapshot(stdout).contains("prompt:> ") +} + +fn prompt_count(stdout: &Arc>) -> usize { + snapshot(stdout).matches("prompt:> ").count() +} + fn response_text_visible(stdout: &Arc>) -> 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>, needle: &str) -> bool { - snapshot(stderr).contains(needle) +fn assert_no_lower_layer_terminal_leaks( + stdout: &Arc>, + stderr: &Arc>, +) -> Result<(), String> { + let leaks = lower_layer_leak_lines("stdout", &snapshot(stdout)) + .into_iter() + .chain(lower_layer_leak_lines("stderr", &snapshot(stderr))) + .collect::>(); + 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 { + 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>) -> String { buf.lock().expect("capture mutex").clone() } diff --git a/xtask/src/main.rs b/xtask/src/main.rs index a535728..2b1274b 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -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 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) -> 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 { + 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) -> 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 { + 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 { + 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 { + 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 }