From 9e0a26d99df8f3913e9d5f3263ea69ce474ffbf0 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Wed, 22 Jul 2026 11:50:53 +0400 Subject: [PATCH] slow but working mvp-chat --- .cargo/config.toml | 1 - Cargo.lock | 5 +- apps/mvp-node/tinygrad_worker.py | 13 + crates/iroh-driver/src/edge_transport.rs | 177 + crates/iroh-driver/src/iroh_driver.rs | 125 +- crates/iroh-driver/src/lib.rs | 6 +- .../mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md | 75 - crates/mvp-system/Cargo.toml | 21 +- .../mvp-system/MVP_NODE_PROVISIONING_SPEC.md | 882 ----- .../mvp-system/{ => specs}/MVP_SYSTEM_SPEC.md | 0 crates/mvp-system/src/actors/mod.rs | 2 +- crates/mvp-system/src/bin/mvp_chat.rs | 559 ++- crates/mvp-system/src/bin/orchestrator.rs | 405 +-- crates/mvp-system/src/bin/worker_node.rs | 254 +- crates/mvp-system/src/config.rs | 2 +- crates/mvp-system/src/lib.rs | 1 + .../mvp-system/tests/gpu_worker_node_e2e.rs | 841 ----- .../tests/gpu_worker_node_e2e/Dockerfile | 23 - .../mvp_tinygrad_worker.py | 228 -- crates/mvp-system/tests/local_e2e_cluster.rs | 356 -- .../tests/local_e2e_cluster/Dockerfile | 24 - .../local_e2e_cluster/tinygrad_cpu_worker.py | 273 -- .../tests/local_unmocked_mvp_e2e.rs | 119 - crates/mvp-system/tests/one_node_chat_e2e.rs | 594 ---- .../mvp-system/tests/support/dumb_worker.rs | 109 - crates/mvp-system/tests/support/local_e2e.rs | 1133 ------ .../tests/support/local_e2e_cluster.rs | 3156 ----------------- xtask/Cargo.toml | 4 + xtask/src/main.rs | 588 ++- 29 files changed, 1603 insertions(+), 8373 deletions(-) create mode 100644 crates/iroh-driver/src/edge_transport.rs delete mode 100644 crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md delete mode 100644 crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md rename crates/mvp-system/{ => specs}/MVP_SYSTEM_SPEC.md (100%) delete mode 100644 crates/mvp-system/tests/gpu_worker_node_e2e.rs delete mode 100644 crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile delete mode 100755 crates/mvp-system/tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py delete mode 100644 crates/mvp-system/tests/local_e2e_cluster.rs delete mode 100644 crates/mvp-system/tests/local_e2e_cluster/Dockerfile delete mode 100755 crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py delete mode 100644 crates/mvp-system/tests/local_unmocked_mvp_e2e.rs delete mode 100644 crates/mvp-system/tests/one_node_chat_e2e.rs delete mode 100644 crates/mvp-system/tests/support/dumb_worker.rs delete mode 100644 crates/mvp-system/tests/support/local_e2e.rs delete mode 100644 crates/mvp-system/tests/support/local_e2e_cluster.rs diff --git a/.cargo/config.toml b/.cargo/config.toml index 22edee3..9cb5b6e 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,4 +1,3 @@ [alias] xtask = "run --package xtask --" 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/Cargo.lock b/Cargo.lock index 1e7fdc6..055b962 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2481,7 +2481,6 @@ dependencies = [ "distribution", "iroh", "iroh-driver", - "iroh-relay", "libc", "parking_lot", "serde", @@ -5736,6 +5735,10 @@ dependencies = [ [[package]] name = "xtask" version = "0.1.0" +dependencies = [ + "libc", + "serde_json", +] [[package]] name = "yasna" diff --git a/apps/mvp-node/tinygrad_worker.py b/apps/mvp-node/tinygrad_worker.py index 25b9c17..7d6330e 100755 --- a/apps/mvp-node/tinygrad_worker.py +++ b/apps/mvp-node/tinygrad_worker.py @@ -11,6 +11,7 @@ import threading import time import struct import traceback +import shutil import urllib.parse import urllib.request from pathlib import Path @@ -159,6 +160,17 @@ def fatal(reason: str, **fields: Any) -> None: def test_mode() -> bool: return os.environ.get("MVP_TINYGRAD_TEST_MODE", "").strip().lower() in {"1", "true", "yes", "on"} +def configure_tinygrad_cuda_compiler(device: str) -> None: + if device.split(":", 1)[0].upper() != "CUDA": + return + if os.environ.get("CUDA_PTX") or os.environ.get("CUDA_CC"): + return + if shutil.which("nvcc") is not None: + return + os.environ["CUDA_PTX"] = "1" + control(type="TinygradCudaCompilerSelected", requested_device=device, compiler="PTX", reason="nvcc_not_found") + + def initialize(cmd: dict[str, Any]) -> None: global Tensor, dtypes, arena @@ -182,6 +194,7 @@ def initialize(cmd: dict[str, Any]) -> None: elapsed_ms=int((time.monotonic() - started) * 1000), ) return + configure_tinygrad_cuda_compiler(device) control(type="TinygradImportStarted", requested_device=device, env_DEV=os.environ.get("DEV")) from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes diff --git a/crates/iroh-driver/src/edge_transport.rs b/crates/iroh-driver/src/edge_transport.rs new file mode 100644 index 0000000..dc36399 --- /dev/null +++ b/crates/iroh-driver/src/edge_transport.rs @@ -0,0 +1,177 @@ +//! Driver-owned byte transport for ring-backed MVP edge protocols. +//! +//! This module deliberately owns only transport framing: one edge id preamble per +//! unidirectional stream, followed by opaque byte chunks. Object-record parsing, +//! ring ownership, and stage semantics stay in the MVP/dataplane crates. + +use std::sync::Arc; + +use distribution::types::NodeId; +use iroh::endpoint::Connection; +use iroh::{Endpoint, EndpointAddr}; +use parking_lot::Mutex; +use tokio::io::AsyncWriteExt; +use tokio::runtime::Handle; +use tokio::sync::mpsc as tokio_mpsc; + +pub const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0"; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum EdgeTransportEvent { + StreamArrived { + peer: NodeId, + edge_id: u64, + stream_id: u64, + }, + BytesRead { + peer: NodeId, + edge_id: u64, + stream_id: u64, + bytes: Vec, + }, + StreamEnded { + peer: NodeId, + edge_id: u64, + stream_id: u64, + }, + StreamFault { + peer: NodeId, + edge_id: Option, + stream_id: Option, + reason: EdgeTransportFault, + }, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum EdgeTransportFault { + ReadError, + WriteError, + ProtocolError, +} + +#[derive(Clone)] +pub struct EdgeSendHandle { + tx: tokio_mpsc::UnboundedSender>, +} + +impl EdgeSendHandle { + pub fn send(&self, bytes: Vec) -> Result<(), String> { + self.tx + .send(bytes) + .map_err(|_| "edge sender task stopped".to_owned()) + } +} + +pub(crate) fn spawn_edge_send_pump( + handle: Handle, + endpoint: Endpoint, + peer: EndpointAddr, + edge_id: u64, +) -> Result { + let (tx, mut rx) = tokio_mpsc::unbounded_channel::>(); + let (ready_tx, ready_rx) = std::sync::mpsc::channel::>(); + handle.spawn(async move { + let result: Result<(), String> = async { + let conn = endpoint + .connect(peer, EDGE_ALPN) + .await + .map_err(|e| format!("connect edge {edge_id}: {e}"))?; + let mut send = conn + .open_uni() + .await + .map_err(|e| format!("open edge stream {edge_id}: {e}"))?; + send.write_all(&encode_edge_preamble(edge_id)) + .await + .map_err(|e| format!("write edge preamble {edge_id}: {e}"))?; + send.flush() + .await + .map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?; + let _ = ready_tx.send(Ok(())); + while let Some(record) = rx.recv().await { + send.write_all(&record) + .await + .map_err(|e| format!("write edge record {edge_id}: {e}"))?; + send.flush() + .await + .map_err(|e| format!("flush edge record {edge_id}: {e}"))?; + } + send.finish() + .map_err(|e| format!("finish edge stream {edge_id}: {e}"))?; + Ok(()) + } + .await; + if let Err(error) = result { + let _ = ready_tx.send(Err(error)); + } + }); + ready_rx + .recv() + .map_err(|e| format!("edge {edge_id} sender startup channel closed: {e}"))??; + Ok(EdgeSendHandle { tx }) +} + +pub(crate) fn spawn_edge_recv_pump( + handle: Handle, + conn: Connection, + peer: NodeId, + events: Arc>>, + stream_group: u64, +) { + handle.spawn(async move { + let mut next_uni_stream_id = stream_group << 32; + while let Ok(mut recv) = conn.accept_uni().await { + next_uni_stream_id = next_uni_stream_id.saturating_add(1); + let current_stream_id = next_uni_stream_id; + let mut preamble = [0u8; 8]; + if recv.read_exact(&mut preamble).await.is_err() { + events.lock().push(EdgeTransportEvent::StreamFault { + peer, + edge_id: None, + stream_id: Some(current_stream_id), + reason: EdgeTransportFault::ProtocolError, + }); + continue; + } + let edge_id = u64::from_le_bytes(preamble); + events.lock().push(EdgeTransportEvent::StreamArrived { + peer, + edge_id, + stream_id: current_stream_id, + }); + let mut chunk = vec![0u8; 4096]; + loop { + match recv.read(&mut chunk).await { + Ok(Some(0)) | Ok(None) => { + events.lock().push(EdgeTransportEvent::StreamEnded { + peer, + edge_id, + stream_id: current_stream_id, + }); + break; + } + Ok(Some(n)) => { + events.lock().push(EdgeTransportEvent::BytesRead { + peer, + edge_id, + stream_id: current_stream_id, + bytes: chunk[..n].to_vec(), + }); + } + Err(_) => { + events.lock().push(EdgeTransportEvent::StreamFault { + peer, + edge_id: Some(edge_id), + stream_id: Some(current_stream_id), + reason: EdgeTransportFault::ReadError, + }); + break; + } + } + } + } + }); +} + +fn encode_edge_preamble(edge_id: u64) -> [u8; 8] { + edge_id.to_le_bytes() +} diff --git a/crates/iroh-driver/src/iroh_driver.rs b/crates/iroh-driver/src/iroh_driver.rs index b17b877..1b7b4bb 100644 --- a/crates/iroh-driver/src/iroh_driver.rs +++ b/crates/iroh-driver/src/iroh_driver.rs @@ -30,7 +30,14 @@ use distribution::swim::actor::SwimIn; use distribution::transport_bridge::{OutFrame, Outbox, RelayMirror, RouteView, peer_addr}; use distribution::types::NodeId; -use crate::datastream_transport::DATASTREAM_ALPN; +use crate::datastream_transport::{ + DATASTREAM_ALPN, DatastreamQuicHeader, DatastreamQuicRead, read_events_from_stream, + spawn_subscription_writer, +}; +use crate::edge_transport::{ + EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, spawn_edge_recv_pump, + spawn_edge_send_pump as spawn_edge_sender_task, +}; use swactor::actor::ActorAddress; use swactor::runtime::Runtime; use swactor_transport::CodecRegistry; @@ -207,6 +214,34 @@ pub struct JoinStatus { // ─── Driver ───────────────────────────────────────────────────────────────── +/// Cloneable logical datastream publisher transport. It hides the raw iroh +/// endpoint and Tokio task handle from callers while leaving datastream +/// subscription/catalog semantics in the datastream crate. +#[derive(Clone)] +pub struct DatastreamPublishHandle { + rt: Handle, + endpoint: Endpoint, +} + +impl DatastreamPublishHandle { + pub fn publish_subscription( + &self, + peer: EndpointAddr, + header: DatastreamQuicHeader, + subscription: datastream::DatastreamSubscription, + idle_sleep: Duration, + ) { + let _ = spawn_subscription_writer( + &self.rt, + self.endpoint.clone(), + peer, + header, + subscription, + idle_sleep, + ); + } +} + /// iroh P2P network transport bridge. /// /// Bridges the actorized distribution protocol (running on a swactor runtime) @@ -236,8 +271,13 @@ pub struct IrohDriver { dialing: Arc>>, /// Connections accepted by the background accept loop (SWIM ALPN). accepted_conns: Arc>>, - /// Connections accepted on non-SWIM ALPNs (streams, datastream, etc.). + /// Connections accepted on non-SWIM ALPNs before driver-owned adapters claim them. other_accepted_conns: Arc, Connection)>>>, + /// Completed datastream QUIC reads from driver-owned DATASTREAM_ALPN adapters. + datastream_reads: Arc>>, + /// Logical edge events emitted by driver-owned EDGE_ALPN byte pumps. + edge_events: Arc>>, + next_edge_stream_group: u64, /// Frames read by per-connection reader tasks, drained synchronously by /// `recv()` / `pump_inbound_to_actors()`. This decouples network reads from the /// state machine so `recv()`/`tick()` are pure-sync (no `block_on`) and can run @@ -395,6 +435,9 @@ impl IrohDriver { Arc::new(Mutex::new(Vec::new())); let other_accepted_conns: Arc, Connection)>>> = Arc::new(Mutex::new(Vec::new())); + let datastream_reads: Arc>> = + Arc::new(Mutex::new(Vec::new())); + let edge_events: Arc>> = Arc::new(Mutex::new(Vec::new())); { let ep = endpoint.clone(); let peer_auth = config.peer_auth.clone(); @@ -443,6 +486,9 @@ impl IrohDriver { dialing: Arc::new(Mutex::new(HashSet::new())), accepted_conns, other_accepted_conns, + datastream_reads, + edge_events, + next_edge_stream_group: 1, incoming: Arc::new(Mutex::new(Vec::new())), evict: Arc::new(Mutex::new(Vec::new())), peer_relay_urls: HashMap::new(), @@ -501,6 +547,81 @@ impl IrohDriver { drained } + /// Claim accepted datastream connections and read them inside driver-owned tasks. + pub fn pump_datastream_ingress(&mut self) { + for (_node, conn) in self.drain_accepted_for_alpn(DATASTREAM_ALPN) { + let reads = Arc::clone(&self.datastream_reads); + self.rt.spawn(async move { + while let Ok(recv) = conn.accept_uni().await { + match read_events_from_stream(recv).await { + Ok(read) => reads.lock().push(read), + Err(_) => break, + } + } + }); + } + } + + /// Drain decoded datastream QUIC reads emitted by driver-owned adapter tasks. + pub fn drain_datastream_reads(&self) -> Vec { + self.datastream_reads.lock().drain(..).collect() + } + + /// Start a driver-owned datastream subscription writer task. + pub fn publish_datastream_subscription( + &self, + peer: EndpointAddr, + header: DatastreamQuicHeader, + subscription: datastream::DatastreamSubscription, + idle_sleep: Duration, + ) { + let _ = spawn_subscription_writer( + &self.rt, + self.endpoint.clone(), + peer, + header, + subscription, + idle_sleep, + ); + } + + /// Return a cloneable logical datastream transport handle for publisher actors. + pub fn datastream_publish_handle(&self) -> DatastreamPublishHandle { + DatastreamPublishHandle { + rt: self.rt.clone(), + endpoint: self.endpoint.clone(), + } + } + + /// Claim accepted MVP edge connections and read opaque edge bytes inside the driver. + pub fn pump_edge_ingress(&mut self) { + for (node, conn) in self.drain_accepted_for_alpn(EDGE_ALPN) { + let stream_group = self.next_edge_stream_group; + self.next_edge_stream_group = self.next_edge_stream_group.saturating_add(1).max(1); + spawn_edge_recv_pump( + self.rt.clone(), + conn, + node, + Arc::clone(&self.edge_events), + stream_group, + ); + } + } + + /// Drain logical edge transport events emitted by driver-owned byte pumps. + pub fn drain_edge_events(&self) -> Vec { + self.edge_events.lock().drain(..).collect() + } + + /// Start a driver-owned EDGE_ALPN send pump and return its logical byte input handle. + pub fn spawn_edge_send_pump( + &self, + peer: EndpointAddr, + edge_id: u64, + ) -> Result { + spawn_edge_sender_task(self.rt.clone(), self.endpoint.clone(), peer, edge_id) + } + /// The node's identity. pub fn node_id(&self) -> NodeId { self.keypair.node_id() diff --git a/crates/iroh-driver/src/lib.rs b/crates/iroh-driver/src/lib.rs index dfba366..4990671 100644 --- a/crates/iroh-driver/src/lib.rs +++ b/crates/iroh-driver/src/lib.rs @@ -5,12 +5,16 @@ //! wire message definitions. pub mod datastream_transport; +pub mod edge_transport; pub mod iroh_driver; pub use iroh_driver::{ - ConnType, IrohDriver, IrohDriverConfig, JoinPhase, JoinStatus, conn_type_of, discover_lan_ips, + ConnType, DatastreamPublishHandle, IrohDriver, IrohDriverConfig, JoinPhase, JoinStatus, + conn_type_of, discover_lan_ips, }; +pub use edge_transport::{EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, EdgeTransportFault}; + pub use datastream_transport::{ DATASTREAM_ALPN, DatastreamQuicHeader, DatastreamQuicRead, DatastreamQuicWriteStats, read_events_from_stream, read_next_event, read_next_uni_from_connection, read_stream_header, diff --git a/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md b/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md deleted file mode 100644 index ab5ac11..0000000 --- a/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md +++ /dev/null @@ -1,75 +0,0 @@ -# Actor Control Audit Ideas - -**status**: early draft - -Grounding from `crates/mvp-system`: the specs already give a useful audit line. `MVP_SYSTEM_SPEC.md` says the orchestrator is run authority, swactor owns the control plane, actors establish/observe/tear down components, and tensor bytes are explicitly not actor-mailbox traffic. `MVP_NODE_PROVISIONING_SPEC.md` also gives a key exception: provider I/O and SSH bootstrap are temporary pre-swactor paths; after convergence, swactor is the live control path. - -Suggested somewhat-deterministic identification passes: - -1. **Execution-boundary denylist scan** -- Yes, and the inverse, any code not called from an Actor::handle(...) needs inspection. - AST-scan for `std::thread::spawn`, `tokio::spawn`, `Handle::spawn`, `spawn_blocking`, `Command::new(...).spawn`, `Runtime::new`, and `block_on`. Anything not inside an actor, runtime bootstrap, hot-path byte pump, or pre-swactor bootstrap allowlist is a candidate. - -2. **Process ownership audit** - Find every `std::process::Child`, `ChildStdin`, `ChildStdout`, `ChildStderr`, and `Command::new`. Require each long-lived child to have an actor owner, stop message, exit observation path, and teardown report; otherwise it is likely imperative supervision. - -3. **Network listener audit** - Scan for `TcpListener`, `UnixListener`, `UnixStream`, `UnixDatagram`, `accept`, and per-connection threads/tasks. A listener is acceptable if it immediately decodes ingress into actor messages; if it owns request state or invokes domain operations directly, flag it. - -4. **Channel-as-shadow-mailbox audit** - Scan for `std::sync::mpsc`, `tokio::sync::mpsc`, `oneshot`, `watch`, `broadcast`, and custom queues. Channels outside actor shells often mean a parallel control surface; classify each as actor ingress adapter, data hot-path helper, test harness, or suspect. - - -7. **Actor reachability taint analysis** -- Yes see my comments on 1 - Treat `impl ActorInterface::handle` and actor constructors as roots, then build a call graph. Side-effectful functions reachable only from bins/tests/background threads but not actor roots become candidates for migration. - - -8. **Side-effect import layering rule** - Flag `std::process`, `std::net`, `tokio::net`, `std::fs`, Docker/VastAI/SSH clients, driver joins, and datastream emitters in modules that are supposed to be pure domain state machines. Pure cores should emit commands/events, not perform effects. - - -10. **Runtime creation inventory** -- If this happens at all, massive red flag. - Enumerate every `tokio::runtime::Runtime::new` and `swactor::runtime::Runtime::new`. Runtime creation should cluster at process/runtime-stack boundaries and tests; nested or ad-hoc runtimes usually indicate imperative islands. - - -11. **Post-handoff control-path check** -- All bootstrap monitoring should be owned by an actor, no exceptions. - Encode the provisioning spec as an audit rule: after `swactor` convergence/handoff, SSH/bootstrap/provider code may not remain the live node control path. Scan for SSH or bootstrap-session methods that can act after convergence without going through a node actor. - - -13. **Datastream emission provenance check** - Find direct calls that emit provisioning/readiness/fault/teardown telemetry. Control-plane telemetry should be derived from actor-observed events or actor-owned adapters; direct emission from random loops can hide imperative authority. - - -18. **External API client audit** - Identify VastAI, Docker, SSH, git, and filesystem operations. Provider plugins can perform provider I/O, but they should be stateless with respect to run authority; any retained run/node state inside the client/plugin is suspect. - - -19. **Ownership matrix by resource** -- Yes, but let us be careful about resource definition to catch these. - Build a table: resource type -> owning actor -> allowed non-actor adapter -> teardown message. Missing owner for processes, sockets, rings, leases, workers, or node records is a concrete migration target. - - -20. **Control-plane exception registry** -- How about a critical section boundary, so that any unactorized code gets flagged - Maintain a small checked-in allowlist: pure core, hot tensor byte path, startup bootstrap, pre-swactor SSH bootstrap, provider I/O adapter, test harness. Every denylist hit must match one exception or be filed as non-actor control code. - - - -22. **Backtrace-based audit mode** -- Yes, but not with a 'registry', and only certain critical datastructures - Wrap side-effect APIs behind crate-local helpers and, in audit builds, record a lightweight backtrace/source tag. During e2e runs, fail or report when control-plane effects happen without an actor frame or registered bootstrap exception. - - - -24. **Spawn wrapper migration** -- Interesting idea, consider later. Eventually want to migrate task/thread behavior to swactor runtime, but that is currently deferred to post-alpha. - Replace direct `thread::spawn`, `tokio::spawn`, and `Command::spawn` with crate-local wrappers like `spawn_actor_adapter`, `spawn_byte_pump`, `spawn_pre_swactor_bootstrap`, `spawn_test_helper`. The wrapper name forces classification and makes unclassified spawns easy to detect. - - - -25. **Shadow-runtime detector** -- Multiple runtimes should be considered always wrong until future notice. - Flag ad-hoc Tokio runtimes or swactor runtimes not created by the runtime stack/binary bootstrap. Multiple runtimes are not always wrong, but they often correlate with code escaping the actor scheduler/control surface. - - -28. **Readiness/fault/teardown vocabulary scan** - Search emitted JSON/log labels and enum variants containing `ready`, `live`, `failed`, `fault`, `stopped`, `exited`, `teardown`, `destroyed`. These are control-plane facts; require actor observation/provenance. - - -33. **Test-harness exclusion rule** - Keep tests out of the main migration signal unless they define production-like support code reused by binaries. The crate has many e2e helpers with threads/processes; classify those separately to avoid noisy false positives. - diff --git a/crates/mvp-system/Cargo.toml b/crates/mvp-system/Cargo.toml index bfb5f0d..07d634e 100644 --- a/crates/mvp-system/Cargo.toml +++ b/crates/mvp-system/Cargo.toml @@ -7,6 +7,7 @@ autobins = false [features] default = [] +dashboard = [] [dependencies] datastream = { path = "../datastream" } @@ -28,9 +29,6 @@ toml = "0.8" libc = "0.2" signal-hook = "0.3" -[dev-dependencies] -iroh-relay = { version = "0.98", features = ["server", "test-utils"] } - [[bin]] name = "mvp-worker-node" path = "src/bin/worker_node.rs" @@ -43,23 +41,6 @@ path = "src/bin/orchestrator.rs" name = "mvp-chat" path = "src/bin/mvp_chat.rs" -[[test]] -name = "local_unmocked_mvp_e2e" -path = "tests/local_unmocked_mvp_e2e.rs" -harness = false -required-features = ["local-e2e"] - -[[test]] -name = "gpu_worker_node_e2e" -path = "tests/gpu_worker_node_e2e.rs" -required-features = ["local-e2e"] - -[[test]] -name = "local-e2e-cluster" -path = "tests/local_e2e_cluster.rs" -harness = false -required-features = ["local-e2e"] - [[test]] name = "mvp_chat_mock" path = "tests/mvp_chat_mock.rs" diff --git a/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md b/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md deleted file mode 100644 index 9a15972..0000000 --- a/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md +++ /dev/null @@ -1,882 +0,0 @@ -# MVP Node Provisioning Specification - ***STALE! FOR HISTORICAL REFERENCE ONLY*** - -**Status:**draft node-provisioning specification. - -This document defines the MVP path from a static runplan node requirement to a -remote swactor runtime joined to the orchestrator-side swarm. It covers provider -leasing, SSH bootstrap, stdout/stderr collection, handoff, and known-lease -teardown. - -It intentionally does not define general run management, automatic replacement, -provider recovery, or a second post-handoff health system. - ---- - -## 1. Purpose - -The MVP needs to rent GPU nodes, bring them to the point where swactor can manage -them, and then stop managing them through SSH. - -The intended path is: - -```text -static runplan - -> logical node specs - -> one NodeManager actor per logical node - -> provider plugin creates a lease - -> BootstrapSession holds SSH until swactor convergence - -> stdout/stderr flows into datastream - -> remote swactor joins - -> BootstrapSession closes SSH and exits - -> NodeManager becomes dormant and keeps lease state for teardown -``` - -The actor system owns state transitions. Provider plugins perform provider I/O. -SSH bootstrap is a temporary pre-swactor transport, not long-term node -management. - ---- - -## 2. Scope - -In scope: - -- static runplan node-group shape -- expansion of node groups into logical node specs -- one `NodeManager` actor per logical node -- node-local inventory/ledger owned by `NodeManager` -- readiness as a `NodeManager` state flag -- stateless provider plugin boundary for Vast.ai -- transient `BootstrapSession` for SSH, remote boot observation, and swactor - startup -- stdout/stderr forwarding from bootstrap SSH to datastream -- handoff from SSH bootstrap to swactor control -- teardown of leases already known to `NodeManager` - -Out of scope: - -- automatic replacement after lease, bootstrap, or runtime failure -- provider-label recovery or hidden provider scans -- post-handoff heartbeat layer outside swactor -- bidding/account/billing policy beyond selecting and destroying leases -- repairing a node after it disappears from swactor - ---- - -## 3. Design Commitments - -`NodeManager` owns one node. It owns both the node finite-state machine and that -node's inventory record. - -There is no central actor that owns the run. A short-lived bootstrap procedure may -expand a runplan and spawn node actors, but it does not retain authority over -node state. - -Readiness is node-local. A node is ready only when its `NodeManager` has recorded -successful swactor handoff. - -The provider plugin is stateless with respect to the run. It maps desired node -shape to provider API calls and maps known lease handles to destroy calls. - -`BootstrapSession` owns SSH and early stdout/stderr. It exists only between -provider endpoint availability and swactor convergence. - -After swactor convergence, swactor is the live control path. The MVP does not add -another liveness or heartbeat system. - -Known lease state is explicit. Teardown uses only lease handles already recorded -by `NodeManager`. - ---- - -## 4. Identifiers - -Identifier types are schematic. Concrete Rust APIs may wrap these as newtypes. - -```rust -struct RunId(u64); -struct LogicalNodeId(String); // e.g. "workers-0" -struct NodeGroupId(String); // e.g. "workers" -struct RoleId(String); // e.g. "worker" -struct ProviderLeaseId(String); // e.g. "vastai:123456" -struct SwactorId(String); -struct BootstrapSessionId(u64); -struct DatastreamStreamId(String); -``` - -`LogicalNodeId` is stable for the run. It is assigned before provisioning and is -used to correlate provider lease, SSH bootstrap logs, and swactor identity. - -`ProviderLeaseId` names the external billing/lease resource. For Vast.ai it wraps -the contract id. - -`SwactorId` is not known until the remote runtime joins. - ---- - -## 5. Static Runplan Node Shape - -The runplan describes desired node groups. It does not describe provider API -steps or SSH polling details. - -```rust -struct RunNodeGroupSpec { - run_id: RunId, - group_id: NodeGroupId, - role: RoleId, - count: u32, - provider: ProviderKind, - shape: DesiredNodeShape, - boot: BootSpec, - swarm_join: SwarmJoinSpec, -} -``` - -Provider-neutral desired shape: - -```rust -struct DesiredNodeShape { - image: String, - disk_gb: u32, - gpu_name: Option, - min_gpu_ram_mb: Option, - min_down_mbps: Option, - min_up_mbps: Option, - min_reliability: Option, - require_verified: bool, - provider_labels: BTreeMap, -} -``` - -Remote boot specification: - -```rust -struct BootSpec { - ssh_user: String, - verify_commands: Vec, - start_swactor_command: String, - stdout_sources: Vec, - stderr_sources: Vec, - timeout_policy: BootstrapTimeoutPolicy, -} -``` - -Swarm join material: - -```rust -struct SwarmJoinSpec { - orch_swactor_addr: String, - join_token_ref: String, - expected_logical_node_id: LogicalNodeId, -} -``` - -A bootstrap procedure expands each group into logical specs: - -```text -workers count=3 - -> workers-0 - -> workers-1 - -> workers-2 -``` - -Each expanded logical spec starts one `NodeManager` actor. - ---- - -## 6. Runtime Topology - - -The orchestrator host runs the swactor actor runtime and a datastream producer. - -```text -orchestrator host - Swactor actor runtime - NodeManager(workers-0) - BootstrapSession(workers-0) while pre-handoff - NodeManager(workers-1) - BootstrapSession(workers-1) while pre-handoff - - Orchestrator control endpoint - receives remote runtime joins - owns post-handoff actor communication - - Provider plugins - VastAiPlugin, called by NodeManager - - Datastream - receives bootstrap stdout/stderr records -``` - ---- - -## 7. NodeManager Actor - -### 7.1 Responsibility - -`NodeManager` owns one logical node's state and lifecycle. - -It: - -- stores desired node spec -- requests a provider lease -- records lease facts -- waits for provider endpoint facts when needed -- starts a `BootstrapSession` -- records compact bootstrap observations -- records swactor identity on join -- sets `ready = true` after handoff -- keeps known lease state while dormant -- releases its known lease on `Destroy` - -It does not: - -- aggregate run readiness -- replace failed nodes -- poll post-handoff liveness -- own provider search state after a plugin call returns -- store full stdout/stderr logs - -### 7.2 Node Record - -```rust -struct NodeRecord { - logical_node_id: LogicalNodeId, - run_id: RunId, - group_id: NodeGroupId, - role: RoleId, - - desired: LogicalNodeSpec, - stage: NodeStage, - ready: bool, - - lease: Option, - connection: Option, - bootstrap: Option, - swactor: Option, - - failed_reason: Option, - destroyed_at: Option, -} -``` - -Provider lease facts: - -```rust -struct LeaseFacts { - provider: ProviderKind, - lease_id: ProviderLeaseId, - provider_contract_id: String, - offer_id: Option, - destroy_handle: DestroyHandle, - provider_metadata: BTreeMap, -} -``` - -SSH endpoint: - -```rust -struct SshEndpoint { - host: String, - port: u16, - user: String, - auth_ref: String, -} -``` - -Bootstrap facts are compact. Full logs belong to datastream. - -```rust -struct BootstrapFacts { - session_id: BootstrapSessionId, - last_stage: BootstrapStage, - last_stdout_seq: Option, - last_stderr_seq: Option, - last_observed_at: SystemTime, -} -``` - -Swactor facts: - -```rust -struct SwactorFacts { - swactor_id: SwactorId, - joined_at: SystemTime, - handed_off_at: Option, -} -``` - -### 7.3 Node Stages - -```text -New - -> LeaseRequested - -> LeaseCreated - -> EndpointKnown - -> BootstrapRunning - -> SwactorJoined - -> HandedOff - -> Dormant -``` - -Terminal stages: - -```text -Failed -Destroyed -``` - -`ready = true` only after handoff has completed. `Dormant` means the actor keeps -state for query and teardown but performs no polling, heartbeating, or repair. - -### 7.4 Inbound Messages - -Messages are logical actor signals. Some implementations may deliver provider -results as awaited futures and then enqueue the equivalent event to the actor FSM. - -```rust -enum NodeManagerMsg { - Start(LogicalNodeSpec), - - LeaseCreated(LeaseFacts, Option), - LeaseFailed(String), - EndpointKnown(SshEndpoint), - EndpointFailed(String), - - BootstrapObserved(BootstrapObservation), - BootstrapFailed(String), - BootstrapClosed, - - SwactorJoined { swactor_id: SwactorId }, - HandoffComplete { swactor_id: SwactorId }, - - Destroy, - GetStatus { reply_to: ActorAddress }, - GetRecord { reply_to: ActorAddress }, -} -``` - -### 7.5 Outbound Effects - -`NodeManager` may perform these effects: - -```text -ProviderPlugin.create_lease(shape) -ProviderPlugin.lookup_endpoint(lease) -spawn BootstrapSession(spec) -BootstrapSession.ConvergenceObserved(swactor_id) -ProviderPlugin.destroy_lease(destroy_handle) -reply with node status or record -``` - -It does not send full logs. `BootstrapSession` writes logs directly to -datastream. - ---- - -## 8. NodeManager FSM Behavior - -### 8.1 Start - -On `Start(spec)`: - -```text -record.desired = spec -record.stage = New -record.ready = false -record.failed_reason = None -``` - -Then: - -```text -record.stage = LeaseRequested -call ProviderPlugin.create_lease(spec.shape) -``` - -### 8.2 Lease Result - -On `LeaseCreated(lease, endpoint)`: - -```text -record.lease = lease -record.stage = LeaseCreated -``` - -If `endpoint` is present: - -```text -record.connection = endpoint -record.stage = EndpointKnown -spawn BootstrapSession -record.stage = BootstrapRunning -``` - -If `endpoint` is absent: - -```text -call ProviderPlugin.lookup_endpoint(lease) until endpoint timeout or success -``` - -On `LeaseFailed(reason)`: - -```text -record.stage = Failed -record.ready = false -record.failed_reason = reason -``` - -### 8.3 Endpoint Result - -On `EndpointKnown(endpoint)`: - -```text -record.connection = endpoint -record.stage = EndpointKnown -spawn BootstrapSession -record.stage = BootstrapRunning -``` - -On `EndpointFailed(reason)`: - -```text -record.stage = Failed -record.ready = false -record.failed_reason = reason -``` - -A lease may still exist after endpoint failure. It is destroyed only when the -actor later receives `Destroy`. - -### 8.4 Bootstrap Observations - -On `BootstrapObserved(obs)`: - -```text -record.bootstrap.last_stage = obs.stage -record.bootstrap.last_observed_at = now -record.bootstrap.last_stdout_seq = obs.last_stdout_seq if present -record.bootstrap.last_stderr_seq = obs.last_stderr_seq if present -``` - -The node stage remains `BootstrapRunning` until swactor join. Optional UI views -may display the finer bootstrap stage from `record.bootstrap.last_stage`. - -On `BootstrapFailed(reason)`: - -```text -record.stage = Failed -record.ready = false -record.failed_reason = reason -``` - -### 8.5 Swactor Join And Handoff - -On `SwactorJoined { swactor_id }`: - -```text -record.swactor.swactor_id = swactor_id -record.swactor.joined_at = now -record.stage = SwactorJoined -send BootstrapSession.ConvergenceObserved(swactor_id) -``` - -On `BootstrapClosed` after swactor join: - -```text -record.swactor.handed_off_at = now -record.stage = HandedOff -record.ready = true -record.stage = Dormant -``` - -The SSH handle must be closed before `ready` becomes true. - -### 8.6 Destroy - -On `Destroy`: - -```text -if BootstrapSession active: - cancel BootstrapSession - -if lease exists and not destroyed: - call ProviderPlugin.destroy_lease(lease.destroy_handle) - -record.stage = Destroyed on success -record.destroyed_at = now -record.ready = false -``` - -Destroy uses only the lease stored in `NodeRecord`. There is no provider scan. - ---- - -## 9. Provider Plugin Boundary - -Provider plugins are adapters. They are not run supervisors. - -```rust -trait ProviderPlugin { - async fn create_lease(&self, request: CreateLeaseRequest) - -> Result; - - async fn lookup_endpoint(&self, lease: &LeaseFacts) - -> Result, ProviderError>; - - async fn destroy_lease(&self, handle: &DestroyHandle) - -> Result<(), ProviderError>; -} -``` - -`create_lease` may search, filter, rank, and create a provider lease. For Vast.ai -this maps to offer search and instance creation. - -`lookup_endpoint` may poll provider APIs until SSH endpoint facts are known. It -must not open SSH or inspect remote boot. - -`destroy_lease` destroys a known provider lease. - -Provider plugin output must include enough facts for teardown: - -```rust -struct CreateLeaseResult { - lease: LeaseFacts, - endpoint: Option, -} -``` - -The plugin must not: - -- own `NodeRecord` -- stream stdout/stderr -- start swactor -- infer run readiness -- replace failed nodes -- recover unknown leases by provider label - ---- - -## 10. Vast.ai Plugin Mapping - -For Vast.ai, `CreateLeaseRequest` is derived from `DesiredNodeShape`: - -```text -gpu_name -> SelectionPolicy.gpu_name -min_gpu_ram_mb -> SelectionPolicy.min_gpu_ram_mb -min_down_mbps -> SelectionPolicy.min_down_mbps -min_up_mbps -> SelectionPolicy.min_up_mbps -min_reliability -> SelectionPolicy.min_reliability -require_verified -> SelectionPolicy.require_verified -image -> CreateInstanceRequest.image -disk_gb -> CreateInstanceRequest.disk_gb -provider labels -> CreateInstanceRequest.label / env labels as needed -``` - -The plugin may wait until Vast.ai exposes a usable SSH endpoint. Once that -endpoint is returned, provider provisioning is complete from the plugin's -perspective. - -The plugin does not determine whether the remote image booted correctly. That is -`BootstrapSession` work. - ---- - -## 11. BootstrapSession - -### 11.1 Responsibility - -`BootstrapSession` is a transient child of `NodeManager`. - -It owns: - -- SSH connection attempts -- SSH handle -- remote bootstrap command handles -- stdout/stderr collection before swactor handoff -- boot verification commands -- swactor start command -- waiting for convergence acknowledgement -- closing SSH after handoff - -It exits after either convergence or failure. - -### 11.2 Input - -```rust -struct BootstrapSessionSpec { - run_id: RunId, - logical_node_id: LogicalNodeId, - lease_id: ProviderLeaseId, - ssh: SshEndpoint, - boot: BootSpec, - swarm_join: SwarmJoinSpec, - datastream: DatastreamStreamId, - timeout_policy: BootstrapTimeoutPolicy, -} -``` - -### 11.3 Stages - -```text -Created - -> SshConnecting - -> SshReady - -> StdoutStreaming - -> BootChecking - -> SwactorStarting - -> WaitingForSwactorJoin - -> Converged - -> Closed -``` - -Failure stages: - -```text -SshTimeout -BootCheckFailed -StartFailed -JoinTimeout -StreamError -Cancelled -``` - -### 11.4 Behavior - -1. Connect SSH until timeout. -2. Prove the machine is touchable by running a small command and reading output. -3. Start stdout/stderr capture for configured sources. -4. Emit full log records to datastream. -5. Emit compact observations to `NodeManager`. -6. Run boot verification commands. -7. Run or verify the swactor start command with the join spec. -8. Wait for convergence acknowledgement. -9. Flush datastream writes. -10. Close SSH. -11. Notify `NodeManager` with `BootstrapClosed`. - -### 11.5 Datastream Records - -Bootstrap logs use a stable stream per logical node. - -```rust -struct BootstrapLogRecord { - run_id: RunId, - logical_node_id: LogicalNodeId, - lease_id: ProviderLeaseId, - source: BootstrapLogSource, // ssh-bootstrap - stream: BootstrapLogStream, // stdout | stderr - seq: u64, - timestamp: SystemTime, - line: String, -} -``` - -`NodeManager` stores only sequence numbers and the latest compact observation. -It does not retain log bodies. - -### 11.6 Convergence Signal - -Preferred signal path: - -```text -remote swactor runtime -> orchestrator control endpoint -> NodeManager.SwactorJoined(swactor_id) -NodeManager -> BootstrapSession.ConvergenceObserved(swactor_id) -BootstrapSession -> NodeManager.BootstrapClosed -``` - -This keeps swactor membership authoritative while still letting -`BootstrapSession` close the SSH transport. - ---- - -## 12. Handoff Contract - -Handoff is complete only when all are true: - -- the remote swactor runtime has joined the orchestrator-side actor runtime -- the actor runtime can address the remote by `SwactorId` -- `NodeManager` has recorded `SwactorFacts` -- bootstrap stdout/stderr records have been flushed -- SSH has been closed -- `NodeManager.ready == true` -- `NodeManager.stage == Dormant` - -After handoff: - -- `BootstrapSession` is gone -- `NodeManager` does not poll the node -- the swactor actor runtime owns live communication -- the node is considered usable by the MVP run - ---- - -## 13. Failure Behavior - -Failures before handoff are terminal for the logical node. - -```text -LeaseFailed -EndpointFailed -BootstrapFailed -JoinTimeout -``` - -Terminal behavior: - -```text -record.stage = Failed -record.ready = false -record.failed_reason = reason -``` - -The MVP does not create a replacement lease. - -A failed node with a known lease is still eligible for explicit teardown through -`Destroy`. - -Failures after handoff are handled by actor-runtime behavior. `NodeManager` is -dormant and does not repair the node. If the runtime loses the remote node, the -run stalls or fails according to existing swactor behavior. - ---- - -## 14. Teardown Behavior - -Teardown targets `NodeManager` actors. - -```text -Teardown caller -> NodeManager.Destroy -NodeManager -> BootstrapSession.Cancel if active -NodeManager -> ProviderPlugin.destroy_lease if lease known -NodeManager records Destroyed -``` - -Rules: - -- only known leases are destroyed -- no provider label sweep -- no hidden recovery of missing state -- destroy failure is recorded as node-local failure state - -If the process lost all `NodeManager` state, this MVP spec does not define an -automatic cleanup path. Operator/provider-side cleanup remains manual for that -case. - ---- - -## 15. Single End-To-End Worked Example - -Input runplan node group: - -```text -run_id: 42 -group: workers -role: worker -count: 2 -provider: vastai -shape: - image: ghcr.io/acme/mvp-worker:sha123 - disk_gb: 80 - gpu_name: RTX 4090 - min_gpu_ram_mb: 20000 -boot: - ssh_user: root - verify_commands: - - test -x /opt/mvp/swactor - start_swactor_command: - /opt/mvp/swactor-node --join ${ORCH_ADDR} --node ${LOGICAL_NODE_ID} -swarm_join: - orch_swactor_addr: quic://orch.example:9443 - join_token_ref: secret://run-42-join-token -``` - -Expansion: - -```text -workers-0 -workers-1 -``` - -For `workers-0`: - -1. Bootstrap procedure spawns `NodeManager(workers-0)` with its logical spec. -2. `NodeManager` records `New`, then `LeaseRequested`. -3. `NodeManager` calls `VastAiPlugin.create_lease`. -4. `VastAiPlugin` searches offers, creates a Vast.ai instance, and returns: - -```text -lease_id: vastai:123 -contract_id: 123 -offer_id: 9001 -ssh: root@203.0.113.10:22001 -``` - -5. `NodeManager` records `LeaseCreated`, `EndpointKnown`, then spawns - `BootstrapSession(workers-0)` and records `BootstrapRunning`. -6. `BootstrapSession` connects over SSH, runs a probe command, and sends: - -```text -BootstrapObserved(stage=SshReady) -``` - -7. `BootstrapSession` streams bootstrap stdout/stderr into datastream: - -```text -run=42 node=workers-0 stream=stdout seq=1 line="container boot entered" -run=42 node=workers-0 stream=stdout seq=2 line="swactor binary found" -``` - -8. `BootstrapSession` runs `test -x /opt/mvp/swactor`, then runs the swactor - start command with `LOGICAL_NODE_ID=workers-0` and the run join material. -9. Remote swactor runtime joins the orchestrator-side actor runtime. -10. The orchestrator control endpoint sends: - -```text -NodeManager(workers-0).SwactorJoined(swactor_id=swactor-a7) -``` - -11. `NodeManager` records `SwactorJoined` and tells the bootstrap session that - convergence was observed. -12. `BootstrapSession` flushes datastream writes, closes SSH, and sends - `BootstrapClosed`. -13. `NodeManager` records: - -```text -stage = Dormant -ready = true -swactor_id = swactor-a7 -lease_id = vastai:123 -``` - -The same sequence runs independently for `workers-1`. - -The run bootstrap caller can determine node readiness by querying both -`NodeManager` actors: - -```text -workers-0.ready == true -workers-1.ready == true -``` - -At run completion, teardown sends `Destroy` to both node managers. Each manager -destroys only its recorded Vast.ai contract and records `Destroyed`. - ---- - -## 16. Implementation Boundaries - -The code should preserve these boundaries even if local test plugins combine -steps for convenience: - -- provider lease creation is not SSH bootstrap -- SSH bootstrap is not post-handoff supervision -- `NodeManager` state is the per-node ledger -- datastream owns log bodies -- the swactor actor runtime owns live communication after handoff -- teardown uses known lease handles only - -A local Docker test provider may emit lease, endpoint, bootstrap, and swactor -join observations quickly, but the observations should still map onto the same -FSM stages. This keeps local tests aligned with Vast.ai behavior. diff --git a/crates/mvp-system/MVP_SYSTEM_SPEC.md b/crates/mvp-system/specs/MVP_SYSTEM_SPEC.md similarity index 100% rename from crates/mvp-system/MVP_SYSTEM_SPEC.md rename to crates/mvp-system/specs/MVP_SYSTEM_SPEC.md diff --git a/crates/mvp-system/src/actors/mod.rs b/crates/mvp-system/src/actors/mod.rs index a8dc9ad..9d96f5a 100644 --- a/crates/mvp-system/src/actors/mod.rs +++ b/crates/mvp-system/src/actors/mod.rs @@ -1,4 +1,4 @@ -//! swactor actor shells for the MVP system local E2E stack. +//! swactor actor shells for the MVP system runtime. //! //! Each actor module owns its message type. Pure state machines remain in the //! existing domain modules; actors translate mailbox messages into those cores and diff --git a/crates/mvp-system/src/bin/mvp_chat.rs b/crates/mvp-system/src/bin/mvp_chat.rs index 371547e..06b6f92 100644 --- a/crates/mvp-system/src/bin/mvp_chat.rs +++ b/crates/mvp-system/src/bin/mvp_chat.rs @@ -1,4 +1,5 @@ -use std::fs; +use std::collections::BTreeMap; +use std::fs::{self, File, OpenOptions}; use std::io::{self, BufRead, BufReader, IsTerminal, Write}; use std::net::{Shutdown, TcpStream}; #[cfg(all(target_os = "linux", not(test)))] @@ -10,7 +11,12 @@ use std::sync::{Mutex, mpsc}; use std::thread; use std::time::{Duration, Instant}; +use datastream::{ + ChannelContent, ChannelId, DatastreamEndpoint, DatastreamProducer, Frame, Lifetime, NodeId, + StreamDescriptor, StreamId, StreamOrigin, +}; use serde::Deserialize; +use serde_json::{Value, json}; #[cfg(target_os = "linux")] use signal_hook::consts::signal::{SIGINT, SIGTERM}; #[cfg(target_os = "linux")] @@ -30,6 +36,10 @@ const REPO_MODEL_CACHE_DIR: &str = ".model-cache"; const DEFAULT_MAX_TOKENS: u32 = 64; const ORCH_SHUTDOWN_GRACE_MS: u64 = 5_000; const ORCH_SHUTDOWN_POLL_MS: u64 = 50; +const CHAT_LIFECYCLE_CHANNEL: &str = "mvp.chat.lifecycle"; +const CHAT_RUNTIME_CHANNEL: &str = "mvp.chat.runtime"; +const CHAT_PROMPT_CHANNEL: &str = "mvp.chat.prompt"; +const CHAT_COMPONENT_CHANNEL: &str = "mvp.chat.component"; #[derive(Debug)] enum PromptInput { @@ -63,19 +73,122 @@ where I: IntoIterator, { let config = Config::from_args(args)?; + let mut progress = ChatDatastream::new(1, config.datastream_frame_log.clone())?; + progress.emit( + CHAT_LIFECYCLE_CHANNEL, + "config", + "ready", + json!({ + "provider": config.provider.as_str(), + "pipeline_stages": config.pipeline_stages, + "max_tokens": config.max_tokens, + "cached_model": config.cached_model.as_ref().map(|model| model.host_path.to_string_lossy().to_string()), + "dump_logs": config.datastream_frame_log.as_ref().map(|path| path.to_string_lossy().to_string()), + }), + ); confirm_vastai_if_needed(&config)?; - let image_ref = prepare_runtime(&config)?; - let mut orch = OrchChild::spawn(&config, &image_ref)?; + let image_ref = match prepare_runtime(&config) { + Ok(image_ref) => { + progress.emit( + CHAT_RUNTIME_CHANNEL, + "prepare_runtime", + "ready", + json!({"image_ref": image_ref}), + ); + image_ref + } + Err(error) => { + progress.emit( + CHAT_RUNTIME_CHANNEL, + "prepare_runtime", + "failed", + json!({"error": error}), + ); + progress.archive_pending()?; + return Err(error); + } + }; + let mut orch = match OrchChild::spawn(&config, &image_ref) { + Ok(orch) => { + progress.emit( + CHAT_COMPONENT_CHANNEL, + "orchestrator_process", + "started", + json!({"binary": config.orch_bin.to_string_lossy()}), + ); + orch + } + Err(error) => { + progress.emit( + CHAT_COMPONENT_CHANNEL, + "orchestrator_process", + "failed", + json!({"error": error}), + ); + progress.archive_pending()?; + return Err(error); + } + }; let rpc_addr = match orch.wait_ready(config.rpc_addr.clone()) { - Ok(addr) => addr, + Ok(addr) => { + progress.emit( + CHAT_RUNTIME_CHANNEL, + "prompt_rpc", + "ready", + json!({"addr": addr}), + ); + addr + } Err(_) if STOP_REQUESTED.load(Ordering::SeqCst) => { + progress.emit( + CHAT_LIFECYCLE_CHANNEL, + "shutdown", + "requested", + json!({"reason": "interrupted_before_ready"}), + ); orch.shutdown(); + progress.emit( + CHAT_COMPONENT_CHANNEL, + "orchestrator_process", + "stopped", + json!({"reason": "interrupted_before_ready"}), + ); + progress.archive_pending()?; return Ok(()); } - Err(error) => return Err(error), + Err(error) => { + progress.emit( + CHAT_RUNTIME_CHANNEL, + "prompt_rpc", + "failed", + json!({"error": error}), + ); + orch.shutdown(); + progress.emit( + CHAT_COMPONENT_CHANNEL, + "orchestrator_process", + "stopped", + json!({"reason": "startup_failed"}), + ); + progress.archive_pending()?; + return Err(error); + } }; - let result = run_chat_loop(&rpc_addr, config.max_tokens); + let result = run_chat_loop_with_progress(&rpc_addr, config.max_tokens, Some(&mut progress)); + progress.emit( + CHAT_LIFECYCLE_CHANNEL, + "shutdown", + "requested", + json!({"reason": "prompt_loop_exited", "ok": result.is_ok()}), + ); orch.shutdown(); + progress.emit( + CHAT_COMPONENT_CHANNEL, + "orchestrator_process", + "stopped", + json!({"reason": "shutdown_requested"}), + ); + progress.archive_pending()?; result } @@ -95,6 +208,174 @@ struct Config { skip_rebuild: bool, } +struct ChatDatastream { + stream: StreamId, + endpoint: DatastreamEndpoint, + producer: DatastreamProducer, + channels: BTreeMap, + channel_names: BTreeMap, + archive_path: Option, + pending: Vec<(String, StreamId, String, Frame)>, +} + +impl ChatDatastream { + fn new(run_id: u64, archive_path: Option) -> Result { + let stream = StreamId::new(NodeId::new("mvp-chat"), Lifetime(run_id)); + let endpoint = DatastreamEndpoint::with_descriptor( + StreamDescriptor { + stream: stream.clone(), + label: Some("mvp chat".to_owned()), + origin: StreamOrigin::Orchestrator, + }, + 1024, + 256, + ); + let producer = endpoint.producer(); + let mut out = Self { + stream, + endpoint, + producer, + channels: BTreeMap::new(), + channel_names: BTreeMap::new(), + archive_path, + pending: Vec::new(), + }; + for name in [ + CHAT_LIFECYCLE_CHANNEL, + CHAT_RUNTIME_CHANNEL, + CHAT_PROMPT_CHANNEL, + CHAT_COMPONENT_CHANNEL, + ] { + out.channel_by_name(name); + } + Ok(out) + } + + fn channel_by_name(&mut self, name: &str) -> ChannelId { + if let Some(id) = self.channels.get(name).copied() { + return id; + } + let id = self.producer.register_channel( + name, + ChannelContent::JsonRecord { + schema: Some(name.to_owned()), + }, + ); + self.channels.insert(name.to_owned(), id); + self.channel_names.insert(id, name.to_owned()); + id + } + + fn emit(&mut self, channel: &str, phase: &str, status: &str, detail: Value) { + let id = self.channel_by_name(channel); + let payload = serde_json::to_vec(&json!({ + "type": "ChatProgress", + "phase": phase, + "status": status, + "detail": detail, + })) + .expect("serialize mvp-chat progress event"); + self.producer.submit_bytes(id, payload); + self.flush(); + } + + fn flush(&mut self) { + let stream = self.stream.clone(); + for frame in self.endpoint.mux().drain() { + let channel = self + .channel_names + .get(&frame.channel) + .cloned() + .unwrap_or_else(|| format!("channel#{}", frame.channel.0)); + self.pending + .push(("mvp-chat".to_owned(), stream.clone(), channel, frame)); + } + } + + fn archive_pending(&mut self) -> Result<(), String> { + let Some(path) = self.archive_path.as_deref() else { + self.pending.clear(); + return Ok(()); + }; + if self.pending.is_empty() { + return Ok(()); + } + let mut archive = ChatFrameArchive::open(path)?; + for (source, stream, channel, frame) in self.pending.drain(..) { + archive.record(&source, &stream, &channel, &frame)?; + } + Ok(()) + } +} + +struct ChatFrameArchive { + file: File, + next_seq: u64, +} + +impl ChatFrameArchive { + fn open(path: &Path) -> Result { + if let Some(parent) = path.parent() + && !parent.as_os_str().is_empty() + { + fs::create_dir_all(parent).map_err(|e| { + format!( + "create mvp-chat datastream frame log dir {}: {e}", + parent.display() + ) + })?; + } + let next_seq = match File::open(path) { + Ok(file) => BufReader::new(file).lines().count() as u64, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => 0, + Err(error) => { + return Err(format!( + "read mvp-chat datastream frame log {}: {error}", + path.display() + )); + } + }; + let file = OpenOptions::new() + .create(true) + .append(true) + .open(path) + .map_err(|e| format!("open mvp-chat datastream frame log {}: {e}", path.display()))?; + Ok(Self { file, next_seq }) + } + + fn record( + &mut self, + source: &str, + stream: &StreamId, + channel: &str, + frame: &Frame, + ) -> Result<(), String> { + 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": channel, + "channel_id": frame.channel.0, + "position": frame.position.0, + "payload": payload, + }); + self.next_seq += 1; + let mut line = serde_json::to_vec(&record) + .map_err(|e| format!("serialize mvp-chat frame log: {e}"))?; + line.push(b'\n'); + self.file + .write_all(&line) + .map_err(|e| format!("write mvp-chat frame log: {e}"))?; + self.file + .flush() + .map_err(|e| format!("flush mvp-chat frame log: {e}")) + } +} + #[derive(Clone, Debug, Default, Deserialize)] #[serde(default, deny_unknown_fields)] struct ChatTomlConfig { @@ -255,6 +536,7 @@ impl Config { self.max_tokens.to_string(), "--pipeline-stages".to_owned(), self.pipeline_stages.to_string(), + "--no-dashboard".to_owned(), ]; if self.provider == ProviderKind::Process { args.extend([ @@ -412,7 +694,7 @@ fn resolve_vastai_config( node_image: &str, ) -> Result { ResolvedVastAiConfig { - api_key: first_non_empty([env_optional("VAST_API_KEY")]).unwrap_or_default(), + api_key: first_non_empty([env_optional("VASTAI_API_KEY")]).unwrap_or_default(), relay_url: first_non_empty([file.relay_url.clone()]).unwrap_or_default(), image: node_image.to_owned(), bootstrap_command: first_non_empty([file.bootstrap_command.clone()]).unwrap_or_default(), @@ -692,42 +974,70 @@ fn stdin_prompt_events() -> mpsc::Receiver { rx } -fn run_chat_loop(addr: &str, max_tokens: u32) -> Result<(), String> { - run_chat_loop_with_input(addr, max_tokens, stdin_prompt_events()) +fn run_chat_loop_with_progress( + addr: &str, + max_tokens: u32, + progress: Option<&mut ChatDatastream>, +) -> Result<(), String> { + run_chat_loop_with_input_and_progress(addr, max_tokens, stdin_prompt_events(), progress) } -fn run_chat_loop_with_input( +fn run_chat_loop_with_input_and_progress( addr: &str, max_tokens: u32, input_rx: mpsc::Receiver, + progress: Option<&mut ChatDatastream>, ) -> Result<(), String> { - let mut stream = - TcpStream::connect(addr).map_err(|e| format!("connect prompt RPC {addr}: {e}"))?; - let reader = BufReader::new( - stream - .try_clone() - .map_err(|e| format!("clone prompt RPC stream: {e}"))?, + let mut progress = progress; + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "prompt_rpc", + "connecting", + json!({"addr": addr}), ); - run_chat_session(&mut stream, reader, input_rx, max_tokens) -} - -fn run_chat_session( - writer: &mut W, - reader: R, - input_rx: mpsc::Receiver, - max_tokens: u32, -) -> Result<(), String> -where - R: BufRead, - W: Write, -{ - let mut output = io::stdout(); - run_chat_session_with_output(writer, reader, input_rx, max_tokens, &mut output) + let mut stream = match TcpStream::connect(addr) { + Ok(stream) => { + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "prompt_rpc", + "connected", + json!({"addr": addr}), + ); + stream + } + Err(error) => { + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "prompt_rpc", + "failed", + json!({"addr": addr, "error": error.to_string()}), + ); + return Err(format!("connect prompt RPC {addr}: {error}")); + } + }; + let reader = match stream.try_clone() { + Ok(stream) => BufReader::new(stream), + Err(error) => { + emit_chat_progress( + &mut progress, + CHAT_RUNTIME_CHANNEL, + "prompt_rpc_clone", + "failed", + json!({"error": error.to_string()}), + ); + return Err(format!("clone prompt RPC stream: {error}")); + } + }; + run_chat_session_with_progress(&mut stream, reader, input_rx, max_tokens, progress) } +#[cfg(test)] fn run_chat_session_with_output( writer: &mut W, - mut reader: R, + reader: R, input_rx: mpsc::Receiver, max_tokens: u32, output: &mut O, @@ -737,17 +1047,101 @@ where W: Write, O: Write, { + run_chat_session_with_output_and_progress(writer, reader, input_rx, max_tokens, output, None) +} + +fn run_chat_session_with_progress( + writer: &mut W, + reader: R, + input_rx: mpsc::Receiver, + max_tokens: u32, + progress: Option<&mut ChatDatastream>, +) -> Result<(), String> +where + R: BufRead, + W: Write, +{ + let mut output = io::stdout(); + run_chat_session_with_output_and_progress( + writer, + reader, + input_rx, + max_tokens, + &mut output, + progress, + ) +} + +fn emit_chat_progress( + progress: &mut Option<&mut ChatDatastream>, + channel: &str, + phase: &str, + status: &str, + detail: Value, +) { + if let Some(progress) = progress.as_deref_mut() { + progress.emit(channel, phase, status, detail); + } +} + +fn run_chat_session_with_output_and_progress( + writer: &mut W, + mut reader: R, + input_rx: mpsc::Receiver, + max_tokens: u32, + output: &mut O, + progress: Option<&mut ChatDatastream>, +) -> Result<(), String> +where + R: BufRead, + W: Write, + O: Write, +{ + let mut progress = progress; let mut next_request_id = 1_u64; loop { if STOP_REQUESTED.load(Ordering::SeqCst) { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_loop", + "exited", + json!({"reason": "stop_requested"}), + ); return Ok(()); } + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "waiting_for_prompt", + "started", + json!({"next_request_id": next_request_id}), + ); write!(output, "prompt:> ").map_err(|e| format!("write prompt: {e}"))?; output.flush().map_err(|e| format!("flush prompt: {e}"))?; let prompt = match input_rx.recv() { Ok(PromptInput::Line(line)) => line.trim_end().to_owned(), - Ok(PromptInput::Closed | PromptInput::StopRequested) | Err(_) => return Ok(()), + Ok(PromptInput::Closed) | Err(_) => { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_loop", + "exited", + json!({"reason": "input_closed"}), + ); + return Ok(()); + } + Ok(PromptInput::StopRequested) => { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_loop", + "exited", + json!({"reason": "stop_requested"}), + ); + return Ok(()); + } }; if prompt.trim().is_empty() { continue; @@ -755,6 +1149,13 @@ where let request_id = next_request_id; next_request_id = next_request_id.wrapping_add(1).max(1); + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_submitted", + "ready", + json!({"request_id": request_id, "prompt_bytes": prompt.len(), "max_tokens": max_tokens}), + ); write_json_line( writer, &SubmitPrompt { @@ -764,22 +1165,72 @@ where }, )?; writeln!(output, "decoding...").map_err(|e| format!("write decoding marker: {e}"))?; + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "decoding", + "started", + json!({"request_id": request_id}), + ); let mut response_started = false; loop { if STOP_REQUESTED.load(Ordering::SeqCst) { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_loop", + "exited", + json!({"reason": "stop_requested"}), + ); return Ok(()); } let mut line = String::new(); match reader.read_line(&mut line) { - Ok(0) => return Err("prompt RPC closed".to_owned()), + Ok(0) => { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_rpc", + "failed", + json!({"request_id": request_id, "error": "prompt RPC closed"}), + ); + return Err("prompt RPC closed".to_owned()); + } Ok(_) => {} - Err(error) => return Err(format!("read prompt RPC event: {error}")), + Err(error) => { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_rpc", + "failed", + json!({"request_id": request_id, "error": error.to_string()}), + ); + return Err(format!("read prompt RPC event: {error}")); + } } - let event = serde_json::from_str::(&line) - .map_err(|e| format!("parse prompt RPC event: {e}"))?; + let event = match serde_json::from_str::(&line) { + Ok(event) => event, + Err(error) => { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_event_parse", + "failed", + json!({"request_id": request_id, "error": error.to_string()}), + ); + return Err(format!("parse prompt RPC event: {error}")); + } + }; let seen = event.request_id(); if seen != request_id { + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "prompt_request_id", + "failed", + json!({"expected": request_id, "observed": seen}), + ); return Err(format!( "prompt RPC protocol error: response request_id {seen} does not match active request_id {request_id}" )); @@ -795,6 +1246,13 @@ where output .flush() .map_err(|e| format!("flush response text: {e}"))?; + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "response_text", + "observed", + json!({"request_id": request_id, "text_bytes": text.len()}), + ); } PromptEvent::Done { .. } => { if response_started { @@ -803,11 +1261,25 @@ where writeln!(output, "Response: ") .map_err(|e| format!("write empty response: {e}"))?; } + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "request_completed", + "ready", + json!({"request_id": request_id, "response_started": response_started}), + ); break; } PromptEvent::Fault { error, .. } => { writeln!(output, "error: {error}") .map_err(|e| format!("write prompt fault: {e}"))?; + emit_chat_progress( + &mut progress, + CHAT_PROMPT_CHANNEL, + "request_faulted", + "ready", + json!({"request_id": request_id, "error": error}), + ); break; } } @@ -1061,8 +1533,11 @@ mod tests { static PROCESS_STATE_LOCK: Mutex<()> = Mutex::new(()); static NEXT_TEMP_ID: AtomicU64 = AtomicU64::new(1); - const PROCESS_ENV_KEYS: &[&str] = - &["VAST_API_KEY", "MVP_PIPELINE_STAGES", "MVP_RUNTIME_CONFIG"]; + const PROCESS_ENV_KEYS: &[&str] = &[ + "VASTAI_API_KEY", + "MVP_PIPELINE_STAGES", + "MVP_RUNTIME_CONFIG", + ]; struct TempDir { path: PathBuf, @@ -1521,7 +1996,7 @@ bootstrap_command = "boot" "#, ); with_process_state( - &[("VAST_API_KEY", Some("secret"))], + &[("VASTAI_API_KEY", Some("secret"))], Some(missing_relay.path()), || { let config_arg = missing_relay_config.to_string_lossy().into_owned(); @@ -1546,7 +2021,7 @@ bootstrap_command = "boot" "#, ); with_process_state( - &[("VAST_API_KEY", Some("secret"))], + &[("VASTAI_API_KEY", Some("secret"))], Some(local_image.path()), || { let config_arg = local_image_config.to_string_lossy().into_owned(); @@ -1571,7 +2046,7 @@ bootstrap_command = "boot" "#, ); with_process_state( - &[("VAST_API_KEY", Some("secret"))], + &[("VASTAI_API_KEY", Some("secret"))], Some(valid.path()), || { let config_arg = valid_config.to_string_lossy().into_owned(); diff --git a/crates/mvp-system/src/bin/orchestrator.rs b/crates/mvp-system/src/bin/orchestrator.rs index c42f273..0a636d5 100644 --- a/crates/mvp-system/src/bin/orchestrator.rs +++ b/crates/mvp-system/src/bin/orchestrator.rs @@ -17,9 +17,9 @@ use datastream::{ }; use distribution::node::DistributedNodeConfig; use distribution::types::{MemberState, NodeId as DistNodeId}; -use iroh::{Endpoint, EndpointAddr}; +use iroh::EndpointAddr; use iroh_driver::{ - DATASTREAM_ALPN, IrohDriver, IrohDriverConfig, read_next_event, read_stream_header, + DATASTREAM_ALPN, EDGE_ALPN, EdgeSendHandle, EdgeTransportEvent, IrohDriver, IrohDriverConfig, }; use mvp_system::actors::node_agent::{ NodeAgentMsg, StageEdgeKindWire, StageInboundEdgeWire, StageObjectSpecWire, @@ -28,7 +28,7 @@ use mvp_system::actors::node_agent::{ use mvp_system::actors::orchestrator::{OrchestratorActor, OrchestratorReport}; use mvp_system::actors::register_mvp_actor_codecs; use mvp_system::config::{DEFAULT_CONFIG_PATH, TomlConfigOverlay}; -#[cfg(feature = "local-e2e")] +#[cfg(feature = "dashboard")] use mvp_system::dashboard_view::MvpClusterDashboardView; use mvp_system::distribution_stack::DistributionRuntimeStack; use mvp_system::gpu_worker_ingress_parser as ingress; @@ -60,7 +60,7 @@ use mvp_system::vastai_provisioning::{ use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::ActorAddress; -use tokio::io::AsyncWriteExt; +#[cfg(test)] use tokio::sync::mpsc as tokio_mpsc; const DEFAULT_IMAGE: &str = "swactor-mvp-node:latest"; @@ -85,7 +85,6 @@ const MVP_STAGE_ROUTE: &str = "mvp.orch.stage_route"; const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator"; const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; -const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0"; fn main() -> ExitCode { match run() { @@ -97,41 +96,6 @@ fn main() -> ExitCode { } } -fn build_pipeline_edge_endpoint( - handle: &tokio::runtime::Handle, - relay: &RelayRuntimeConfig, -) -> Result { - let relay_mode = relay.mode.clone(); - let custom_relay = matches!(&relay_mode, iroh::RelayMode::Custom(_)); - handle - .block_on(async move { - let mut builder = Endpoint::builder(iroh::endpoint::presets::Minimal) - .relay_mode(relay_mode) - .alpns(vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()]); - if custom_relay { - builder = builder.ca_roots_config(iroh::tls::CaRootsConfig::insecure_skip_verify()); - } - builder.bind().await - }) - .map_err(|e| format!("create pipeline edge endpoint: {e}")) -} - -fn pipeline_edge_endpoint_addr( - endpoint: &Endpoint, - relay: &RelayRuntimeConfig, -) -> Result { - let mut addr = endpoint.addr(); - if addr.relay_urls().next().is_none() { - if let Some(url) = &relay.url { - addr = addr.with_relay_url( - url.parse() - .map_err(|e| format!("parse pipeline edge relay URL {url:?}: {e}"))?, - ); - } - } - Ok(addr) -} - fn run() -> Result<(), String> { let mut config = Config::from_defaults_toml_env_args(std::env::args().skip(1))?; config.prepare_vastai_ssh_key()?; @@ -429,29 +393,8 @@ fn run() -> Result<(), String> { let sink = PluginSink::new(Arc::new(ChannelObservationSink { tx: Mutex::new(obs_tx), })); - let pipeline_edge_endpoint = if pipeline_plan.is_some() { - let endpoint = build_pipeline_edge_endpoint(tokio.handle(), &config.relay)?; - let addr = pipeline_edge_endpoint_addr(&endpoint, &config.relay)?; - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "pipeline_edge_endpoint", - "ready", - json!({"endpoint":addr}), - ); - Some(endpoint) - } else { - None - }; - let pipeline_token_ingress = pipeline_edge_endpoint - .as_ref() - .map(|endpoint| PipelineTokenIngress::start(tokio.handle().clone(), endpoint.clone())); let coordinator_endpoint = driver.endpoint_addr(); - let pipeline_coordinator_endpoint = match &pipeline_edge_endpoint { - Some(endpoint) => pipeline_edge_endpoint_addr(endpoint, &config.relay)?, - None => coordinator_endpoint.clone(), - }; + let pipeline_coordinator_endpoint = coordinator_endpoint.clone(); let (mut provisioned_nodes, ready) = start_and_provision_workers( provisioner, &config, @@ -539,8 +482,6 @@ fn run() -> Result<(), String> { tokenizer_reply_actor, config.provider, pipeline_plan.as_ref(), - pipeline_edge_endpoint.as_ref(), - pipeline_token_ingress, ready.first_stage.endpoint.clone(), ); if let Err(error) = &result { @@ -3105,71 +3046,60 @@ struct CollectedDatastreamFrame { } fn drain_datastream_connections( - driver: &IrohDriver, + driver: &mut IrohDriver, frame_tx: &mpsc::Sender, ) { - for (_node, conn) in driver.drain_accepted_for_alpn(DATASTREAM_ALPN) { - let tx = frame_tx.clone(); - driver.runtime_handle().spawn(async move { - while let Ok(mut recv) = conn.accept_uni().await { - let Ok(header) = read_stream_header(&mut recv).await else { - break; - }; - let mut channels = header - .channels - .iter() - .map(|descriptor| { - ( - ChannelRef { - stream: descriptor.stream.clone(), - channel: descriptor.id, - }, - descriptor.name.clone(), - ) - }) - .collect::>(); - loop { - let event = match read_next_event(&mut recv, &header.stream).await { - Ok(Some(event)) => event, - Ok(None) => break, - Err(_) => break, - }; - match event { - DatastreamEvent::ChannelDeclared(descriptor) => { - channels.insert( - ChannelRef { - stream: descriptor.stream.clone(), - channel: descriptor.id, - }, - descriptor.name, - ); - } - DatastreamEvent::Frame(delivery) => { - let channel_name = - channels.get(&delivery.channel).cloned().unwrap_or_else(|| { - format!("channel#{}", delivery.channel.channel.0) - }); - let frame = Frame::new( - delivery.channel.channel, - delivery.position, - delivery.payload, - ); - if tx - .send(CollectedDatastreamFrame { - stream: delivery.channel.stream, - channel_name, - frame, - }) - .is_err() - { - break; - } - } - DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {} + driver.pump_datastream_ingress(); + for read in driver.drain_datastream_reads() { + let mut channels = read + .header + .channels + .iter() + .map(|descriptor| { + ( + ChannelRef { + stream: descriptor.stream.clone(), + channel: descriptor.id, + }, + descriptor.name.clone(), + ) + }) + .collect::>(); + for event in read.events { + match event { + DatastreamEvent::ChannelDeclared(descriptor) => { + channels.insert( + ChannelRef { + stream: descriptor.stream.clone(), + channel: descriptor.id, + }, + descriptor.name, + ); + } + DatastreamEvent::Frame(delivery) => { + let channel_name = channels + .get(&delivery.channel) + .cloned() + .unwrap_or_else(|| format!("channel#{}", delivery.channel.channel.0)); + let frame = Frame::new( + delivery.channel.channel, + delivery.position, + delivery.payload, + ); + if frame_tx + .send(CollectedDatastreamFrame { + stream: delivery.channel.stream, + channel_name, + frame, + }) + .is_err() + { + return; } } + DatastreamEvent::StreamDeclared(_) | DatastreamEvent::StreamEnded(_) => {} } - }); + } } } @@ -3485,12 +3415,13 @@ fn drain_orch_stdio_capture( } } -#[cfg(feature = "local-e2e")] +#[cfg(feature = "dashboard")] struct DashboardSupport { handle: dashboard::DashboardHandle, + _runtime: tokio::runtime::Runtime, } -#[cfg(feature = "local-e2e")] +#[cfg(feature = "dashboard")] impl DashboardSupport { fn start(enabled: bool) -> Result, String> { if !enabled { @@ -3502,10 +3433,17 @@ impl DashboardSupport { .parse::() .map_err(|e| format!("invalid MVP_DASHBOARD_PORT={port:?}: {e}"))?; } - let handle = dashboard::start_dashboard(config); + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| format!("dashboard runtime: {e}"))?; + let handle = dashboard::DashboardHandle::new(config); handle.register_view(Arc::new(MvpClusterDashboardView::new())); - handle.start_http_standalone(); - Ok(Some(Self { handle })) + handle.spawn_http(runtime.handle()); + Ok(Some(Self { + handle, + _runtime: runtime, + })) } fn publish_frame(&self, stream: &StreamId, channel: &str, frame: &Frame) { @@ -3521,15 +3459,15 @@ impl DashboardSupport { } } -#[cfg(not(feature = "local-e2e"))] +#[cfg(not(feature = "dashboard"))] struct DashboardSupport; -#[cfg(not(feature = "local-e2e"))] +#[cfg(not(feature = "dashboard"))] impl DashboardSupport { fn start(enabled: bool) -> Result, String> { if enabled { return Err( - "MVP_DASHBOARD requires building mvp-system with feature local-e2e".to_owned(), + "MVP_DASHBOARD requires building mvp-system with feature dashboard".to_owned(), ); } Ok(None) @@ -3830,31 +3768,23 @@ struct PipelineTokenRecord { eos: bool, } -struct PipelineSendHandle { - tx: tokio_mpsc::UnboundedSender>, +enum PipelineSendHandle { + Driver(EdgeSendHandle), + #[cfg(test)] + Channel(tokio_mpsc::UnboundedSender>), } impl PipelineSendHandle { fn send(&self, bytes: Vec) -> Result<(), String> { - self.tx - .send(bytes) - .map_err(|_| "pipeline token-in sender stopped".to_owned()) + match self { + Self::Driver(handle) => handle.send(bytes), + #[cfg(test)] + Self::Channel(tx) => tx + .send(bytes) + .map_err(|_| "pipeline token-in sender stopped".to_owned()), + } } } - -struct PipelineTokenIngress { - recv_rx: mpsc::Receiver>, - recv_tx: mpsc::Sender>, -} - -impl PipelineTokenIngress { - fn start(handle: tokio::runtime::Handle, endpoint: Endpoint) -> Self { - let (recv_tx, recv_rx) = mpsc::channel(); - spawn_pipeline_token_acceptor(handle, endpoint, recv_tx.clone()); - Self { recv_rx, recv_tx } - } -} - struct PendingEncode { request_id: u64, } @@ -3889,11 +3819,9 @@ struct PipelinePromptRuntime { impl PipelinePromptRuntime { fn new( - handle: tokio::runtime::Handle, - endpoint: Endpoint, + driver: &IrohDriver, plan: &run_plan::RunPlan, first_stage_endpoint: EndpointAddr, - ingress: PipelineTokenIngress, tokenizer_encode_actor: ActorAddress, tokenizer_decode_actor: ActorAddress, tokenizer_reply_to: ActorAddress, @@ -3908,19 +3836,17 @@ impl PipelinePromptRuntime { .iter() .find(|edge| edge.kind == run_plan::EdgeKind::TokenOut) .ok_or_else(|| "pipeline plan missing token-out edge".to_owned())?; + let (recv_tx, recv_rx) = mpsc::channel(); Ok(Self { token_in_edge_id: token_in_edge.edge_id.0, token_out_edge_id: token_out_edge.edge_id.0, token_spec: token_in_edge.object_spec, token_out_spec: token_out_edge.object_spec, - token_in_sender: spawn_pipeline_token_sender( - handle, - endpoint, - first_stage_endpoint, - token_in_edge.edge_id.0, - )?, - recv_rx: ingress.recv_rx, - recv_tx: ingress.recv_tx, + token_in_sender: PipelineSendHandle::Driver( + driver.spawn_edge_send_pump(first_stage_endpoint, token_in_edge.edge_id.0)?, + ), + recv_rx, + recv_tx, tokenizer_encode_actor, tokenizer_decode_actor, tokenizer_reply_to, @@ -3950,7 +3876,6 @@ impl PipelinePromptRuntime { node_id: u64, ) -> Result<(), String> { let request_id = request.request_id; - self.next_sequence = 0; self.generated_tokens.clear(); self.final_text.clear(); self.recv_buffer.clear(); @@ -4048,6 +3973,7 @@ impl PipelinePromptRuntime { "ready", json!({"node_actor":self.tokenizer_encode_actor,"reply_to":self.tokenizer_reply_to,"tokens":tokens.len()}), ); + let sequence = self.next_sequence; orch_datastream.emit_prompt( dashboard, run_id, @@ -4055,9 +3981,9 @@ impl PipelinePromptRuntime { request_id, "pipeline_token_in", "started", - json!({"edge_id":self.token_in_edge_id,"sequence":0,"tokens":tokens.len()}), + json!({"edge_id":self.token_in_edge_id,"sequence":sequence,"tokens":tokens.len()}), ); - self.send_token_in(0, &tokens)?; + self.send_token_in(sequence, &tokens)?; orch_datastream.emit_prompt( dashboard, run_id, @@ -4065,7 +3991,7 @@ impl PipelinePromptRuntime { request_id, "pipeline_token_in", "ready", - json!({"edge_id":self.token_in_edge_id,"sequence":0}), + json!({"edge_id":self.token_in_edge_id,"sequence":sequence}), ); Ok(()) } @@ -4175,9 +4101,31 @@ impl PipelinePromptRuntime { self.pending_decode = None; } - fn poll_driver(&mut self, driver: &IrohDriver) { - for (_node, conn) in driver.drain_other_connections() { - spawn_pipeline_token_receiver(driver.tokio_handle(), conn, self.recv_tx.clone()); + fn poll_driver(&mut self, driver: &mut IrohDriver) { + driver.pump_edge_ingress(); + for event in driver.drain_edge_events() { + match event { + EdgeTransportEvent::BytesRead { edge_id, bytes, .. } + if edge_id == self.token_out_edge_id => + { + let _ = self.recv_tx.send(bytes); + } + EdgeTransportEvent::StreamFault { + edge_id: Some(edge_id), + reason, + .. + } if edge_id == self.token_out_edge_id => { + if let Some(request_id) = + self.active.as_ref().map(|active| active.request.request_id) + { + self.fault_active( + request_id, + format!("pipeline token-out stream fault: {reason:?}"), + ); + } + } + _ => {} + } } } @@ -4308,96 +4256,6 @@ fn take_pipeline_token_record( Ok(Some(out)) } -fn spawn_pipeline_token_sender( - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, - peer: EndpointAddr, - edge_id: u64, -) -> Result { - let (tx, mut rx) = tokio_mpsc::unbounded_channel::>(); - let (ready_tx, ready_rx) = mpsc::channel::>(); - handle.spawn(async move { - let result: Result<(), String> = async { - let conn = endpoint - .connect(peer, EDGE_ALPN) - .await - .map_err(|e| format!("connect token-in edge {edge_id}: {e}"))?; - let mut send = conn - .open_uni() - .await - .map_err(|e| format!("open token-in stream {edge_id}: {e}"))?; - send.write_all(&edge_id.to_le_bytes()) - .await - .map_err(|e| format!("write token-in preamble {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush token-in preamble {edge_id}: {e}"))?; - let _ = ready_tx.send(Ok(())); - while let Some(record) = rx.recv().await { - send.write_all(&record) - .await - .map_err(|e| format!("write token-in record {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush token-in record {edge_id}: {e}"))?; - } - send.finish() - .map_err(|e| format!("finish token-in stream {edge_id}: {e}"))?; - Ok(()) - } - .await; - if let Err(error) = result { - let _ = ready_tx.send(Err(error)); - } - }); - ready_rx - .recv() - .map_err(|e| format!("token-in sender startup channel closed: {e}"))??; - Ok(PipelineSendHandle { tx }) -} - -fn spawn_pipeline_token_receiver( - handle: tokio::runtime::Handle, - conn: iroh::endpoint::Connection, - tx: mpsc::Sender>, -) { - handle.spawn(async move { - while let Ok(mut recv) = conn.accept_uni().await { - let mut preamble = [0u8; 8]; - if recv.read_exact(&mut preamble).await.is_err() { - continue; - } - let mut chunk = vec![0u8; 4096]; - loop { - match recv.read(&mut chunk).await { - Ok(Some(0)) | Ok(None) => break, - Ok(Some(n)) => { - if tx.send(chunk[..n].to_vec()).is_err() { - break; - } - } - Err(_) => break, - } - } - } - }); -} - -fn spawn_pipeline_token_acceptor( - handle: tokio::runtime::Handle, - endpoint: Endpoint, - tx: mpsc::Sender>, -) { - let accept_handle = handle.clone(); - handle.spawn(async move { - while let Some(incoming) = endpoint.accept().await { - if let Ok(conn) = incoming.await { - spawn_pipeline_token_receiver(accept_handle.clone(), conn, tx.clone()); - } - } - }); -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum PromptRuntimeMode { DirectInferPrompt, @@ -4434,28 +4292,17 @@ fn serve_prompts( tokenizer_reply_to: ActorAddress, provider: ProviderKind, pipeline_plan: Option<&run_plan::RunPlan>, - pipeline_edge_endpoint: Option<&Endpoint>, - pipeline_token_ingress: Option, prompt_endpoint: EndpointAddr, ) -> Result<(), String> { let mut pipeline_runtime = match prompt_runtime_mode(pipeline_plan) { - PromptRuntimeMode::PipelineTokenEdges => { - let endpoint = pipeline_edge_endpoint - .ok_or_else(|| "pipeline mode requires an edge endpoint".to_owned())? - .clone(); - let ingress = pipeline_token_ingress - .ok_or_else(|| "pipeline mode requires edge ingress".to_owned())?; - Some(PipelinePromptRuntime::new( - driver.tokio_handle(), - endpoint, - pipeline_plan.expect("pipeline mode requires plan"), - prompt_endpoint, - ingress, - tokenizer_encode_actor, - tokenizer_decode_actor, - tokenizer_reply_to, - )?) - } + PromptRuntimeMode::PipelineTokenEdges => Some(PipelinePromptRuntime::new( + driver, + pipeline_plan.expect("pipeline mode requires plan"), + prompt_endpoint, + tokenizer_encode_actor, + tokenizer_decode_actor, + tokenizer_reply_to, + )?), PromptRuntimeMode::DirectInferPrompt => None, }; let mut active: Option = None; @@ -5469,7 +5316,7 @@ mod tests { token_out_edge_id, token_spec, token_out_spec, - token_in_sender: PipelineSendHandle { tx: token_in_tx }, + token_in_sender: PipelineSendHandle::Channel(token_in_tx), recv_rx, recv_tx, recv_buffer: Vec::new(), diff --git a/crates/mvp-system/src/bin/worker_node.rs b/crates/mvp-system/src/bin/worker_node.rs index 97db979..aabb288 100644 --- a/crates/mvp-system/src/bin/worker_node.rs +++ b/crates/mvp-system/src/bin/worker_node.rs @@ -20,7 +20,8 @@ use distribution::node::DistributedNodeConfig; use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; use iroh_driver::{ - DATASTREAM_ALPN, DatastreamQuicHeader, IrohDriver, IrohDriverConfig, spawn_subscription_writer, + DATASTREAM_ALPN, DatastreamPublishHandle, DatastreamQuicHeader, EDGE_ALPN, EdgeSendHandle, + EdgeTransportEvent, IrohDriver, IrohDriverConfig, }; use mvp_system::actors::node_agent::{ NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageInboundEdgeWire, @@ -40,7 +41,6 @@ use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::ActorAddress; use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; -use tokio::sync::mpsc as tokio_mpsc; const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py"; const DEFAULT_DEVICE: &str = "CUDA"; @@ -58,7 +58,6 @@ 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"; -const EDGE_ALPN: &[u8] = b"mvp/pipeline-edge/0"; fn node_event_payload( config: &DeploymentConfig, @@ -509,65 +508,6 @@ fn spawn_arena_sampler( }); } -#[derive(Clone, Debug)] -enum DriverIngressEvent { - StreamArrived { - edge_id: u64, - stream_id: u64, - }, - BytesRead { - edge_id: u64, - stream_id: u64, - bytes: Vec, - }, -} - -#[derive(Clone)] -struct SendPumpHandle { - tx: tokio_mpsc::UnboundedSender>, -} - -impl SendPumpHandle { - fn send(&self, record: Vec) -> Result<(), String> { - self.tx - .send(record) - .map_err(|_| "edge sender task stopped".to_owned()) - } -} - -struct DriverRuntime { - tx: mpsc::Sender, - rx: mpsc::Receiver, - next_stream_id: u64, -} - -impl DriverRuntime { - fn new() -> Self { - let (tx, rx) = mpsc::channel(); - Self { - tx, - rx, - next_stream_id: 1, - } - } - - fn poll_iroh(&mut self, driver: &IrohDriver) { - for (_node, conn) in driver.drain_other_connections() { - spawn_recv_pump( - driver.tokio_handle(), - conn, - self.tx.clone(), - self.next_stream_id, - ); - self.next_stream_id = self.next_stream_id.saturating_add(1); - } - } - - fn try_recv(&self) -> Option { - self.rx.try_recv().ok() - } -} - #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] struct ObjectKey { edge_id: u64, @@ -585,7 +525,6 @@ struct LoadedObject { struct WorkerEdgeRuntime { establisher: edge::EdgeEstablisher, driver_model: driver_model::Driver, - driver_runtime: DriverRuntime, edge_command_cursor: usize, edge_event_cursor: usize, driver_event_cursor: usize, @@ -593,7 +532,7 @@ struct WorkerEdgeRuntime { outbound_edge: Option, inbound_ring_id: Option, outbound_ring_id: Option, - outbound_sender: Option, + outbound_sender: Option, next_output_object_id: u64, object_handles: BTreeMap, ingress_streams: BTreeMap>, @@ -607,7 +546,6 @@ impl WorkerEdgeRuntime { local_node_id: driver_model::NodeId(local_node_id), alpn: driver_model::Alpn(String::from_utf8_lossy(EDGE_ALPN).into_owned()), }), - driver_runtime: DriverRuntime::new(), edge_command_cursor: 0, edge_event_cursor: 0, driver_event_cursor: 0, @@ -625,7 +563,7 @@ impl WorkerEdgeRuntime { #[allow(clippy::too_many_arguments)] fn poll_iroh( &mut self, - driver: &IrohDriver, + driver: &mut IrohDriver, stack: &DistributionRuntimeStack, node_actor: ActorAddress, worker: &mut TinygradWorker, @@ -633,10 +571,12 @@ impl WorkerEdgeRuntime { config: &DeploymentConfig, datastream: &mut NodeDatastream, ) -> Result<(), String> { - self.driver_runtime.poll_iroh(driver); - while let Some(event) = self.driver_runtime.try_recv() { + driver.pump_edge_ingress(); + for event in driver.drain_edge_events() { match event { - DriverIngressEvent::StreamArrived { edge_id, stream_id } => { + EdgeTransportEvent::StreamArrived { + edge_id, stream_id, .. + } => { self.driver_model .observe(driver_model::DriverEvent::IncomingUniStream { edge_id: driver_model::EdgeId(edge_id), @@ -649,14 +589,14 @@ impl WorkerEdgeRuntime { arena_manager, config, datastream, - driver.tokio_handle(), - driver.endpoint().clone(), + driver, )?; } - DriverIngressEvent::BytesRead { + EdgeTransportEvent::BytesRead { edge_id, stream_id, bytes, + .. } => { self.ingest_stream_bytes( edge_id, @@ -668,10 +608,29 @@ impl WorkerEdgeRuntime { arena_manager, config, datastream, - driver.tokio_handle(), - driver.endpoint().clone(), + driver, )?; } + EdgeTransportEvent::StreamEnded { .. } => {} + EdgeTransportEvent::StreamFault { + edge_id: Some(edge_id), + .. + } => { + self.driver_model + .observe(driver_model::DriverEvent::ReadError { + edge_id: driver_model::EdgeId(edge_id), + }); + self.drive_edge_workflow( + stack, + node_actor, + worker, + arena_manager, + config, + datastream, + driver, + )?; + } + EdgeTransportEvent::StreamFault { edge_id: None, .. } => {} } } Ok(()) @@ -687,8 +646,7 @@ impl WorkerEdgeRuntime { arena_manager: &Arc>, config: &DeploymentConfig, datastream: &mut NodeDatastream, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, + driver: &mut IrohDriver, ) -> Result<(), String> { self.inbound_edge = Some(edge.clone()); self.establisher @@ -706,8 +664,7 @@ impl WorkerEdgeRuntime { arena_manager, config, datastream, - handle, - endpoint, + driver, ) } @@ -721,8 +678,7 @@ impl WorkerEdgeRuntime { arena_manager: &Arc>, config: &DeploymentConfig, datastream: &mut NodeDatastream, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, + driver: &mut IrohDriver, ) -> Result<(), String> { if edge.consumer_endpoint.is_none() { stack @@ -754,8 +710,7 @@ impl WorkerEdgeRuntime { arena_manager, config, datastream, - handle, - endpoint, + driver, ) } @@ -872,8 +827,7 @@ impl WorkerEdgeRuntime { arena_manager: &Arc>, config: &DeploymentConfig, datastream: &mut NodeDatastream, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, + driver: &mut IrohDriver, ) -> Result<(), String> { let Some(inbound) = self.inbound_edge.clone() else { return Ok(()); @@ -937,8 +891,7 @@ impl WorkerEdgeRuntime { arena_manager, config, datastream, - handle, - endpoint, + driver, ) } @@ -951,8 +904,7 @@ impl WorkerEdgeRuntime { arena_manager: &Arc>, config: &DeploymentConfig, datastream: &mut NodeDatastream, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, + driver: &mut IrohDriver, ) -> Result<(), String> { loop { let mut progressed = false; @@ -1099,12 +1051,7 @@ impl WorkerEdgeRuntime { }, }, )); - self.outbound_sender = Some(spawn_send_pump( - handle.clone(), - endpoint.clone(), - peer, - edge_id.0, - )?); + self.outbound_sender = Some(driver.spawn_edge_send_pump(peer, edge_id.0)?); } edge::EdgeCommand::EstablishRecv { edge_id, .. } => { let record = self @@ -1291,104 +1238,6 @@ fn take_complete_ingress_record( Ok(Some(buffer.drain(..record.total_len).collect())) } -fn spawn_send_pump( - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, - peer: EndpointAddr, - edge_id: u64, -) -> Result { - let (tx, mut rx) = tokio_mpsc::unbounded_channel::>(); - let (ready_tx, ready_rx) = mpsc::channel::>(); - handle.spawn(async move { - let result: Result<(), String> = async { - let conn = endpoint - .connect(peer, EDGE_ALPN) - .await - .map_err(|e| format!("connect edge {edge_id}: {e}"))?; - let mut send = conn - .open_uni() - .await - .map_err(|e| format!("open edge stream {edge_id}: {e}"))?; - send.write_all(&driver_model::encode_edge_preamble(driver_model::EdgeId( - edge_id, - ))) - .await - .map_err(|e| format!("write edge preamble {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush edge preamble {edge_id}: {e}"))?; - let _ = ready_tx.send(Ok(())); - while let Some(record) = rx.recv().await { - send.write_all(&record) - .await - .map_err(|e| format!("write edge record {edge_id}: {e}"))?; - send.flush() - .await - .map_err(|e| format!("flush edge record {edge_id}: {e}"))?; - } - send.finish() - .map_err(|e| format!("finish edge stream {edge_id}: {e}"))?; - Ok(()) - } - .await; - if let Err(error) = result { - let _ = ready_tx.send(Err(error)); - } - }); - ready_rx - .recv() - .map_err(|e| format!("edge {edge_id} sender startup channel closed: {e}"))??; - Ok(SendPumpHandle { tx }) -} - -fn spawn_recv_pump( - handle: tokio::runtime::Handle, - conn: iroh::endpoint::Connection, - tx: mpsc::Sender, - stream_id: u64, -) { - handle.spawn(async move { - let mut next_uni_stream_id = stream_id << 32; - while let Ok(mut recv) = conn.accept_uni().await { - next_uni_stream_id = next_uni_stream_id.saturating_add(1); - let current_stream_id = next_uni_stream_id; - let mut preamble = [0u8; 8]; - if recv.read_exact(&mut preamble).await.is_err() { - continue; - } - let edge_id = u64::from_le_bytes(preamble); - if tx - .send(DriverIngressEvent::StreamArrived { - edge_id, - stream_id: current_stream_id, - }) - .is_err() - { - break; - } - let mut chunk = vec![0u8; 4096]; - loop { - match recv.read(&mut chunk).await { - Ok(Some(0)) | Ok(None) => break, - Ok(Some(n)) => { - if tx - .send(DriverIngressEvent::BytesRead { - edge_id, - stream_id: current_stream_id, - bytes: chunk[..n].to_vec(), - }) - .is_err() - { - break; - } - } - Err(_) => break, - } - } - } - }); -} - fn value_u64(value: &Value, field: &str) -> Result { value .get(field) @@ -1580,9 +1429,10 @@ fn run() -> Result<(), String> { let arena_fd = arena_manager.lock().arena_fd(); let mut datastream = node_datastream(&config); + let datastream_transport = driver.datastream_publish_handle(); let datastream_publisher = match stack .runtime - .spawn(datastream.publisher_actor(tokio.handle().clone(), driver.endpoint())) + .spawn(datastream.publisher_actor(datastream_transport)) { Ok(actor) => actor, Err(error) => { @@ -1859,7 +1709,7 @@ fn run() -> Result<(), String> { datastream.tick(); worker.drain_stderr(&config, &mut datastream); edge_runtime.poll_iroh( - &driver, + &mut driver, &stack, node_actor, &mut worker, @@ -2154,11 +2004,7 @@ impl NodeDatastream { } } - fn publisher_actor( - &self, - tokio: tokio::runtime::Handle, - iroh_endpoint: iroh::Endpoint, - ) -> DatastreamPublisherActor { + fn publisher_actor(&self, transport: DatastreamPublishHandle) -> DatastreamPublisherActor { DatastreamPublisherActor::new( Arc::clone(&self.endpoint), move |subscribe: DatastreamSubscribe, subscription: DatastreamSubscription| { @@ -2169,9 +2015,7 @@ impl NodeDatastream { ) else { return; }; - let _ = spawn_subscription_writer( - &tokio, - iroh_endpoint.clone(), + transport.publish_subscription( subscribe.collector, header, subscription, @@ -2920,8 +2764,7 @@ fn handle_stage_command( arena_manager, config, datastream, - driver.tokio_handle(), - driver.endpoint().clone(), + driver, )?; emit_node_event( datastream, @@ -2950,8 +2793,7 @@ fn handle_stage_command( arena_manager, config, datastream, - driver.tokio_handle(), - driver.endpoint().clone(), + driver, )?; emit_node_event( datastream, diff --git a/crates/mvp-system/src/config.rs b/crates/mvp-system/src/config.rs index 008f14b..64b0cc9 100644 --- a/crates/mvp-system/src/config.rs +++ b/crates/mvp-system/src/config.rs @@ -174,7 +174,7 @@ impl TomlConfigOverlay { impl ResolvedVastAiConfig { pub fn validate(self) -> Result { - require_non_empty("VAST_API_KEY", &self.api_key)?; + require_non_empty("VASTAI_API_KEY", &self.api_key)?; require_non_empty("relay.url", &self.relay_url)?; require_non_empty("vastai.image", &self.image)?; require_non_empty("vastai.bootstrap_command", &self.bootstrap_command)?; diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 682e814..4b0422b 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -5,6 +5,7 @@ pub mod actors; pub mod arena_manager; pub mod bootstrap_datastream; pub mod config; +pub mod dashboard_view; pub mod device_bridge; pub mod distribution_stack; pub mod docker_cluster_provisioning; diff --git a/crates/mvp-system/tests/gpu_worker_node_e2e.rs b/crates/mvp-system/tests/gpu_worker_node_e2e.rs deleted file mode 100644 index d4f615f..0000000 --- a/crates/mvp-system/tests/gpu_worker_node_e2e.rs +++ /dev/null @@ -1,841 +0,0 @@ -use std::ffi::CString; -use std::io::{BufRead, BufReader, Write}; -use std::os::fd::RawFd; -use std::os::unix::process::CommandExt; -use std::path::{Path, PathBuf}; -use std::process::{Command, Stdio}; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::mpsc; -use std::thread; -use std::time::{Duration, Instant}; - -use datastream::emit::{DatastreamEmitter, EmitterConfig, FrameSink}; -use datastream::{Frame, StreamId}; -use serde_json::{Value, json}; -use swactor::actor::{ActorAddress, ActorInterface}; -use swactor::runtime::{Ctx, ExternalSender, Runtime, RuntimeConfig}; - -const IMAGE: &str = "swactor-mvp-gpu-worker-node-e2e:latest"; -const WORKER_EVENTS_CHANNEL: &str = "mvp.worker.events"; -const ARENA_BYTES: usize = 8192; -const INGRESS_RING_ID: u64 = 8001; -const EGRESS_RING_ID: u64 = 8002; -const INGRESS_EDGE_ID: u64 = 7001; -const EGRESS_EDGE_ID: u64 = 7002; -const INGRESS_BASE: usize = 0; -const EGRESS_BASE: usize = 4096; -const RING_BYTES: usize = 1024; -const HEADER_LEN: usize = 48; -const PREFLIGHT_WATCHDOG: Duration = Duration::from_secs(45); -const EVENT_WATCHDOG: Duration = Duration::from_secs(60); - -#[test] -fn gpu_worker_node_e2e_cuda() { - if std::env::var_os("MVP_SYSTEM_CUDA_E2E_IN_CONTAINER").is_some() { - run_integrated_node_harness(); - } else if std::env::var_os("MVP_SYSTEM_CUDA_E2E").is_some() { - build_and_run_docker_fixture(); - } else { - eprintln!("skipping; set MVP_SYSTEM_CUDA_E2E=1 to run docker CUDA e2e"); - } -} - -fn build_and_run_docker_fixture() { - let crate_dir = Path::new(env!("CARGO_MANIFEST_DIR")); - let workspace = crate_dir - .parent() - .and_then(Path::parent) - .expect("workspace root") - .canonicalize() - .expect("canonical workspace root"); - let context = - std::env::temp_dir().join(format!("mvp-system-docker-context-{}", std::process::id())); - let _ = std::fs::remove_dir_all(&context); - copy_workspace_context(&workspace, &context); - let dockerfile = context.join("crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile"); - - phase("building CUDA Docker fixture image"); - let build = Command::new("docker") - .args(["build", "-f"]) - .arg(&dockerfile) - .args(["-t", IMAGE]) - .arg(&context) - .status() - .expect("run docker build"); - assert!(build.success(), "docker build failed with status {build}"); - - phase("running CUDA Docker fixture"); - let run = Command::new("docker") - .args(["run", "--rm", "--gpus"]) - .arg(std::env::var("MVP_CUDA_GPUS").unwrap_or_else(|_| "all".to_owned())) - .args(["-e", "MVP_SYSTEM_CUDA_E2E_IN_CONTAINER=1"]) - .args(["-e", "CARGO_TARGET_DIR=/tmp/mvp-system-target"]) - .arg(IMAGE) - .status() - .expect("run docker CUDA fixture"); - assert!( - run.success(), - "docker CUDA fixture failed with status {run}" - ); -} - -fn copy_workspace_context(source: &Path, dest: &Path) { - std::fs::create_dir_all(dest).expect("create docker context"); - for entry in std::fs::read_dir(source).expect("read workspace") { - let entry = entry.expect("read workspace entry"); - let name = entry.file_name(); - let name = name.to_string_lossy(); - if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") { - continue; - } - copy_context_entry(&entry.path(), &dest.join(name.as_ref())); - } -} - -fn copy_context_entry(source: &Path, dest: &Path) { - let metadata = std::fs::symlink_metadata(source).expect("context metadata"); - if metadata.file_type().is_symlink() { - return; - } - if metadata.is_dir() { - std::fs::create_dir_all(dest).expect("create context dir"); - for entry in std::fs::read_dir(source).expect("read context dir") { - let entry = entry.expect("read context entry"); - let name = entry.file_name(); - let name = name.to_string_lossy(); - if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") { - continue; - } - copy_context_entry(&entry.path(), &dest.join(name.as_ref())); - } - } else if metadata.is_file() { - if let Some(parent) = dest.parent() { - std::fs::create_dir_all(parent).expect("create context parent"); - } - std::fs::copy(source, dest).expect("copy context file"); - } -} - -fn phase(message: &str) { - eprintln!("gpu-worker-node-e2e: {message}"); -} - -fn run_cuda_preflight() { - let script = r#" -import os -print("cuda preflight: importing tinygrad", flush=True) -from tinygrad import Tensor, dtypes -print(f"cuda preflight: DEV={os.environ.get('DEV')}", flush=True) -print("cuda preflight: realizing Tensor([1])", flush=True) -value = Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() -print(f"cuda preflight: ok {value}", flush=True) -"#; - let child = Command::new("python3") - .arg("-c") - .arg(script) - .env("DEV", "CUDA") - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .expect("spawn CUDA preflight"); - let pid = child.id(); - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - let _ = tx.send(child.wait_with_output()); - }); - - match rx.recv_timeout(PREFLIGHT_WATCHDOG) { - Ok(output) => { - let output = output.expect("wait CUDA preflight"); - eprintln!( - "gpu-worker-node-e2e: CUDA preflight stdout:\n{}", - String::from_utf8_lossy(&output.stdout) - ); - assert!( - output.status.success(), - "CUDA preflight failed with status {}\nstdout:\n{}\nstderr:\n{}", - output.status, - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - } - Err(mpsc::RecvTimeoutError::Timeout) => { - unsafe { - libc::kill(pid as libc::pid_t, libc::SIGKILL); - } - let output = rx - .recv_timeout(Duration::from_secs(5)) - .ok() - .and_then(Result::ok); - let (stdout, stderr) = output - .as_ref() - .map(|output| { - ( - String::from_utf8_lossy(&output.stdout).into_owned(), - String::from_utf8_lossy(&output.stderr).into_owned(), - ) - }) - .unwrap_or_else(|| ("".to_owned(), "".to_owned())); - panic!( - "CUDA preflight test watchdog after {:?}; killed pid {pid}\nstdout:\n{stdout}\nstderr:\n{stderr}", - PREFLIGHT_WATCHDOG - ); - } - Err(mpsc::RecvTimeoutError::Disconnected) => panic!("CUDA preflight waiter disconnected"), - } -} - -fn run_integrated_node_harness() { - phase("running CUDA tinygrad preflight"); - run_cuda_preflight(); - - phase("creating shared arena and telemetry socket"); - let arena_fd = create_arena(ARENA_BYTES); - let worker_path = Path::new(env!("CARGO_MANIFEST_DIR")) - .join("tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py"); - let socket_path = - std::env::temp_dir().join(format!("mvp-worker-events-{}.sock", std::process::id())); - let _ = std::fs::remove_file(&socket_path); - - let (frame_tx, frame_rx) = mpsc::channel(); - let mut emitter = DatastreamEmitter::new( - EmitterConfig { - node_hex: "node-11".to_owned(), - life: 1, - mux_capacity: 256, - }, - Box::new(ChannelFrameSink { tx: frame_tx }), - ); - let ingest_alive = Arc::new(AtomicBool::new(true)); - let ingest_thread = spawn_worker_event_ingest( - socket_path.clone(), - emitter.event_sink(), - Arc::clone(&ingest_alive), - ); - - phase("spawning Rust worker node actor"); - let rt = Runtime::new(RuntimeConfig::default()); - let reports = rt.new_inbox::().expect("report inbox"); - let node = rt - .spawn(GpuWorkerNodeActor::new(rt.create_sender(), *reports.addr())) - .expect("spawn gpu worker node actor"); - - phase("spawning tinygrad worker process"); - rt.send_to( - node, - NodeMsg::Start(StartWorker { - worker_path, - socket_path: socket_path.clone(), - arena_fd, - arena_bytes: ARENA_BYTES as u64, - }), - ) - .expect("send start"); - wait_for_report( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "worker process start", - |report| matches!(report, HarnessReport::ProcessStarted), - ); - - phase("initializing CUDA backend"); - send_command( - &rt, - node, - json!({"type":"InitializeWorker","helper_abi_version":1}), - ); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "WorkerReady", - "worker ready", - ); - - phase("installing ingress and egress rings"); - send_command( - &rt, - node, - install_ring_command( - INGRESS_RING_ID, - INGRESS_EDGE_ID, - "in", - "ingress", - INGRESS_BASE, - ), - ); - send_command( - &rt, - node, - install_ring_command(EGRESS_RING_ID, EGRESS_EDGE_ID, "out", "egress", EGRESS_BASE), - ); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "RingInstalled", - "ingress ring installed", - ); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "RingInstalled", - "egress ring installed", - ); - - phase("configuring worker role"); - send_command(&rt, node, json!({"type":"ConfigureRole","role_id":1})); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "RoleLoaded", - "role loaded", - ); - - phase("copying ingress object into CUDA tensor"); - let input_record = object_record(9000, 0, &[1, 2, 3, 4]); - pwrite_all(arena_fd, INGRESS_BASE, &input_record); - send_command( - &rt, - node, - json!({"type":"RingReadable","ring_id":INGRESS_RING_ID,"committed_bytes":input_record.len()}), - ); - let loaded = wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "ObjectLoaded", - "object loaded", - ); - let handle = loaded["handle"]["id"].as_u64().expect("device handle id"); - - phase("executing CUDA step and writing egress object"); - send_command( - &rt, - node, - json!({ - "type":"ExecuteStep", - "step_id":9001, - "input_handle":handle, - "egress_ring_id":EGRESS_RING_ID, - "output_object_id":9001 - }), - ); - let produced = wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "ObjectProduced", - "object produced", - ); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "StepCompleted", - "step completed", - ); - let committed = produced["committed_bytes"] - .as_u64() - .expect("committed bytes") as usize; - let mut egress = vec![0u8; committed]; - pread_exact(arena_fd, EGRESS_BASE, &mut egress); - assert_eq!(decode_payload_words(&egress), vec![2, 4, 6, 8]); - - phase("releasing device object"); - send_command( - &rt, - node, - json!({"type":"ReleaseDeviceObject","handle":handle}), - ); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "DeviceObjectReleased", - "device object released", - ); - - phase("shutting down worker process"); - send_command(&rt, node, json!({"type":"ShutdownWorker"})); - wait_for_control( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "WorkerStopped", - "worker stopped", - ); - wait_for_report( - &rt, - node, - &mut emitter, - &frame_rx, - &reports, - "worker process exit", - |report| matches!(report, HarnessReport::ProcessExited(0)), - ); - - phase("asserting worker telemetry"); - let telemetry = collect_telemetry(&mut emitter, &frame_rx, Duration::from_secs(1)); - assert_has_worker_event(&telemetry, "importing_tinygrad"); - assert_has_worker_event(&telemetry, "tinygrad_imported"); - assert_has_worker_event(&telemetry, "realizing_cuda_probe"); - assert_has_worker_event(&telemetry, "backend_initialized"); - assert_has_worker_event(&telemetry, "worker_ready"); - assert_has_worker_event(&telemetry, "ring_installed"); - assert_has_worker_event(&telemetry, "role_loaded"); - assert_has_worker_event(&telemetry, "object_copy_started"); - assert_has_worker_event_with(&telemetry, "object_loaded", |event| { - event["object_id"] == 9000 && event["sequence"] == 0 && event["device_sum"] == 10 - }); - assert_has_worker_event(&telemetry, "execute_step_started"); - assert_has_worker_event_with(&telemetry, "object_produced", |event| { - event["object_id"] == 9001 && event["sequence"] == 0 && event["device_sum"] == 20 - }); - assert_has_worker_event(&telemetry, "step_completed"); - assert_has_worker_event(&telemetry, "device_object_released"); - assert_has_worker_event(&telemetry, "worker_stopped"); - - ingest_alive.store(false, Ordering::SeqCst); - let _ = ingest_thread.join(); - let _ = std::fs::remove_file(&socket_path); - unsafe { - libc::close(arena_fd); - } -} - -struct ChannelFrameSink { - tx: mpsc::Sender, -} - -impl FrameSink for ChannelFrameSink { - fn ship(&mut self, _stream: &StreamId, frame: &Frame) { - self.tx.send(frame.clone()).expect("ship datastream frame"); - } -} - -#[derive(Clone)] -struct StartWorker { - worker_path: PathBuf, - socket_path: PathBuf, - arena_fd: RawFd, - arena_bytes: u64, -} - -#[derive(Clone)] -enum NodeMsg { - Start(StartWorker), - SendCommand(Value), - ControlLine(String), - StderrLine(String), - ProcessExited(i32), - KillWorker(String), -} - -#[derive(Clone, Debug)] -enum HarnessReport { - ProcessStarted, - ControlEvent(Value), - StderrLine(String), - ProcessExited(i32), -} - -struct GpuWorkerNodeActor { - sender: ExternalSender, - report_to: ActorAddress, - stdin: Option, - child_pid: Option, -} - -impl GpuWorkerNodeActor { - fn new(sender: ExternalSender, report_to: ActorAddress) -> Self { - Self { - sender, - report_to, - stdin: None, - child_pid: None, - } - } -} - -impl ActorInterface for GpuWorkerNodeActor { - type Incoming = NodeMsg; - type Response = (); - - fn handle(&mut self, ctx: &Ctx, msg: NodeMsg) { - match msg { - NodeMsg::Start(start) => self.start_worker(ctx, start), - NodeMsg::SendCommand(value) => { - let stdin = self.stdin.as_mut().expect("worker stdin"); - writeln!(stdin, "{}", value).expect("write worker command"); - stdin.flush().expect("flush worker command"); - } - NodeMsg::ControlLine(line) => { - let value: Value = serde_json::from_str(&line).expect("control JSON"); - ctx.send(self.report_to, HarnessReport::ControlEvent(value)) - .expect("send control report"); - } - NodeMsg::StderrLine(line) => { - ctx.send(self.report_to, HarnessReport::StderrLine(line)) - .expect("send stderr report"); - } - NodeMsg::ProcessExited(code) => { - self.child_pid = None; - ctx.send(self.report_to, HarnessReport::ProcessExited(code)) - .expect("send exit report"); - } - NodeMsg::KillWorker(reason) => { - self.kill_worker(&reason); - } - } - } -} - -impl GpuWorkerNodeActor { - fn start_worker(&mut self, ctx: &Ctx, start: StartWorker) { - let mut command = Command::new(start.worker_path); - command - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .env("SWACTOR_ARENA_FD", start.arena_fd.to_string()) - .env("SWACTOR_ARENA_BYTES", start.arena_bytes.to_string()) - .env("SWACTOR_WORKER_EVENT_SOCK", &start.socket_path) - .env("SWACTOR_NODE_ID", "11") - .env("SWACTOR_RUN_ID", "77") - .env("SWACTOR_STAGE_INDEX", "0") - .env("DEV", "CUDA"); - unsafe { - command.pre_exec(|| Ok(())); - } - let mut child = command.spawn().expect("spawn tinygrad worker"); - let stdout = child.stdout.take().expect("worker stdout"); - let stderr = child.stderr.take().expect("worker stderr"); - self.stdin = Some(child.stdin.take().expect("worker stdin")); - self.child_pid = Some(child.id()); - - let target = ctx.self_addr(); - let sender = self.sender.clone(); - thread::spawn(move || { - for line in BufReader::new(stdout).lines() { - match line { - Ok(line) => { - if sender.send_to(target, NodeMsg::ControlLine(line)).is_err() { - break; - } - } - Err(_) => break, - } - } - }); - - let target = ctx.self_addr(); - let sender = self.sender.clone(); - thread::spawn(move || { - for line in BufReader::new(stderr).lines() { - match line { - Ok(line) => { - eprintln!("gpu-worker-node-e2e worker stderr: {line}"); - if sender.send_to(target, NodeMsg::StderrLine(line)).is_err() { - break; - } - } - Err(_) => break, - } - } - }); - - let target = ctx.self_addr(); - let sender = self.sender.clone(); - thread::spawn(move || { - let code = child - .wait() - .ok() - .and_then(|status| status.code()) - .unwrap_or(-1); - let _ = sender.send_to(target, NodeMsg::ProcessExited(code)); - }); - - ctx.send(self.report_to, HarnessReport::ProcessStarted) - .expect("send started report"); - } - - fn kill_worker(&mut self, reason: &str) { - if let Some(pid) = self.child_pid.take() { - eprintln!("gpu-worker-node-e2e: killing worker pid {pid}: {reason}"); - unsafe { - libc::kill(pid as libc::pid_t, libc::SIGKILL); - } - } - } -} - -fn spawn_worker_event_ingest( - socket_path: PathBuf, - sink: datastream::emit::DatastreamEventSink, - alive: Arc, -) -> thread::JoinHandle<()> { - let (ready_tx, ready_rx) = mpsc::channel(); - let handle = thread::spawn(move || { - let socket = std::os::unix::net::UnixDatagram::bind(&socket_path).expect("bind UDS ingest"); - ready_tx.send(()).expect("signal UDS ingest ready"); - socket - .set_read_timeout(Some(Duration::from_millis(50))) - .expect("set UDS read watchdog"); - let mut buf = vec![0u8; 8192]; - while alive.load(Ordering::SeqCst) { - match socket.recv(&mut buf) { - Ok(len) => { - sink.submit_bytes(WORKER_EVENTS_CHANNEL, buf[..len].to_vec()); - } - Err(error) - if error.kind() == std::io::ErrorKind::WouldBlock - || error.kind() == std::io::ErrorKind::TimedOut => {} - Err(_) => break, - } - } - }); - ready_rx - .recv_timeout(Duration::from_secs(5)) - .expect("UDS ingest socket ready"); - handle -} - -fn send_command(rt: &Runtime, node: ActorAddress, command: Value) { - rt.send_to(node, NodeMsg::SendCommand(command)) - .expect("send node command"); -} - -fn wait_for_control( - rt: &Runtime, - node: ActorAddress, - emitter: &mut DatastreamEmitter, - frame_rx: &mpsc::Receiver, - reports: &swactor::runtime::Inbox, - kind: &str, - phase_name: &str, -) -> Value { - match wait_for_report( - rt, - node, - emitter, - frame_rx, - reports, - phase_name, - |report| matches!(report, HarnessReport::ControlEvent(value) if value["type"] == kind), - ) { - HarnessReport::ControlEvent(value) => value, - other => panic!("unexpected report for {kind}: {other:?}"), - } -} - -fn wait_for_report( - rt: &Runtime, - node: ActorAddress, - emitter: &mut DatastreamEmitter, - _frame_rx: &mpsc::Receiver, - reports: &swactor::runtime::Inbox, - phase_name: &str, - mut predicate: impl FnMut(&HarnessReport) -> bool, -) -> HarnessReport { - let started = Instant::now(); - let mut stderr_lines = Vec::new(); - while started.elapsed() < EVENT_WATCHDOG { - rt.tick(); - emitter.tick(); - while let Some(report) = reports.try_recv() { - if predicate(&report) { - return report; - } - match &report { - HarnessReport::StderrLine(line) => stderr_lines.push(line.clone()), - HarnessReport::ControlEvent(value) if value["type"] == "WorkerFatal" => { - rt.send_to( - node, - NodeMsg::KillWorker(format!("fatal while waiting for {phase_name}")), - ) - .expect("send kill after fatal"); - rt.tick(); - panic!( - "worker fatal while waiting for {phase_name}: {value}\nstderr={stderr_lines:?}" - ); - } - HarnessReport::ProcessExited(code) if *code != 0 => { - panic!( - "worker exited with {code} while waiting for {phase_name}; stderr={stderr_lines:?}" - ); - } - _ => {} - } - } - thread::sleep(Duration::from_millis(5)); - } - panic!( - "test watchdog after {:?} waiting for {phase_name}; stderr={stderr_lines:?}", - EVENT_WATCHDOG - ); -} - -fn collect_telemetry( - emitter: &mut DatastreamEmitter, - frame_rx: &mpsc::Receiver, - duration: Duration, -) -> Vec { - let started = Instant::now(); - let mut frames = Vec::new(); - while started.elapsed() < duration { - emitter.tick(); - while let Ok(frame) = frame_rx.try_recv() { - frames.push(frame); - } - thread::sleep(Duration::from_millis(10)); - } - frames - .into_iter() - .filter(|frame| frame.channel.as_str() == WORKER_EVENTS_CHANNEL) - .map(|frame| serde_json::from_slice::(&frame.payload).expect("telemetry JSON")) - .collect() -} - -fn assert_has_worker_event(events: &[Value], kind: &str) { - assert_has_worker_event_with(events, kind, |_| true); -} - -fn assert_has_worker_event_with(events: &[Value], kind: &str, extra: impl Fn(&Value) -> bool) { - assert!( - events.iter().any(|event| { - event["schema"] == "mvp.worker.event.v1" - && event["kind"] == kind - && event["node_id"] == 11 - && event["run_id"] == 77 - && event["stage_index"] == 0 - && event["worker_generation"] == 1 - && extra(event) - }), - "missing worker event {kind}; events={events:#?}" - ); -} - -fn install_ring_command( - ring_id: u64, - edge_id: u64, - port_id: &str, - direction: &str, - base: usize, -) -> Value { - json!({ - "type":"InstallRing", - "ring_id": ring_id, - "edge_id": edge_id, - "port_id": port_id, - "direction": direction, - "base": base, - "bytes": RING_BYTES, - "object_spec": { - "max_extent": 16, - "alignment": 4, - "layout": "token" - } - }) -} - -fn object_record(object_id: u64, sequence: u64, words: &[i32]) -> Vec { - let mut record = vec![0u8; HEADER_LEN]; - record[0..4].copy_from_slice(b"MO01"); - record[4] = 1; - record[5] = HEADER_LEN as u8; - record[8..16].copy_from_slice(&object_id.to_le_bytes()); - record[16..24].copy_from_slice(&sequence.to_le_bytes()); - record[24..32].copy_from_slice(&((words.len() * 4) as u64).to_le_bytes()); - record[32..40].copy_from_slice(&16u64.to_le_bytes()); - record[40..48].copy_from_slice(&4u64.to_le_bytes()); - for word in words { - record.extend_from_slice(&word.to_le_bytes()); - } - record -} - -fn decode_payload_words(record: &[u8]) -> Vec { - assert_eq!(&record[0..4], b"MO01"); - let extent = u64::from_le_bytes(record[24..32].try_into().unwrap()) as usize; - record[HEADER_LEN..HEADER_LEN + extent] - .chunks_exact(4) - .map(|chunk| i32::from_le_bytes(chunk.try_into().unwrap())) - .collect() -} - -fn create_arena(bytes: usize) -> RawFd { - let name = CString::new("mvp-system-gpu-worker-node-e2e").expect("memfd name"); - let fd = unsafe { libc::memfd_create(name.as_ptr(), 0) }; - assert!( - fd >= 0, - "memfd_create failed: {}", - std::io::Error::last_os_error() - ); - let truncate = unsafe { libc::ftruncate(fd, bytes as libc::off_t) }; - assert_eq!( - truncate, - 0, - "ftruncate failed: {}", - std::io::Error::last_os_error() - ); - fd -} - -fn pwrite_all(fd: RawFd, offset: usize, bytes: &[u8]) { - let written = unsafe { - libc::pwrite( - fd, - bytes.as_ptr().cast(), - bytes.len(), - offset as libc::off_t, - ) - }; - assert_eq!( - written, - bytes.len() as isize, - "pwrite failed: {}", - std::io::Error::last_os_error() - ); -} - -fn pread_exact(fd: RawFd, offset: usize, bytes: &mut [u8]) { - let read = unsafe { - libc::pread( - fd, - bytes.as_mut_ptr().cast(), - bytes.len(), - offset as libc::off_t, - ) - }; - assert_eq!( - read, - bytes.len() as isize, - "pread failed: {}", - std::io::Error::last_os_error() - ); -} diff --git a/crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile b/crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile deleted file mode 100644 index af04281..0000000 --- a/crates/mvp-system/tests/gpu_worker_node_e2e/Dockerfile +++ /dev/null @@ -1,23 +0,0 @@ -FROM nvidia/cuda:12.6.3-runtime-ubuntu24.04 - -RUN apt-get update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - ca-certificates \ - curl \ - build-essential \ - pkg-config \ - python3 \ - python3-pip \ - cuda-cudart-dev-12-6 \ - cuda-nvrtc-12-6 && \ - python3 -m pip install --no-cache-dir --break-system-packages tinygrad==0.12.0 numpy && \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* - -ENV PATH=/root/.cargo/bin:$PATH -ENV DEV=CUDA -ENV PYTHONDONTWRITEBYTECODE=1 -COPY . /workspace -WORKDIR /workspace -CMD ["sh", "-lc", "python3 -c 'import os; print(\"cuda preflight: importing tinygrad\", flush=True); from tinygrad import Tensor,dtypes; print(\"cuda preflight: DEV=\" + str(os.environ.get(\"DEV\")), flush=True); print(\"cuda preflight: realizing Tensor([1])\", flush=True); print(Tensor([1], dtype=dtypes.int32).realize().numpy().tolist(), flush=True)' && cargo test -p mvp-system --features local-e2e --test gpu_worker_node_e2e -- --nocapture"] diff --git a/crates/mvp-system/tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py b/crates/mvp-system/tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py deleted file mode 100755 index 278647a..0000000 --- a/crates/mvp-system/tests/gpu_worker_node_e2e/mvp_tinygrad_worker.py +++ /dev/null @@ -1,228 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import json -import mmap -import os -import socket -import struct -import sys -import time -from typing import Any - - -HEADER_LEN = 48 -GENERATION = 1 - -arena: mmap.mmap | None = None -telemetry: socket.socket | None = None -telemetry_path = os.environ["SWACTOR_WORKER_EVENT_SOCK"] -rings: dict[int, dict[str, Any]] = {} -objects: dict[int, dict[str, Any]] = {} -role_configured = False -Tensor: Any = None -dtypes: Any = None - -next_handle = 42 - - -def control(**event: Any) -> None: - print(json.dumps(event, separators=(",", ":")), flush=True) - - -def log(message: str) -> None: - print(f"mvp_tinygrad_worker: {message}", file=sys.stderr, flush=True) - - -def observe(kind: str, **fields: Any) -> None: - global telemetry - if telemetry is None: - telemetry = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) - event = { - "schema": "mvp.worker.event.v1", - "kind": kind, - "node_id": int(os.environ["SWACTOR_NODE_ID"]), - "run_id": int(os.environ["SWACTOR_RUN_ID"]), - "stage_index": int(os.environ["SWACTOR_STAGE_INDEX"]), - "worker_generation": GENERATION, - "ts_ns": time.time_ns(), - } - event.update(fields) - telemetry.sendto(json.dumps(event, separators=(",", ":")).encode(), telemetry_path) - - -def fatal(reason: str, **fields: Any) -> None: - observe("worker_fatal", reason=reason, **fields) - control(type="WorkerFatal", reason=reason, **fields) - raise SystemExit(1) - - -def require_arena() -> mmap.mmap: - if arena is None: - fatal("ArenaNotMapped") - return arena - - -def require_tinygrad() -> tuple[Any, Any]: - if Tensor is None or dtypes is None: - fatal("BackendNotInitialized") - return Tensor, dtypes - - -def initialize(cmd: dict[str, Any]) -> None: - global arena, Tensor, dtypes - if int(cmd["helper_abi_version"]) != 1: - fatal("UnsupportedHelperAbi", helper_abi_version=cmd["helper_abi_version"]) - fd = int(os.environ["SWACTOR_ARENA_FD"]) - size = int(os.environ["SWACTOR_ARENA_BYTES"]) - arena = mmap.mmap(fd, size) - - log("importing tinygrad") - observe("importing_tinygrad") - import_start = time.monotonic() - from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes - - Tensor = TinyTensor - dtypes = tiny_dtypes - observe("tinygrad_imported", elapsed_ms=int((time.monotonic() - import_start) * 1000)) - - log(f"realizing CUDA probe with DEV={os.environ.get('DEV')}") - observe("realizing_cuda_probe", dev=os.environ.get("DEV")) - probe_start = time.monotonic() - Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() - observe("backend_initialized", elapsed_ms=int((time.monotonic() - probe_start) * 1000)) - observe("worker_ready") - control(type="WorkerReady", generation=GENERATION) - - -def install_ring(cmd: dict[str, Any]) -> None: - ring_id = int(cmd["ring_id"]) - rings[ring_id] = { - "ring_id": ring_id, - "edge_id": int(cmd["edge_id"]), - "port_id": cmd["port_id"], - "direction": cmd["direction"], - "base": int(cmd["base"]), - "bytes": int(cmd["bytes"]), - "max_extent": int(cmd["object_spec"]["max_extent"]), - "alignment": int(cmd["object_spec"]["alignment"]), - } - observe("ring_installed", ring_id=ring_id, edge_id=rings[ring_id]["edge_id"], direction=rings[ring_id]["direction"]) - control(type="RingInstalled", ring_id=ring_id) - - -def configure_role(cmd: dict[str, Any]) -> None: - global role_configured - role_configured = True - observe("role_loaded", role_id=int(cmd["role_id"])) - control(type="RoleLoaded", role_id=int(cmd["role_id"])) - - -def parse_record(base: int, committed_bytes: int, ring: dict[str, Any]) -> tuple[int, int, int, bytes]: - view = require_arena() - if committed_bytes < HEADER_LEN: - fatal("MalformedHeaderLength", ring_id=ring["ring_id"]) - header = view[base : base + HEADER_LEN] - if header[0:4] != b"MO01" or header[4] != 1 or header[5] != HEADER_LEN: - fatal("InvalidObjectHeader", ring_id=ring["ring_id"]) - object_id = struct.unpack_from(" ring["max_extent"] or (ring["alignment"] and extent % ring["alignment"]): - fatal("ObjectExtentInvalid", ring_id=ring["ring_id"], object_id=object_id, extent=extent) - total = HEADER_LEN + extent - if committed_bytes < total: - fatal("EofBeforeFullPayload", ring_id=ring["ring_id"], object_id=object_id) - payload = bytes(view[base + HEADER_LEN : base + total]) - return object_id, sequence, extent, payload - - -def ring_readable(cmd: dict[str, Any]) -> None: - global next_handle - Tensor, dtypes = require_tinygrad() - ring_id = int(cmd["ring_id"]) - ring = rings[ring_id] - if ring["direction"] != "ingress": - fatal("WrongRingDirection", ring_id=ring_id) - object_id, sequence, extent, payload = parse_record(ring["base"], int(cmd["committed_bytes"]), ring) - observe("object_copy_started", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent) - values = list(struct.unpack(f"<{extent // 4}i", payload)) - tensor = Tensor(values, dtype=dtypes.int32).realize() - handle = next_handle - next_handle += 1 - device_sum = int(tensor.sum().item()) - objects[handle] = {"object_id": object_id, "sequence": sequence, "tensor": tensor, "extent": extent} - observe("object_loaded", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent, handle=handle, device_sum=device_sum) - control(type="ObjectLoaded", ring_id=ring_id, edge_id=ring["edge_id"], object_id=object_id, sequence=sequence, extent=extent, handle={"generation": GENERATION, "id": handle}) - - -def write_record(ring: dict[str, Any], object_id: int, sequence: int, words: list[int]) -> int: - payload = b"".join(struct.pack(" ring["max_extent"]: - fatal("OutputExtentInvalid", ring_id=ring["ring_id"], extent=extent) - header = bytearray(HEADER_LEN) - header[0:4] = b"MO01" - header[4] = 1 - header[5] = HEADER_LEN - struct.pack_into(" None: - if not role_configured: - fatal("RoleNotConfigured") - handle = int(cmd["input_handle"]) - step_id = int(cmd["step_id"]) - output_object_id = int(cmd["output_object_id"]) - egress_ring_id = int(cmd["egress_ring_id"]) - obj = objects[handle] - observe("execute_step_started", step_id=step_id, object_id=obj["object_id"], sequence=obj["sequence"], handle=handle) - output = (obj["tensor"] * 2).realize() - words = [int(value) for value in output.numpy().tolist()] - ring = rings[egress_ring_id] - committed = write_record(ring, output_object_id, int(obj["sequence"]), words) - output_sum = sum(words) - observe("object_produced", ring_id=egress_ring_id, edge_id=ring["edge_id"], object_id=output_object_id, sequence=obj["sequence"], extent=len(words) * 4, device_sum=output_sum, committed_bytes=committed) - observe("step_completed", step_id=step_id) - control(type="ObjectProduced", ring_id=egress_ring_id, object_id=output_object_id, sequence=obj["sequence"], committed_bytes=committed) - control(type="StepCompleted", step_id=step_id) - - -def release_device_object(cmd: dict[str, Any]) -> None: - handle = int(cmd["handle"]) - objects.pop(handle, None) - observe("device_object_released", handle=handle) - control(type="DeviceObjectReleased", handle=handle) - - -handlers = { - "InitializeWorker": initialize, - "InstallRing": install_ring, - "ConfigureRole": configure_role, - "RingReadable": ring_readable, - "ExecuteStep": execute_step, - "ReleaseDeviceObject": release_device_object, -} - -for raw in sys.stdin: - if not raw.strip(): - continue - command = json.loads(raw) - if command["type"] == "ShutdownWorker": - observe("worker_stopped") - control(type="WorkerStopped", generation=GENERATION) - break - handlers[command["type"]](command) diff --git a/crates/mvp-system/tests/local_e2e_cluster.rs b/crates/mvp-system/tests/local_e2e_cluster.rs deleted file mode 100644 index b9694af..0000000 --- a/crates/mvp-system/tests/local_e2e_cluster.rs +++ /dev/null @@ -1,356 +0,0 @@ -#![recursion_limit = "256"] - -use std::path::Path; -use std::process::{Command, ExitCode}; -use std::time::{Duration, Instant}; - -#[path = "support/local_e2e_cluster.rs"] -mod local_e2e_cluster; - -const IMAGE: &str = "swactor-mvp-local-e2e-cluster:latest"; -const SKIP_BUILD_ENV: &str = "MVP_LOCAL_E2E_CLUSTER_SKIP_BUILD"; -const BUILD_ONLY_ENV: &str = "MVP_LOCAL_E2E_CLUSTER_BUILD_ONLY"; - -fn main() -> ExitCode { - let args = std::env::args().collect::>(); - match std::env::var("MVP_TEST_ROLE").ok().as_deref() { - Some("cluster-supervisor" | "cluster-relay") => return local_e2e_cluster::run_main(), - Some(role) => { - eprintln!("unknown MVP_TEST_ROLE={role}"); - return ExitCode::from(2); - } - None => {} - } - - if args.iter().any(|arg| arg == "--role=node") { - return local_e2e_cluster::run_main(); - } - - local_e2e_cluster_docker_cpu_pipeline_prompt(); - ExitCode::SUCCESS -} - -fn local_e2e_cluster_docker_cpu_pipeline_prompt() { - if std::env::var_os("MVP_SYSTEM_LOCAL_E2E_CLUSTER").is_none() { - eprintln!("skipping; set MVP_SYSTEM_LOCAL_E2E_CLUSTER=1 to run Docker CPU cluster e2e"); - return; - } - if !Path::new("/var/run/docker.sock").exists() { - eprintln!( - "skipping; /var/run/docker.sock is required for the relay-only Docker cluster e2e" - ); - return; - } - - build_docker_fixture(); - if std::env::var_os(BUILD_ONLY_ENV).is_some() { - return; - } - - let docker = DockerRelayFixture::start(); - let output = docker.run_supervisor("ping"); - - assert!( - output.status.success(), - "mvp-local-e2e-cluster failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - - let value: serde_json::Value = serde_json::from_slice(&output.stdout).expect("json stdout"); - assert_eq!(value["ok"], true); - assert_eq!(value["actor_plane"], "iroh-swactor"); - assert_eq!(value["data_plane"], "iroh-quic-persistent-edge-streams"); - assert_eq!( - value["edge_protocol"], - "edge-id-preamble-mo01-object-records" - ); - assert_eq!( - value["node_local_data_plane"], - "arena-backed-rings-json-metadata-only" - ); - assert_eq!( - value["worker_processes"], - "docker-tinygrad-cpu-worker-per-node" - ); - assert_eq!(value["tinygrad_device"], "CPU"); - assert_eq!(value["prompt_text"], "ping"); - assert_eq!(value["response_text"], "pong"); - assert_eq!(value["response_tokens"].as_array().map(Vec::len), Some(1)); - assert_eq!( - value["engine_builder_pattern"], - "host-coordinator-static-topology-docker-workers" - ); - assert_eq!(value["engine_builder_node_count"], 3); - assert_eq!(value["engine_builder_stage_assignments"], 2); - assert_eq!(value["injected_prompt_observed"], true); - assert_eq!(value["token_received_observed"], true); - assert_eq!(value["run_completed_observed"], true); - assert_eq!(value["run_torn_down_observed"], true); - assert_eq!(value["stop_sent_to_all_nodes"], true); - assert_eq!(value["stage_ready_stdout_count"], 2); - assert_eq!(value["provisioned_node_count"], 2); - assert_eq!(value["provision_node_live_count"], 2); - assert!( - value["provision_stdout_line_count"] - .as_u64() - .is_some_and(|count| count >= 2), - "{value}" - ); - assert!( - value["provision_stderr_line_count"] - .as_u64() - .is_some_and(|count| count >= 2), - "{value}" - ); - assert_eq!(value["provision_nodes_stopped"], true); - assert_eq!(value["relay_only"], true, "{value}"); - assert_eq!(value["relay_url"], "http://relay:7843/"); - assert_eq!(value["orchestrator_endpoint_has_relay"], true, "{value}"); - assert_eq!( - value["orchestrator_endpoint_relay_url"], - "http://relay:7843/" - ); - assert_eq!(value["node0_endpoint_has_relay"], true, "{value}"); - assert_eq!(value["node0_endpoint_relay_url"], "http://relay:7843/"); - assert_eq!(value["node1_endpoint_has_relay"], true, "{value}"); - assert_eq!(value["node1_endpoint_relay_url"], "http://relay:7843/"); - assert!( - value["node0_endpoint"]["addrs"] - .as_array() - .is_some_and(|addrs| !addrs.is_empty()), - "{value}" - ); - assert!( - value["node1_endpoint"]["addrs"] - .as_array() - .is_some_and(|addrs| !addrs.is_empty()), - "{value}" - ); -} -struct DockerRelayFixture { - relay_container: String, - supervisor_network: String, - node0_network: String, - node1_network: String, -} - -impl DockerRelayFixture { - fn start() -> Self { - let suffix = format!("{}-{}", std::process::id(), unique_nanos()); - let relay_container = format!("mvp-local-e2e-relay-{suffix}"); - let supervisor_network = format!("mvp-local-e2e-supervisor-{suffix}"); - let node0_network = format!("mvp-local-e2e-node0-{suffix}"); - let node1_network = format!("mvp-local-e2e-node1-{suffix}"); - for network in [&supervisor_network, &node0_network, &node1_network] { - docker_status( - ["network", "create", network], - "create relay-only Docker network", - ); - } - docker_status( - [ - "run", - "-d", - "--rm", - "--name", - &relay_container, - "--network", - &supervisor_network, - "--network-alias", - "relay", - "-e", - "MVP_TEST_ROLE=cluster-relay", - "-e", - "MVP_LOCAL_E2E_RELAY_LISTEN=0.0.0.0:7843", - IMAGE, - ], - "start relay sidecar", - ); - docker_status( - [ - "network", - "connect", - "--alias", - "relay", - &node0_network, - &relay_container, - ], - "attach relay to node0 network", - ); - docker_status( - [ - "network", - "connect", - "--alias", - "relay", - &node1_network, - &relay_container, - ], - "attach relay to node1 network", - ); - let fixture = Self { - relay_container, - supervisor_network, - node0_network, - node1_network, - }; - fixture.wait_for_relay(); - fixture - } - - fn run_supervisor(&self, prompt: &str) -> std::process::Output { - Command::new("docker") - .args([ - "run", - "--rm", - "--name", - &format!("mvp-local-e2e-supervisor-{}", unique_nanos()), - "--network", - &self.supervisor_network, - "--network-alias", - "supervisor", - "-v", - "/var/run/docker.sock:/var/run/docker.sock", - "-e", - "MVP_TEST_ROLE=cluster-supervisor", - "-e", - &format!("MVP_LOCAL_E2E_CLUSTER_IMAGE={IMAGE}"), - "-e", - &format!("MVP_LOCAL_E2E_DOCKER_NETWORK_NODE0={}", self.node0_network), - "-e", - &format!("MVP_LOCAL_E2E_DOCKER_NETWORK_NODE1={}", self.node1_network), - "-e", - "MVP_IROH_RELAY_MODE=default", - "-e", - "MVP_IROH_RELAY_URL=http://relay:7843/", - IMAGE, - "--prompt", - prompt, - ]) - .output() - .expect("run relay-only local e2e cluster supervisor") - } - - fn wait_for_relay(&self) { - let started = Instant::now(); - while started.elapsed() < Duration::from_secs(20) { - let logs = Command::new("docker") - .args(["logs", &self.relay_container]) - .output() - .expect("read relay sidecar logs"); - let stdout = String::from_utf8_lossy(&logs.stdout); - let stderr = String::from_utf8_lossy(&logs.stderr); - if stdout.contains("relay ready") || stderr.contains("relay ready") { - return; - } - std::thread::sleep(Duration::from_millis(100)); - } - panic!("relay sidecar did not report ready within 20s"); - } -} - -impl Drop for DockerRelayFixture { - fn drop(&mut self) { - let _ = Command::new("docker") - .args(["stop", "-t", "2", &self.relay_container]) - .status(); - for network in [ - &self.node1_network, - &self.node0_network, - &self.supervisor_network, - ] { - let _ = Command::new("docker") - .args(["network", "rm", network]) - .status(); - } - } -} - -fn docker_status(args: [&str; N], action: &str) { - let status = Command::new("docker") - .args(args) - .status() - .unwrap_or_else(|error| panic!("{action}: {error}")); - assert!(status.success(), "{action} failed with status {status}"); -} - -fn unique_nanos() -> u128 { - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .expect("system time after epoch") - .as_nanos() -} - -fn build_docker_fixture() { - if std::env::var_os(SKIP_BUILD_ENV).is_some() { - phase("using existing Docker CPU cluster fixture image"); - return; - } - - let crate_dir = Path::new(env!("CARGO_MANIFEST_DIR")); - let workspace = crate_dir - .parent() - .and_then(Path::parent) - .expect("workspace root") - .canonicalize() - .expect("canonical workspace root"); - let context = std::env::temp_dir().join(format!( - "mvp-system-local-e2e-cluster-context-{}", - std::process::id() - )); - let _ = std::fs::remove_dir_all(&context); - copy_workspace_context(&workspace, &context); - let dockerfile = context.join("crates/mvp-system/tests/local_e2e_cluster/Dockerfile"); - - phase("building Docker CPU cluster fixture image"); - let build = Command::new("docker") - .args(["build", "-f"]) - .arg(&dockerfile) - .args(["-t", IMAGE]) - .arg(&context) - .status() - .expect("run docker build"); - assert!(build.success(), "docker build failed with status {build}"); -} - -fn copy_workspace_context(source: &Path, dest: &Path) { - std::fs::create_dir_all(dest).expect("create docker context"); - for entry in std::fs::read_dir(source).expect("read workspace") { - let entry = entry.expect("read workspace entry"); - let name = entry.file_name(); - let name = name.to_string_lossy(); - if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") { - continue; - } - copy_context_entry(&entry.path(), &dest.join(name.as_ref())); - } -} - -fn copy_context_entry(source: &Path, dest: &Path) { - let metadata = std::fs::symlink_metadata(source).expect("context metadata"); - if metadata.file_type().is_symlink() { - return; - } - if metadata.is_dir() { - std::fs::create_dir_all(dest).expect("create context dir"); - for entry in std::fs::read_dir(source).expect("read context dir") { - let entry = entry.expect("read context entry"); - let name = entry.file_name(); - let name = name.to_string_lossy(); - if matches!(name.as_ref(), ".git" | "target" | ".dockerignore") { - continue; - } - copy_context_entry(&entry.path(), &dest.join(name.as_ref())); - } - } else if metadata.is_file() { - if let Some(parent) = dest.parent() { - std::fs::create_dir_all(parent).expect("create context parent"); - } - std::fs::copy(source, dest).expect("copy context file"); - } -} - -fn phase(message: &str) { - eprintln!("local-e2e-cluster: {message}"); -} diff --git a/crates/mvp-system/tests/local_e2e_cluster/Dockerfile b/crates/mvp-system/tests/local_e2e_cluster/Dockerfile deleted file mode 100644 index ddaadfd..0000000 --- a/crates/mvp-system/tests/local_e2e_cluster/Dockerfile +++ /dev/null @@ -1,24 +0,0 @@ -FROM rust:1-bookworm - -RUN apt-get update && \ - DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \ - ca-certificates \ - docker.io \ - pkg-config \ - python3 \ - python3-pip && \ - python3 -m pip install --no-cache-dir --break-system-packages tinygrad==0.12.0 numpy && \ - apt-get clean && \ - rm -rf /var/lib/apt/lists/* - -ENV DEV=CPU -ENV PYTHONDONTWRITEBYTECODE=1 -ENV CARGO_TARGET_DIR=/tmp/mvp-local-e2e-target -COPY . /workspace -WORKDIR /workspace -RUN cargo test -p mvp-system --features local-e2e --test local-e2e-cluster --no-run && \ - test_bin="$(find /tmp/mvp-local-e2e-target/debug/deps -maxdepth 1 -type f -perm /111 \( -name 'local_e2e_cluster-*' -o -name 'local-e2e-cluster-*' \) | head -n 1)" && \ - cp "$test_bin" /usr/local/bin/local-e2e-cluster && \ - rm -rf /tmp/mvp-local-e2e-target - -ENTRYPOINT ["/usr/local/bin/local-e2e-cluster"] diff --git a/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py b/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py deleted file mode 100755 index ca3ecdf..0000000 --- a/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py +++ /dev/null @@ -1,273 +0,0 @@ -#!/usr/bin/env python3 -from __future__ import annotations - -import json -import mmap -import os -import struct -import sys -from typing import Any - -HEADER_LEN = 40 -worker_generation = 0 -arena: mmap.mmap | None = None -rings: dict[int, dict[str, Any]] = {} -objects: dict[int, dict[str, Any]] = {} -role: dict[str, Any] = {} -Tensor: Any = None -dtypes: Any = None -next_handle = 42 - - -def control(**event: Any) -> None: - print(json.dumps(event, separators=(",", ":")), flush=True) - - -def fatal(reason: str, **fields: Any) -> None: - control(type="WorkerFatal", reason=reason, **fields) - raise SystemExit(1) - - -def require_arena() -> mmap.mmap: - if arena is None: - fatal("ArenaNotMapped") - return arena - - -def require_tinygrad() -> tuple[Any, Any]: - if Tensor is None or dtypes is None: - fatal("BackendNotInitialized") - return Tensor, dtypes - - -def initialize(cmd: dict[str, Any]) -> None: - global arena, Tensor, dtypes, worker_generation - if int(cmd["required_ring_helper_abi"]) != 1: - fatal("UnsupportedHelperAbi", required_ring_helper_abi=cmd["required_ring_helper_abi"]) - worker_generation = int(cmd["worker_generation"]) - fd = int(os.environ["SWACTOR_ARENA_FD"]) - size = int(cmd.get("arena_ceiling", os.environ["SWACTOR_ARENA_BYTES"])) - arena = mmap.mmap(fd, size) - os.environ.setdefault("DEV", cmd.get("backend", {}).get("device", "CPU")) - from tinygrad import Tensor as TinyTensor, dtypes as tiny_dtypes - - Tensor = TinyTensor - dtypes = tiny_dtypes - Tensor([1], dtype=dtypes.int32).realize().numpy().tolist() - control( - type="WorkerReady", - pid=os.getpid(), - worker_generation=worker_generation, - ring_helper_abi=1, - backend={"device": os.environ.get("DEV", "CPU")}, - ) - - -def install_ring(cmd: dict[str, Any]) -> None: - ring_id = int(cmd["ring_id"]) - layout = cmd["layout"] - spec = cmd["object_spec"] - rings[ring_id] = { - "ring_id": ring_id, - "edge_id": int(cmd["edge_id"]), - "port_id": cmd["port_id"], - "direction": cmd["direction"], - "data_offset": int(layout["data_offset"]), - "data_capacity": int(layout["data_capacity"]), - "max_extent": int(spec["max_extent"]), - "alignment": int(spec["alignment"]), - "next_sequence": 0, - } - control(type="RingInstalled", ring_id=ring_id, edge_id=rings[ring_id]["edge_id"], port_id=cmd["port_id"]) - - -def configure_role(cmd: dict[str, Any]) -> None: - config = cmd["config"] - role.clear() - role.update( - role_id=int(cmd["role_id"]), - run_id=int(config["run_id"]), - stage_index=int(config["stage_index"]), - layer_start=int(config["layer_start"]), - layer_end_exclusive=int(config["layer_end_exclusive"]), - ) - control(type="RoleConfigured", role_id=role["role_id"]) - - -def load_weights(cmd: dict[str, Any]) -> None: - if not role: - fatal("RoleNotConfigured") - role["model_id"] = cmd["model_id"] - role["gguf_source"] = cmd["gguf_source"] - role["tokenizer"] = cmd["tokenizer"] - role["weight_layer_start"] = int(cmd["layer_start"]) - role["weight_layer_end_exclusive"] = int(cmd["layer_end_exclusive"]) - control( - type="WeightsLoaded", - model_id=role["model_id"], - layer_start=role["weight_layer_start"], - layer_end_exclusive=role["weight_layer_end_exclusive"], - ) - - -def parse_record(ring: dict[str, Any]) -> tuple[int, int, int, bytes]: - view = require_arena() - base = ring["data_offset"] - header = view[base : base + HEADER_LEN] - version = struct.unpack_from(" ring["max_extent"] or (ring["alignment"] and extent % ring["alignment"]): - fatal("ObjectExtentInvalid", ring_id=ring["ring_id"], object_id=object_id, extent=extent) - if sequence != ring["next_sequence"]: - fatal("SequenceViolation", ring_id=ring["ring_id"], expected=ring["next_sequence"], actual=sequence) - payload = bytes(view[base + HEADER_LEN : base + HEADER_LEN + extent]) - ring["next_sequence"] += 1 - return object_id, sequence, extent, payload - - -def ring_readable(cmd: dict[str, Any]) -> None: - global next_handle - tensor, dtype_mod = require_tinygrad() - ring_id = int(cmd["ring_id"]) - ring = rings[ring_id] - if ring["direction"] != "ingress": - fatal("WrongRingDirection", ring_id=ring_id) - object_id, sequence, extent, payload = parse_record(ring) - values = list(struct.unpack(f"<{extent // 4}i", payload)) - loaded = tensor(values, dtype=dtype_mod.int32).realize() - handle = next_handle - next_handle += 1 - objects[handle] = { - "object_id": object_id, - "sequence": sequence, - "tensor": loaded, - "extent": extent, - } - control( - type="ObjectLoaded", - ring_id=ring_id, - edge_id=ring["edge_id"], - port_id=ring["port_id"], - object_id=object_id, - sequence=sequence, - extent=extent, - device_handle={"worker_generation": worker_generation, "id": handle}, - ) - - -def write_record(ring: dict[str, Any], object_id: int, sequence: int, words: list[int], flags: int) -> int: - payload = b"".join(struct.pack(" ring["max_extent"]: - fatal("OutputExtentInvalid", ring_id=ring["ring_id"], extent=extent) - header = bytearray(HEADER_LEN) - header[0:4] = b"MO01" - struct.pack_into(" None: - if not role: - fatal("RoleNotConfigured") - tensor, _ = require_tinygrad() - del tensor - role_id = int(cmd["role_id"]) - if role_id != role["role_id"]: - fatal("RoleMismatch", expected=role["role_id"], actual=role_id) - step_id = int(cmd["step_id"]) - input_binding = cmd["inputs"][0] - output_binding = cmd["outputs"][0] - device_handle = input_binding["device_handle"] - if int(device_handle["worker_generation"]) != worker_generation: - fatal("OldGenerationHandle", handle=device_handle) - handle = int(device_handle["id"]) - obj = objects[handle] - if int(input_binding["object_id"]) != obj["object_id"] or int(input_binding["sequence"]) != obj["sequence"]: - fatal("InputBindingMismatch", step_id=step_id) - transformed = int(obj["tensor"].sum().item()) + role["layer_start"] + role["layer_end_exclusive"] + role["stage_index"] - if bool(cmd.get("runtime", {}).get("final_stage")): - words = [6 if transformed % 2 == 1 else 8] - else: - words = [transformed if transformed > 0 else 1] - ring = rings[int(output_binding["ring_id"])] - if ring["direction"] != "egress": - fatal("WrongRingDirection", ring_id=ring["ring_id"]) - committed = write_record( - ring, - int(output_binding["object_id"]), - int(output_binding["sequence"]), - words, - int(output_binding.get("flags", 0)), - ) - control( - type="ObjectProduced", - ring_id=ring["ring_id"], - edge_id=ring["edge_id"], - port_id=ring["port_id"], - object_id=int(output_binding["object_id"]), - sequence=int(output_binding["sequence"]), - committed_bytes=committed, - ) - if bool(cmd.get("release_inputs_after")): - objects.pop(handle, None) - control(type="StepCompleted", role_id=role_id, step_id=step_id) - - -def release_device_object(cmd: dict[str, Any]) -> None: - handle = int(cmd["device_handle"]["id"]) - objects.pop(handle, None) - control(type="DeviceObjectReleased", device_handle=cmd["device_handle"]) - - -def uninstall_ring(cmd: dict[str, Any]) -> None: - ring_id = int(cmd["ring_id"]) - rings.pop(ring_id, None) - control(type="RingQuiesced", ring_id=ring_id) - - -def shutdown_worker(_: dict[str, Any]) -> None: - control(type="WorkerStopped", reason="Graceful") - raise SystemExit(0) - - -HANDLERS = { - "InitializeWorker": initialize, - "InstallRing": install_ring, - "ConfigureRole": configure_role, - "LoadWeights": load_weights, - "RingReadable": ring_readable, - "ExecuteStep": execute_step, - "ReleaseDeviceObject": release_device_object, - "UninstallRing": uninstall_ring, - "ShutdownWorker": shutdown_worker, -} - -for raw in sys.stdin: - if not raw.strip(): - continue - try: - command = json.loads(raw) - except json.JSONDecodeError as exc: - fatal("InvalidJson", error=str(exc)) - handler = HANDLERS.get(command.get("type")) - if handler is None: - fatal("UnknownCommand", command=command.get("type")) - handler(command) diff --git a/crates/mvp-system/tests/local_unmocked_mvp_e2e.rs b/crates/mvp-system/tests/local_unmocked_mvp_e2e.rs deleted file mode 100644 index af38b47..0000000 --- a/crates/mvp-system/tests/local_unmocked_mvp_e2e.rs +++ /dev/null @@ -1,119 +0,0 @@ -use std::io::Write; -use std::path::PathBuf; -use std::process::{Command, ExitCode, Stdio}; - -#[path = "support/dumb_worker.rs"] -mod dumb_worker; -#[path = "support/local_e2e.rs"] -mod local_e2e; - -fn main() -> ExitCode { - let args = std::env::args().collect::>(); - match std::env::var("MVP_TEST_ROLE").ok().as_deref() { - Some("dumb-worker") => return dumb_worker::run_main(), - Some("local-e2e") => return local_e2e::run_main(), - Some(role) => { - eprintln!("unknown MVP_TEST_ROLE={role}"); - return ExitCode::from(2); - } - None => {} - } - - if args.iter().any(|arg| arg == "--role=node") { - return local_e2e::run_main(); - } - - local_e2e_binary_drives_real_local_process_deployment(); - dumb_worker_is_a_real_child_process_protocol_endpoint(); - ExitCode::SUCCESS -} - -fn local_e2e_binary_drives_real_local_process_deployment() { - let output = Command::new(current_test_exe()) - .env("MVP_TEST_ROLE", "local-e2e") - .output() - .expect("run local e2e supervisor"); - - assert!( - output.status.success(), - "mvp-local-e2e failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - - let value: serde_json::Value = serde_json::from_slice(&output.stdout).expect("json stdout"); - assert_eq!(value["ok"], true); - assert_eq!(value["actor_plane"], "iroh-swactor"); - assert_eq!(value["data_plane"], "tcp-loopback-streams"); - assert_eq!(value["worker_processes"], "mvp-dumb-worker-per-node"); - assert_eq!( - value["engine_builder_pattern"], - "pool-first-static-launcher" - ); - assert_eq!(value["engine_builder_node_count"], 3); - assert_eq!(value["engine_builder_stage_assignments"], 2); - assert!( - value["engine_builder_event_count"] - .as_u64() - .is_some_and(|count| count >= 10), - "{value}" - ); - assert_eq!(value["injected_prompt_observed"], true); - assert_eq!(value["token_received_observed"], true); - assert_eq!(value["run_completed_observed"], true); - assert_eq!(value["run_torn_down_observed"], true); - assert_eq!(value["stop_sent_to_all_nodes"], true); - assert_eq!(value["stage_ready_stdout_count"], 2); - assert!(value["processes"]["node0"].as_u64().is_some(), "{value}"); - assert!(value["processes"]["node1"].as_u64().is_some(), "{value}"); - assert!( - value["node0_endpoint"]["addrs"] - .as_array() - .is_some_and(|addrs| !addrs.is_empty()), - "{value}" - ); - assert!( - value["node1_endpoint"]["addrs"] - .as_array() - .is_some_and(|addrs| !addrs.is_empty()), - "{value}" - ); -} - -fn dumb_worker_is_a_real_child_process_protocol_endpoint() { - let mut child = Command::new(current_test_exe()) - .env("MVP_TEST_ROLE", "dumb-worker") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .expect("spawn dumb worker test role"); - - { - let stdin = child.stdin.as_mut().expect("worker stdin"); - writeln!( - stdin, - "{{\"type\":\"InitializeWorker\",\"helper_abi_version\":1}}" - ) - .expect("write initialize"); - writeln!(stdin, "{{\"type\":\"ExecuteStep\",\"step_id\":7}}").expect("write execute"); - writeln!(stdin, "{{\"type\":\"ShutdownWorker\"}}").expect("write shutdown"); - } - - let output = child.wait_with_output().expect("worker output"); - assert!( - output.status.success(), - "worker failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - let stdout = String::from_utf8(output.stdout).expect("utf8 stdout"); - assert!(stdout.contains("\"type\":\"WorkerReady\""), "{stdout}"); - assert!(stdout.contains("\"type\":\"StepCompleted\""), "{stdout}"); - assert!(stdout.contains("\"step_id\":7"), "{stdout}"); - assert!(stdout.contains("\"type\":\"WorkerStopped\""), "{stdout}"); -} - -fn current_test_exe() -> PathBuf { - std::env::current_exe().expect("current test exe") -} diff --git a/crates/mvp-system/tests/one_node_chat_e2e.rs b/crates/mvp-system/tests/one_node_chat_e2e.rs deleted file mode 100644 index 4eddd9b..0000000 --- a/crates/mvp-system/tests/one_node_chat_e2e.rs +++ /dev/null @@ -1,594 +0,0 @@ -use std::io::{Read, Write}; -use std::net::TcpStream; -use std::process::{Child, Command, Stdio}; -use std::sync::{Arc, Mutex}; -use std::thread; -use std::time::{Duration, Instant}; - -use serde_json::Value; - -#[cfg(target_os = "linux")] -use std::os::unix::process::CommandExt; - -const TEST_WATCHDOG: Duration = Duration::from_secs(1_800); -const PROMPT_WATCHDOG: Duration = Duration::from_secs(600); -const SHUTDOWN_WATCHDOG: Duration = Duration::from_secs(60); -const DASHBOARD_ADDR: &str = "127.0.0.1:9090"; -const DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; - -#[test] -fn one_node_chat_docker_cuda_e2e() { - let root = workspace_root(); - require_docker(&root); - let container_prefix = format!("mvp-orchestrator-e2e-{}", std::process::id()); - - let mut command = Command::new("cargo"); - command - .current_dir(&root) - .args(["mvp-chat", "--docker"]) - .env("MVP_RUNTIME_CONFIG", "local") - .env("MVP_IROH_RELAY_MODE", "disabled") - .env(DOCKER_CONTAINER_PREFIX_ENV, &container_prefix) - .env("MVP_TINYGRAD_TEST_MODE", "1") - .env("MVP_LAYER_END_EXCLUSIVE", "1") - .env("MVP_PROMPT_MAX_TOKENS", "3") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()); - #[cfg(target_os = "linux")] - unsafe { - command.pre_exec(|| { - if libc::setpgid(0, 0) == 0 { - Ok(()) - } else { - Err(std::io::Error::last_os_error()) - } - }); - } - - let mut child = command.spawn().expect("spawn cargo mvp-chat"); - let mut stdin = child.stdin.take().expect("cargo mvp-chat stdin"); - let stdout = Arc::new(Mutex::new(String::new())); - let stderr = Arc::new(Mutex::new(String::new())); - let stdout_reader = spawn_capture(child.stdout.take().expect("stdout"), Arc::clone(&stdout)); - let stderr_reader = spawn_capture(child.stderr.take().expect("stderr"), Arc::clone(&stderr)); - - let mut result = run_full_flow( - &mut child, - &mut stdin, - &stdout, - "mvp.provisioning.logs", - "mvp-entrypoint", - dashboard_has_worker_prompt_completed, - ); - if result.is_err() { - request_child_interrupt(&child); - let _ = wait_child(&mut child, SHUTDOWN_WATCHDOG); - let _ = child.kill(); - let _ = child.wait(); - } - let _ = stdout_reader.join(); - let _ = stderr_reader.join(); - if result.is_ok() { - result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr); - } - assert_containers_with_prefix_removed(&root, &container_prefix); - - if let Err(error) = result { - panic!( - "{error}\nstdout:\n{}\nstderr:\n{}\ndashboard frames:\n{}", - snapshot(&stdout), - snapshot(&stderr), - dashboard_snapshot().unwrap_or_else(|err| format!("")) - ); - } -} - -#[test] -fn one_node_chat_process_cached_model_e2e() { - let root = workspace_root(); - let cached_model = root - .join(".model-cache") - .join("SmolLM2-135M-Instruct.Q4_0.gguf"); - assert!( - cached_model.is_file(), - "cached model fixture is required: {}", - cached_model.display() - ); - let container_prefix = format!("mvp-orchestrator-process-e2e-{}", std::process::id()); - - let mut command = Command::new("cargo"); - command - .current_dir(&root) - .args(["mvp-chat", "--cached-model"]) - .args(["--pipeline-stages", "2", "--dump-logs"]) - .env("MVP_RUNTIME_CONFIG", "local") - .env("MVP_IROH_RELAY_MODE", "disabled") - .env(DOCKER_CONTAINER_PREFIX_ENV, &container_prefix) - .env("MVP_TINYGRAD_TEST_MODE", "1") - .env("MVP_LAYER_END_EXCLUSIVE", "1") - .env("MVP_PROMPT_MAX_TOKENS", "3") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()); - #[cfg(target_os = "linux")] - unsafe { - command.pre_exec(|| { - if libc::setpgid(0, 0) == 0 { - Ok(()) - } else { - Err(std::io::Error::last_os_error()) - } - }); - } - - let mut child = command.spawn().expect("spawn cargo mvp-chat"); - let mut stdin = child.stdin.take().expect("cargo mvp-chat stdin"); - let stdout = Arc::new(Mutex::new(String::new())); - let stderr = Arc::new(Mutex::new(String::new())); - let stdout_reader = spawn_capture(child.stdout.take().expect("stdout"), Arc::clone(&stdout)); - let stderr_reader = spawn_capture(child.stderr.take().expect("stderr"), Arc::clone(&stderr)); - - let mut result = run_full_flow( - &mut child, - &mut stdin, - &stdout, - "mvp.node.bootstrap", - "runtime_ready_local", - dashboard_has_pipeline_prompt_output, - ); - if result.is_err() { - request_child_interrupt(&child); - let _ = wait_child(&mut child, SHUTDOWN_WATCHDOG); - let _ = child.kill(); - let _ = child.wait(); - } - let _ = stdout_reader.join(); - let _ = stderr_reader.join(); - if result.is_ok() { - result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr); - } - assert_no_containers_with_prefix_if_docker_available(&root, &container_prefix); - - if let Err(error) = result { - panic!( - "{error}\nstdout:\n{}\nstderr:\n{}\ndashboard frames:\n{}", - snapshot(&stdout), - snapshot(&stderr), - dashboard_snapshot().unwrap_or_else(|err| format!("")) - ); - } -} - -fn run_full_flow( - child: &mut Child, - stdin: &mut impl Write, - stdout: &Arc>, - provisioning_channel_substr: &str, - provisioning_payload_substr: &str, - prompt_dashboard: fn() -> bool, -) -> Result<(), String> { - wait_for_child_or(TEST_WATCHDOG, child, dashboard_responding) - .map_err(|e| format!("dashboard API not live: {e}"))?; - wait_for_child_or(TEST_WATCHDOG, child, || { - dashboard_has_channel_or_payload(provisioning_channel_substr, provisioning_payload_substr) - }) - .map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?; - wait_for_child_or(TEST_WATCHDOG, child, || { - dashboard_has_frame("mvp.worker.weights", "LoadWeightsStarted") - }) - .map_err(|e| format!("weight loading start not visible in dashboard: {e}"))?; - wait_for_child_or(TEST_WATCHDOG, child, || { - dashboard_has_frame("mvp.worker.weights", "WeightsLoaded") - }) - .map_err(|e| format!("WeightsLoaded not visible in dashboard: {e}"))?; - wait_for_child_or(TEST_WATCHDOG, 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_WATCHDOG, child, || { - stdout_contains(stdout, "decoding...") - }) - .map_err(|e| format!("prompt was not submitted to chat loop: {e}"))?; - wait_for_child_or(PROMPT_WATCHDOG, child, || response_text_visible(stdout)) - .map_err(|e| format!("decoded response text not visible: {e}"))?; - wait_for_child_or(PROMPT_WATCHDOG, child, prompt_dashboard) - .map_err(|e| format!("prompt result not visible in dashboard: {e}"))?; - wait_for_child_or(PROMPT_WATCHDOG, child, dashboard_has_orch_prompt_lifecycle) - .map_err(|e| format!("orchestrator prompt lifecycle not visible in dashboard: {e}"))?; - wait_for_child_or(PROMPT_WATCHDOG, child, || prompt_count(stdout) >= 2) - .map_err(|e| format!("chat prompt did not return after response: {e}"))?; - writeln!(stdin, "/exit").map_err(|e| format!("write exit command: {e}"))?; - stdin - .flush() - .map_err(|e| format!("flush exit command: {e}"))?; - let status = wait_child(child, SHUTDOWN_WATCHDOG) - .ok_or_else(|| "cargo mvp-chat did not exit after /exit".to_owned())?; - if status.success() { - Ok(()) - } else { - Err(format!("cargo mvp-chat exited with {status}")) - } -} - -fn wait_for_child_or( - watchdog: Duration, - child: &mut Child, - mut predicate: impl FnMut() -> bool, -) -> Result<(), String> { - let start = Instant::now(); - while start.elapsed() < watchdog { - if predicate() { - return Ok(()); - } - if let Some(status) = child.try_wait().map_err(|e| format!("poll child: {e}"))? { - return Err(format!("child exited early with {status}")); - } - thread::sleep(Duration::from_millis(250)); - } - Err("test watchdog".to_owned()) -} - -fn spawn_capture( - mut reader: impl Read + Send + 'static, - out: Arc>, -) -> thread::JoinHandle<()> { - thread::spawn(move || { - let mut buf = [0_u8; 8192]; - loop { - match reader.read(&mut buf) { - Ok(0) => break, - Ok(n) => out - .lock() - .expect("capture mutex") - .push_str(&String::from_utf8_lossy(&buf[..n])), - Err(_) => break, - } - } - }) -} - -fn dashboard_responding() -> bool { - dashboard_frames().is_ok() -} - -fn dashboard_has_channel_or_payload(channel_substr: &str, payload_substr: &str) -> bool { - dashboard_frames() - .map(|frames| { - frames.iter().any(|frame| { - frame.channel.contains(channel_substr) || frame.payload.contains(payload_substr) - }) - }) - .unwrap_or(false) -} - -fn dashboard_has_frame(channel_substr: &str, payload_substr: &str) -> bool { - dashboard_frames() - .map(|frames| { - frames.iter().any(|frame| { - frame.channel.contains(channel_substr) && frame.payload.contains(payload_substr) - }) - }) - .unwrap_or(false) -} - -fn dashboard_has_worker_prompt_completed() -> bool { - dashboard_has_frame("mvp.worker.prompt", "PromptCompleted") -} - -fn dashboard_has_pipeline_prompt_output() -> bool { - dashboard_has_frame("mvp.orch.prompt", "pipeline_token_out") -} - -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" - && (direct_orch_prompt_lifecycle(&events, prompt_work) - || pipeline_orch_prompt_lifecycle(&events, prompt_work)) - }) - }) - .unwrap_or(false) -} - -fn direct_orch_prompt_lifecycle( - events: &[OrchPromptObservation], - prompt_work: &OrchPromptObservation, -) -> bool { - 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") - }) -} - -fn pipeline_orch_prompt_lifecycle( - events: &[OrchPromptObservation], - prompt_work: &OrchPromptObservation, -) -> bool { - events.iter().any(|event| { - same_prompt(prompt_work, event) - && event.phase == "pipeline_tokenizer_encode" - && event.status == "ready" - }) && events.iter().any(|event| { - same_prompt(prompt_work, event) - && event.phase == "pipeline_token_in" - && event.status == "ready" - }) && events.iter().any(|event| { - same_prompt(prompt_work, event) - && event.phase == "pipeline_token_out" - && event.status == "observed" - }) && events.iter().any(|event| { - same_prompt(prompt_work, event) - && event.phase == "pipeline_tokenizer_decode" - && event.status == "ready" - }) -} - -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, - payload: String, -} - -fn dashboard_frames() -> Result, String> { - let response = http_get("/api/frames")?; - let (_, body) = response - .split_once("\r\n\r\n") - .ok_or_else(|| "HTTP response missing body".to_owned())?; - let values = serde_json::from_str::>(body) - .map_err(|e| format!("parse dashboard frames JSON: {e}; body={body:?}"))?; - Ok(values - .into_iter() - .map(|value| SeenFrame { - channel: value - .get("channel") - .and_then(Value::as_str) - .unwrap_or_default() - .to_owned(), - payload: decode_payload(value.get("payload")).unwrap_or_default(), - }) - .collect()) -} - -fn dashboard_snapshot() -> Result { - let mut frames = dashboard_frames()?; - let keep = frames.len().saturating_sub(40); - frames.drain(0..keep); - Ok(frames - .into_iter() - .map(|frame| format!("{} {}", frame.channel, frame.payload)) - .collect::>() - .join("\n")) -} - -fn decode_payload(value: Option<&Value>) -> Option { - let bytes = value? - .as_array()? - .iter() - .map(|byte| byte.as_u64().map(|n| n as u8)) - .collect::>>()?; - Some(String::from_utf8_lossy(&bytes).to_string()) -} - -fn http_get(path: &str) -> Result { - let mut stream = TcpStream::connect(DASHBOARD_ADDR) - .map_err(|e| format!("connect dashboard {DASHBOARD_ADDR}: {e}"))?; - stream - .set_read_timeout(Some(Duration::from_secs(2))) - .map_err(|e| format!("set read watchdog: {e}"))?; - write!( - stream, - "GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n" - ) - .map_err(|e| format!("write HTTP request: {e}"))?; - let mut response = String::new(); - stream - .read_to_string(&mut response) - .map_err(|e| format!("read HTTP response: {e}"))?; - if response.starts_with("HTTP/1.1 200") { - Ok(response) - } else { - Err(format!("non-200 dashboard response: {response:?}")) - } -} - -fn stdout_contains(stdout: &Arc>, needle: &str) -> bool { - snapshot(stdout).contains(needle) -} - -fn 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 { - snapshot(stdout) - .split("Response: ") - .skip(1) - .any(|text| !text.lines().next().unwrap_or_default().trim().is_empty()) -} - -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-orchestrator:") - || trimmed.starts_with("mvp-worker-node:") - || trimmed.starts_with("mvp_tinygrad_worker:"); - leaked.then(|| format!("{stream}: {line}")) - }) - .collect() -} -fn snapshot(buf: &Arc>) -> String { - buf.lock().expect("capture mutex").clone() -} - -fn wait_child(child: &mut Child, watchdog: Duration) -> Option { - let start = Instant::now(); - while start.elapsed() < watchdog { - if let Some(status) = child.try_wait().expect("poll child") { - return Some(status); - } - thread::sleep(Duration::from_millis(100)); - } - None -} - -fn request_child_interrupt(child: &Child) { - #[cfg(target_os = "linux")] - unsafe { - let _ = libc::kill(-(child.id() as libc::pid_t), libc::SIGINT); - } - #[cfg(not(target_os = "linux"))] - { - let _ = child; - } -} - -fn require_docker(root: &std::path::Path) { - let version = Command::new("docker") - .current_dir(root) - .arg("version") - .output() - .expect("run docker version"); - assert!( - version.status.success(), - "docker is not available\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&version.stdout), - String::from_utf8_lossy(&version.stderr) - ); -} - -fn assert_no_containers_with_prefix_if_docker_available(root: &std::path::Path, prefix: &str) { - let docker_available = Command::new("docker") - .current_dir(root) - .arg("version") - .output() - .map(|output| output.status.success()) - .unwrap_or(false); - if docker_available { - assert_containers_with_prefix_removed(root, prefix); - } -} -fn assert_containers_with_prefix_removed(root: &std::path::Path, prefix: &str) { - let start = Instant::now(); - let mut containers = String::new(); - while start.elapsed() < SHUTDOWN_WATCHDOG { - let output = Command::new("docker") - .current_dir(root) - .args([ - "ps", - "-a", - "--filter", - &format!("name=^{prefix}-"), - "--format", - "{{.Names}}", - ]) - .output() - .expect("run docker ps"); - assert!( - output.status.success(), - "docker ps failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - containers = String::from_utf8_lossy(&output.stdout).into_owned(); - if containers.trim().is_empty() { - return; - } - thread::sleep(Duration::from_millis(100)); - } - panic!("containers with prefix {prefix} still exist:\n{containers}"); -} - -fn workspace_root() -> std::path::PathBuf { - std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .parent() - .and_then(std::path::Path::parent) - .expect("workspace root") - .to_path_buf() -} diff --git a/crates/mvp-system/tests/support/dumb_worker.rs b/crates/mvp-system/tests/support/dumb_worker.rs deleted file mode 100644 index a400914..0000000 --- a/crates/mvp-system/tests/support/dumb_worker.rs +++ /dev/null @@ -1,109 +0,0 @@ -use std::io::{self, BufRead, Write}; -use std::process::ExitCode; - -use serde_json::json; - -pub fn run_main() -> ExitCode { - let stdin = io::stdin(); - let mut stdout = io::stdout(); - - for line in stdin.lock().lines() { - let line = match line { - Ok(line) => line, - Err(error) => { - let _ = writeln!( - stdout, - "{}", - json!({"type":"WorkerFatal","reason":format!("stdin:{error}")}) - ); - return ExitCode::from(1); - } - }; - if line.trim().is_empty() { - continue; - } - let value: serde_json::Value = match serde_json::from_str(&line) { - Ok(value) => value, - Err(error) => { - let _ = writeln!( - stdout, - "{}", - json!({"type":"ProcessFault","reason":format!("invalid-json:{error}")}) - ); - continue; - } - }; - let Some(kind) = value.get("type").and_then(|value| value.as_str()) else { - let _ = writeln!( - stdout, - "{}", - json!({"type":"ProcessFault","reason":"missing-type"}) - ); - continue; - }; - match kind { - "InitializeWorker" => { - let _ = writeln!(stdout, "{}", json!({"type":"WorkerReady","generation":1})); - } - "InstallRing" => { - let ring_id = value - .get("ring_id") - .and_then(|value| value.as_u64()) - .unwrap_or(0); - let _ = writeln!( - stdout, - "{}", - json!({"type":"RingInstalled","ring_id":ring_id}) - ); - } - "ExecuteStep" => { - let step_id = value - .get("step_id") - .and_then(|value| value.as_u64()) - .unwrap_or(0); - let _ = writeln!( - stdout, - "{}", - json!({"type":"StepCompleted","step_id":step_id}) - ); - } - "ReleaseDeviceObject" => { - let handle = value - .get("handle") - .and_then(|value| value.as_u64()) - .unwrap_or(0); - let _ = writeln!( - stdout, - "{}", - json!({"type":"DeviceObjectReleased","handle":handle}) - ); - } - "RingReadable" => { - let ring_id = value - .get("ring_id") - .and_then(|value| value.as_u64()) - .unwrap_or(0); - let _ = writeln!( - stdout, - "{}", - json!({"type":"RingReadableAck","ring_id":ring_id}) - ); - } - "ShutdownWorker" => { - let _ = writeln!(stdout, "{}", json!({"type":"WorkerStopped"})); - let _ = stdout.flush(); - return ExitCode::SUCCESS; - } - _ => { - let _ = writeln!( - stdout, - "{}", - json!({"type":"ProcessFault","reason":"unknown-command","command":kind}) - ); - } - } - let _ = stdout.flush(); - } - - ExitCode::SUCCESS -} diff --git a/crates/mvp-system/tests/support/local_e2e.rs b/crates/mvp-system/tests/support/local_e2e.rs deleted file mode 100644 index e1f4de5..0000000 --- a/crates/mvp-system/tests/support/local_e2e.rs +++ /dev/null @@ -1,1133 +0,0 @@ -use std::collections::VecDeque; -use std::io::{BufRead, BufReader, Write}; -use std::net::{SocketAddr, TcpListener, TcpStream}; -use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::mpsc::{self, Receiver}; -use std::thread; -use std::time::{Duration, Instant}; - -use distribution::node::DistributedNodeConfig; -use iroh::EndpointAddr; -use iroh_driver::{IrohDriver, IrohDriverConfig}; -use mvp_system::actors::node_agent::{ - NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageProvisionWire, -}; -use mvp_system::actors::orchestrator::{ - LifecycleEventWire, OrchestratorActor, OrchestratorMsg, OrchestratorReport, RunCommandWire, - StageRefWire, -}; -use mvp_system::actors::register_mvp_actor_codecs; -use mvp_system::distribution_stack::DistributionRuntimeStack; -use mvp_system::driver_pumps; -use mvp_system::engine_builder as engine; -use mvp_system::orchestrator_run_fsm as fsm; -use mvp_system::run_plan as plan; -use mvp_system::stage_controller as stage; -use mvp_system::tx_rx_edge_actor as edge_actor; -use serde::{Deserialize, Serialize}; -use serde_json::json; -use swactor::actor::ActorAddress; -use swactor::runtime::ExternalSender; - -const RUN_ID: u64 = 77; -const ORCHESTRATOR_LOGICAL_NODE_ID: u64 = 900; -const NODE0_LOGICAL_ID: u64 = 11; -const NODE1_LOGICAL_ID: u64 = 12; -const MAX_TOKENS: u64 = 1; - -static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); - -pub fn run_main() -> ExitCode { - install_signal_handlers(); - let args = std::env::args().collect::>(); - let result = if args.iter().any(|arg| arg == "--role=node") { - run_node_role(&args) - } else { - run_supervisor_once(RUN_ID, true) - }; - - match result { - Ok(()) => ExitCode::SUCCESS, - Err(error) => { - eprintln!("mvp-local-e2e: {error}"); - ExitCode::from(1) - } - } -} - -#[derive(Clone, Debug, Serialize, Deserialize)] -struct EdgeFrame { - edge_id: u64, - object_id: u64, - sequence: u64, - kind: String, - token_id: Option, - eos: bool, -} - -#[derive(Clone, Debug, Deserialize)] -struct NodeReady { - #[serde(rename = "type")] - kind: String, - endpoint: EndpointAddr, - node_actor: ActorAddress, - token_in_addr: SocketAddr, - logical_node_id: u64, - stage_index: u32, -} - -#[derive(Clone, Debug, Deserialize)] -struct NodeStdoutLine { - #[serde(rename = "type")] - kind: String, - event: Option, -} - -fn run_supervisor_once(run_id: u64, print_summary: bool) -> Result<(), String> { - let _tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; - let mut driver = new_driver(_tokio.handle().clone())?; - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - DistributedNodeConfig::default(), - register_mvp_actor_codecs, - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - - let orchestrator_report = stack - .runtime - .new_inbox::() - .map_err(|e| format!("orchestrator report inbox: {e}"))?; - let orchestrator_addr = stack - .runtime - .spawn(OrchestratorActor::new( - fsm::RunConfig { - run_id: fsm::RunId(run_id), - max_tokens: MAX_TOKENS, - prompt: vec![1, 2, 3], - }, - Some(*orchestrator_report.addr()), - )) - .map_err(|e| format!("spawn orchestrator actor: {e}"))?; - stack.register_local_actor(driver.register_actor(orchestrator_addr, 1)); - - let token_out_listener = TcpListener::bind("127.0.0.1:0") - .map_err(|e| format!("bind orchestrator token-out listener: {e}"))?; - let token_out_addr = token_out_listener - .local_addr() - .map_err(|e| format!("token-out local addr: {e}"))?; - spawn_token_out_reader( - token_out_listener, - stack.runtime.create_sender(), - orchestrator_addr, - ); - - let topology = build_local_engine_topology(run_id)?; - let run_plan = topology.role_plan.run_plan.clone(); - let stage0 = run_plan - .stages - .iter() - .find(|stage| stage.stage_index == 0) - .cloned() - .ok_or_else(|| "engine builder did not assign stage 0".to_owned())?; - let stage1 = run_plan - .stages - .iter() - .find(|stage| stage.stage_index == 1) - .cloned() - .ok_or_else(|| "engine builder did not assign stage 1".to_owned())?; - - let self_endpoint_json = serde_json::to_string(&driver.endpoint_addr()) - .map_err(|e| format!("serialize endpoint addr: {e}"))?; - let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr) - .map_err(|e| format!("serialize orchestrator actor: {e}"))?; - - let mut node1 = spawn_node_process( - stage1.node_id.0, - stage1.stage_index, - stage1.inbound_edge.0, - token_out_addr, - &self_endpoint_json, - &orchestrator_actor_json, - )?; - let mut node0 = spawn_node_process( - stage0.node_id.0, - stage0.stage_index, - stage0.inbound_edge.0, - node1.ready.token_in_addr, - &self_endpoint_json, - &orchestrator_actor_json, - )?; - - wait_for_routes( - &mut driver, - &stack, - &[node0.ready.node_actor, node1.ready.node_actor], - )?; - - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObservePoolReady { - nodes: run_plan - .stages - .iter() - .map(|stage| stage.node_id.0) - .collect(), - }, - ) - .map_err(|e| format!("observe pool ready: {e}"))?; - - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObservePlanAvailable { - run_id, - stages: run_plan - .stages - .iter() - .map(|stage| StageRefWire { - stage_index: stage.stage_index, - node_id: stage.node_id.0, - }) - .collect(), - }, - ) - .map_err(|e| format!("observe plan: {e}"))?; - - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenInEndpointReady, - ) - .map_err(|e| format!("observe token-in endpoint: {e}"))?; - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenOutEndpointReady, - ) - .map_err(|e| format!("observe token-out endpoint: {e}"))?; - - let mut injected = false; - let mut completed = false; - let mut torn_down = false; - let mut stage_ready_count = 0usize; - let mut token_received = false; - let mut sent_stop_to_node0 = false; - let mut sent_stop_to_node1 = false; - - let mut token_in_object_allocator = - edge_actor::ObjectIdAllocator::new(edge_actor::EdgeId(stage0.inbound_edge.0)); - while !STOP_REQUESTED.load(Ordering::SeqCst) { - pump_network(&mut driver, &stack); - stage_ready_count += drain_node_stdout(&node0.stdout_rx); - stage_ready_count += drain_node_stdout(&node1.stdout_rx); - - while let Some(report) = orchestrator_report.try_recv() { - match report { - OrchestratorReport::Command(command) => match command { - RunCommandWire::ProvisionStage { stage_index, .. } => { - let provision = stage_provision_wire(&run_plan, stage_index)?; - let target = if stage_index == 0 { - node0.ready.node_actor - } else { - node1.ready.node_actor - }; - stack - .runtime - .send_to(target, NodeAgentMsg::ProvisionStage(provision)) - .map_err(|e| format!("send provision to stage {stage_index}: {e}"))?; - } - RunCommandWire::InjectTokenObject { - sequence, payload, .. - } => { - injected = true; - - let token_id = match payload { - mvp_system::actors::orchestrator::TokenObjectPayloadWire::Prompt { - tokens, - } => tokens.first().copied(), - mvp_system::actors::orchestrator::TokenObjectPayloadWire::Decode { - token_id, - .. - } => Some(token_id), - }; - let prompt_object = token_in_object_allocator.alloc(); - write_edge_frame( - node0.ready.token_in_addr, - EdgeFrame { - edge_id: stage0.inbound_edge.0, - object_id: prompt_object.object_id.0, - sequence, - kind: "token".to_owned(), - token_id, - eos: false, - }, - )?; - } - RunCommandWire::StopRun { - stage_index, - run_id, - } => { - let target = if stage_index == 0 { - sent_stop_to_node0 = true; - node0.ready.node_actor - } else { - sent_stop_to_node1 = true; - node1.ready.node_actor - }; - stack - .runtime - .send_to(target, NodeAgentMsg::StopRun { run_id }) - .map_err(|e| format!("send stop to stage {stage_index}: {e}"))?; - } - RunCommandWire::TearDownTokenEndpoints { .. } => { - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenEndpointsStopped, - ) - .map_err(|e| format!("observe token endpoints stopped: {e}"))?; - } - RunCommandWire::CreateTokenInEndpoint { .. } - | RunCommandWire::CreateTokenOutEndpoint { .. } => {} - }, - OrchestratorReport::Lifecycle(event) => match event { - LifecycleEventWire::RunCompleted { .. } => { - token_received = true; - completed = true; - } - LifecycleEventWire::RunTornDown { .. } => { - torn_down = true; - } - LifecycleEventWire::RunRejected { .. } - | LifecycleEventWire::RunFaulted { .. } - | LifecycleEventWire::RunOperatorStopped { .. } => { - return Err(format!("run failed: {event:?}")); - } - }, - OrchestratorReport::StageFault { - run_id, - stage_index, - } => return Err(format!("stage {stage_index} faulted in run {run_id}")), - OrchestratorReport::NodeRuntimeReady { .. } - | OrchestratorReport::NodeRuntimeReadyAck { .. } - | OrchestratorReport::WeightsReady { .. } - | OrchestratorReport::StageReady { .. } - | OrchestratorReport::Snapshot { .. } => {} - } - } - - token_received |= completed; - - if injected && completed && torn_down && sent_stop_to_node0 && sent_stop_to_node1 { - shutdown_node(&mut node0); - shutdown_node(&mut node1); - let builder_stage_assignments = topology - .events - .iter() - .filter(|event| { - matches!( - event, - engine::EngineEvent::RoleAssigned { - role: engine::RoleKind::StageWorker { .. }, - .. - } - ) - }) - .count(); - let summary = json!({ - "ok": true, - "actor_plane": "iroh-swactor", - "processes": { - "orchestrator": std::process::id(), - "node0": node0.child.id(), - "node1": node1.child.id(), - }, - "node0_endpoint": node0.ready.endpoint, - "node1_endpoint": node1.ready.endpoint, - "node0_logical_id": node0.ready.logical_node_id, - "node1_logical_id": node1.ready.logical_node_id, - "node0_stage_index": node0.ready.stage_index, - "node1_stage_index": node1.ready.stage_index, - "data_plane": "tcp-loopback-streams", - "worker_processes": "mvp-dumb-worker-per-node", - "injected_prompt_observed": injected, - "token_received_observed": token_received, - "run_completed_observed": completed, - "run_torn_down_observed": torn_down, - "stop_sent_to_all_nodes": sent_stop_to_node0 && sent_stop_to_node1, - "engine_builder_pattern": "pool-first-static-launcher", - "engine_builder_event_count": topology.events.len(), - "engine_builder_node_count": topology.node_summaries.len(), - "engine_builder_stage_assignments": builder_stage_assignments, - "node_route_count": stack.route_view.read().map(|view| view.len()).unwrap_or_default(), - "stage_ready_stdout_count": stage_ready_count, - }); - if print_summary { - println!("{summary}"); - } else { - eprintln!("mvp-local-e2e: run {run_id} summary {summary}"); - } - return Ok(()); - } - - thread::sleep(Duration::from_millis(10)); - } - - shutdown_node(&mut node0); - shutdown_node(&mut node1); - Err("interrupted".to_owned()) -} - -fn run_node_role(args: &[String]) -> Result<(), String> { - let logical_node_id = parse_arg(args, "--logical-node-id")? - .parse::() - .map_err(|e| format!("logical node id: {e}"))?; - let stage_index = parse_arg(args, "--stage-index")? - .parse::() - .map_err(|e| format!("stage index: {e}"))?; - let inbound_edge_id = parse_arg(args, "--inbound-edge-id")? - .parse::() - .map_err(|e| format!("inbound edge id: {e}"))?; - let outbound_addr = parse_arg(args, "--outbound-addr")? - .parse::() - .map_err(|e| format!("outbound addr: {e}"))?; - let coordinator: EndpointAddr = - serde_json::from_str(parse_arg(args, "--coordinator-endpoint")?) - .map_err(|e| format!("coordinator endpoint json: {e}"))?; - let orchestrator_addr: ActorAddress = - serde_json::from_str(parse_arg(args, "--orchestrator-actor")?) - .map_err(|e| format!("orchestrator actor json: {e}"))?; - - let _tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; - let mut driver = new_driver(_tokio.handle().clone())?; - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - DistributedNodeConfig::default(), - register_mvp_actor_codecs, - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - driver.join(&[coordinator]); - - let node_report = stack - .runtime - .new_inbox::() - .map_err(|e| format!("node report inbox: {e}"))?; - let node_actor = stack - .runtime - .spawn(NodeAgentActor::new( - stage::NodeId(logical_node_id), - orchestrator_addr, - Some(*node_report.addr()), - )) - .map_err(|e| format!("spawn node actor: {e}"))?; - stack.register_local_actor(driver.register_actor(node_actor, 1)); - - let listener = TcpListener::bind("127.0.0.1:0") - .map_err(|e| format!("bind node inbound edge listener: {e}"))?; - let token_in_addr = listener - .local_addr() - .map_err(|e| format!("node inbound local addr: {e}"))?; - spawn_inbound_edge_reader( - listener, - inbound_edge_id, - stack.runtime.create_sender(), - node_actor, - ); - - println!( - "{}", - json!({ - "type": "ready", - "role": "node", - "endpoint": driver.endpoint_addr(), - "node_actor": node_actor, - "token_in_addr": token_in_addr, - "logical_node_id": logical_node_id, - "stage_index": stage_index, - }) - ); - std::io::stdout() - .flush() - .map_err(|e| format!("flush ready: {e}"))?; - - let (stdin_tx, stdin_rx) = mpsc::channel::(); - thread::spawn(move || { - let stdin = std::io::stdin(); - for line in stdin.lock().lines().map_while(Result::ok) { - if let Ok(value) = serde_json::from_str::(&line) { - let _ = stdin_tx.send(value); - } - } - }); - - let mut worker = WorkerProc::spawn()?; - let mut pending_commands = VecDeque::new(); - let mut outbound_stream: Option = None; - let mut outbound_object_allocator: Option = None; - - while !STOP_REQUESTED.load(Ordering::SeqCst) { - if let Ok(value) = stdin_rx.try_recv() { - if value.get("type").and_then(|value| value.as_str()) == Some("shutdown") { - let _ = worker.shutdown(); - return Ok(()); - } - } - - pump_network(&mut driver, &stack); - - while let Some(report) = node_report.try_recv() { - match report { - NodeAgentReport::Command(command) => pending_commands.push_back(command), - NodeAgentReport::Lifecycle(event) => { - println!( - "{}", - json!({ - "type": "node_lifecycle", - "stage_index": stage_index, - "event": format!("{event:?}"), - }) - ); - let _ = std::io::stdout().flush(); - } - NodeAgentReport::PromptRequested { .. } | NodeAgentReport::Snapshot { .. } => {} - } - } - - let route_to_orchestrator_ready = stack - .route_view - .read() - .map(|view| view.contains_key(&orchestrator_addr)) - .unwrap_or(false); - let mut deferred = VecDeque::new(); - while let Some(command) = pending_commands.pop_front() { - if !route_to_orchestrator_ready && readiness_command(&command) { - deferred.push_back(command); - continue; - } - handle_node_command( - command, - &stack.runtime, - node_actor, - &mut worker, - outbound_addr, - &mut outbound_stream, - &mut outbound_object_allocator, - stage_index, - )?; - } - pending_commands = deferred; - - thread::sleep(Duration::from_millis(10)); - } - let _ = worker.shutdown(); - Err("interrupted".to_owned()) -} - -fn new_driver(handle: tokio::runtime::Handle) -> Result { - IrohDriver::with_handle( - handle, - IrohDriverConfig { - secret_key: None, - relay_mode: iroh::RelayMode::Disabled, - node: DistributedNodeConfig::default(), - peer_auth: None, - additional_alpns: vec![], - }, - ) - .map_err(|e| format!("create iroh driver: {e}")) -} - -fn pump_network(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) { - stack.tick_protocol_actors(Instant::now()); - driver.pump_inbound_to_actors(); - stack.pump_runtime_once(); - driver.drain_outbox(&stack.outbox); -} - -fn wait_for_routes( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - actors: &[ActorAddress], -) -> Result<(), String> { - loop { - pump_network(driver, stack); - let ready = stack - .route_view - .read() - .map(|view| actors.iter().all(|actor| view.contains_key(actor))) - .unwrap_or(false); - if ready { - return Ok(()); - } - thread::sleep(Duration::from_millis(20)); - } -} - -struct LocalEngineTopology { - role_plan: engine::RoleAssignmentPlan, - events: Vec, - node_summaries: Vec, -} - -fn build_local_engine_topology(run_id: u64) -> Result { - let cluster = engine::ClusterBuilder::new( - "local-process-e2e", - engine::ModelSpec::pipelined_causal_llm( - "local-e2e-fixture", - engine::ModelArtifact::TestTinyLlm { - path: "local-process://local-e2e-fixture".to_owned(), - }, - 4, - 8, - engine::DTypeFamily::BFloat, - 2, - 8, - 99, - plan::TokenizerSource::EmbeddedGguf, - ), - ) - .run_id(run_id) - .image( - engine::NodeImageSpec::new("mvp-local-e2e") - .worker_runtime(engine::WorkerRuntimeSpec::DumbProcess), - ) - .pool_provider(engine::StaticPoolProvider::new(vec![ - engine::NodeLease::new( - "orchestrator", - engine::NodeId(ORCHESTRATOR_LOGICAL_NODE_ID), - [engine::NodeCapability::Coordinator], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - engine::NodeLease::new( - "node0", - engine::NodeId(NODE0_LOGICAL_ID), - [engine::NodeCapability::Worker], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - engine::NodeLease::new( - "node1", - engine::NodeId(NODE1_LOGICAL_ID), - [engine::NodeCapability::Worker], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - ])) - .launcher(engine::StaticNodeLauncher) - .planner( - engine::FixedLinearPipelinePlanner::new(2).runtime(plan::RuntimeConfig { - max_tokens: MAX_TOKENS as u32, - prompt: plan::PromptSource::Inline("1 2 3".to_owned()), - sampling: plan::SamplingPolicy { - temperature_millis: 0, - top_k: 1, - }, - token_output_policy: plan::TokenOutputPolicy::EmitAll, - }), - ) - .launch() - .map_err(|e| format!("local engine builder launch: {e}"))?; - let role_plan = cluster.role_plan().clone(); - let events = cluster.events().to_vec(); - let node_summaries = cluster.node_summaries(); - cluster - .shutdown() - .map_err(|e| format!("local engine builder shutdown: {e}"))?; - Ok(LocalEngineTopology { - role_plan, - events, - node_summaries, - }) -} - -fn stage_provision_wire( - plan: &plan::RunPlan, - stage_index: u32, -) -> Result { - let provision = plan::derive_stage_provision(plan, stage_index) - .map_err(|e| format!("derive stage provision {stage_index}: {e:?}"))?; - Ok(StageProvisionWire { - run_id: provision.run_id.0, - authorized_orchestrator: ORCHESTRATOR_LOGICAL_NODE_ID, - node_id: provision.node_id.0, - stage_index: provision.stage_index, - stage_count: provision.stage_count, - layer_start: provision.layer_start, - layer_end_exclusive: provision.layer_end_exclusive, - inbound_edge_id: provision.inbound.edge_id.0, - outbound_edge_id: provision.outbound.edge_id.0, - model_id: provision.model.model_id, - gguf_source: provision.gguf_source, - tokenizer: provision.tokenizer, - }) -} - -fn write_edge_frame(addr: SocketAddr, frame: EdgeFrame) -> Result<(), String> { - let mut stream = TcpStream::connect(addr).map_err(|e| format!("connect edge {addr}: {e}"))?; - stream - .write_all(&driver_pumps::encode_edge_preamble(driver_pumps::EdgeId( - frame.edge_id, - ))) - .map_err(|e| format!("write edge preamble: {e}"))?; - write_json_frame(&mut stream, &frame) -} - -fn write_json_frame(stream: &mut TcpStream, frame: &EdgeFrame) -> Result<(), String> { - serde_json::to_writer(&mut *stream, frame).map_err(|e| format!("serialize edge frame: {e}"))?; - stream - .write_all(b"\n") - .map_err(|e| format!("write edge frame newline: {e}"))?; - stream.flush().map_err(|e| format!("flush edge frame: {e}")) -} - -fn spawn_inbound_edge_reader( - listener: TcpListener, - expected_edge_id: u64, - sender: ExternalSender, - node_actor: ActorAddress, -) { - thread::spawn(move || { - for incoming in listener.incoming() { - let Ok(mut stream) = incoming else { continue }; - let mut preamble = [0u8; 8]; - if std::io::Read::read_exact(&mut stream, &mut preamble).is_err() { - continue; - } - if preamble != expected_edge_id.to_le_bytes() { - continue; - } - let mut reader = BufReader::new(stream); - loop { - let mut line = String::new(); - match reader.read_line(&mut line) { - Ok(0) => break, - Ok(_) => { - if let Ok(frame) = serde_json::from_str::(&line) { - let _ = sender.send_to( - node_actor, - NodeAgentMsg::ObjectLoaded { - edge_id: frame.edge_id, - object_id: frame.object_id, - sequence: frame.sequence, - handle_generation: 1, - handle_id: frame.object_id, - }, - ); - } - } - Err(_) => break, - } - } - } - }); -} - -fn spawn_token_out_reader( - listener: TcpListener, - sender: ExternalSender, - orchestrator_addr: ActorAddress, -) { - thread::spawn(move || { - for incoming in listener.incoming() { - let Ok(mut stream) = incoming else { continue }; - let mut preamble = [0u8; 8]; - if std::io::Read::read_exact(&mut stream, &mut preamble).is_err() { - continue; - } - let mut reader = BufReader::new(stream); - let mut line = String::new(); - if reader.read_line(&mut line).is_ok() - && let Ok(frame) = serde_json::from_str::(&line) - { - let _ = sender.send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenReceived { - sequence: frame.sequence, - token_id: frame.token_id.unwrap_or(99), - eos: frame.eos, - }, - ); - } - } - }); -} - -fn handle_node_command( - command: StageCommandWire, - runtime: &swactor::runtime::Runtime, - node_actor: ActorAddress, - worker: &mut WorkerProc, - outbound_addr: SocketAddr, - outbound_stream: &mut Option, - outbound_object_allocator: &mut Option, - local_stage_index: u32, -) -> Result<(), String> { - match command { - StageCommandWire::EstablishInboundEdge { edge_id } => runtime - .send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id }) - .map_err(|e| format!("mark inbound ready: {e}")), - StageCommandWire::EstablishOutboundEdge { edge_id } => { - let mut stream = TcpStream::connect(outbound_addr) - .map_err(|e| format!("connect outbound edge {edge_id} to {outbound_addr}: {e}"))?; - stream - .write_all(&driver_pumps::encode_edge_preamble(driver_pumps::EdgeId( - edge_id, - ))) - .map_err(|e| format!("write outbound preamble: {e}"))?; - *outbound_stream = Some(stream); - *outbound_object_allocator = Some(edge_actor::ObjectIdAllocator::new( - edge_actor::EdgeId(edge_id), - )); - runtime - .send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id }) - .map_err(|e| format!("mark outbound ready: {e}")) - } - StageCommandWire::ConfigureWorkerRole { .. } => { - worker.initialize()?; - runtime - .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) - .map_err(|e| format!("mark worker ready: {e}")) - } - StageCommandWire::LoadWeights { .. } => runtime - .send_to( - node_actor, - NodeAgentMsg::MarkWeightsReady { - run_id: RUN_ID, - node_id: 0, - stage_index: local_stage_index, - }, - ) - .map_err(|e| format!("mark weights ready: {e}")), - StageCommandWire::ExecuteStep { - step_id, - sequence, - output_edge_ids, - .. - } => { - worker.execute_step(step_id)?; - runtime - .send_to(node_actor, NodeAgentMsg::StepCompleted { step_id }) - .map_err(|e| format!("mark step completed: {e}"))?; - let output_edge_id = output_edge_ids.first().copied().unwrap_or(0); - let output_key = outbound_object_allocator - .as_mut() - .ok_or_else(|| "outbound object allocator missing for ExecuteStep".to_owned())? - .alloc(); - let stream = outbound_stream - .as_mut() - .ok_or_else(|| "outbound stream missing for ExecuteStep".to_owned())?; - let final_stage = local_stage_index == 1; - write_json_frame( - stream, - &EdgeFrame { - edge_id: output_edge_id, - object_id: output_key.object_id.0, - sequence, - kind: if final_stage { - "token".to_owned() - } else { - "activation".to_owned() - }, - token_id: final_stage.then_some(99), - eos: final_stage, - }, - ) - } - StageCommandWire::ReleaseInputHandle { .. } | StageCommandWire::RewireEdge { .. } => Ok(()), - StageCommandWire::StopLocalEdges { run_id } => { - runtime - .send_to(node_actor, NodeAgentMsg::LocalEdgesStopped { run_id }) - .map_err(|e| format!("mark local edges stopped: {e}"))?; - runtime - .send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id }) - .map_err(|e| format!("mark worker rings quiesced: {e}")) - } - StageCommandWire::ReleaseRunDeviceObjects { run_id } => { - runtime - .send_to(node_actor, NodeAgentMsg::DeviceObjectsReleased { run_id }) - .map_err(|e| format!("mark device objects released: {e}"))?; - runtime - .send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id }) - .map_err(|e| format!("mark worker role reset: {e}")) - } - } -} - -fn readiness_command(command: &StageCommandWire) -> bool { - matches!( - command, - StageCommandWire::EstablishInboundEdge { .. } - | StageCommandWire::EstablishOutboundEdge { .. } - | StageCommandWire::ConfigureWorkerRole { .. } - | StageCommandWire::LoadWeights { .. } - ) -} - -struct WorkerProc { - child: Child, - stdin: ChildStdin, - stdout: BufReader, - cleaned: bool, -} - -impl WorkerProc { - fn spawn() -> Result { - let mut child = Command::new(current_test_exe()?) - .env("MVP_TEST_ROLE", "dumb-worker") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::null()) - .spawn() - .map_err(|e| format!("spawn dumb worker: {e}"))?; - let stdin = match child.stdin.take() { - Some(stdin) => stdin, - None => { - kill_child(&mut child); - return Err("worker stdin missing".to_owned()); - } - }; - let stdout = match child.stdout.take() { - Some(stdout) => stdout, - None => { - kill_child(&mut child); - return Err("worker stdout missing".to_owned()); - } - }; - Ok(Self { - child, - stdin, - stdout: BufReader::new(stdout), - cleaned: false, - }) - } - - fn initialize(&mut self) -> Result<(), String> { - self.command( - json!({"type":"InitializeWorker","helper_abi_version":1}), - "WorkerReady", - ) - } - - fn execute_step(&mut self, step_id: u64) -> Result<(), String> { - self.command( - json!({"type":"ExecuteStep","step_id":step_id}), - "StepCompleted", - ) - } - - fn shutdown(&mut self) -> Result<(), String> { - let result = self.command(json!({"type":"ShutdownWorker"}), "WorkerStopped"); - self.terminate(); - result - } - - fn terminate(&mut self) { - if self.cleaned { - return; - } - self.cleaned = true; - kill_child(&mut self.child); - } - - fn command(&mut self, command: serde_json::Value, expected: &str) -> Result<(), String> { - writeln!(self.stdin, "{command}").map_err(|e| format!("write worker command: {e}"))?; - self.stdin - .flush() - .map_err(|e| format!("flush worker stdin: {e}"))?; - let mut line = String::new(); - self.stdout - .read_line(&mut line) - .map_err(|e| format!("read worker stdout: {e}"))?; - let value: serde_json::Value = serde_json::from_str(&line) - .map_err(|e| format!("parse worker stdout {line:?}: {e}"))?; - if value.get("type").and_then(|value| value.as_str()) == Some(expected) { - Ok(()) - } else { - Err(format!("worker emitted {value}, expected {expected}")) - } - } -} - -impl Drop for WorkerProc { - fn drop(&mut self) { - self.terminate(); - } -} - -struct NodeChild { - child: Child, - stdin: ChildStdin, - stdout_rx: Receiver, - ready: NodeReady, - cleaned: bool, -} - -fn spawn_node_process( - logical_node_id: u64, - stage_index: u32, - inbound_edge_id: u64, - outbound_addr: SocketAddr, - coordinator_endpoint_json: &str, - orchestrator_actor_json: &str, -) -> Result { - let mut child = Command::new(std::env::current_exe().map_err(|e| format!("current exe: {e}"))?) - .arg("--role=node") - .arg("--logical-node-id") - .arg(logical_node_id.to_string()) - .arg("--stage-index") - .arg(stage_index.to_string()) - .arg("--inbound-edge-id") - .arg(inbound_edge_id.to_string()) - .arg("--outbound-addr") - .arg(outbound_addr.to_string()) - .arg("--coordinator-endpoint") - .arg(coordinator_endpoint_json) - .arg("--orchestrator-actor") - .arg(orchestrator_actor_json) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .map_err(|e| format!("spawn node {stage_index}: {e}"))?; - - let stdin = match child.stdin.take() { - Some(stdin) => stdin, - None => { - kill_child(&mut child); - return Err("node stdin missing".to_owned()); - } - }; - let stdout = match child.stdout.take() { - Some(stdout) => stdout, - None => { - kill_child(&mut child); - return Err("node stdout missing".to_owned()); - } - }; - let mut reader = BufReader::new(stdout); - let mut ready_line = String::new(); - if let Err(error) = reader.read_line(&mut ready_line) { - kill_child(&mut child); - return Err(format!("read node ready: {error}")); - } - let ready: NodeReady = match serde_json::from_str(&ready_line) { - Ok(ready) => ready, - Err(error) => { - kill_child(&mut child); - return Err(format!("parse node ready {ready_line:?}: {error}")); - } - }; - if ready.kind != "ready" { - kill_child(&mut child); - return Err(format!("node first line was not ready: {ready_line}")); - } - - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - for line in reader.lines().map_while(Result::ok) { - if let Ok(value) = serde_json::from_str::(&line) { - let _ = tx.send(value); - } - } - }); - - Ok(NodeChild { - child, - stdin, - stdout_rx: rx, - ready, - cleaned: false, - }) -} - -impl Drop for NodeChild { - fn drop(&mut self) { - shutdown_node(self); - } -} - -fn shutdown_node(node: &mut NodeChild) { - if node.cleaned { - return; - } - node.cleaned = true; - let _ = writeln!(node.stdin, "{}", json!({"type":"shutdown"})); - let _ = node.stdin.flush(); - let started = Instant::now(); - while started.elapsed() < Duration::from_secs(3) { - if matches!(node.child.try_wait(), Ok(Some(_))) { - return; - } - thread::sleep(Duration::from_millis(20)); - } - kill_child(&mut node.child); -} - -fn drain_node_stdout(rx: &Receiver) -> usize { - let mut stage_ready_count = 0; - while let Ok(line) = rx.try_recv() { - if line.kind != "node_lifecycle" { - continue; - } - let Some(event) = line.event.as_deref() else { - continue; - }; - if event.contains("StageReady") { - stage_ready_count += 1; - } - } - stage_ready_count -} - -fn kill_child(child: &mut Child) { - if !matches!(child.try_wait(), Ok(Some(_))) { - let _ = child.kill(); - } - let _ = child.wait(); -} - -extern "C" fn request_stop(_: libc::c_int) { - STOP_REQUESTED.store(true, Ordering::SeqCst); -} - -fn install_signal_handlers() { - #[cfg(target_os = "linux")] - unsafe { - libc::signal(libc::SIGINT, request_stop as *const () as usize); - libc::signal(libc::SIGTERM, request_stop as *const () as usize); - libc::signal(libc::SIGHUP, request_stop as *const () as usize); - } -} - -fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> { - let index = args - .iter() - .position(|arg| arg == name) - .ok_or_else(|| format!("missing {name}"))?; - args.get(index + 1) - .map(String::as_str) - .ok_or_else(|| format!("missing value for {name}")) -} - -fn current_test_exe() -> Result { - std::env::current_exe().map_err(|e| format!("current test exe: {e}")) -} diff --git a/crates/mvp-system/tests/support/local_e2e_cluster.rs b/crates/mvp-system/tests/support/local_e2e_cluster.rs deleted file mode 100644 index 0218ba6..0000000 --- a/crates/mvp-system/tests/support/local_e2e_cluster.rs +++ /dev/null @@ -1,3156 +0,0 @@ -use std::collections::{BTreeMap, HashMap, VecDeque}; -use std::io::{BufRead, BufReader, Write}; -use std::net::SocketAddr; -use std::path::PathBuf; -use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; -use std::sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - mpsc::{self, Receiver, Sender}, -}; -use std::thread; -use std::time::{Duration, Instant}; - -use dashboard::swactor::{RUNTIME_ACTORS, RUNTIME_STATS, RUNTIME_WORKERS}; -use datastream::frame::{ChannelId, Frame, Position}; -use datastream::{ - DatastreamEndpoint, DatastreamProducer, DatastreamSubscription, Lifetime, NodeId, StreamId, -}; -use distribution::node::DistributedNodeConfig; -use iroh::EndpointAddr; -use iroh_driver::{DATASTREAM_ALPN, IrohDriver, IrohDriverConfig}; -use mvp_system::actors::node_agent::{ - NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageProvisionWire, -}; -use mvp_system::actors::orchestrator::{ - LifecycleEventWire, OrchestratorActor, OrchestratorMsg, OrchestratorReport, RunCommandWire, - StageRefWire, -}; -use mvp_system::actors::register_mvp_actor_codecs; -use mvp_system::arena_manager as arena; -use mvp_system::dashboard_view::MvpClusterDashboardView; -use mvp_system::distribution_stack::DistributionRuntimeStack; -use mvp_system::docker_cluster_provisioning as docker_provision; -use mvp_system::driver_pumps as driver_model; -use mvp_system::edge_establisher as edge; -use mvp_system::engine_builder as engine; -use mvp_system::gpu_worker_ctl as worker_ctl; -use mvp_system::gpu_worker_ingress_parser as ingress; -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, MVP_IROH_RELAY_URL_ENV, RelayProvider, RelayProvisionRequest, - RelayPurpose, relay_runtime_config_from_env, -}; -use mvp_system::run_plan as plan; -use mvp_system::stage_controller as stage; -use mvp_system::tx_rx_edge_actor as edge_actor; -use serde::Deserialize; -use serde_json::{Value, json}; -use swactor::actor::ActorAddress; -use tokio::sync::mpsc as tokio_mpsc; - -const RUN_ID: u64 = 77; -const ORCHESTRATOR_LOGICAL_NODE_ID: u64 = 900; -const NODE0_LOGICAL_ID: u64 = 11; -const NODE1_LOGICAL_ID: u64 = 12; -const MAX_TOKENS: u64 = 1; -const DEFAULT_PROMPT: &str = "ping"; -const DEFAULT_DOCKER_IMAGE: &str = "swactor-mvp-local-e2e-cluster:latest"; -const EDGE_ALPN: &[u8] = b"swactor/edge/1"; -const OBJECT_MAX_EXTENT: u64 = 16; -const OBJECT_ALIGNMENT: u64 = 4; -const ARENA_BYTES: usize = 16 * 1024; -const RING_BYTES: usize = 4096; -const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500); -const LOCAL_E2E_ROUTE_TIMEOUT: Duration = Duration::from_secs(60); -const LOCAL_E2E_WORKFLOW_TIMEOUT: Duration = Duration::from_secs(120); -const LOCAL_E2E_RELAY_LISTEN_ENV: &str = "MVP_LOCAL_E2E_RELAY_LISTEN"; -const LOCAL_E2E_DOCKER_NETWORK_NODE0_ENV: &str = "MVP_LOCAL_E2E_DOCKER_NETWORK_NODE0"; -const LOCAL_E2E_DOCKER_NETWORK_NODE1_ENV: &str = "MVP_LOCAL_E2E_DOCKER_NETWORK_NODE1"; -static STOP_REQUESTED: AtomicBool = AtomicBool::new(false); - -struct MvpDashboard { - url: String, - endpoint: DatastreamEndpoint, - producer: DatastreamProducer, - subscription: DatastreamSubscription, - handle: dashboard::DashboardHandle, - runtime_position: u64, - runtime_snapshot_interval: Duration, - last_runtime_snapshot: Option, -} - -impl MvpDashboard { - fn start_from_env() -> Result { - let mut config = dashboard::DashboardConfig::default(); - if let Some(port) = std::env::var_os("MVP_DASHBOARD_PORT") { - let port = port - .to_string_lossy() - .parse::() - .map_err(|e| format!("invalid MVP_DASHBOARD_PORT: {e}"))?; - config.port = port; - } - let runtime_snapshot_interval = runtime_snapshot_interval_from_env()?; - let url = format!("http://127.0.0.1:{}/view/datastream/live", config.port); - let handle = dashboard::start_dashboard(config.clone()); - handle.register_view(Arc::new(MvpClusterDashboardView::new())); - handle.start_http_standalone(); - let endpoint = DatastreamEndpoint::with_capacity( - StreamId::new(NodeId::new("mvp-system-orch"), Lifetime(1)), - 4096, - config.frame_buffer, - ); - let producer = endpoint.producer(); - let subscription = endpoint.subscribe_all("dashboard"); - Ok(Self { - url, - endpoint, - producer, - subscription, - handle, - runtime_position: 0, - runtime_snapshot_interval, - last_runtime_snapshot: None, - }) - } - - fn url(&self) -> &str { - &self.url - } - - fn drain(&mut self) { - self.endpoint.tick(); - for delivery in self.subscription.drain_available() { - self.handle.ingest(&delivery.stream, &delivery.frame); - } - } - - fn record_event(&mut self, event: obs::Event) { - let record = mvp_system::telemetry::MvpLifecycleRecord::new(event); - self.producer.submit_record(&record); - self.drain(); - } - - fn record_provision_log(&mut self, line: mvp_system::provisioning::ProvisionLogLine) { - let channel = mvp_system::telemetry::mvp_provision_log_channel(line.node_id, line.stream); - let record = mvp_system::telemetry::MvpProvisionLogRecord::new(line); - let payload = serde_json::to_vec(&record).expect("serialize provisioning log record"); - self.producer.submit_bytes(channel, payload); - self.drain(); - } - - fn publish_runtime_snapshot_throttled(&mut self, stack: &DistributionRuntimeStack) { - let due = self.last_runtime_snapshot.map_or(true, |last| { - last.elapsed() >= self.runtime_snapshot_interval - }); - if self.runtime_snapshot_interval.is_zero() || due { - self.publish_runtime_snapshot(stack); - } - } - - fn publish_runtime_snapshot(&mut self, stack: &DistributionRuntimeStack) { - let stats = stack.runtime.stats(); - let actors = stats - .actors - .iter() - .map(|(address, worker_id)| json!([address.to_string(), worker_id])) - .collect::>(); - let total_mailbox_depth = stats - .workers - .iter() - .map(|worker| worker.mailbox_depth) - .sum::(); - - self.ingest_runtime_json( - RUNTIME_STATS, - json!({ - "num_workers": stats.num_workers, - "uptime_ms": stats.uptime_ms, - "actors_live": stats.actors.len(), - "mailbox_depth": total_mailbox_depth, - "actors": actors, - "workers": &stats.workers, - "actor_details": &stats.actor_details, - "tick_timings": &stats.tick_timings, - }), - ); - self.ingest_runtime_json(RUNTIME_WORKERS, json!({ "workers": &stats.workers })); - if !stats.actor_details.is_empty() { - self.ingest_runtime_json(RUNTIME_ACTORS, json!({ "actors": &stats.actor_details })); - } - self.last_runtime_snapshot = Some(Instant::now()); - } - - fn ingest_runtime_json(&mut self, channel: &str, value: Value) { - let payload = serde_json::to_vec(&value).expect("serialize runtime dashboard frame"); - let frame = Frame::new( - ChannelId::new(channel), - Position(self.runtime_position), - payload, - ); - self.runtime_position = self.runtime_position.wrapping_add(1); - self.handle.ingest(self.endpoint.stream_id(), &frame); - } -} - -extern "C" fn request_stop(_: libc::c_int) { - STOP_REQUESTED.store(true, Ordering::SeqCst); -} - -fn install_signal_handlers() { - #[cfg(target_os = "linux")] - unsafe { - libc::signal(libc::SIGINT, request_stop as *const () as usize); - libc::signal(libc::SIGTERM, request_stop as *const () as usize); - libc::signal(libc::SIGHUP, request_stop as *const () as usize); - } -} - -fn runtime_snapshot_interval_from_env() -> Result { - let Some(value) = std::env::var_os("MVP_RUNTIME_SNAPSHOT_MS") else { - return Ok(DEFAULT_RUNTIME_SNAPSHOT_INTERVAL); - }; - let millis = value - .to_string_lossy() - .parse::() - .map_err(|e| format!("invalid MVP_RUNTIME_SNAPSHOT_MS: {e}"))?; - Ok(Duration::from_millis(millis)) -} - -pub fn run_main() -> ExitCode { - install_signal_handlers(); - let args = std::env::args().collect::>(); - let result = if std::env::var("MVP_TEST_ROLE").ok().as_deref() == Some("cluster-relay") { - run_relay_role() - } else if args.iter().any(|arg| arg == "--role=node") { - run_node_role(&args) - } else if std::env::var_os("MVP_DASHBOARD").is_some() { - run_supervisor_dashboard_loop() - } else { - let prompt = parse_optional_arg(&args, "--prompt") - .unwrap_or(DEFAULT_PROMPT) - .to_owned(); - run_supervisor_once(RUN_ID, None, true, prompt) - }; - - match result { - Ok(()) => ExitCode::SUCCESS, - Err(error) => { - eprintln!("mvp-local-e2e-cluster: {error}"); - ExitCode::from(1) - } - } -} - -struct LocalRelayGuard { - _server: iroh_relay::server::Server, - _rt: tokio::runtime::Runtime, -} - -fn run_relay_role() -> Result<(), String> { - let listen = - std::env::var(LOCAL_E2E_RELAY_LISTEN_ENV).unwrap_or_else(|_| "0.0.0.0:7843".to_owned()); - let _relay = spawn_local_relay(&listen)?; - eprintln!("mvp-local-e2e-cluster: relay ready on {listen}"); - while !STOP_REQUESTED.load(Ordering::SeqCst) { - thread::sleep(Duration::from_millis(100)); - } - Ok(()) -} - -fn spawn_local_relay(listen: &str) -> Result { - let bind_addr = listen - .parse::() - .map_err(|e| format!("invalid {LOCAL_E2E_RELAY_LISTEN_ENV}={listen:?}: {e}"))?; - let rt = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .worker_threads(2) - .build() - .map_err(|e| format!("relay tokio runtime: {e}"))?; - let server = rt - .block_on(async { - iroh_relay::server::Server::spawn(iroh_relay::server::ServerConfig::<(), ()> { - relay: Some(iroh_relay::server::RelayConfig { - http_bind_addr: bind_addr, - tls: None, - limits: Default::default(), - key_cache_capacity: Some(256), - access: iroh_relay::server::AccessConfig::Everyone, - }), - quic: None, - metrics_addr: None, - }) - .await - }) - .map_err(|e| format!("spawn relay server: {e}"))?; - Ok(LocalRelayGuard { - _server: server, - _rt: rt, - }) -} - -#[derive(Clone, Debug, Deserialize)] -struct NodeStdoutLine { - #[serde(rename = "type")] - kind: String, - stage_index: Option, - event: Option, - endpoint: Option, - node_actor: Option, - logical_node_id: Option, -} - -type LocalDockerNodeProvisioner = docker_provision::DockerNodeProvisioner< - LocalE2eDockerCli, - LocalE2eBootstrapFactory, - LocalE2eBootstrapDatastream, ->; - -struct ProvisionedDockerNode { - node_id: u64, - stage_index: u32, - endpoint: EndpointAddr, - node_actor: ActorAddress, - provider_process_id: Option, - provisioner: LocalDockerNodeProvisioner, - events: Receiver, - cleaned: bool, -} - -impl ProvisionedDockerNode { - fn stop(&mut self) -> Result<(), String> { - if self.cleaned { - return Ok(()); - } - self.provisioner - .stop() - .map_err(|e| format!("stop node {}: {e:?}", self.node_id))?; - self.cleaned = true; - Ok(()) - } -} - -impl Drop for ProvisionedDockerNode { - fn drop(&mut self) { - let _ = self.stop(); - } -} - -#[derive(Default)] -struct ProvisionStats { - node_live_count: usize, - stdout_line_count: usize, - stderr_line_count: usize, -} - -#[derive(Clone, Debug)] -enum LocalDockerNodeEvent { - Stdout(String), - Stderr(String), - Exited(Option), -} - -#[derive(Clone, Debug)] -enum DriverIngressEvent { - StreamArrived { - edge_id: u64, - stream_id: u64, - }, - BytesRead { - edge_id: u64, - stream_id: u64, - bytes: Vec, - }, -} - -#[derive(Clone)] -struct SendPumpHandle { - tx: tokio_mpsc::UnboundedSender>, -} - -impl SendPumpHandle { - fn send(&self, record: Vec) -> Result<(), String> { - self.tx - .send(record) - .map_err(|_| "edge sender task stopped".to_owned()) - } -} - -struct DriverRuntime { - tx: Sender, - rx: Receiver, - next_stream_id: u64, - connection_count: usize, -} - -impl DriverRuntime { - fn new() -> Self { - let (tx, rx) = mpsc::channel(); - Self { - tx, - rx, - next_stream_id: 1, - connection_count: 0, - } - } - - fn poll_iroh(&mut self, driver: &IrohDriver) { - for (_node, conn) in driver.drain_other_connections() { - self.connection_count += 1; - spawn_recv_pump( - driver.tokio_handle(), - conn, - self.tx.clone(), - self.next_stream_id, - ); - self.next_stream_id += 1; - } - } - - fn try_recv(&self) -> Option { - self.rx.try_recv().ok() - } -} - -fn record_dashboard_event(dashboard: &mut Option<&mut MvpDashboard>, event: obs::Event) { - if let Some(dashboard) = dashboard.as_deref_mut() { - dashboard.record_event(event); - } -} -fn record_dashboard_provision_log( - dashboard: &mut Option<&mut MvpDashboard>, - run_id: u64, - node_id: u64, - stream: ProvisionLogStream, - line: &str, -) { - if let Some(dashboard) = dashboard.as_deref_mut() { - dashboard.record_provision_log(mvp_system::provisioning::ProvisionLogLine { - run_id, - node_id, - stream, - line: line.to_owned(), - }); - } -} - -fn drain_dashboard(dashboard: &mut Option<&mut MvpDashboard>) { - if let Some(dashboard) = dashboard.as_deref_mut() { - dashboard.drain(); - } -} - -fn publish_runtime_snapshot( - dashboard: &mut Option<&mut MvpDashboard>, - stack: &DistributionRuntimeStack, -) { - if let Some(dashboard) = dashboard.as_deref_mut() { - dashboard.publish_runtime_snapshot(stack); - } -} - -fn publish_runtime_snapshot_throttled( - dashboard: &mut Option<&mut MvpDashboard>, - stack: &DistributionRuntimeStack, -) { - if let Some(dashboard) = dashboard.as_deref_mut() { - dashboard.publish_runtime_snapshot_throttled(stack); - } -} - -fn run_event(run_id: u64, kind: obs::EventKind) -> obs::Event { - obs::Event::RunScoped { - kind, - run_id: obs::RunId(run_id), - reason: None, - component: obs::Component::Orchestrator, - } -} - -fn run_fault_event(run_id: u64) -> obs::Event { - obs::Event::RunScoped { - kind: obs::EventKind::RunFaulted, - run_id: obs::RunId(run_id), - reason: None, - component: obs::Component::Orchestrator, - } -} - -fn node_event(node_id: u64, kind: obs::EventKind) -> obs::Event { - obs::Event::NodeScoped { - kind, - node_id: obs::NodeId(node_id), - component: obs::Component::NodeBoot, - } -} - -fn stage_event(run_id: u64, stage_index: u32, kind: obs::EventKind) -> obs::Event { - obs::Event::StageScoped { - kind, - run_id: obs::RunId(run_id), - stage_index: obs::StageIndex(stage_index), - reason: None, - component: obs::Component::StageController, - } -} - -fn object_event(run_id: u64, object_id: u64, sequence: u64, kind: obs::EventKind) -> obs::Event { - let _ = run_id; - obs::Event::ObjectScoped { - kind, - object_id: obs::ObjectId(object_id), - sequence: obs::Sequence(sequence), - component: obs::Component::TokenEndpoint, - } -} - -fn run_supervisor_dashboard_loop() -> Result<(), String> { - let mut dashboard = MvpDashboard::start_from_env()?; - eprintln!("mvp-local-e2e-cluster: dashboard {}", dashboard.url()); - eprintln!( - "mvp-local-e2e-cluster: MVP_DASHBOARD=1, repeating Docker cluster scenario until Ctrl+C" - ); - let mut run_id = RUN_ID; - while !STOP_REQUESTED.load(Ordering::SeqCst) { - match run_supervisor_once( - run_id, - Some(&mut dashboard), - false, - DEFAULT_PROMPT.to_owned(), - ) { - Ok(()) => eprintln!("mvp-local-e2e-cluster: run {run_id} ok"), - Err(error) => eprintln!("mvp-local-e2e-cluster: run {run_id} failed: {error}"), - } - run_id = run_id.saturating_add(1); - thread::sleep(Duration::from_secs(1)); - } - Ok(()) -} - -fn run_supervisor_once( - run_id: u64, - mut dashboard: Option<&mut MvpDashboard>, - print_summary: bool, - prompt_text: String, -) -> Result<(), String> { - let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?; - let mut driver = new_driver(tokio.handle().clone())?; - let mut driver_runtime = DriverRuntime::new(); - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - DistributedNodeConfig::default(), - register_mvp_actor_codecs, - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - - let prompt_tokens = tokenize_prompt(&prompt_text); - let orchestrator_report = stack - .runtime - .new_inbox::() - .map_err(|e| format!("orchestrator report inbox: {e}"))?; - let orchestrator_addr = stack - .runtime - .spawn(OrchestratorActor::new( - fsm::RunConfig { - run_id: fsm::RunId(run_id), - max_tokens: MAX_TOKENS, - prompt: prompt_tokens.clone(), - }, - Some(*orchestrator_report.addr()), - )) - .map_err(|e| format!("spawn orchestrator actor: {e}"))?; - stack.register_local_actor(driver.register_actor(orchestrator_addr, 1)); - publish_runtime_snapshot(&mut dashboard, &stack); - let mut provision_stats = ProvisionStats::default(); - - let topology = build_local_engine_topology(run_id)?; - let run_plan = topology.role_plan.run_plan.clone(); - let stage0 = run_plan - .stages - .iter() - .find(|stage| stage.stage_index == 0) - .cloned() - .ok_or_else(|| "engine builder did not assign stage 0".to_owned())?; - let stage1 = run_plan - .stages - .iter() - .find(|stage| stage.stage_index == 1) - .cloned() - .ok_or_else(|| "engine builder did not assign stage 1".to_owned())?; - - let self_endpoint = driver.endpoint_addr(); - let self_endpoint_json = serde_json::to_string(&self_endpoint) - .map_err(|e| format!("serialize endpoint addr: {e}"))?; - let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr) - .map_err(|e| format!("serialize orchestrator actor: {e}"))?; - - let mut node1 = provision_local_docker_node( - &mut driver, - &stack, - &mut dashboard, - &mut provision_stats, - local_docker_spec( - run_id, - stage1.node_id.0, - stage1.stage_index, - &self_endpoint_json, - &self_endpoint_json, - &orchestrator_actor_json, - ), - )?; - let node1_endpoint_json = serde_json::to_string(&node1.endpoint) - .map_err(|e| format!("serialize node1 endpoint: {e}"))?; - let mut node0 = match provision_local_docker_node( - &mut driver, - &stack, - &mut dashboard, - &mut provision_stats, - local_docker_spec( - run_id, - stage0.node_id.0, - stage0.stage_index, - &node1_endpoint_json, - &self_endpoint_json, - &orchestrator_actor_json, - ), - ) { - Ok(node) => node, - Err(error) => { - let _ = stop_provisioned_nodes(&mut [&mut node1], &mut driver, &stack, &mut dashboard); - return Err(error); - } - }; - - for stage in [&stage1, &stage0] { - record_dashboard_event( - &mut dashboard, - node_event(stage.node_id.0, obs::EventKind::NodeStarted), - ); - record_dashboard_event( - &mut dashboard, - node_event(stage.node_id.0, obs::EventKind::NodeAvailable), - ); - } - - wait_for_routes(&mut driver, &stack, &[node0.node_actor, node1.node_actor])?; - - let token_in_sender = spawn_send_pump( - driver.tokio_handle(), - driver.endpoint().clone(), - node0.endpoint.clone(), - stage0.inbound_edge.0, - )?; - let mut token_in_object_allocator = - edge_actor::ObjectIdAllocator::new(edge_actor::EdgeId(stage0.inbound_edge.0)); - - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObservePoolReady { - nodes: run_plan - .stages - .iter() - .map(|stage| stage.node_id.0) - .collect(), - }, - ) - .map_err(|e| format!("observe pool ready: {e}"))?; - record_dashboard_event(&mut dashboard, run_event(run_id, obs::EventKind::PoolReady)); - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObservePlanAvailable { - run_id, - stages: run_plan - .stages - .iter() - .map(|stage| StageRefWire { - stage_index: stage.stage_index, - node_id: stage.node_id.0, - }) - .collect(), - }, - ) - .map_err(|e| format!("observe plan: {e}"))?; - record_dashboard_event( - &mut dashboard, - run_event(run_id, obs::EventKind::RunPlanned), - ); - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenInEndpointReady, - ) - .map_err(|e| format!("observe token-in endpoint: {e}"))?; - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenOutEndpointReady, - ) - .map_err(|e| format!("observe token-out endpoint: {e}"))?; - record_dashboard_event( - &mut dashboard, - run_event(run_id, obs::EventKind::ReadinessBarrierPassed), - ); - - let mut injected = false; - let mut completed = false; - let mut torn_down = false; - let mut stage_ready_count = 0usize; - let mut token_received = false; - let mut sent_stop_to_node0 = false; - let mut sent_stop_to_node1 = false; - let mut response_tokens = Vec::::new(); - let mut edge_stream_count = 0usize; - let mut token_out_streams = HashMap::>::new(); - let workflow_started = Instant::now(); - - while !STOP_REQUESTED.load(Ordering::SeqCst) { - if workflow_started.elapsed() >= LOCAL_E2E_WORKFLOW_TIMEOUT { - let _ = stop_provisioned_nodes( - &mut [&mut node0, &mut node1], - &mut driver, - &stack, - &mut dashboard, - ); - return Err(format!( - "workflow timed out after {:?}: injected={injected}, token_received={token_received}, completed={completed}, torn_down={torn_down}, stop_node0={sent_stop_to_node0}, stop_node1={sent_stop_to_node1}, stage_ready_count={stage_ready_count}, edge_stream_count={edge_stream_count}", - LOCAL_E2E_WORKFLOW_TIMEOUT - )); - } - - pump_network(&mut driver, &stack); - driver_runtime.poll_iroh(&driver); - while let Some(event) = driver_runtime.try_recv() { - match event { - DriverIngressEvent::StreamArrived { edge_id, .. } => { - if edge_id == stage1.outbound_edge.0 { - edge_stream_count += 1; - } - } - DriverIngressEvent::BytesRead { - edge_id, - stream_id, - bytes, - } => { - if edge_id != stage1.outbound_edge.0 { - continue; - } - let records = { - let buffer = token_out_streams.entry(stream_id).or_default(); - buffer.extend_from_slice(&bytes); - let mut records = Vec::new(); - while let Some(record) = take_complete_ingress_record(buffer)? { - records.push(record); - } - records - }; - for record in records { - let metadata = decode_ingress_record(&record)?; - let words = object_record_words(&metadata, &record)?; - if let Some(token_id) = words.first().copied() { - response_tokens.push(token_id); - token_received = true; - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenReceived { - sequence: metadata.sequence, - token_id, - eos: true, - }, - ) - .map_err(|e| format!("observe token received: {e}"))?; - record_dashboard_event( - &mut dashboard, - object_event( - run_id, - metadata.object_id.0, - metadata.sequence, - obs::EventKind::TokenReceived, - ), - ); - } - } - } - } - } - stage_ready_count += drain_provisioned_node_events( - &mut [&mut node0, &mut node1], - run_id, - &mut dashboard, - &mut provision_stats, - )?; - drain_dashboard(&mut dashboard); - publish_runtime_snapshot_throttled(&mut dashboard, &stack); - - while let Some(report) = orchestrator_report.try_recv() { - match report { - OrchestratorReport::Command(command) => match command { - RunCommandWire::ProvisionStage { stage_index, .. } => { - let provision = stage_provision_wire(&run_plan, stage_index)?; - let target = if stage_index == 0 { - node0.node_actor - } else { - node1.node_actor - }; - stack - .runtime - .send_to(target, NodeAgentMsg::ProvisionStage(provision)) - .map_err(|e| format!("send provision to stage {stage_index}: {e}"))?; - record_dashboard_event( - &mut dashboard, - stage_event(run_id, stage_index, obs::EventKind::StageProvisionStarted), - ); - } - RunCommandWire::InjectTokenObject { - sequence, payload, .. - } => { - injected = true; - let tokens = match payload { - mvp_system::actors::orchestrator::TokenObjectPayloadWire::Prompt { - tokens, - } => tokens, - mvp_system::actors::orchestrator::TokenObjectPayloadWire::Decode { - token_id, - .. - } => vec![token_id], - }; - let prompt_object = token_in_object_allocator.alloc(); - let record = object_record(prompt_object.object_id.0, sequence, &tokens); - token_in_sender.send(record)?; - record_dashboard_event( - &mut dashboard, - object_event( - run_id, - prompt_object.object_id.0, - sequence, - obs::EventKind::PromptInjected, - ), - ); - } - RunCommandWire::StopRun { - stage_index, - run_id, - } => { - let target = if stage_index == 0 { - sent_stop_to_node0 = true; - node0.node_actor - } else { - sent_stop_to_node1 = true; - node1.node_actor - }; - stack - .runtime - .send_to(target, NodeAgentMsg::StopRun { run_id }) - .map_err(|e| format!("send stop to stage {stage_index}: {e}"))?; - record_dashboard_event( - &mut dashboard, - stage_event(run_id, stage_index, obs::EventKind::StopRunSent), - ); - } - RunCommandWire::TearDownTokenEndpoints { .. } => { - stack - .runtime - .send_to( - orchestrator_addr, - OrchestratorMsg::ObserveTokenEndpointsStopped, - ) - .map_err(|e| format!("observe token endpoints stopped: {e}"))?; - } - RunCommandWire::CreateTokenInEndpoint { .. } - | RunCommandWire::CreateTokenOutEndpoint { .. } => {} - }, - OrchestratorReport::Lifecycle(event) => match event { - LifecycleEventWire::RunCompleted { run_id } => { - record_dashboard_event( - &mut dashboard, - run_event(run_id, obs::EventKind::RunCompleted), - ); - completed = true; - } - LifecycleEventWire::RunTornDown { run_id } => { - record_dashboard_event( - &mut dashboard, - run_event(run_id, obs::EventKind::RunTornDown), - ); - torn_down = true; - } - LifecycleEventWire::RunRejected { run_id } - | LifecycleEventWire::RunFaulted { run_id } - | LifecycleEventWire::RunOperatorStopped { run_id } => { - record_dashboard_event(&mut dashboard, run_fault_event(run_id)); - return Err(format!("run failed: {event:?}")); - } - }, - OrchestratorReport::StageFault { - run_id, - stage_index, - } => { - record_dashboard_event(&mut dashboard, run_fault_event(run_id)); - return Err(format!("stage {stage_index} faulted in run {run_id}")); - } - OrchestratorReport::NodeRuntimeReady { .. } - | OrchestratorReport::NodeRuntimeReadyAck { .. } - | OrchestratorReport::WeightsReady { .. } - | OrchestratorReport::StageReady { .. } - | OrchestratorReport::Snapshot { .. } => {} - } - } - - if injected - && token_received - && completed - && torn_down - && sent_stop_to_node0 - && sent_stop_to_node1 - { - stop_provisioned_nodes( - &mut [&mut node0, &mut node1], - &mut driver, - &stack, - &mut dashboard, - )?; - let builder_stage_assignments = topology - .events - .iter() - .filter(|event| { - matches!( - event, - engine::EngineEvent::RoleAssigned { - role: engine::RoleKind::StageWorker { .. }, - .. - } - ) - }) - .count(); - let response_text = detokenize_response(&response_tokens); - let relay_url = relay_url_from_env(); - let orchestrator_endpoint_relay_url = endpoint_relay_url(&self_endpoint); - let node0_endpoint_relay_url = endpoint_relay_url(&node0.endpoint); - let node1_endpoint_relay_url = endpoint_relay_url(&node1.endpoint); - let summary = json!({ - "ok": true, - "actor_plane": "iroh-swactor", - "data_plane": "iroh-quic-persistent-edge-streams", - "edge_protocol": "edge-id-preamble-mo01-object-records", - "node_local_data_plane": "arena-backed-rings-json-metadata-only", - "processes": { - "orchestrator": std::process::id(), - "node0": node0.provider_process_id, - "node1": node1.provider_process_id, - }, - "node0_endpoint": node0.endpoint, - "node1_endpoint": node1.endpoint, - "orchestrator_endpoint": self_endpoint, - "relay_only": relay_url.is_some(), - "relay_url": relay_url, - "orchestrator_endpoint_has_relay": orchestrator_endpoint_relay_url.is_some(), - "orchestrator_endpoint_relay_url": orchestrator_endpoint_relay_url, - "node0_endpoint_has_relay": node0_endpoint_relay_url.is_some(), - "node0_endpoint_relay_url": node0_endpoint_relay_url, - "node1_endpoint_has_relay": node1_endpoint_relay_url.is_some(), - "node1_endpoint_relay_url": node1_endpoint_relay_url, - "node0_logical_id": node0.node_id, - "node1_logical_id": node1.node_id, - "node0_stage_index": node0.stage_index, - "node1_stage_index": node1.stage_index, - "worker_processes": "docker-tinygrad-cpu-worker-per-node", - "docker_image": docker_image(), - "prompt_text": prompt_text, - "prompt_tokens": prompt_tokens, - "response_tokens": response_tokens, - "response_text": response_text, - "tinygrad_device": "CPU", - "injected_prompt_observed": injected, - "token_received_observed": token_received, - "run_completed_observed": completed, - "run_torn_down_observed": torn_down, - "stop_sent_to_all_nodes": sent_stop_to_node0 && sent_stop_to_node1, - "engine_builder_pattern": "host-coordinator-static-topology-docker-workers", - "engine_builder_event_count": topology.events.len(), - "engine_builder_node_count": topology.node_summaries.len(), - "engine_builder_stage_assignments": builder_stage_assignments, - "node_route_count": stack.route_view.read().map(|view| view.len()).unwrap_or_default(), - "stage_ready_stdout_count": stage_ready_count, - "provisioned_node_count": 2, - "provision_node_live_count": provision_stats.node_live_count, - "provision_stdout_line_count": provision_stats.stdout_line_count, - "provision_stderr_line_count": provision_stats.stderr_line_count, - "provision_nodes_stopped": true, - "edge_stream_object_count": edge_stream_count, - }); - if print_summary { - println!("{summary}"); - } else { - eprintln!("mvp-local-e2e-cluster: run {run_id} summary {summary}"); - } - return Ok(()); - } - - thread::sleep(Duration::from_millis(10)); - } - - record_dashboard_event(&mut dashboard, run_fault_event(run_id)); - let _ = stop_provisioned_nodes( - &mut [&mut node0, &mut node1], - &mut driver, - &stack, - &mut dashboard, - ); - Err("interrupted".to_owned()) -} - -fn run_node_role(args: &[String]) -> Result<(), String> { - let logical_node_id = parse_arg(args, "--logical-node-id")? - .parse::() - .map_err(|e| format!("logical node id: {e}"))?; - let stage_index = parse_arg(args, "--stage-index")? - .parse::() - .map_err(|e| format!("stage index: {e}"))?; - let coordinator: EndpointAddr = - serde_json::from_str(parse_arg(args, "--coordinator-endpoint")?) - .map_err(|e| format!("coordinator endpoint json: {e}"))?; - let downstream_endpoint: EndpointAddr = - serde_json::from_str(parse_arg(args, "--outbound-endpoint")?) - .map_err(|e| format!("outbound endpoint json: {e}"))?; - let orchestrator_addr: ActorAddress = - serde_json::from_str(parse_arg(args, "--orchestrator-actor")?) - .map_err(|e| format!("orchestrator actor json: {e}"))?; - - let tokio = tokio::runtime::Runtime::new().map_err(|e| format!("node tokio runtime: {e}"))?; - let mut driver = new_driver(tokio.handle().clone())?; - driver.join(&[coordinator]); - let mut driver_runtime = DriverRuntime::new(); - let stack = DistributionRuntimeStack::new_with_codecs( - driver.node_id(), - DistributedNodeConfig::default(), - register_mvp_actor_codecs, - ); - driver.enable_actor_bridge( - stack.runtime.clone(), - stack.codec.clone(), - stack.actor_bridge_routes(), - stack.actors.swim, - stack.relay_mirror.clone(), - stack.route_view.clone(), - ); - - let node_report = stack - .runtime - .new_inbox::() - .map_err(|e| format!("node report inbox: {e}"))?; - let node_actor = stack - .runtime - .spawn(NodeAgentActor::new( - stage::NodeId(logical_node_id), - orchestrator_addr, - Some(*node_report.addr()), - )) - .map_err(|e| format!("spawn node agent: {e}"))?; - stack.register_local_actor(driver.register_actor(node_actor, 1)); - - let mut arena_manager = arena::ArenaManager::boot(arena::ArenaConfig { - node_id: arena::NodeId(logical_node_id), - reservation_ceiling: ARENA_BYTES as u64, - base_alignment: 64, - }) - .map_err(|e| format!("boot ArenaManager: {e:?}"))?; - let mut worker = GpuWorkerRuntime::spawn(stage_index, arena_manager.arena_fd())?; - let mut edge_establisher = edge::EdgeEstablisher::new(edge::NodeId(logical_node_id)); - let mut tx_edge_actor: Option = None; - let mut rx_edge_actor: Option = None; - let mut driver_fsm = driver_model::Driver::new(driver_model::DriverConfig { - local_node_id: driver_model::NodeId(logical_node_id), - alpn: driver_model::Alpn(String::from_utf8_lossy(EDGE_ALPN).into_owned()), - }); - let mut edge_command_cursor = 0usize; - let mut edge_event_cursor = 0usize; - let mut driver_event_cursor = 0usize; - eprintln!( - "{}", - json!({ - "type": "boot_stream", - "stream": "stderr", - "logical_node_id": logical_node_id, - "stage_index": stage_index, - }) - ); - println!( - "{}", - json!({ - "type": "ready", - "role": "node", - "endpoint": driver.endpoint_addr(), - "node_actor": node_actor, - "logical_node_id": logical_node_id, - "stage_index": stage_index, - }) - ); - std::io::stdout() - .flush() - .map_err(|e| format!("flush node ready: {e}"))?; - - let mut pending_commands = VecDeque::new(); - let mut object_handles = HashMap::::new(); - let mut ingress_streams = HashMap::>::new(); - let mut outbound_sender: Option = None; - let mut outbound_object_allocator: Option = None; - let mut inbound_edge_id = None; - let mut outbound_edge_id = None; - let shutdown_rx = spawn_shutdown_listener(); - let mut inbound_ring_id = None; - let mut outbound_ring_id = None; - - while !STOP_REQUESTED.load(Ordering::SeqCst) { - pump_network(&mut driver, &stack); - driver_runtime.poll_iroh(&driver); - while let Some(event) = driver_runtime.try_recv() { - let (edge_id, stream_id, bytes) = match event { - DriverIngressEvent::StreamArrived { edge_id, stream_id } => { - driver_fsm.observe(driver_model::DriverEvent::IncomingUniStream { - edge_id: driver_model::EdgeId(edge_id), - stream_id: driver_model::StreamId(stream_id), - }); - drive_edge_workflow( - &mut edge_establisher, - &mut edge_command_cursor, - &mut edge_event_cursor, - &mut driver_event_cursor, - &mut arena_manager, - &mut worker, - &mut driver_fsm, - &mut tx_edge_actor, - &mut rx_edge_actor, - &mut outbound_sender, - &mut inbound_edge_id, - &mut outbound_edge_id, - &mut inbound_ring_id, - &mut outbound_ring_id, - &stack.runtime, - node_actor, - driver.tokio_handle(), - driver.endpoint().clone(), - downstream_endpoint.clone(), - )?; - continue; - } - DriverIngressEvent::BytesRead { - edge_id, - stream_id, - bytes, - } => (edge_id, stream_id, bytes), - }; - if Some(edge_id) != inbound_edge_id { - continue; - } - let records = { - let buffer = ingress_streams.entry(stream_id).or_default(); - buffer.extend_from_slice(&bytes); - let mut records = Vec::new(); - while let Some(record) = take_complete_ingress_record(buffer)? { - records.push(record); - } - records - }; - for record in records { - let ring_id = inbound_ring_id.ok_or_else(|| "inbound ring missing".to_owned())?; - let lease = arena_manager - .lookup_lease(arena::RingId(ring_id)) - .ok_or_else(|| format!("inbound ring {ring_id} lease missing"))?; - arena_manager - .write_arena(lease.layout.data_offset, &record) - .map_err(|e| format!("write ingress ring: {e}"))?; - let loaded = worker.ring_readable(ring_id, edge_id)?; - let object_key = edge_actor::ObjectKey::new( - edge_actor::EdgeId(edge_id), - edge_actor::ObjectId(loaded.object_id), - ); - object_handles.insert(object_key, loaded.clone()); - stack - .runtime - .send_to( - node_actor, - NodeAgentMsg::ObjectLoaded { - edge_id, - object_id: loaded.object_id, - sequence: loaded.sequence, - handle_generation: loaded.handle_generation, - handle_id: loaded.handle_id, - }, - ) - .map_err(|e| format!("report object loaded: {e}"))?; - println!( - "{}", - json!({ - "type":"node_lifecycle", - "stage_index":stage_index, - "event":"ObjectLoaded", - "object_id":loaded.object_id, - "sequence":loaded.sequence, - }) - ); - std::io::stdout() - .flush() - .map_err(|e| format!("flush object loaded stdout: {e}"))?; - } - } - - while let Some(report) = node_report.try_recv() { - match report { - NodeAgentReport::Command(command) => pending_commands.push_back(command), - NodeAgentReport::Lifecycle(event) => { - println!( - "{}", - json!({ - "type":"node_lifecycle", - "stage_index":stage_index, - "event":format!("{event:?}"), - }) - ); - std::io::stdout() - .flush() - .map_err(|e| format!("flush lifecycle stdout: {e}"))?; - } - NodeAgentReport::PromptRequested { .. } - | NodeAgentReport::RuntimeReadyAck { .. } - | NodeAgentReport::Snapshot { .. } => {} - } - } - - if let Some(index) = pending_commands.iter().position(readiness_command) { - let command = pending_commands.remove(index).unwrap(); - handle_node_command( - command, - &stack.runtime, - node_actor, - &mut worker, - &mut outbound_sender, - &mut outbound_object_allocator, - &mut inbound_edge_id, - &mut outbound_edge_id, - &mut object_handles, - &mut inbound_ring_id, - &mut outbound_ring_id, - &mut edge_establisher, - &mut edge_command_cursor, - &mut edge_event_cursor, - &mut driver_event_cursor, - &mut arena_manager, - &mut driver_fsm, - &mut tx_edge_actor, - &mut rx_edge_actor, - stage_index, - driver.tokio_handle(), - driver.endpoint().clone(), - downstream_endpoint.clone(), - )?; - } else if let Some(command) = pending_commands.pop_front() { - handle_node_command( - command, - &stack.runtime, - node_actor, - &mut worker, - &mut outbound_sender, - &mut outbound_object_allocator, - &mut inbound_edge_id, - &mut outbound_edge_id, - &mut object_handles, - &mut inbound_ring_id, - &mut outbound_ring_id, - &mut edge_establisher, - &mut edge_command_cursor, - &mut edge_event_cursor, - &mut driver_event_cursor, - &mut arena_manager, - &mut driver_fsm, - &mut tx_edge_actor, - &mut rx_edge_actor, - stage_index, - driver.tokio_handle(), - driver.endpoint().clone(), - downstream_endpoint.clone(), - )?; - } - - if shutdown_rx.try_recv().is_ok() || STOP_REQUESTED.load(Ordering::SeqCst) { - worker.shutdown().ok(); - return Ok(()); - } - - thread::sleep(Duration::from_millis(10)); - } - worker.shutdown().ok(); - Ok(()) -} - -fn new_driver(handle: tokio::runtime::Handle) -> Result { - let relay_mode = if relay_url_from_env().is_some() { - relay_runtime_config_from_env(RUN_ID)?.mode - } else { - let mut relay_provider = LocalShimRelayProvider; - let relay = relay_provider.provision_relay(RelayProvisionRequest { - run_id: RUN_ID, - purpose: RelayPurpose::Combined, - })?; - relay_provider.relay_mode(&relay)? - }; - IrohDriver::with_handle( - handle, - IrohDriverConfig { - secret_key: None, - relay_mode, - node: DistributedNodeConfig::default(), - peer_auth: None, - additional_alpns: vec![EDGE_ALPN.to_vec(), DATASTREAM_ALPN.to_vec()], - }, - ) - .map_err(|e| format!("create iroh driver: {e}")) -} - -fn pump_network(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) { - stack.tick_protocol_actors(Instant::now()); - driver.pump_inbound_to_actors(); - stack.pump_runtime_once(); - driver.drain_outbox(&stack.outbox); -} - -fn wait_for_routes( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - actors: &[ActorAddress], -) -> Result<(), String> { - let started = Instant::now(); - loop { - pump_network(driver, stack); - let route_count = stack - .route_view - .read() - .map(|view| view.len()) - .unwrap_or_default(); - let ready = stack - .route_view - .read() - .map(|view| actors.iter().all(|actor| view.contains_key(actor))) - .unwrap_or(false); - if ready { - return Ok(()); - } - if started.elapsed() >= LOCAL_E2E_ROUTE_TIMEOUT { - return Err(format!( - "routes not ready after {:?}: expected {} actor routes, observed {route_count}", - LOCAL_E2E_ROUTE_TIMEOUT, - actors.len() - )); - } - thread::sleep(Duration::from_millis(20)); - } -} - -struct LocalEngineTopology { - role_plan: engine::RoleAssignmentPlan, - events: Vec, - node_summaries: Vec, -} - -fn build_local_engine_topology(run_id: u64) -> Result { - let cluster = engine::ClusterBuilder::new( - "local-e2e-cluster", - engine::ModelSpec::pipelined_causal_llm( - "local-e2e-cluster-tinygrad-cpu-fixture", - engine::ModelArtifact::TestTinyLlm { - path: "docker-cpu://local-e2e-cluster-tinygrad-cpu-fixture".to_owned(), - }, - 4, - 8, - engine::DTypeFamily::BFloat, - 2, - 8, - 99, - plan::TokenizerSource::EmbeddedGguf, - ), - ) - .run_id(run_id) - .image(engine::NodeImageSpec::new(docker_image()).worker_runtime( - engine::WorkerRuntimeSpec::External { - name: "tinygrad-cpu".to_owned(), - }, - )) - .pool_provider(engine::StaticPoolProvider::new(vec![ - engine::NodeLease::new( - "orchestrator", - engine::NodeId(ORCHESTRATOR_LOGICAL_NODE_ID), - [engine::NodeCapability::Coordinator], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - engine::NodeLease::new( - "node0", - engine::NodeId(NODE0_LOGICAL_ID), - [engine::NodeCapability::Worker], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - engine::NodeLease::new( - "node1", - engine::NodeId(NODE1_LOGICAL_ID), - [engine::NodeCapability::Worker], - ) - .resources(engine::ResourceFacts::cpu_only(2, 2 << 30)), - ])) - .launcher(engine::StaticNodeLauncher) - .planner( - engine::FixedLinearPipelinePlanner::new(2).runtime(plan::RuntimeConfig { - max_tokens: MAX_TOKENS as u32, - prompt: plan::PromptSource::Inline(DEFAULT_PROMPT.to_owned()), - sampling: plan::SamplingPolicy { - temperature_millis: 0, - top_k: 1, - }, - token_output_policy: plan::TokenOutputPolicy::EmitAll, - }), - ) - .launch() - .map_err(|e| format!("local engine builder launch: {e}"))?; - let role_plan = cluster.role_plan().clone(); - let events = cluster.events().to_vec(); - let node_summaries = cluster.node_summaries(); - cluster - .shutdown() - .map_err(|e| format!("local engine builder shutdown: {e}"))?; - Ok(LocalEngineTopology { - role_plan, - events, - node_summaries, - }) -} - -fn stage_provision_wire( - plan: &plan::RunPlan, - stage_index: u32, -) -> Result { - let provision = plan::derive_stage_provision(plan, stage_index) - .map_err(|e| format!("derive stage provision {stage_index}: {e:?}"))?; - Ok(StageProvisionWire { - run_id: provision.run_id.0, - authorized_orchestrator: ORCHESTRATOR_LOGICAL_NODE_ID, - node_id: provision.node_id.0, - stage_index: provision.stage_index, - stage_count: provision.stage_count, - layer_start: provision.layer_start, - layer_end_exclusive: provision.layer_end_exclusive, - inbound_edge_id: provision.inbound.edge_id.0, - outbound_edge_id: provision.outbound.edge_id.0, - model_id: provision.model.model_id, - gguf_source: provision.gguf_source, - tokenizer: provision.tokenizer, - }) -} - -fn handle_node_command( - command: StageCommandWire, - runtime: &swactor::runtime::Runtime, - node_actor: ActorAddress, - worker: &mut GpuWorkerRuntime, - outbound_sender: &mut Option, - outbound_object_allocator: &mut Option, - inbound_edge_id: &mut Option, - outbound_edge_id: &mut Option, - object_handles: &mut HashMap, - inbound_ring_id: &mut Option, - outbound_ring_id: &mut Option, - edge_establisher: &mut edge::EdgeEstablisher, - edge_command_cursor: &mut usize, - edge_event_cursor: &mut usize, - driver_event_cursor: &mut usize, - arena_manager: &mut arena::ArenaManager, - driver_fsm: &mut driver_model::Driver, - tx_edge_actor: &mut Option, - rx_edge_actor: &mut Option, - local_stage_index: u32, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, - downstream_endpoint: EndpointAddr, -) -> Result<(), String> { - match command { - StageCommandWire::EstablishInboundEdge { edge_id } => { - *inbound_edge_id = Some(edge_id); - *rx_edge_actor = Some(edge_actor::RxEdgeActor::new(edge_actor::RxConfig { - edge_id: edge_actor::EdgeId(edge_id), - role_port: edge_actor::PortId("input".to_owned()), - })); - edge_establisher.observe(edge::EdgeEvent::ProvisionRx(edge::ProvisionRx { - run_id: edge::RunId(RUN_ID), - edge_id: edge::EdgeId(edge_id), - local_node_id: edge::NodeId(u64::from(local_stage_index) + 11), - object_spec: edge::ObjectSpec::test_activation(), - ring_spec: edge_ring_spec(), - })); - drive_edge_workflow( - edge_establisher, - edge_command_cursor, - edge_event_cursor, - driver_event_cursor, - arena_manager, - worker, - driver_fsm, - tx_edge_actor, - rx_edge_actor, - outbound_sender, - inbound_edge_id, - outbound_edge_id, - inbound_ring_id, - outbound_ring_id, - runtime, - node_actor, - handle, - endpoint, - downstream_endpoint, - ) - } - StageCommandWire::EstablishOutboundEdge { edge_id } => { - *outbound_edge_id = Some(edge_id); - *outbound_object_allocator = Some(edge_actor::ObjectIdAllocator::new( - edge_actor::EdgeId(edge_id), - )); - *tx_edge_actor = Some(edge_actor::TxEdgeActor::new(edge_actor::TxConfig { - edge_id: edge_actor::EdgeId(edge_id), - role_port: edge_actor::PortId("output".to_owned()), - })); - edge_establisher.observe(edge::EdgeEvent::ProvisionTx(edge::ProvisionTx { - run_id: edge::RunId(RUN_ID), - edge_id: edge::EdgeId(edge_id), - local_node_id: edge::NodeId(u64::from(local_stage_index) + 11), - consumer_node_id: edge::NodeId(if local_stage_index == 0 { - NODE1_LOGICAL_ID - } else { - ORCHESTRATOR_LOGICAL_NODE_ID - }), - object_spec: edge::ObjectSpec::test_activation(), - ring_spec: edge_ring_spec(), - })); - drive_edge_workflow( - edge_establisher, - edge_command_cursor, - edge_event_cursor, - driver_event_cursor, - arena_manager, - worker, - driver_fsm, - tx_edge_actor, - rx_edge_actor, - outbound_sender, - inbound_edge_id, - outbound_edge_id, - inbound_ring_id, - outbound_ring_id, - runtime, - node_actor, - handle, - endpoint, - downstream_endpoint, - ) - } - StageCommandWire::ConfigureWorkerRole { - run_id, - stage_index, - layer_start, - layer_end_exclusive, - } => { - worker.configure_role(run_id, stage_index, layer_start, layer_end_exclusive)?; - runtime - .send_to(node_actor, NodeAgentMsg::MarkWorkerReady) - .map_err(|e| format!("mark worker ready: {e}")) - } - StageCommandWire::LoadWeights { - model_id, - gguf_source, - tokenizer, - layer_start, - layer_end_exclusive, - } => { - worker.load_weights( - model_id, - gguf_source, - tokenizer, - layer_start, - layer_end_exclusive, - )?; - runtime - .send_to( - node_actor, - NodeAgentMsg::MarkWeightsReady { - run_id: RUN_ID, - node_id: u64::from(local_stage_index) + 11, - stage_index: local_stage_index, - }, - ) - .map_err(|e| format!("mark weights ready: {e}")) - } - StageCommandWire::ExecuteStep { - step_id, - input_edge_id, - object_id, - sequence, - .. - } => { - let input_key = edge_actor::ObjectKey::new( - edge_actor::EdgeId(input_edge_id), - edge_actor::ObjectId(object_id), - ); - let loaded = object_handles - .get(&input_key) - .cloned() - .ok_or_else(|| format!("object {input_key:?} has no loaded device handle"))?; - if loaded.sequence != sequence { - return Err(format!( - "object {input_key:?} sequence {} does not match command sequence {sequence}", - loaded.sequence - )); - } - if loaded.edge_id != input_edge_id { - return Err(format!( - "object {input_key:?} was loaded from edge {}, not command edge {input_edge_id}", - loaded.edge_id - )); - } - let output_ring = outbound_ring_id.ok_or_else(|| "outbound ring missing".to_owned())?; - let output_key = outbound_object_allocator - .as_mut() - .ok_or_else(|| "outbound object allocator missing".to_owned())? - .alloc(); - let output_object_id = output_key.object_id.0; - let committed_bytes = worker.execute_step( - u64::from(local_stage_index) + 1, - step_id, - object_id, - loaded.sequence, - loaded.handle_id, - output_ring, - output_object_id, - loaded.sequence, - local_stage_index == 1, - )?; - let lease = arena_manager - .lookup_lease(arena::RingId(output_ring)) - .ok_or_else(|| format!("outbound ring {output_ring} lease missing"))?; - let record = arena_manager - .read_arena(lease.layout.data_offset, committed_bytes) - .map_err(|e| format!("read egress ring: {e}"))?; - driver_fsm.observe(driver_model::DriverEvent::EgressBytesCommitted { - edge_id: driver_model::EdgeId( - outbound_edge_id.ok_or_else(|| "outbound edge missing".to_owned())?, - ), - bytes: record.clone(), - }); - driver_fsm.observe(driver_model::DriverEvent::RingReadable { - ring_id: driver_model::RingId(output_ring), - }); - if let Some(tx_actor) = tx_edge_actor.as_mut() { - tx_actor.observe(edge_actor::TxEvent::ObjectProduced { - edge_id: edge_actor::EdgeId( - outbound_edge_id.ok_or_else(|| "outbound edge missing".to_owned())?, - ), - object_id: output_key.object_id, - sequence: loaded.sequence, - }); - } - let sender = outbound_sender - .as_ref() - .ok_or_else(|| "outbound edge sender missing".to_owned())?; - sender.send(record)?; - runtime - .send_to(node_actor, NodeAgentMsg::StepCompleted { step_id }) - .map_err(|e| format!("mark step completed: {e}")) - } - StageCommandWire::ReleaseInputHandle { handle_id, .. } => { - worker.release_device_object(handle_id) - } - StageCommandWire::StopLocalEdges { run_id } => { - runtime - .send_to(node_actor, NodeAgentMsg::LocalEdgesStopped { run_id }) - .map_err(|e| format!("mark local edges stopped: {e}"))?; - runtime - .send_to(node_actor, NodeAgentMsg::WorkerRingsQuiesced { run_id }) - .map_err(|e| format!("mark worker rings quiesced: {e}")) - } - StageCommandWire::ReleaseRunDeviceObjects { run_id } => { - runtime - .send_to(node_actor, NodeAgentMsg::DeviceObjectsReleased { run_id }) - .map_err(|e| format!("mark device objects released: {e}"))?; - runtime - .send_to(node_actor, NodeAgentMsg::WorkerRoleReset { run_id }) - .map_err(|e| format!("mark worker role reset: {e}")) - } - StageCommandWire::RewireEdge { .. } => Ok(()), - } -} - -fn readiness_command(command: &StageCommandWire) -> bool { - matches!( - command, - StageCommandWire::EstablishInboundEdge { .. } - | StageCommandWire::EstablishOutboundEdge { .. } - | StageCommandWire::ConfigureWorkerRole { .. } - | StageCommandWire::LoadWeights { .. } - ) -} - -#[derive(Clone, Debug)] -struct LoadedObject { - object_id: u64, - edge_id: u64, - sequence: u64, - handle_generation: u64, - handle_id: u64, -} -fn edge_ring_spec() -> edge::RingSpec { - edge::RingSpec { - header_bytes: 128, - data_bytes: RING_BYTES as u64, - alignment: 64, - } -} - -fn drive_edge_workflow( - edge_establisher: &mut edge::EdgeEstablisher, - edge_command_cursor: &mut usize, - edge_event_cursor: &mut usize, - driver_event_cursor: &mut usize, - arena_manager: &mut arena::ArenaManager, - worker: &mut GpuWorkerRuntime, - driver_fsm: &mut driver_model::Driver, - tx_edge_actor: &mut Option, - rx_edge_actor: &mut Option, - outbound_sender: &mut Option, - inbound_edge_id: &mut Option, - outbound_edge_id: &mut Option, - inbound_ring_id: &mut Option, - outbound_ring_id: &mut Option, - runtime: &swactor::runtime::Runtime, - node_actor: ActorAddress, - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, - downstream_endpoint: EndpointAddr, -) -> Result<(), String> { - loop { - let mut progressed = false; - - while *edge_command_cursor < edge_establisher.commands().len() { - let command = edge_establisher.commands()[*edge_command_cursor].clone(); - *edge_command_cursor += 1; - progressed = true; - match command { - edge::EdgeCommand::LeaseRing { - request_id, - direction, - ring_spec, - .. - } => { - let events = - arena_manager.request(arena::ArenaRequest::LeaseRing(arena::LeaseRing { - request_id: arena::LeaseRequestId(request_id.0), - ring_spec: arena::RingSpec { - header_bytes: ring_spec.header_bytes, - data_bytes: ring_spec.data_bytes, - alignment: ring_spec.alignment, - }, - })); - for event in events { - match event { - arena::ArenaEvent::RingLeased { lease } => { - let edge_layout = edge::RingLayout { - start_offset: lease.layout.start_offset, - header_offset: lease.layout.header_offset, - data_offset: lease.layout.data_offset, - end_offset: lease.layout.end_offset, - data_bytes: lease.layout.data_bytes, - alignment: lease.layout.alignment, - }; - edge_establisher.observe(edge::EdgeEvent::RingLeased { - request_id: edge::LeaseRequestId(lease.request_id.0), - ring_id: edge::RingId(lease.ring_id.0), - layout: edge_layout, - }); - } - arena::ArenaEvent::RingLeaseRejected { request_id, reason } => { - let reason = match reason { - arena::RingLeaseRejection::CannotFitWithinCeiling => { - edge::RingLeaseRejection::CannotFit - } - arena::RingLeaseRejection::ArenaShuttingDown => { - edge::RingLeaseRejection::ArenaShuttingDown - } - }; - edge_establisher.observe(edge::EdgeEvent::RingLeaseRejected { - request_id: edge::LeaseRequestId(request_id.0), - reason, - }); - } - arena::ArenaEvent::RingLeaseQueued { .. } - | arena::ArenaEvent::RingReleased { .. } - | arena::ArenaEvent::RingReleaseRejected { .. } - | arena::ArenaEvent::CancelledFreshLeaseReleased { .. } => {} - } - } - if matches!(direction, edge::RingDirection::Ingress) { - // The eventual ring id is learned from the install command. - } - } - edge::EdgeCommand::InstallWorkerRing { - edge_id, - ring_id, - direction, - .. - } => { - let lease = arena_manager - .lookup_lease(arena::RingId(ring_id.0)) - .ok_or_else(|| format!("ring {} lease missing", ring_id.0))? - .clone(); - let (port, direction_name) = match direction { - edge::RingDirection::Ingress => { - *inbound_ring_id = Some(ring_id.0); - ("input", "ingress") - } - edge::RingDirection::Egress => { - *outbound_ring_id = Some(ring_id.0); - ("output", "egress") - } - }; - worker.install_ring( - ring_id.0, - edge_id.0, - port, - direction_name, - lease.layout, - )?; - edge_establisher.observe(edge::EdgeEvent::RingInstalled { edge_id, ring_id }); - } - edge::EdgeCommand::EstablishSend { - edge_id, - consumer_node_id, - .. - } => { - let record = edge_establisher - .local_record(edge_id) - .ok_or_else(|| format!("edge {} record missing", edge_id.0))?; - let ring_id = record - .ring_id - .ok_or_else(|| format!("edge {} ring missing", edge_id.0))?; - driver_fsm.observe(driver_model::DriverEvent::EstablishSend( - driver_model::EstablishSend { - edge_id: driver_model::EdgeId(edge_id.0), - peer_node_id: driver_model::NodeId(consumer_node_id.0), - layout: driver_model::RingLayout { - ring_id: driver_model::RingId(ring_id.0), - byte_capacity: RING_BYTES, - direction: driver_model::RingDirection::Egress, - }, - }, - )); - *outbound_sender = Some(spawn_send_pump( - handle.clone(), - endpoint.clone(), - downstream_endpoint.clone(), - edge_id.0, - )?); - } - edge::EdgeCommand::EstablishRecv { edge_id, .. } => { - let record = edge_establisher - .local_record(edge_id) - .ok_or_else(|| format!("edge {} record missing", edge_id.0))?; - let ring_id = record - .ring_id - .ok_or_else(|| format!("edge {} ring missing", edge_id.0))?; - driver_fsm.observe(driver_model::DriverEvent::EstablishRecv( - driver_model::EstablishRecv { - edge_id: driver_model::EdgeId(edge_id.0), - layout: driver_model::RingLayout { - ring_id: driver_model::RingId(ring_id.0), - byte_capacity: RING_BYTES, - direction: driver_model::RingDirection::Ingress, - }, - }, - )); - } - edge::EdgeCommand::CancelQueuedLease { request_id, .. } => { - let _ = arena_manager.request(arena::ArenaRequest::CancelLease { - request_id: arena::LeaseRequestId(request_id.0), - }); - } - edge::EdgeCommand::StopPump { edge_id, .. } => { - driver_fsm.observe(driver_model::DriverEvent::StopEdge { - edge_id: driver_model::EdgeId(edge_id.0), - }); - } - edge::EdgeCommand::UninstallWorkerRing { ring_id, .. } => { - edge_establisher.observe(edge::EdgeEvent::RingQuiesced { ring_id }); - } - edge::EdgeCommand::ReleaseArenaLease { ring_id, proof } => { - let proof = if proof == edge::QuiescenceProof::verified() { - arena::QuiescenceProof::verified() - } else { - arena::QuiescenceProof::missing() - }; - for event in arena_manager.request(arena::ArenaRequest::ReleaseRing { - ring_id: arena::RingId(ring_id.0), - proof, - }) { - if let arena::ArenaEvent::RingReleased { .. } = event { - edge_establisher.observe(edge::EdgeEvent::Stopped { - edge_id: edge::EdgeId(0), - }); - } - } - } - } - } - - while *driver_event_cursor < driver_fsm.events().len() { - let event = driver_fsm.events()[*driver_event_cursor].clone(); - *driver_event_cursor += 1; - progressed = true; - match event { - driver_model::DriverEventOut::DriverEdgeReady { edge_id } => { - edge_establisher.observe(edge::EdgeEvent::DriverEdgeReady { - edge_id: edge::EdgeId(edge_id.0), - }); - } - driver_model::DriverEventOut::StreamFault { edge_id, reason } => { - let reason = match reason { - driver_model::StreamFaultReason::ReadError => { - edge::StreamFaultReason::ReadError - } - driver_model::StreamFaultReason::WriteError => { - edge::StreamFaultReason::WriteError - } - driver_model::StreamFaultReason::ProtocolError => { - edge::StreamFaultReason::ProtocolError - } - }; - edge_establisher.observe(edge::EdgeEvent::StreamFault { - edge_id: edge::EdgeId(edge_id.0), - reason, - }); - } - driver_model::DriverEventOut::PumpStopped { edge_id, ring_id } => { - edge_establisher.observe(edge::EdgeEvent::PumpStopped { - edge_id: edge::EdgeId(edge_id.0), - ring_id: edge::RingId(ring_id.0), - }); - } - driver_model::DriverEventOut::StreamClosed { .. } => {} - } - } - - while *edge_event_cursor < edge_establisher.events().len() { - let event = edge_establisher.events()[*edge_event_cursor].clone(); - *edge_event_cursor += 1; - progressed = true; - match event { - edge::EdgeLifecycleEvent::EdgeReady { edge_id, .. } => { - if Some(edge_id.0) == *inbound_edge_id { - if let Some(rx_actor) = rx_edge_actor.as_mut() { - rx_actor.observe(edge_actor::RxEvent::EdgeReady { - edge_id: edge_actor::EdgeId(edge_id.0), - }); - } - runtime - .send_to( - node_actor, - NodeAgentMsg::MarkInboundEdgeReady { edge_id: edge_id.0 }, - ) - .map_err(|e| format!("mark inbound ready: {e}"))?; - } - if Some(edge_id.0) == *outbound_edge_id { - if let Some(tx_actor) = tx_edge_actor.as_mut() { - tx_actor.observe(edge_actor::TxEvent::EdgeReady { - edge_id: edge_actor::EdgeId(edge_id.0), - }); - } - runtime - .send_to( - node_actor, - NodeAgentMsg::MarkOutboundEdgeReady { edge_id: edge_id.0 }, - ) - .map_err(|e| format!("mark outbound ready: {e}"))?; - } - } - edge::EdgeLifecycleEvent::EdgeFaulted { .. } - | edge::EdgeLifecycleEvent::EdgeStopped { .. } => {} - } - } - - if !progressed { - break; - } - } - Ok(()) -} - -struct GpuWorkerRuntime { - ctl: worker_ctl::GpuWorkerCtl, - child: Child, - stdin: ChildStdin, - stdout: BufReader, - generation: u64, -} - -impl GpuWorkerRuntime { - fn spawn(stage_index: u32, arena_fd: i32) -> Result { - let mut ctl = worker_ctl::GpuWorkerCtl::new(worker_ctl::WorkerConfig { - node_id: worker_ctl::NodeId(u64::from(stage_index) + 1), - arena_env: worker_ctl::ArenaEnv { - arena_fd, - arena_bytes: ARENA_BYTES as u64, - }, - }); - ctl.observe(worker_ctl::WorkerCtlEvent::StartWorker); - - let mut child = Command::new("python3") - .arg(tinygrad_worker_path()) - .env("SWACTOR_ARENA_FD", arena_fd.to_string()) - .env("SWACTOR_ARENA_BYTES", ARENA_BYTES.to_string()) - .env("SWACTOR_STAGE_INDEX", stage_index.to_string()) - .env("DEV", "CPU") - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::inherit()) - .spawn() - .map_err(|e| format!("spawn tinygrad CPU worker: {e}"))?; - let process_id = worker_ctl::ProcessId(u64::from(child.id())); - ctl.observe(worker_ctl::WorkerCtlEvent::ProcessStarted { pid: process_id }); - let stdin = child - .stdin - .take() - .ok_or_else(|| "worker stdin missing".to_owned())?; - let stdout = child - .stdout - .take() - .ok_or_else(|| "worker stdout missing".to_owned())?; - let mut worker = Self { - ctl, - child, - stdin, - stdout: BufReader::new(stdout), - generation: 1, - }; - worker.initialize()?; - Ok(worker) - } - - fn initialize(&mut self) -> Result<(), String> { - let value = self.command( - json!({ - "type":"InitializeWorker", - "worker_generation": self.generation, - "arena_ceiling": ARENA_BYTES, - "required_ring_helper_abi": 1, - "backend": {"device":"CPU"}, - }), - "WorkerReady", - )?; - let generation = value - .get("worker_generation") - .or_else(|| value.get("generation")) - .and_then(|value| value.as_u64()) - .unwrap_or(self.generation); - self.generation = generation; - self.ctl.observe(worker_ctl::WorkerCtlEvent::WorkerReady { - generation: worker_ctl::WorkerGeneration(generation), - }); - Ok(()) - } - - fn install_ring( - &mut self, - ring_id: u64, - edge_id: u64, - port_id: &str, - direction: &str, - layout: arena::RingLayout, - ) -> Result<(), String> { - self.ctl.observe(worker_ctl::WorkerCtlEvent::ActorCommand( - worker_ctl::ActorCommand::InstallRing { - generation: worker_ctl::WorkerGeneration(self.generation), - ring_id: worker_ctl::RingId(ring_id), - }, - )); - self.command( - json!({ - "type":"InstallRing", - "ring_id":ring_id, - "edge_id":edge_id, - "port_id":port_id, - "direction":direction, - "layout": { - "ring_id": ring_id, - "arena_offset": layout.start_offset, - "header_offset": layout.header_offset, - "data_offset": layout.data_offset, - "data_capacity": layout.data_bytes, - "alignment": layout.alignment, - }, - "object_spec": { - "kind":"token_or_activation", - "max_extent": OBJECT_MAX_EXTENT, - "dtype_family":"i32", - "dtype_width_bytes":4, - "layout":"linear", - "alignment": OBJECT_ALIGNMENT, - "sequence_policy":"strict_increasing", - } - }), - "RingInstalled", - )?; - self.ctl.observe(worker_ctl::WorkerCtlEvent::StdoutEvent( - worker_ctl::WorkerEvent::RingInstalled { - ring_id: worker_ctl::RingId(ring_id), - }, - )); - Ok(()) - } - - fn configure_role( - &mut self, - run_id: u64, - stage_index: u32, - layer_start: u32, - layer_end_exclusive: u32, - ) -> Result<(), String> { - self.command( - json!({ - "type":"ConfigureRole", - "role_id":stage_index + 1, - "config": { - "run_id":run_id, - "stage_index":stage_index, - "layer_start":layer_start, - "layer_end_exclusive":layer_end_exclusive, - } - }), - "RoleConfigured", - ) - .map(|_| ()) - } - - fn load_weights( - &mut self, - model_id: String, - gguf_source: plan::GgufSource, - tokenizer: plan::TokenizerSource, - layer_start: u32, - layer_end_exclusive: u32, - ) -> Result<(), String> { - self.command( - json!({ - "type":"LoadWeights", - "model_id":model_id, - "gguf_source":gguf_source, - "tokenizer":tokenizer, - "layer_start":layer_start, - "layer_end_exclusive":layer_end_exclusive, - }), - "WeightsLoaded", - ) - .map(|_| ()) - } - - fn ring_readable(&mut self, ring_id: u64, edge_id: u64) -> Result { - let value = self.command( - json!({"type":"RingReadable","ring_id":ring_id}), - "ObjectLoaded", - )?; - let loaded = LoadedObject { - object_id: value - .get("object_id") - .and_then(|value| value.as_u64()) - .ok_or_else(|| format!("ObjectLoaded missing object_id: {value}"))?, - edge_id, - sequence: value - .get("sequence") - .and_then(|value| value.as_u64()) - .ok_or_else(|| format!("ObjectLoaded missing sequence: {value}"))?, - handle_generation: value - .get("device_handle") - .or_else(|| value.get("handle")) - .and_then(|value| { - value - .get("worker_generation") - .or_else(|| value.get("generation")) - }) - .and_then(|value| value.as_u64()) - .ok_or_else(|| format!("ObjectLoaded missing handle generation: {value}"))?, - handle_id: value - .get("device_handle") - .or_else(|| value.get("handle")) - .and_then(|value| value.get("id")) - .and_then(|value| value.as_u64()) - .ok_or_else(|| format!("ObjectLoaded missing handle id: {value}"))?, - }; - self.ctl.observe(worker_ctl::WorkerCtlEvent::StdoutEvent( - worker_ctl::WorkerEvent::ObjectLoaded { - object_id: worker_ctl::ObjectId(loaded.object_id), - sequence: loaded.sequence, - }, - )); - Ok(loaded) - } - - fn execute_step( - &mut self, - role_id: u64, - step_id: u64, - input_object_id: u64, - input_sequence: u64, - input_handle: u64, - output_ring_id: u64, - output_object_id: u64, - output_sequence: u64, - final_stage: bool, - ) -> Result { - self.ctl.observe(worker_ctl::WorkerCtlEvent::ActorCommand( - worker_ctl::ActorCommand::ExecuteStep { - generation: worker_ctl::WorkerGeneration(self.generation), - step_id: worker_ctl::StepId(step_id), - input: worker_ctl::DeviceHandle { - generation: worker_ctl::WorkerGeneration(self.generation), - id: input_handle, - }, - }, - )); - let produced = self.command( - json!({ - "type":"ExecuteStep", - "role_id":role_id, - "step_id":step_id, - "inputs":[{ - "port_id":"input", - "object_id":input_object_id, - "sequence":input_sequence, - "device_handle":{"worker_generation":self.generation,"id":input_handle}, - }], - "outputs":[{ - "port_id":"output", - "ring_id":output_ring_id, - "object_id":output_object_id, - "sequence":output_sequence, - "extent":4, - "flags": if final_stage { 1 } else { 0 }, - }], - "runtime":{"final_stage":final_stage}, - "release_inputs_after":false, - }), - "ObjectProduced", - )?; - let produced_object_id = produced - .get("object_id") - .and_then(|value| value.as_u64()) - .unwrap_or(output_object_id); - let produced_sequence = produced - .get("sequence") - .and_then(|value| value.as_u64()) - .unwrap_or(output_sequence); - self.ctl.observe(worker_ctl::WorkerCtlEvent::StdoutEvent( - worker_ctl::WorkerEvent::ObjectProduced { - object_id: worker_ctl::ObjectId(produced_object_id), - sequence: produced_sequence, - }, - )); - self.expect_event("StepCompleted")?; - self.ctl.observe(worker_ctl::WorkerCtlEvent::StdoutEvent( - worker_ctl::WorkerEvent::StepCompleted { - step_id: worker_ctl::StepId(step_id), - }, - )); - produced - .get("committed_bytes") - .and_then(|value| value.as_u64()) - .map(|value| value as usize) - .ok_or_else(|| format!("ObjectProduced missing committed_bytes: {produced}")) - } - - fn release_device_object(&mut self, handle: u64) -> Result<(), String> { - self.ctl.observe(worker_ctl::WorkerCtlEvent::ActorCommand( - worker_ctl::ActorCommand::ReleaseDeviceObject { - generation: worker_ctl::WorkerGeneration(self.generation), - handle: worker_ctl::DeviceHandle { - generation: worker_ctl::WorkerGeneration(self.generation), - id: handle, - }, - }, - )); - self.command( - json!({"type":"ReleaseDeviceObject","device_handle":{"worker_generation":self.generation,"id":handle}}), - "DeviceObjectReleased", - ) - .map(|_| ()) - } - - fn shutdown(&mut self) -> Result<(), String> { - self.ctl - .observe(worker_ctl::WorkerCtlEvent::ShutdownRequested); - self.command( - json!({"type":"ShutdownWorker","mode":"Graceful"}), - "WorkerStopped", - ) - .map(|_| ()) - } - - fn expect_event(&mut self, expected: &str) -> Result { - let mut line = String::new(); - self.stdout - .read_line(&mut line) - .map_err(|e| format!("read worker stdout: {e}"))?; - let value: serde_json::Value = serde_json::from_str(&line) - .map_err(|e| format!("parse worker stdout {line:?}: {e}"))?; - if value.get("type").and_then(|value| value.as_str()) == Some(expected) { - Ok(value) - } else { - Err(format!("worker emitted {value}, expected {expected}")) - } - } - - fn command( - &mut self, - command: serde_json::Value, - expected: &str, - ) -> Result { - if command.get("payload").is_some() - || command.get("bytes").is_some() - || command.get("data").is_some() - { - return Err(format!( - "worker command illegally carried payload bytes: {command}" - )); - } - writeln!(self.stdin, "{command}").map_err(|e| format!("write worker command: {e}"))?; - self.stdin - .flush() - .map_err(|e| format!("flush worker stdin: {e}"))?; - self.expect_event(expected) - } -} - -impl Drop for GpuWorkerRuntime { - fn drop(&mut self) { - let _ = self.child.kill(); - let _ = self.child.wait(); - } -} - -fn local_docker_spec( - run_id: u64, - node_id: u64, - stage_index: u32, - outbound_endpoint_json: &str, - coordinator_endpoint_json: &str, - orchestrator_actor_json: &str, -) -> NodeProvisionSpec { - let mut env = vec![ - ("DEV".to_owned(), "CPU".to_owned()), - ("PYTHONDONTWRITEBYTECODE".to_owned(), "1".to_owned()), - ( - "MVP_TINYGRAD_WORKER".to_owned(), - "/workspace/crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py" - .to_owned(), - ), - ]; - if let Some(relay_url) = relay_url_from_env() { - env.push(("MVP_IROH_RELAY_MODE".to_owned(), "default".to_owned())); - env.push((MVP_IROH_RELAY_URL_ENV.to_owned(), relay_url)); - } - - NodeProvisionSpec { - run_id, - node_id, - stage_index: Some(stage_index), - image: docker_image(), - env, - args: vec![ - "--role=node".to_owned(), - "--logical-node-id".to_owned(), - node_id.to_string(), - "--stage-index".to_owned(), - stage_index.to_string(), - "--outbound-endpoint".to_owned(), - outbound_endpoint_json.to_owned(), - "--coordinator-endpoint".to_owned(), - coordinator_endpoint_json.to_owned(), - "--orchestrator-actor".to_owned(), - orchestrator_actor_json.to_owned(), - ], - mounts: Vec::new(), - } -} - -struct LocalE2eContainer { - container_name: String, - stdin: ChildStdin, -} - -struct LocalE2eDockerCli { - spec: NodeProvisionSpec, - events: Sender, - nodes: HashMap, - provider_process_id: Option, -} - -impl LocalE2eDockerCli { - fn new(spec: NodeProvisionSpec, events: Sender) -> Self { - Self { - spec, - events, - nodes: HashMap::new(), - provider_process_id: None, - } - } - - fn provider_process_id(&self) -> Option { - self.provider_process_id - } -} - -impl docker_provision::DockerCli for LocalE2eDockerCli { - fn run_container( - &mut self, - request: docker_provision::DockerRunRequest, - ) -> Result { - let mut env = request.env.clone(); - for (key, value) in &self.spec.env { - env.insert(key.clone(), value.clone()); - } - - let mut command = Command::new("docker"); - command.arg("run").arg("--rm"); - if let Some(network) = docker_network_for_node(self.spec.node_id) { - command.arg("--network").arg(network); - } else { - command - .arg("--add-host") - .arg("host.docker.internal:host-gateway"); - } - command.arg("--name").arg(&request.container_name).arg("-i"); - for (key, value) in &request.labels { - command.arg("--label").arg(format!("{key}={value}")); - } - for (key, value) in &env { - command.arg("-e").arg(format!("{key}={value}")); - } - command.arg(&self.spec.image); - for arg in &self.spec.args { - command.arg(arg); - } - - let mut child = command - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .spawn() - .map_err(|e| { - docker_provision::DockerCliError::new(format!( - "spawn Docker node {}: {e}", - self.spec.node_id - )) - })?; - - let provider_process_id = child.id(); - self.provider_process_id = Some(provider_process_id); - let stdin = child.stdin.take().ok_or_else(|| { - docker_provision::DockerCliError::new(format!( - "Docker node {} stdin missing", - self.spec.node_id - )) - })?; - let stdout = child.stdout.take().ok_or_else(|| { - docker_provision::DockerCliError::new(format!( - "Docker node {} stdout missing", - self.spec.node_id - )) - })?; - let stderr = child.stderr.take().ok_or_else(|| { - docker_provision::DockerCliError::new(format!( - "Docker node {} stderr missing", - self.spec.node_id - )) - })?; - - self.nodes.insert( - request.container_name.clone(), - LocalE2eContainer { - container_name: request.container_name.clone(), - stdin, - }, - ); - - spawn_local_docker_reader(stdout, self.events.clone(), true); - spawn_local_docker_reader(stderr, self.events.clone(), false); - let events = self.events.clone(); - thread::spawn(move || match child.wait() { - Ok(status) => { - let _ = events.send(LocalDockerNodeEvent::Exited(status.code())); - } - Err(error) => { - let _ = events.send(LocalDockerNodeEvent::Stderr(format!( - "wait Docker node: {error}" - ))); - let _ = events.send(LocalDockerNodeEvent::Exited(None)); - } - }); - - Ok(docker_provision::DockerRunResult { - container_id: request.container_name, - }) - } - - fn inspect_ssh_endpoint( - &mut self, - container_id: &str, - ) -> Result, docker_provision::DockerCliError> { - Ok(Some(node_provision::SshEndpoint { - host: "127.0.0.1".to_owned(), - port: 0, - user: "local-e2e".to_owned(), - auth_ref: format!("local-docker:{container_id}"), - })) - } - - fn remove_force(&mut self, container_id: &str) -> Result<(), docker_provision::DockerCliError> { - let Some(mut node) = self.nodes.remove(container_id) else { - return Ok(()); - }; - let _ = writeln!(node.stdin, "shutdown"); - let _ = node.stdin.flush(); - let status = Command::new("docker") - .arg("stop") - .arg("-t") - .arg("2") - .arg(&node.container_name) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - .status() - .map_err(|e| { - docker_provision::DockerCliError::new(format!( - "docker stop {}: {e}", - node.container_name - )) - })?; - if status.success() { - Ok(()) - } else { - Err(docker_provision::DockerCliError::new(format!( - "docker stop {} exited with {status}", - node.container_name - ))) - } - } -} - -#[derive(Default)] -struct LocalE2eBootstrapFactory; - -impl docker_provision::SshBootstrapClientFactory for LocalE2eBootstrapFactory { - type Client = LocalE2eBootstrapClient; - - fn client_for(&mut self, _spec: &node_provision::BootstrapSessionSpec) -> Self::Client { - LocalE2eBootstrapClient - } -} - -struct LocalE2eBootstrapClient; - -impl docker_provision::BootstrapSshClient for LocalE2eBootstrapClient { - fn connect( - &mut self, - _endpoint: &node_provision::SshEndpoint, - ) -> Result<(), docker_provision::BootstrapSshError> { - Ok(()) - } - - fn probe_stdout(&mut self) -> Result<(), docker_provision::BootstrapSshError> { - Ok(()) - } - - fn read_bootstrap_logs( - &mut self, - _stdout_sources: &[String], - _stderr_sources: &[String], - ) -> Result< - Vec<(node_provision::BootstrapLogStream, String)>, - docker_provision::BootstrapSshError, - > { - Ok(Vec::new()) - } - - fn run_verify_commands( - &mut self, - _commands: &[String], - ) -> Result<(), docker_provision::BootstrapSshError> { - Ok(()) - } - - fn start_swactor( - &mut self, - _command: &str, - _join: &node_provision::SwarmJoinSpec, - ) -> Result<(), docker_provision::BootstrapSshError> { - Ok(()) - } - - fn close(&mut self) {} -} - -#[derive(Default)] -struct LocalE2eBootstrapDatastream; - -impl node_provision::BootstrapDatastreamSink for LocalE2eBootstrapDatastream { - fn record(&mut self, _record: node_provision::BootstrapLogRecord) {} - - fn flush(&mut self) {} -} - -fn spawn_local_docker_reader( - stream: impl std::io::Read + Send + 'static, - events: Sender, - stdout: bool, -) { - thread::spawn(move || { - let reader = BufReader::new(stream); - for next in reader.lines() { - let Ok(line) = next else { - break; - }; - let event = if stdout { - LocalDockerNodeEvent::Stdout(line) - } else { - LocalDockerNodeEvent::Stderr(line) - }; - if events.send(event).is_err() { - break; - } - } - }); -} - -fn logical_node_spec_from_local(spec: &NodeProvisionSpec) -> node_provision::LogicalNodeSpec { - let logical_node_id = node_provision::LogicalNodeId(spec.node_id.to_string()); - node_provision::LogicalNodeSpec { - run_id: node_provision::RunId(spec.run_id), - logical_node_id: logical_node_id.clone(), - group_id: node_provision::NodeGroupId("local-e2e".to_owned()), - role: node_provision::RoleId("stage-worker".to_owned()), - provider: node_provision::ProviderKind::Docker, - shape: node_provision::DesiredNodeShape { - image: spec.image.clone(), - disk_gb: 0, - gpu_name: None, - min_gpu_ram_mb: None, - min_down_mbps: None, - min_up_mbps: None, - min_reliability: None, - require_verified: false, - provider_labels: BTreeMap::from([ - ("mvp.local_e2e".to_owned(), "true".to_owned()), - ("mvp.node_id".to_owned(), spec.node_id.to_string()), - ]), - }, - boot: node_provision::BootSpec { - ssh_user: "local-e2e".to_owned(), - verify_commands: Vec::new(), - start_swactor_command: spec.args.join(" "), - stdout_sources: Vec::new(), - stderr_sources: Vec::new(), - }, - swarm_join: node_provision::SwarmJoinSpec { - orch_swactor_addr: "local-e2e".to_owned(), - join_token_ref: "local-e2e".to_owned(), - expected_logical_node_id: logical_node_id, - }, - } -} - -fn provision_local_docker_node( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - dashboard: &mut Option<&mut MvpDashboard>, - stats: &mut ProvisionStats, - spec: NodeProvisionSpec, -) -> Result { - let expected_run_id = spec.run_id; - let expected_node_id = spec.node_id; - let expected_stage_index = spec.stage_index.unwrap_or_default(); - let logical_spec = logical_node_spec_from_local(&spec); - let (events_tx, events_rx) = mpsc::channel(); - let cli = LocalE2eDockerCli::new(spec, events_tx); - let provider = docker_provision::DockerProvider::new(cli); - let mut provisioner = docker_provision::DockerNodeProvisioner::new( - provider, - LocalE2eBootstrapFactory, - LocalE2eBootstrapDatastream, - ); - provisioner - .start(logical_spec) - .map_err(|e| format!("start Docker node {expected_node_id}: {e:?}"))?; - let provider_process_id = provisioner.provider().cli().provider_process_id(); - - while !STOP_REQUESTED.load(Ordering::SeqCst) { - pump_network(driver, stack); - drain_dashboard(dashboard); - publish_runtime_snapshot_throttled(dashboard, stack); - while let Ok(event) = events_rx.try_recv() { - match event { - LocalDockerNodeEvent::Stdout(line) => { - handle_local_node_stdout( - expected_run_id, - expected_node_id, - &line, - dashboard, - stats, - ); - if let Some((endpoint, node_actor, stage_index)) = - ready_node_from_stdout(expected_node_id, expected_stage_index, &line)? - { - provisioner - .observe_swactor_join(node_provision::SwactorId(format!( - "local-e2e-node-{expected_node_id}" - ))) - .map_err(|e| { - format!("complete Docker node {expected_node_id} handoff: {e:?}") - })?; - stats.node_live_count += 1; - return Ok(ProvisionedDockerNode { - node_id: expected_node_id, - stage_index, - endpoint, - node_actor, - provider_process_id, - provisioner, - events: events_rx, - cleaned: false, - }); - } - } - LocalDockerNodeEvent::Stderr(line) => { - handle_local_node_stderr( - expected_run_id, - expected_node_id, - &line, - dashboard, - stats, - ); - } - LocalDockerNodeEvent::Exited(status) => { - return Err(format!( - "node {expected_node_id} process exited before ready: {status:?}" - )); - } - } - } - thread::sleep(Duration::from_millis(10)); - } - Err(format!("interrupted provisioning node {expected_node_id}")) -} - -fn ready_node_from_stdout( - expected_node_id: u64, - default_stage_index: u32, - line: &str, -) -> Result, String> { - let Ok(line) = serde_json::from_str::(line) else { - return Ok(None); - }; - if line.kind != "ready" { - return Ok(None); - } - if line.logical_node_id != Some(expected_node_id) { - return Ok(None); - } - let endpoint = line - .endpoint - .ok_or_else(|| format!("ready line for node {expected_node_id} missing endpoint"))?; - let node_actor = line - .node_actor - .ok_or_else(|| format!("ready line for node {expected_node_id} missing node actor"))?; - Ok(Some(( - endpoint, - node_actor, - line.stage_index.unwrap_or(default_stage_index), - ))) -} - -fn drain_provisioned_node_events( - nodes: &mut [&mut ProvisionedDockerNode], - run_id: u64, - dashboard: &mut Option<&mut MvpDashboard>, - stats: &mut ProvisionStats, -) -> Result { - let mut stage_ready_count = 0usize; - for node in nodes.iter_mut() { - while let Ok(event) = node.events.try_recv() { - match event { - LocalDockerNodeEvent::Stdout(line) => { - stage_ready_count += - handle_local_node_stdout(run_id, node.node_id, &line, dashboard, stats); - } - LocalDockerNodeEvent::Stderr(line) => { - handle_local_node_stderr(run_id, node.node_id, &line, dashboard, stats); - } - LocalDockerNodeEvent::Exited(status) => { - return Err(format!( - "node {} process exited before stop: {status:?}", - node.node_id - )); - } - } - } - } - Ok(stage_ready_count) -} - -fn handle_local_node_stdout( - run_id: u64, - node_id: u64, - line: &str, - dashboard: &mut Option<&mut MvpDashboard>, - stats: &mut ProvisionStats, -) -> usize { - stats.stdout_line_count += 1; - record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stdout, line); - record_stage_ready_from_stdout(run_id, ProvisionLogStream::Stdout, line, dashboard) -} - -fn handle_local_node_stderr( - run_id: u64, - node_id: u64, - line: &str, - dashboard: &mut Option<&mut MvpDashboard>, - stats: &mut ProvisionStats, -) { - stats.stderr_line_count += 1; - record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stderr, line); -} - -fn record_stage_ready_from_stdout( - run_id: u64, - stream: ProvisionLogStream, - line: &str, - dashboard: &mut Option<&mut MvpDashboard>, -) -> usize { - if stream != ProvisionLogStream::Stdout { - return 0; - } - let Ok(line) = serde_json::from_str::(line) else { - return 0; - }; - if line.kind != "node_lifecycle" { - return 0; - } - let Some(stage_index) = line.stage_index else { - return 0; - }; - if line - .event - .as_deref() - .is_some_and(|event| event.contains("StageReady")) - { - record_dashboard_event( - dashboard, - stage_event(run_id, stage_index, obs::EventKind::StageReady), - ); - 1 - } else { - 0 - } -} - -fn stop_provisioned_nodes( - nodes: &mut [&mut ProvisionedDockerNode], - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - dashboard: &mut Option<&mut MvpDashboard>, -) -> Result<(), String> { - let mut first_error = None; - for node in nodes.iter_mut().rev() { - if first_error.is_none() { - if let Err(error) = node.stop() { - first_error = Some(error); - } - } else { - let _ = node.stop(); - } - } - pump_network(driver, stack); - drain_dashboard(dashboard); - publish_runtime_snapshot(dashboard, stack); - if let Some(error) = first_error { - Err(error) - } else { - Ok(()) - } -} - -fn spawn_send_pump( - handle: tokio::runtime::Handle, - endpoint: iroh::Endpoint, - peer: EndpointAddr, - edge_id: u64, -) -> Result { - let (tx, mut rx) = tokio_mpsc::unbounded_channel::>(); - let (ready_tx, ready_rx) = mpsc::channel::>(); - handle.spawn(async move { - let result: Result<(), String> = async { - let conn = endpoint - .connect(peer, EDGE_ALPN) - .await - .map_err(|e| format!("connect edge {edge_id}: {e}"))?; - let mut send = conn - .open_uni() - .await - .map_err(|e| format!("open edge stream {edge_id}: {e}"))?; - send.write_all(&driver_model::encode_edge_preamble(driver_model::EdgeId( - edge_id, - ))) - .await - .map_err(|e| format!("write edge preamble {edge_id}: {e}"))?; - let _ = ready_tx.send(Ok(())); - while let Some(record) = rx.recv().await { - send.write_all(&record) - .await - .map_err(|e| format!("write edge record {edge_id}: {e}"))?; - } - send.finish() - .map_err(|e| format!("finish edge stream {edge_id}: {e}"))?; - Ok(()) - } - .await; - if let Err(error) = result { - let _ = ready_tx.send(Err(error)); - } - }); - ready_rx - .recv() - .map_err(|e| format!("edge {edge_id} sender startup channel closed: {e}"))??; - Ok(SendPumpHandle { tx }) -} - -fn spawn_recv_pump( - handle: tokio::runtime::Handle, - conn: iroh::endpoint::Connection, - tx: Sender, - stream_id: u64, -) { - handle.spawn(async move { - let mut next_uni_stream_id = stream_id << 32; - while let Ok(mut recv) = conn.accept_uni().await { - next_uni_stream_id += 1; - let current_stream_id = next_uni_stream_id; - let mut preamble = [0u8; 8]; - if recv.read_exact(&mut preamble).await.is_err() { - continue; - } - let edge_id = u64::from_le_bytes(preamble); - if tx - .send(DriverIngressEvent::StreamArrived { - edge_id, - stream_id: current_stream_id, - }) - .is_err() - { - break; - } - let mut chunk = vec![0u8; 4096]; - loop { - match recv.read(&mut chunk).await { - Ok(Some(0)) | Ok(None) => break, - Ok(Some(n)) => { - if tx - .send(DriverIngressEvent::BytesRead { - edge_id, - stream_id: current_stream_id, - bytes: chunk[..n].to_vec(), - }) - .is_err() - { - break; - } - } - Err(_) => break, - } - } - } - }); -} - -fn object_record_spec() -> ingress::ObjectSpec { - ingress::ObjectSpec { - max_extent: OBJECT_MAX_EXTENT, - alignment: OBJECT_ALIGNMENT, - layout: ingress::ObjectLayout::Token, - } -} - -fn object_record(object_id: u64, sequence: u64, words: &[u32]) -> Vec { - let mut payload = Vec::with_capacity(words.len() * 4); - for word in words { - payload.extend_from_slice(&(*word as i32).to_le_bytes()); - } - ingress::ObjectRecordBuilder::new(object_record_spec()) - .object_id(ingress::ObjectId(object_id)) - .sequence(sequence) - .payload(payload) - .encode() -} - -fn take_complete_ingress_record(buffer: &mut Vec) -> Result>, String> { - let record = match ingress::read_object_record(buffer, object_record_spec(), false) - .map_err(|reason| format!("invalid object record: {reason:?}"))? - { - ingress::ObjectRecordRead::Incomplete => return Ok(None), - ingress::ObjectRecordRead::Complete(record) => record, - }; - Ok(Some(buffer.drain(..record.total_len).collect())) -} - -fn decode_ingress_record(record: &[u8]) -> Result { - match ingress::read_object_record(record, object_record_spec(), true) - .map_err(|reason| format!("invalid object record: {reason:?}"))? - { - ingress::ObjectRecordRead::Incomplete => Err("object record incomplete".to_owned()), - ingress::ObjectRecordRead::Complete(record) => Ok(record), - } -} - -fn object_record_words( - metadata: &ingress::ObjectRecord, - record: &[u8], -) -> Result, String> { - let payload = metadata - .payload(record) - .ok_or_else(|| format!("object {} payload truncated", metadata.object_id.0))?; - if payload.len() % 4 != 0 { - return Err(format!( - "object {} payload length {} is not word-aligned", - metadata.object_id.0, - payload.len() - )); - } - Ok(payload - .chunks_exact(4) - .map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap())) - .collect()) -} - -fn spawn_shutdown_listener() -> Receiver<()> { - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - let stdin = std::io::stdin(); - for line in stdin.lock().lines().map_while(Result::ok) { - if line.trim() == "shutdown" { - let _ = tx.send(()); - break; - } - } - }); - rx -} - -fn relay_url_from_env() -> Option { - std::env::var(MVP_IROH_RELAY_URL_ENV) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -fn endpoint_relay_url(endpoint: &EndpointAddr) -> Option { - endpoint.relay_urls().next().map(|url| url.to_string()) -} - -fn docker_network_for_node(node_id: u64) -> Option { - let env_name = match node_id { - NODE0_LOGICAL_ID => LOCAL_E2E_DOCKER_NETWORK_NODE0_ENV, - NODE1_LOGICAL_ID => LOCAL_E2E_DOCKER_NETWORK_NODE1_ENV, - _ => return None, - }; - std::env::var(env_name) - .ok() - .map(|value| value.trim().to_owned()) - .filter(|value| !value.is_empty()) -} - -fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> { - let index = args - .iter() - .position(|arg| arg == name) - .ok_or_else(|| format!("missing {name}"))?; - args.get(index + 1) - .map(String::as_str) - .ok_or_else(|| format!("missing value for {name}")) -} - -fn parse_optional_arg<'a>(args: &'a [String], name: &str) -> Option<&'a str> { - let index = args.iter().position(|arg| arg == name)?; - args.get(index + 1).map(String::as_str) -} - -fn docker_image() -> String { - std::env::var("MVP_LOCAL_E2E_CLUSTER_IMAGE").unwrap_or_else(|_| DEFAULT_DOCKER_IMAGE.to_owned()) -} - -fn tinygrad_worker_path() -> PathBuf { - std::env::var("MVP_TINYGRAD_WORKER") - .map(PathBuf::from) - .unwrap_or_else(|_| { - PathBuf::from("crates/mvp-system/tests/local_e2e_cluster/tinygrad_cpu_worker.py") - }) -} - -fn tokenize_prompt(prompt: &str) -> Vec { - let mut tokens = prompt - .split_whitespace() - .map(|word| match word.to_ascii_lowercase().as_str() { - "ping" => 2, - "hello" => 3, - "local" => 4, - "cluster" => 5, - "pong" => 6, - "world" => 7, - "ok" => 8, - other => 20 + (other.bytes().fold(0u32, |acc, byte| acc + u32::from(byte)) % 100), - }) - .collect::>(); - if tokens.is_empty() { - tokens.push(3); - } - tokens -} - -fn detokenize_response(tokens: &[u32]) -> String { - tokens - .iter() - .map(|token| match *token { - 1 => "".to_owned(), - 2 => "ping".to_owned(), - 3 => "hello".to_owned(), - 4 => "local".to_owned(), - 5 => "cluster".to_owned(), - 6 => "pong".to_owned(), - 7 => "world".to_owned(), - 8 => "ok".to_owned(), - other => format!(""), - }) - .collect::>() - .join(" ") -} diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index e240aec..a91dfdf 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -4,3 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] +serde_json = "1" + +[target.'cfg(target_os = "linux")'.dependencies] +libc = "0.2" diff --git a/xtask/src/main.rs b/xtask/src/main.rs index d53c0b5..5311bcc 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -1,11 +1,38 @@ -use std::process::{Command, ExitCode}; -use std::time::Instant; +use std::fs; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, ExitCode, ExitStatus, Stdio}; +use std::thread; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +#[cfg(target_os = "linux")] +use std::os::unix::process::CommandExt; + +use serde_json::Value; struct TestStep { label: &'static str, args: &'static [&'static str], } +const MVP_CHAT_CHECK_TIMEOUT_SECS: u64 = 900; +const MVP_CHAT_CHECK_POLL_MS: u64 = 100; +const MVP_CHAT_CHECK_TERM_GRACE_MS: u64 = 2_000; +const MVP_CHAT_CHECK_PROMPTS: &[u8] = b"ping\nsecond prompt\n"; + +struct MvpChatCheckPaths { + root: PathBuf, + dump_log: PathBuf, +} + +struct MvpChatCheckOutput { + status: ExitStatus, + stdout: String, + stderr: String, + timed_out: bool, + stdin_error: Option, +} + const BASIC_TESTS: &[TestStep] = &[ TestStep { label: "root crate", @@ -62,9 +89,9 @@ USAGE: cargo xtask COMMANDS: mvp-chat [--process|--docker|--vastai] [--pipeline-stages n] [--cached-model] [-- args...] Run the human chat wrapper against the real orchestrator/worker bins. - 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." + mvp-chat-check Run real cargo mvp-chat acceptance check. + test Run the basic non-binding test barrier: root crate plus each + non-binding repository package with `cargo test -p`." ); } @@ -84,6 +111,10 @@ fn run_step(step: &TestStep) -> bool { fn run_tests() -> ExitCode { let start = Instant::now(); + let check = run_mvp_chat_check(); + if check != ExitCode::SUCCESS { + return check; + } for (index, step) in BASIC_TESTS.iter().enumerate() { if !run_step(step) { @@ -106,7 +137,12 @@ fn run_tests() -> ExitCode { fn run_mvp_chat(args: Vec) -> ExitCode { let mut command = Command::new(cargo_bin()); command.args(["run", "--package", "mvp-system", "--bin", "mvp-chat", "--"]); - command.args(args); + let forwarded = if args.first().is_some_and(|arg| arg == "--") { + args[1..].to_vec() + } else { + args + }; + command.args(forwarded); match command.status() { Ok(status) if status.success() => ExitCode::SUCCESS, @@ -123,10 +159,550 @@ fn run_mvp_chat(args: Vec) -> ExitCode { } } +fn workspace_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("xtask manifest dir has a parent") + .to_path_buf() +} + +fn unique_temp_dir(prefix: &str) -> PathBuf { + let pid = std::process::id(); + for attempt in 0..100 { + let timestamp_nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + let root = + std::env::temp_dir().join(format!("{prefix}-{pid}-{timestamp_nanos}-{attempt}")); + match fs::create_dir(&root) { + Ok(()) => return root, + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => continue, + Err(error) => panic!("mvp-chat-check: create temp dir {}: {error}", root.display()), + } + } + panic!("mvp-chat-check: could not allocate unique temp dir for prefix {prefix}"); +} + +fn write_mvp_chat_check_paths(root: &Path) -> Result { + if !root.is_dir() { + return Err(format!( + "mvp-chat-check: temp root {} is not a directory", + root.display() + )); + } + let dump_log = root.join("mvp-chat.ndjson"); + if dump_log.exists() { + return Err(format!( + "mvp-chat-check: dump log path already exists: {}", + dump_log.display() + )); + } + Ok(MvpChatCheckPaths { + root: root.to_path_buf(), + dump_log, + }) +} + +fn run_mvp_chat_check() -> ExitCode { + let workspace = workspace_root(); + let temp_root = unique_temp_dir("mvp-chat-check"); + let paths = match write_mvp_chat_check_paths(&temp_root) { + Ok(paths) => paths, + Err(error) => { + eprintln!("{error}"); + eprintln!( + "mvp-chat-check: temp directory kept at {}", + temp_root.display() + ); + return ExitCode::from(1); + } + }; + + let output = match run_mvp_chat_check_process(&workspace, &paths) { + Ok(output) => output, + Err(error) => return fail_mvp_chat_check(&error, &paths, "", "", None), + }; + + if let Some(error) = &output.stdin_error { + return fail_mvp_chat_check( + error, + &paths, + &output.stdout, + &output.stderr, + Some(&output.status), + ); + } + if output.timed_out { + let reason = format!("timeout after {MVP_CHAT_CHECK_TIMEOUT_SECS} seconds"); + return fail_mvp_chat_check( + &reason, + &paths, + &output.stdout, + &output.stderr, + Some(&output.status), + ); + } + if !output.status.success() { + return fail_mvp_chat_check( + "child exited nonzero", + &paths, + &output.stdout, + &output.stderr, + Some(&output.status), + ); + } + + let responses = match assert_stdout_contains_two_prompt_cycles(&output.stdout) { + Ok(responses) => responses, + Err(error) => { + return fail_mvp_chat_check( + &error, + &paths, + &output.stdout, + &output.stderr, + Some(&output.status), + ); + } + }; + if let Err(error) = assert_dump_log_facts(&paths.dump_log) { + return fail_mvp_chat_check( + &error, + &paths, + &output.stdout, + &output.stderr, + Some(&output.status), + ); + } + + if let Err(error) = fs::remove_dir_all(&paths.root) { + eprintln!( + "mvp-chat-check: remove temp directory {}: {error}", + paths.root.display() + ); + return ExitCode::from(1); + } + + println!("mvp-chat-check: ok"); + for (index, response) in responses.iter().enumerate() { + println!("mvp-chat-check: response {}: {}", index + 1, response); + } + ExitCode::SUCCESS +} + +fn run_mvp_chat_check_process( + workspace: &Path, + paths: &MvpChatCheckPaths, +) -> Result { + let mut command = Command::new(cargo_bin()); + command + .current_dir(workspace) + .args(["mvp-chat", "--", "--cached-model"]) + .arg(format!("--dump-logs={}", paths.dump_log.display())) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()); + + #[cfg(target_os = "linux")] + unsafe { + command.pre_exec(|| { + let result = libc::setpgid(0, 0); + if result == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } + }); + } + + let mut child = command + .spawn() + .map_err(|e| format!("mvp-chat-check: spawn cargo mvp-chat: {e}"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "mvp-chat-check: child stdout was not piped".to_owned())?; + let stderr = child + .stderr + .take() + .ok_or_else(|| "mvp-chat-check: child stderr was not piped".to_owned())?; + let stdout_reader = thread::spawn(move || read_pipe_to_string(stdout, "stdout")); + let stderr_reader = thread::spawn(move || read_pipe_to_string(stderr, "stderr")); + + let stdin_error = match child.stdin.take() { + Some(mut stdin) => { + let result = stdin.write_all(MVP_CHAT_CHECK_PROMPTS); + drop(stdin); + result + .err() + .map(|error| format!("mvp-chat-check: write child stdin: {error}")) + } + None => Some("mvp-chat-check: child stdin was not piped".to_owned()), + }; + + let (status, timed_out) = if stdin_error.is_some() { + (terminate_mvp_chat_child(&mut child)?, false) + } else { + wait_mvp_chat_check_child(&mut child)? + }; + + let stdout = join_reader(stdout_reader, "stdout")?; + let stderr = join_reader(stderr_reader, "stderr")?; + Ok(MvpChatCheckOutput { + status, + stdout, + stderr, + timed_out, + stdin_error, + }) +} + +fn wait_mvp_chat_check_child(child: &mut Child) -> Result<(ExitStatus, bool), String> { + let timeout = Duration::from_secs(MVP_CHAT_CHECK_TIMEOUT_SECS); + let poll = Duration::from_millis(MVP_CHAT_CHECK_POLL_MS); + let deadline = Instant::now() + timeout; + loop { + match child.try_wait() { + Ok(Some(status)) => return Ok((status, false)), + Ok(None) if Instant::now() >= deadline => { + return terminate_mvp_chat_child(child).map(|status| (status, true)); + } + Ok(None) => thread::sleep(poll), + Err(error) => return Err(format!("mvp-chat-check: poll child status: {error}")), + } + } +} + +fn terminate_mvp_chat_child(child: &mut Child) -> Result { + #[cfg(target_os = "linux")] + { + signal_mvp_chat_process_group(child, libc::SIGTERM); + let grace_polls = MVP_CHAT_CHECK_TERM_GRACE_MS / MVP_CHAT_CHECK_POLL_MS; + for _ in 0..grace_polls { + match child.try_wait() { + Ok(Some(status)) => return Ok(status), + Ok(None) => thread::sleep(Duration::from_millis(MVP_CHAT_CHECK_POLL_MS)), + Err(error) => { + return Err(format!( + "mvp-chat-check: poll child after SIGTERM: {error}" + )); + } + } + } + signal_mvp_chat_process_group(child, libc::SIGKILL); + } + + #[cfg(not(target_os = "linux"))] + { + let _ = child.kill(); + } + + child + .wait() + .map_err(|e| format!("mvp-chat-check: wait for terminated child: {e}")) +} + +#[cfg(target_os = "linux")] +fn signal_mvp_chat_process_group(child: &Child, signal: libc::c_int) { + let process_group = -(child.id() as libc::pid_t); + let _ = unsafe { libc::kill(process_group, signal) }; +} + +fn read_pipe_to_string(mut reader: R, label: &'static str) -> Result { + let mut text = String::new(); + reader + .read_to_string(&mut text) + .map_err(|e| format!("mvp-chat-check: read child {label}: {e}"))?; + Ok(text) +} + +fn join_reader( + handle: thread::JoinHandle>, + label: &str, +) -> Result { + handle + .join() + .map_err(|_| format!("mvp-chat-check: child {label} reader panicked"))? +} + +fn fail_mvp_chat_check( + reason: &str, + paths: &MvpChatCheckPaths, + stdout: &str, + stderr: &str, + status: Option<&ExitStatus>, +) -> ExitCode { + eprintln!("mvp-chat-check: failed: {reason}"); + if let Some(status) = status { + eprintln!("mvp-chat-check: child exit status: {status}"); + } + eprintln!( + "mvp-chat-check: temp directory kept at {}", + paths.root.display() + ); + eprintln!("--- captured stdout ---"); + if stdout.is_empty() { + eprintln!(""); + } else { + eprint!("{stdout}"); + if !stdout.ends_with('\n') { + eprintln!(); + } + } + eprintln!("--- captured stderr ---"); + if stderr.is_empty() { + eprintln!(""); + } else { + eprint!("{stderr}"); + if !stderr.ends_with('\n') { + eprintln!(); + } + } + ExitCode::from(1) +} + +fn assert_stdout_contains_two_prompt_cycles(stdout: &str) -> Result, String> { + let decoding_count = stdout.matches("decoding...").count(); + if decoding_count < 2 { + return Err(format!( + "mvp-chat-check: expected at least two decoding... markers, found {decoding_count}" + )); + } + let response_count = stdout.matches("Response: ").count(); + if response_count < 2 { + return Err(format!( + "mvp-chat-check: expected at least two Response: prefixes, found {response_count}" + )); + } + + let mut cursor = 0; + let mut responses = Vec::with_capacity(2); + for cycle in 1..=2 { + let prompt_at = find_stdout_marker(stdout, "prompt:>", cursor, cycle, "prompt")?; + let decoding_at = find_stdout_marker( + stdout, + "decoding...", + prompt_at + "prompt:>".len(), + cycle, + "decoding", + )?; + let response_at = find_stdout_marker( + stdout, + "Response: ", + decoding_at + "decoding...".len(), + cycle, + "response", + )?; + let response_start = response_at + "Response: ".len(); + let response_end = stdout[response_start..] + .find('\n') + .map_or(stdout.len(), |offset| response_start + offset); + let response = &stdout[response_start..response_end]; + if !response.chars().any(|ch| !ch.is_whitespace()) { + return Err(format!( + "mvp-chat-check: empty Response text for prompt cycle {cycle}" + )); + } + responses.push(response.to_owned()); + cursor = response_end; + } + Ok(responses) +} + +fn find_stdout_marker( + stdout: &str, + marker: &str, + start: usize, + cycle: usize, + label: &str, +) -> Result { + stdout[start..] + .find(marker) + .map(|offset| start + offset) + .ok_or_else(|| { + format!("mvp-chat-check: missing {label} marker for prompt cycle {cycle}") + }) +} + +#[derive(Default)] +struct DumpLogFacts { + chat_config_ready: bool, + prepare_runtime_ready: bool, + prompt_rpc_ready: bool, + orch_iroh_driver_ready: bool, + node_iroh_driver_ready: bool, + node_worker_initialize_ready: bool, + orch_weights_loaded_ready: bool, + response_text_1: bool, + request_completed_1: bool, + response_text_2: bool, + request_completed_2: bool, + shutdown_requested: bool, + orchestrator_stopped: bool, +} + +fn assert_dump_log_facts(path: &Path) -> Result<(), String> { + let content = fs::read_to_string(path) + .map_err(|e| format!("mvp-chat-check: read dump log {}: {e}", path.display()))?; + if content.lines().next().is_none() { + return Err(format!( + "mvp-chat-check: dump log {} is empty", + path.display() + )); + } + + let mut facts = DumpLogFacts::default(); + for (line_index, line) in content.lines().enumerate() { + if line.trim().is_empty() { + continue; + } + let outer: Value = serde_json::from_str(line).map_err(|e| { + format!( + "mvp-chat-check: parse dump log {} line {}: {e}", + path.display(), + line_index + 1 + ) + })?; + let channel = outer + .get("channel") + .and_then(Value::as_str) + .ok_or_else(|| { + format!( + "mvp-chat-check: dump log line {} missing channel", + line_index + 1 + ) + })?; + let payload = outer.get("payload").ok_or_else(|| { + format!( + "mvp-chat-check: dump log line {} missing payload", + line_index + 1 + ) + })?; + if payload.get("encoding").and_then(Value::as_str) != Some("utf8") { + continue; + } + let inner_text = payload.get("value").and_then(Value::as_str).ok_or_else(|| { + format!( + "mvp-chat-check: dump log line {} missing utf8 payload value", + line_index + 1 + ) + })?; + let event: Value = serde_json::from_str(inner_text).map_err(|e| { + format!( + "mvp-chat-check: parse inner event on dump log line {}: {e}", + line_index + 1 + ) + })?; + record_dump_log_event(channel, &event, &mut facts)?; + } + + require_dump_log_fact(facts.chat_config_ready, "config ready")?; + require_dump_log_fact(facts.prepare_runtime_ready, "prepare_runtime ready")?; + require_dump_log_fact(facts.prompt_rpc_ready, "prompt_rpc ready")?; + require_dump_log_fact(facts.orch_iroh_driver_ready, "OrchBootstrap iroh_driver ready")?; + require_dump_log_fact(facts.node_iroh_driver_ready, "NodeEvent iroh_driver ready")?; + require_dump_log_fact( + facts.node_worker_initialize_ready, + "NodeEvent worker_initialize ready", + )?; + require_dump_log_fact(facts.orch_weights_loaded_ready, "OrchBootstrap weights_loaded ready")?; + require_dump_log_fact(facts.response_text_1, "response_text request_id=1")?; + require_dump_log_fact(facts.request_completed_1, "request_completed request_id=1")?; + require_dump_log_fact(facts.response_text_2, "response_text request_id=2")?; + require_dump_log_fact(facts.request_completed_2, "request_completed request_id=2")?; + require_dump_log_fact(facts.shutdown_requested, "shutdown requested")?; + require_dump_log_fact(facts.orchestrator_stopped, "orchestrator_process stopped") +} + +fn record_dump_log_event( + channel: &str, + event: &Value, + facts: &mut DumpLogFacts, +) -> Result<(), String> { + let event_type = event.get("type").and_then(Value::as_str); + let phase = event.get("phase").and_then(Value::as_str); + let status = event.get("status").and_then(Value::as_str); + if status == Some("failed") { + return Err(format!( + "mvp-chat-check: failed event channel={channel} type={} phase={} detail={}", + event_type.unwrap_or(""), + phase.unwrap_or(""), + event.get("detail").unwrap_or(&Value::Null) + )); + } + + match (channel, event_type, phase, status) { + ("mvp.chat.lifecycle", Some("ChatProgress"), Some("config"), Some("ready")) => { + facts.chat_config_ready = true; + } + ("mvp.chat.runtime", Some("ChatProgress"), Some("prepare_runtime"), Some("ready")) => { + facts.prepare_runtime_ready = true; + } + ("mvp.chat.runtime", Some("ChatProgress"), Some("prompt_rpc"), Some("ready")) => { + facts.prompt_rpc_ready = true; + } + (_, Some("OrchBootstrap"), Some("iroh_driver"), Some("ready")) => { + facts.orch_iroh_driver_ready = true; + } + (_, Some("NodeEvent"), Some("iroh_driver"), Some("ready")) => { + facts.node_iroh_driver_ready = true; + } + (_, Some("NodeEvent"), Some("worker_initialize"), Some("ready")) => { + facts.node_worker_initialize_ready = true; + } + (_, Some("OrchBootstrap"), Some("weights_loaded"), Some("ready")) => { + facts.orch_weights_loaded_ready = true; + } + ("mvp.chat.prompt", Some("ChatProgress"), Some("response_text"), Some("observed")) => { + match dump_log_request_id(event) { + Some(1) => facts.response_text_1 = true, + Some(2) => facts.response_text_2 = true, + _ => {} + } + } + ("mvp.chat.prompt", Some("ChatProgress"), Some("request_completed"), Some("ready")) => { + match dump_log_request_id(event) { + Some(1) => facts.request_completed_1 = true, + Some(2) => facts.request_completed_2 = true, + _ => {} + } + } + ("mvp.chat.lifecycle", Some("ChatProgress"), Some("shutdown"), Some("requested")) => { + facts.shutdown_requested = true; + } + ( + "mvp.chat.component", + Some("ChatProgress"), + Some("orchestrator_process"), + Some("stopped"), + ) => { + facts.orchestrator_stopped = true; + } + _ => {} + } + Ok(()) +} + +fn dump_log_request_id(event: &Value) -> Option { + event + .get("detail") + .and_then(|detail| detail.get("request_id")) + .and_then(Value::as_u64) +} + +fn require_dump_log_fact(found: bool, fact: &str) -> Result<(), String> { + if found { + Ok(()) + } else { + Err(format!("mvp-chat-check: missing {fact}")) + } +} + fn main() -> ExitCode { let mut args = std::env::args().skip(1); match args.next().as_deref() { Some("test") if args.next().is_none() => run_tests(), + Some("mvp-chat-check") if args.next().is_none() => run_mvp_chat_check(), Some("mvp-chat") => run_mvp_chat(args.collect()), Some("help" | "--help" | "-h") | None => { print_usage();