From 4811564d0f61debdaa497fbb08424127099d37ec Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Thu, 9 Jul 2026 12:53:53 +0400 Subject: [PATCH] feat(mvp-system): actor admin control, iroh relay debug Add node_agent actor and orchestrator/worker_node admin control; iroh-driver relay debugging; ACTOR_CONTROL_AUDIT_IDEAS notes. Signed-off-by: Zachery Aaron Shores-Chmielewski --- Cargo.lock | 1 + crates/distribution/tests/telemetry.rs | 1 + crates/iroh-driver/src/iroh_driver.rs | 47 +- .../iroh-driver/tests/datastream_transport.rs | 1 + crates/iroh-driver/tests/iroh_driver.rs | 29 +- .../mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md | 75 ++ crates/mvp-system/Cargo.toml | 5 +- .../mvp-system/MVP_NODE_PROVISIONING_SPEC.md | 3 +- crates/mvp-system/MVP_SYSTEM_SPEC.md | 1 + crates/mvp-system/src/actors/node_agent.rs | 78 ++ crates/mvp-system/src/actors/orchestrator.rs | 6 + .../mvp-system/src/bin/mvp_one_node_chat.rs | 23 +- crates/mvp-system/src/bin/orchestrator.rs | 318 ++++++- crates/mvp-system/src/bin/worker_node.rs | 828 ++++++++++++++++-- crates/mvp-system/src/distribution_stack.rs | 50 ++ crates/mvp-system/src/vastai_provisioning.rs | 17 +- crates/mvp-system/tests/local_e2e_cluster.rs | 200 ++++- .../tests/local_e2e_cluster/Dockerfile | 10 +- crates/mvp-system/tests/one_node_chat_e2e.rs | 95 +- .../tests/support/local_e2e_cluster.rs | 190 +++- 20 files changed, 1750 insertions(+), 228 deletions(-) create mode 100644 crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md diff --git a/Cargo.lock b/Cargo.lock index c6a6770..f6b55c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2477,6 +2477,7 @@ dependencies = [ "distribution", "iroh", "iroh-driver", + "iroh-relay", "libc", "parking_lot", "serde", diff --git a/crates/distribution/tests/telemetry.rs b/crates/distribution/tests/telemetry.rs index 5b274ca..edec249 100644 --- a/crates/distribution/tests/telemetry.rs +++ b/crates/distribution/tests/telemetry.rs @@ -49,6 +49,7 @@ mod datastream_records { fn distribution_emits_owned_channel_through_datastream_mux() { let stream = StreamId::new(NodeId::new("dist-node"), Lifetime(1)); let mux = Mux::unbounded(stream); + mux.set_frame_timing_enabled(false); let state = DistributionState { registry_size: 9, ..Default::default() diff --git a/crates/iroh-driver/src/iroh_driver.rs b/crates/iroh-driver/src/iroh_driver.rs index 8559902..7926308 100644 --- a/crates/iroh-driver/src/iroh_driver.rs +++ b/crates/iroh-driver/src/iroh_driver.rs @@ -237,8 +237,8 @@ pub struct IrohDriver { peer_relay_urls: HashMap, /// Real-time join status for each peer being joined. join_statuses: Arc>>, - /// Relay URL exposed by the bound endpoint, if any. - relay_url: Option, + /// Relay URL configured or exposed by the bound endpoint, if any. + relay_url: Option, /// Actor-bridge wiring, installed via [`Self::enable_actor_bridge`]. When /// present, the driver decodes inbound frames into actor messages /// ([`Self::pump_inbound_to_actors`]) and writes actor-produced outbound @@ -330,6 +330,11 @@ impl IrohDriver { config: IrohDriverConfig, ) -> Result> { let effective_relay_mode = config.relay_mode; + let configured_relay_url = match &effective_relay_mode { + RelayMode::Custom(relay_map) => relay_map.urls::>().into_iter().next(), + _ => None, + }; + let custom_relay = matches!(&effective_relay_mode, RelayMode::Custom(_)); // A custom relay is operator-controlled (typically `iroh-driver-relay` // on a VPS, serving QUIC Address Discovery with a self-signed cert). @@ -337,8 +342,6 @@ impl IrohDriver { // that, address discovery fails and every connection stays // `conn_type=Relay`, which defeats hole-punching and makes a NAT'd peer // (e.g. a locally-run orchestrator) reachable only over the relay. - let custom_relay = matches!(effective_relay_mode, RelayMode::Custom(_)); - let endpoint = rt.block_on(async { let mut alpns = vec![ALPN.to_vec()]; alpns.extend(config.additional_alpns.iter().cloned()); @@ -362,7 +365,8 @@ impl IrohDriver { .addr() .relay_urls() .next() - .map(|url| url.to_string()); + .cloned() + .or(configured_relay_url); // The driver's signing identity matches the iroh endpoint: both use // ed25519-dalek, so we reconstruct our Keypair from iroh's secret key. @@ -458,16 +462,20 @@ impl IrohDriver { self.keypair.node_id() } - /// The endpoint's full address (public key + direct socket addresses). + /// The endpoint's current relay/direct advertised address. /// - /// Constructs the address from the endpoint's public key and bound - /// sockets. For sockets bound to `0.0.0.0`, emits one address per + /// Starts from Iroh's live endpoint address, which includes the current + /// home relay when one is available, then merges normalized direct socket + /// addresses. For sockets bound to `0.0.0.0`, emits one address per /// discovered LAN IP so that peers on the same network can connect /// directly. IPv6 unspecified is mapped to localhost. pub fn endpoint_addr(&self) -> EndpointAddr { - let key = PublicKey::from_bytes(&self.keypair.node_id().0) - .expect("node_id is a valid public key"); - let mut addr = EndpointAddr::new(key); + let mut addr = self.endpoint.addr(); + if addr.relay_urls().next().is_none() { + if let Some(relay) = self.relay_url.clone() { + addr = addr.with_relay_url(relay); + } + } for sa in self.direct_addresses() { addr = addr.with_ip_addr(sa); } @@ -622,7 +630,7 @@ impl IrohDriver { .peer_relay_urls .get(&seed_node_id) .cloned() - .or_else(|| self.endpoint.addr().relay_urls().next().cloned()) + .or_else(|| self.home_relay_url()) { seed_addr.clone().with_relay_url(relay) } else { @@ -983,7 +991,7 @@ impl IrohDriver { .and_then(|s| s.parse::().ok()) { Some(r) - } else if let Some(r) = self.endpoint.addr().relay_urls().next().cloned() { + } else if let Some(r) = self.home_relay_url() { Some(r) } else { None @@ -1051,14 +1059,19 @@ impl IrohDriver { } } - /// Relay URL exposed by the bound endpoint, if any. + /// Relay URL configured or exposed by the bound endpoint, if any. pub fn relay_url(&self) -> Option<&str> { - self.relay_url.as_deref() + self.relay_url.as_ref().map(|url| url.as_str()) } - /// The endpoint's home relay URL (from RelayMode::Custom), if connected. + /// The endpoint's live or configured home relay URL, if any. pub fn home_relay_url(&self) -> Option { - self.endpoint.addr().relay_urls().next().cloned() + self.endpoint + .addr() + .relay_urls() + .next() + .cloned() + .or_else(|| self.relay_url.clone()) } /// Async teardown for the unified driver loop, which runs on a tokio worker diff --git a/crates/iroh-driver/tests/datastream_transport.rs b/crates/iroh-driver/tests/datastream_transport.rs index 2e529f2..5148081 100644 --- a/crates/iroh-driver/tests/datastream_transport.rs +++ b/crates/iroh-driver/tests/datastream_transport.rs @@ -28,6 +28,7 @@ async fn iroh_datastream_alpn_carries_endpoint_subscription() { let endpoint = DatastreamEndpoint::with_capacity(stream.clone(), 8, 8); let subscription = endpoint.subscribe_all("iroh"); let producer = endpoint.producer(); + producer.set_frame_timing_enabled(false); producer.submit_text("runtime.log", "alpha"); producer.submit_text("runtime.log", "beta"); diff --git a/crates/iroh-driver/tests/iroh_driver.rs b/crates/iroh-driver/tests/iroh_driver.rs index e4b13a7..671aa2e 100644 --- a/crates/iroh-driver/tests/iroh_driver.rs +++ b/crates/iroh-driver/tests/iroh_driver.rs @@ -11,7 +11,7 @@ mod common; use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::time::{Duration, Instant}; use common::iroh::*; use distribution::peer_auth::PeerAllowList; @@ -38,6 +38,33 @@ fn iroh_driver_reports_listen_addr_and_no_routes() { driver.shutdown(); } +#[test] +fn endpoint_addr_includes_home_relay() { + let (relay_url, _relay_guard) = spawn_test_relay(); + let expected_relay_url = relay_url.to_string(); + let mut node = make_driver_with_relay(relay_url.clone()); + let start = Instant::now(); + let (relay_advertised, observed_relay_url) = loop { + let endpoint = node.endpoint_addr(); + let current_relay_url = endpoint.relay_urls().next().map(|url| url.to_string()); + if current_relay_url.as_deref() == Some(expected_relay_url.as_str()) { + break (true, current_relay_url); + } + if start.elapsed() >= Duration::from_secs(5) { + break (false, current_relay_url); + } + pump_one(&mut node); + std::thread::sleep(Duration::from_millis(10)); + }; + + node.shutdown(); + + assert!( + relay_advertised, + "advertised endpoint did not include home relay {expected_relay_url} within timeout; last relay URL: {observed_relay_url:?}" + ); +} + // ─── Join integration tests ───────────────────────────────────────────── #[test] diff --git a/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md b/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md new file mode 100644 index 0000000..ab5ac11 --- /dev/null +++ b/crates/mvp-system/ACTOR_CONTROL_AUDIT_IDEAS.md @@ -0,0 +1,75 @@ +# 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 05234b2..7f87343 100644 --- a/crates/mvp-system/Cargo.toml +++ b/crates/mvp-system/Cargo.toml @@ -19,7 +19,7 @@ swactor-transport = { path = "../transport" } distribution = { path = "../distribution" } iroh-driver = { path = "../iroh-driver" } iroh = "0.98" -tokio = { version = "1", features = ["rt-multi-thread", "macros", "process", "io-util", "sync", "time"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "process", "io-util", "sync", "time", "net"] } swactor-vastai = { path = "../../tools/vastai" } parking_lot = "0.12" blake3 = "1" @@ -28,6 +28,9 @@ toml = "0.8" [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2" +[dev-dependencies] +iroh-relay = { version = "0.98", features = ["server", "test-utils"] } + [[bin]] name = "mvp-worker-node" path = "src/bin/worker_node.rs" diff --git a/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md b/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md index 003d311..9a15972 100644 --- a/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md +++ b/crates/mvp-system/MVP_NODE_PROVISIONING_SPEC.md @@ -1,6 +1,7 @@ # MVP Node Provisioning Specification + ***STALE! FOR HISTORICAL REFERENCE ONLY*** -**Status:** draft node-provisioning specification. +**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 diff --git a/crates/mvp-system/MVP_SYSTEM_SPEC.md b/crates/mvp-system/MVP_SYSTEM_SPEC.md index caa7bac..a0965ae 100644 --- a/crates/mvp-system/MVP_SYSTEM_SPEC.md +++ b/crates/mvp-system/MVP_SYSTEM_SPEC.md @@ -1,4 +1,5 @@ # MVP System Specification + ***STALE! FOR HISTORICAL REFERENCE ONLY*** **Status:** draft consolidated system specification. diff --git a/crates/mvp-system/src/actors/node_agent.rs b/crates/mvp-system/src/actors/node_agent.rs index f5c212f..93677b3 100644 --- a/crates/mvp-system/src/actors/node_agent.rs +++ b/crates/mvp-system/src/actors/node_agent.rs @@ -58,6 +58,13 @@ pub enum NodeAgentMsg { stage_index: u32, endpoint: EndpointAddr, node_actor: ActorAddress, + readiness_id: u64, + }, + RuntimeReadyAck { + run_id: u64, + node_id: u64, + stage_index: u32, + readiness_id: u64, }, MarkWeightsReady, MarkInboundEdgeReady { @@ -184,6 +191,12 @@ pub enum NodeAgentReport { max_tokens: u32, reply_to: ActorAddress, }, + RuntimeReadyAck { + run_id: u64, + node_id: u64, + stage_index: u32, + readiness_id: u64, + }, Snapshot { commands: Vec, events: Vec, @@ -234,6 +247,7 @@ impl NodeAgentActor { stage_index, endpoint, node_actor, + readiness_id, } => { self.core.observe(stage::StageEvent::WorkerReady); let _ = ctx.send( @@ -244,9 +258,29 @@ impl NodeAgentActor { stage_index, endpoint, node_actor, + readiness_id, }, ); } + NodeAgentMsg::RuntimeReadyAck { + run_id, + node_id, + stage_index, + readiness_id, + } => { + if let Some(report_to) = self.report_to { + let _ = ctx.send( + report_to, + NodeAgentReport::RuntimeReadyAck { + run_id, + node_id, + stage_index, + readiness_id, + }, + ); + } + return; + } NodeAgentMsg::MarkWeightsReady => self.core.observe(stage::StageEvent::WeightsReady), NodeAgentMsg::MarkInboundEdgeReady { edge_id } => { self.core.observe(stage::StageEvent::InboundEdgeReady { @@ -534,6 +568,7 @@ mod tests { stage_index: 3, endpoint: endpoint.clone(), node_actor, + readiness_id: 99, }, ) .expect("send runtime loaded"); @@ -547,7 +582,50 @@ mod tests { stage_index: 3, endpoint, node_actor, + readiness_id: 99, }) ); } + + #[test] + fn node_agent_runtime_ready_ack_reports_worker_loop() { + let runtime = Runtime::new(RuntimeConfig::default()); + let orchestrator_inbox = runtime + .new_inbox::() + .expect("orchestrator inbox"); + let reports = runtime + .new_inbox::() + .expect("node report inbox"); + let actor = runtime + .spawn(NodeAgentActor::new( + stage::NodeId(11), + *orchestrator_inbox.addr(), + Some(*reports.addr()), + )) + .expect("spawn node agent"); + + runtime + .send_to( + actor, + NodeAgentMsg::RuntimeReadyAck { + run_id: 7, + node_id: 11, + stage_index: 3, + readiness_id: 99, + }, + ) + .expect("send runtime ready ack"); + runtime.tick(); + + assert_eq!( + reports.try_recv(), + Some(NodeAgentReport::RuntimeReadyAck { + run_id: 7, + node_id: 11, + stage_index: 3, + readiness_id: 99, + }) + ); + assert_eq!(reports.try_recv(), None); + } } diff --git a/crates/mvp-system/src/actors/orchestrator.rs b/crates/mvp-system/src/actors/orchestrator.rs index d6fb52a..88345db 100644 --- a/crates/mvp-system/src/actors/orchestrator.rs +++ b/crates/mvp-system/src/actors/orchestrator.rs @@ -33,6 +33,7 @@ pub enum OrchestratorMsg { stage_index: u32, endpoint: EndpointAddr, node_actor: ActorAddress, + readiness_id: u64, }, ObserveTokenInEndpointReady, ObserveTokenOutEndpointReady, @@ -134,6 +135,7 @@ pub enum OrchestratorReport { stage_index: u32, endpoint: EndpointAddr, node_actor: ActorAddress, + readiness_id: u64, }, Snapshot { commands: Vec, @@ -269,6 +271,7 @@ impl ActorInterface for OrchestratorActor { stage_index, endpoint, node_actor, + readiness_id, } = msg.clone() { if let Some(report_to) = self.report_to { @@ -280,6 +283,7 @@ impl ActorInterface for OrchestratorActor { stage_index, endpoint, node_actor, + readiness_id, }, ); } @@ -421,6 +425,7 @@ mod tests { stage_index: 3, endpoint: endpoint.clone(), node_actor, + readiness_id: 99, }, ) .expect("send runtime ready"); @@ -434,6 +439,7 @@ mod tests { stage_index: 3, endpoint, node_actor, + readiness_id: 99, }) ); assert_eq!(reports.try_recv(), None); diff --git a/crates/mvp-system/src/bin/mvp_one_node_chat.rs b/crates/mvp-system/src/bin/mvp_one_node_chat.rs index 45ed60d..eca9f5d 100644 --- a/crates/mvp-system/src/bin/mvp_one_node_chat.rs +++ b/crates/mvp-system/src/bin/mvp_one_node_chat.rs @@ -1109,6 +1109,22 @@ fn orch_local_e2e_marker(bin: &Path) -> PathBuf { marker } +fn orch_binary_fingerprint(bin: &Path, root: &Path) -> Result { + let display = display_workspace_path(root, bin); + let metadata = fs::metadata(bin).map_err(|e| format!("stat {display}: {e}"))?; + let modified = metadata + .modified() + .map_err(|e| format!("modified time {display}: {e}"))?; + let modified_ns = modified + .duration_since(std::time::UNIX_EPOCH) + .map_err(|e| format!("modified time before Unix epoch for {display}: {e}"))? + .as_nanos(); + Ok(format!( + "local-e2e\nlen={}\nmodified_ns={modified_ns}\n", + metadata.len() + )) +} + fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result { if !bin.is_file() { return Ok(true); @@ -1117,13 +1133,16 @@ fn orch_local_e2e_marker_stale(bin: &Path, root: &Path) -> Result if !marker.is_file() { return Ok(true); } - Ok(modified_time(root, &marker)? < modified_time(root, bin)?) + let expected = orch_binary_fingerprint(bin, root)?; + let actual = fs::read_to_string(&marker).unwrap_or_default(); + Ok(actual != expected) } fn write_orch_local_e2e_marker(bin: &Path, root: &Path) -> Result<(), String> { let marker = orch_local_e2e_marker(bin); let display = display_workspace_path(root, &marker); - fs::write(&marker, b"local-e2e\n").map_err(|e| format!("write {display}: {e}")) + let fingerprint = orch_binary_fingerprint(bin, root)?; + fs::write(&marker, fingerprint).map_err(|e| format!("write {display}: {e}")) } fn latest_mtime(root: &Path, path: &Path) -> Result { diff --git a/crates/mvp-system/src/bin/orchestrator.rs b/crates/mvp-system/src/bin/orchestrator.rs index 4dfc2e4..60cab63 100644 --- a/crates/mvp-system/src/bin/orchestrator.rs +++ b/crates/mvp-system/src/bin/orchestrator.rs @@ -11,6 +11,7 @@ use std::time::{Duration, Instant}; use datastream::{ChannelId, DatastreamSink, Frame, Lifetime, Mux, NodeId, StreamId}; use distribution::node::DistributedNodeConfig; +use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; use iroh_driver::{IrohDriver, IrohDriverConfig}; use mvp_system::actors::node_agent::{NodeAgentMsg, StageProvisionWire}; @@ -60,6 +61,8 @@ const PUMP_INTERVAL: Duration = Duration::from_millis(10); const MVP_ORCH_BOOTSTRAP: &str = "mvp.orch.bootstrap"; const MVP_ORCH_PROMPT: &str = "mvp.orch.prompt"; const DATASTREAM_FRAME_LOG_ENV: &str = "MVP_DATASTREAM_FRAME_LOG"; +const DEFAULT_DOCKER_CONTAINER_PREFIX: &str = "mvp-orchestrator"; +const MVP_DOCKER_CONTAINER_PREFIX_ENV: &str = "MVP_DOCKER_CONTAINER_PREFIX"; fn main() -> ExitCode { match run() { @@ -484,36 +487,38 @@ fn run() -> Result<(), String> { } }; provisioned_node.complete_bootstrap()?; - driver.join(std::slice::from_ref(&ready.endpoint)); + match enqueue_runtime_ready_ack(&stack, &ready, config.run_id, config.node_id) { + Ok(()) => { + driver.drain_outbox(&stack.outbox); + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "runtime_ready_ack", + "ready", + json!({"node_actor":ready.node_actor,"readiness_id":ready.readiness_id}), + ); + } + Err(error) => { + orch_datastream.emit_bootstrap( + dashboard.as_ref(), + config.run_id, + config.node_id, + "runtime_ready_ack", + "failed", + json!({"node_actor":ready.node_actor,"readiness_id":ready.readiness_id,"error":error}), + ); + return Err(error); + } + } orch_datastream.emit_bootstrap( dashboard.as_ref(), config.run_id, config.node_id, "node_join", "ready", - json!({"endpoint":&ready.endpoint}), + json!({"endpoint":&ready.endpoint,"source":"runtime_ready_barrier"}), ); - match wait_for_route(&mut driver, &stack, ready.node_actor, &stop_rx) { - Ok(()) => orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "node_route", - "ready", - json!({"node_actor":ready.node_actor}), - ), - Err(error) => { - orch_datastream.emit_bootstrap( - dashboard.as_ref(), - config.run_id, - config.node_id, - "node_route", - "failed", - json!({"node_actor":ready.node_actor,"error":error}), - ); - return Err(error); - } - } orch_datastream.emit_bootstrap( dashboard.as_ref(), config.run_id, @@ -1609,7 +1614,7 @@ impl Config { bootstrap_runtime: Arc, ) -> Result, String> { match self.provider { - ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new("mvp-orchestrator"))), + ProviderKind::Docker => Ok(Box::new(LocalDockerPlugin::new(docker_container_prefix()))), ProviderKind::VastAi => { let vastai = self.vastai.as_ref().ok_or_else(|| { "VastAI config was not resolved for provider vastai".to_owned() @@ -1805,6 +1810,33 @@ struct RuntimeReady { endpoint: EndpointAddr, node_actor: ActorAddress, stage_index: u32, + readiness_id: u64, + swim_node_id: DistNodeId, +} + +fn runtime_ready_barrier_met(stack: &DistributionRuntimeStack, ready: &RuntimeReady) -> bool { + stack.member_state(ready.swim_node_id) == Some(MemberState::Alive) + && stack.route_owner(ready.node_actor) == Some(ready.swim_node_id) +} + +fn enqueue_runtime_ready_ack( + stack: &DistributionRuntimeStack, + ready: &RuntimeReady, + run_id: u64, + node_id: u64, +) -> Result<(), String> { + stack + .runtime + .send_to( + ready.node_actor, + NodeAgentMsg::RuntimeReadyAck { + run_id, + node_id, + stage_index: ready.stage_index, + readiness_id: ready.readiness_id, + }, + ) + .map_err(|e| format!("send runtime ready ack: {e}")) } struct ProvisionedNodeGuard<'a> { @@ -2294,6 +2326,10 @@ fn wait_for_runtime_ready( node_id: u64, provider: ProviderKind, ) -> Result { + let mut pending_ready: Option = None; + let mut node_swim_started = false; + let mut node_swim_ready = false; + let mut node_route_started = false; loop { pump(driver, stack); drain_frames(frame_rx, dashboard, orch_datastream); @@ -2321,39 +2357,72 @@ fn wait_for_runtime_ready( stage_index, endpoint, node_actor, + readiness_id, } = report { if report_run_id == run_id && report_node_id == node_id { - return Ok(RuntimeReady { + let reset_progress = pending_ready + .as_ref() + .map(|ready| ready.readiness_id != readiness_id) + .unwrap_or(true); + if reset_progress { + node_swim_started = false; + node_swim_ready = false; + node_route_started = false; + } + let swim_node_id = DistNodeId(*endpoint.id.as_bytes()); + pending_ready = Some(RuntimeReady { endpoint, node_actor, stage_index, + readiness_id, + swim_node_id, }); } } } - thread::sleep(PUMP_INTERVAL); - } -} - -fn wait_for_route( - driver: &mut IrohDriver, - stack: &DistributionRuntimeStack, - actor: ActorAddress, - stop_rx: &mpsc::Receiver<()>, -) -> Result<(), String> { - loop { - pump(driver, stack); - if stop_requested(stop_rx) { - return Err("shutdown requested while waiting for node route".to_owned()); - } - let ready = stack - .route_view - .read() - .map(|view| view.contains_key(&actor)) - .unwrap_or(false); - if ready { - return Ok(()); + if let Some(ready) = pending_ready.as_ref() { + if runtime_ready_barrier_met(stack, ready) { + return Ok(ready.clone()); + } + let swim_ready = stack.member_state(ready.swim_node_id) == Some(MemberState::Alive); + let route_ready = stack.route_owner(ready.node_actor) == Some(ready.swim_node_id); + if !swim_ready { + if !node_swim_started { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_swim", + "started", + json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_swim_started = true; + } + } else if !route_ready { + if !node_swim_ready { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_swim", + "ready", + json!({"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_swim_ready = true; + } + if !node_route_started { + orch_datastream.emit_bootstrap( + dashboard, + run_id, + node_id, + "node_route", + "started", + json!({"node_actor":ready.node_actor,"node":format!("{:?}", ready.swim_node_id),"readiness_id":ready.readiness_id}), + ); + node_route_started = true; + } + } } thread::sleep(PUMP_INTERVAL); } @@ -2812,6 +2881,11 @@ fn env_optional(name: &str) -> Option { .filter(|value| !value.is_empty()) } +fn docker_container_prefix() -> String { + env_optional(MVP_DOCKER_CONTAINER_PREFIX_ENV) + .unwrap_or_else(|| DEFAULT_DOCKER_CONTAINER_PREFIX.to_owned()) +} + fn optional_env(name: &str) -> Option<(String, String)> { env_optional(name).map(|value| (name.to_owned(), value)) } @@ -3045,6 +3119,7 @@ mod tests { "MVP_VASTAI_SSH_IDENTITY", "VASTAI_API_KEY", SWACTOR_IROH_RELAY_URL_ENV, + MVP_DOCKER_CONTAINER_PREFIX_ENV, ]; struct RestoreEnv { @@ -3126,6 +3201,16 @@ mod tests { .map(|(_, value)| value.as_str()) } + #[test] + fn docker_container_prefix_defaults_and_trims_env_override() { + with_clean_env(&[], || { + assert_eq!(docker_container_prefix(), DEFAULT_DOCKER_CONTAINER_PREFIX); + }); + with_clean_env(&[(MVP_DOCKER_CONTAINER_PREFIX_ENV, " custom-prefix ")], || { + assert_eq!(docker_container_prefix(), "custom-prefix"); + }); + } + #[test] fn expand_home_path_expands_leading_home_segment() { with_clean_env(&[("HOME", "/tmp/mvp-vastai-home")], || { @@ -3574,6 +3659,39 @@ bootstrap_command = "/run" assert_eq!(env_value(&disabled_env, MVP_IROH_RELAY_URL_ENV), None); } + #[test] + fn node_spec_preserves_relay_transport_in_coordinator_endpoint() { + with_clean_env(&[], || { + let config = Config::from_layers_with_path_and_args(None, std::iter::empty::()) + .expect("config parses"); + let secret = iroh::SecretKey::from_bytes(&[10; 32]); + let coordinator = EndpointAddr::new(secret.public()).with_relay_url( + "http://relay.example.com" + .parse::() + .unwrap(), + ); + let datastream_sink = ActorAddress([19; 32]); + let orchestrator_actor = ActorAddress([20; 32]); + + let spec = config + .node_spec(coordinator, datastream_sink, orchestrator_actor) + .expect("node spec builds"); + let coordinator_endpoint_json = env_value(&spec.env, "MVP_COORDINATOR_ENDPOINT") + .expect("coordinator endpoint env is present"); + let coordinator_endpoint = + serde_json::from_str::(coordinator_endpoint_json) + .expect("coordinator endpoint env deserializes"); + + assert_eq!( + coordinator_endpoint + .relay_urls() + .next() + .map(|url| url.to_string()), + Some("http://relay.example.com/".to_owned()) + ); + }); + } + #[test] fn docker_config_construction_ignores_malformed_vastai_environment() { let config = with_clean_env( @@ -3735,4 +3853,110 @@ bootstrap_command = "/run" assert!(stop_requested(&rx)); assert!(!stop_requested(&rx)); } + + fn endpoint(seed: u8) -> EndpointAddr { + EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public()) + } + + fn insert_route(stack: &DistributionRuntimeStack, actor: ActorAddress, owner: DistNodeId) { + let mut route_view = match stack.route_view.write() { + Ok(route_view) => route_view, + Err(poisoned) => poisoned.into_inner(), + }; + route_view.insert(actor, owner); + } + + fn mark_alive(stack: &DistributionRuntimeStack, node_id: DistNodeId) { + stack + .runtime + .send_to( + stack.actors.membership_fanout, + distribution::swim::actor::MembershipChanged { + node_id, + state: MemberState::Alive, + incarnation: 1, + }, + ) + .expect("send membership change"); + stack.pump_runtime_once(); + } + + #[test] + fn runtime_ready_barrier_waits_for_specific_swim_and_route() { + let remote = DistNodeId([2; 32]); + let node_actor = ActorAddress::new_random(); + let ready = RuntimeReady { + endpoint: endpoint(2), + node_actor, + stage_index: 3, + readiness_id: 99, + swim_node_id: remote, + }; + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + mark_alive(&stack, remote); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + insert_route(&stack, node_actor, remote); + assert!(!runtime_ready_barrier_met(&stack, &ready)); + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + mark_alive(&stack, remote); + insert_route(&stack, node_actor, remote); + assert!(runtime_ready_barrier_met(&stack, &ready)); + } + + #[test] + fn enqueue_runtime_ready_ack_reports_to_node_agent() { + use mvp_system::actors::node_agent::{NodeAgentActor, NodeAgentReport}; + use mvp_system::actors::orchestrator::OrchestratorMsg; + use mvp_system::stage_controller as stage; + + let stack = + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()); + let orchestrator_inbox = stack + .runtime + .new_inbox::() + .expect("orchestrator inbox"); + let reports = stack + .runtime + .new_inbox::() + .expect("node report inbox"); + let node_actor = stack + .runtime + .spawn(NodeAgentActor::new( + stage::NodeId(11), + *orchestrator_inbox.addr(), + Some(*reports.addr()), + )) + .expect("spawn node agent"); + let ready = RuntimeReady { + endpoint: endpoint(9), + node_actor, + stage_index: 3, + readiness_id: 99, + swim_node_id: DistNodeId([2; 32]), + }; + + enqueue_runtime_ready_ack(&stack, &ready, 7, 11).expect("enqueue runtime ready ack"); + stack.pump_runtime_once(); + + assert_eq!( + reports.try_recv(), + Some(NodeAgentReport::RuntimeReadyAck { + run_id: 7, + node_id: 11, + stage_index: ready.stage_index, + readiness_id: 99, + }) + ); + } } diff --git a/crates/mvp-system/src/bin/worker_node.rs b/crates/mvp-system/src/bin/worker_node.rs index 6190981..b5c8a79 100644 --- a/crates/mvp-system/src/bin/worker_node.rs +++ b/crates/mvp-system/src/bin/worker_node.rs @@ -1,5 +1,6 @@ -use std::fs::{File, OpenOptions}; -use std::io::{BufRead, BufReader, Write}; +use std::fs::{self, File, OpenOptions}; +use std::io::{BufRead, BufReader, Read, Write}; +use std::path::PathBuf; use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio}; use std::sync::{ Arc, OnceLock, @@ -14,6 +15,7 @@ use datastream::emit::{ }; use distribution::node::DistributedNodeConfig; +use distribution::types::{MemberState, NodeId as DistNodeId}; use iroh::EndpointAddr; use iroh_driver::{IrohDriver, IrohDriverConfig}; use mvp_system::actors::node_agent::{ @@ -29,6 +31,7 @@ use mvp_system::stage_controller as stage; use parking_lot::Mutex; use serde_json::{Value, json}; use swactor::actor::ActorAddress; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt}; const DEFAULT_WORKER_SCRIPT: &str = "/usr/local/share/mvp/tinygrad_worker.py"; const DEFAULT_DEVICE: &str = "CUDA"; @@ -38,6 +41,8 @@ const DEFAULT_MODEL_ID: &str = "llama-3.2-1b-instruct-q4"; const DEFAULT_ARENA_BYTES: u64 = 64 * 1024 * 1024; const DEFAULT_ARENA_ALIGNMENT: u64 = 64; const PUMP_INTERVAL: Duration = Duration::from_millis(10); +const RUNTIME_READY_RETRY_INITIAL: Duration = Duration::from_millis(100); +const RUNTIME_READY_RETRY_MAX: Duration = Duration::from_secs(2); const NODE_BOOTSTRAP_CHANNEL: &str = "mvp.node.bootstrap"; const NODE_RUNTIME_CHANNEL: &str = "mvp.node.runtime"; const NODE_STAGE_CHANNEL: &str = "mvp.node.stage"; @@ -98,6 +103,298 @@ fn emit_node_event( datastream.tick(); } +#[derive(serde::Deserialize, serde::Serialize)] +#[serde(tag = "type")] +enum DebugJoinRequestWire { + JoinEndpoint { endpoint: EndpointAddr }, +} + +#[derive(serde::Deserialize, serde::Serialize)] +#[serde(tag = "type")] +enum DebugJoinResponseWire { + JoinQueued { + peer_node_id: String, + has_relay: bool, + direct_addr_count: usize, + }, + JoinRejected { + error: String, + detail: String, + }, +} + +enum DebugJoinCommand { + JoinEndpoint { + endpoint: EndpointAddr, + reply: tokio::sync::oneshot::Sender, + }, +} + +enum DebugJoinClientError { + Cli(String), + Runtime(String), +} + +fn debug_join_client_main(args: Vec) -> ExitCode { + match run_debug_join_client(args) { + Ok(response) => { + let queued = matches!(response, DebugJoinResponseWire::JoinQueued { .. }); + match serde_json::to_string(&response) { + Ok(line) => println!("{line}"), + Err(error) => { + eprintln!("mvp-worker-node debug-join: serialize response: {error}"); + return ExitCode::from(1); + } + } + if queued { + ExitCode::SUCCESS + } else { + ExitCode::from(1) + } + } + Err(DebugJoinClientError::Cli(error)) => { + eprintln!("mvp-worker-node debug-join: {error}"); + ExitCode::from(2) + } + Err(DebugJoinClientError::Runtime(error)) => { + eprintln!("mvp-worker-node debug-join: {error}"); + ExitCode::from(1) + } + } +} + +fn run_debug_join_client(args: Vec) -> Result { + let mut socket = None; + let mut endpoint_json = None; + let mut read_endpoint_stdin = false; + let mut iter = args.into_iter(); + while let Some(arg) = iter.next() { + match arg.as_str() { + "--socket" => { + socket = Some(PathBuf::from(iter.next().ok_or_else(|| { + DebugJoinClientError::Cli("--socket requires a path".to_owned()) + })?)); + } + "--endpoint-json" => { + endpoint_json = Some(iter.next().ok_or_else(|| { + DebugJoinClientError::Cli("--endpoint-json requires JSON".to_owned()) + })?); + } + "--endpoint-json-stdin" => read_endpoint_stdin = true, + other => { + return Err(DebugJoinClientError::Cli(format!( + "unknown argument {other:?}; usage: debug-join --socket (--endpoint-json | --endpoint-json-stdin)" + ))); + } + } + } + let socket = socket.ok_or_else(|| { + DebugJoinClientError::Cli( + "missing --socket ; usage: debug-join --socket (--endpoint-json | --endpoint-json-stdin)".to_owned(), + ) + })?; + let endpoint_json = match (endpoint_json, read_endpoint_stdin) { + (Some(_), true) => { + return Err(DebugJoinClientError::Cli( + "use either --endpoint-json or --endpoint-json-stdin, not both".to_owned(), + )); + } + (Some(json), false) => json, + (None, true) => { + let mut json = String::new(); + std::io::stdin() + .read_to_string(&mut json) + .map_err(|e| DebugJoinClientError::Runtime(format!("read endpoint stdin: {e}")))?; + json + } + (None, false) => { + return Err(DebugJoinClientError::Cli( + "missing endpoint JSON; use --endpoint-json or --endpoint-json-stdin" + .to_owned(), + )); + } + }; + let endpoint = serde_json::from_str::(&endpoint_json) + .map_err(|e| DebugJoinClientError::Cli(format!("parse endpoint JSON: {e}")))?; + let request = debug_join_request_line(endpoint).map_err(DebugJoinClientError::Runtime)?; + let mut stream = std::os::unix::net::UnixStream::connect(&socket) + .map_err(|e| DebugJoinClientError::Runtime(format!("connect {}: {e}", socket.display())))?; + stream + .write_all(request.as_bytes()) + .map_err(|e| DebugJoinClientError::Runtime(format!("write request: {e}")))?; + stream + .flush() + .map_err(|e| DebugJoinClientError::Runtime(format!("flush request: {e}")))?; + let mut response_line = String::new(); + BufReader::new(stream) + .read_line(&mut response_line) + .map_err(|e| DebugJoinClientError::Runtime(format!("read response: {e}")))?; + if response_line.trim().is_empty() { + return Err(DebugJoinClientError::Runtime( + "debug join socket closed without response".to_owned(), + )); + } + serde_json::from_str::(&response_line) + .map_err(|e| DebugJoinClientError::Runtime(format!("parse response JSON: {e}"))) +} + +fn debug_join_request_line(endpoint: EndpointAddr) -> Result { + serde_json::to_string(&DebugJoinRequestWire::JoinEndpoint { endpoint }) + .map(|mut line| { + line.push('\n'); + line + }) + .map_err(|e| format!("serialize debug join request: {e}")) +} + +fn spawn_debug_join_listener( + handle: tokio::runtime::Handle, + path: PathBuf, +) -> Result, String> { + use std::os::unix::fs::PermissionsExt; + + match fs::remove_file(&path) { + Ok(()) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(format!( + "remove stale debug join socket {}: {error}", + path.display() + )); + } + } + let listener = { + let _guard = handle.enter(); + tokio::net::UnixListener::bind(&path) + .map_err(|e| format!("bind debug join socket {}: {e}", path.display()))? + }; + fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) + .map_err(|e| format!("chmod debug join socket {}: {e}", path.display()))?; + let (command_tx, command_rx) = tokio::sync::mpsc::unbounded_channel::(); + handle.spawn(async move { + loop { + match listener.accept().await { + Ok((stream, _addr)) => { + let command_tx = command_tx.clone(); + tokio::spawn(async move { + handle_debug_join_stream(stream, command_tx).await; + }); + } + Err(error) => { + eprintln!("mvp-worker-node debug join listener stopped: {error}"); + break; + } + } + } + }); + Ok(command_rx) +} + +async fn handle_debug_join_stream( + stream: tokio::net::UnixStream, + command_tx: tokio::sync::mpsc::UnboundedSender, +) { + let mut reader = tokio::io::BufReader::new(stream); + let mut line = String::new(); + let response = match reader.read_line(&mut line).await { + Ok(0) => DebugJoinResponseWire::JoinRejected { + error: "MalformedCommand".to_owned(), + detail: "empty request".to_owned(), + }, + Ok(_) => match parse_debug_join_request(&line) { + Ok(DebugJoinRequestWire::JoinEndpoint { endpoint }) => { + let (reply, response_rx) = tokio::sync::oneshot::channel(); + if command_tx + .send(DebugJoinCommand::JoinEndpoint { endpoint, reply }) + .is_err() + { + DebugJoinResponseWire::JoinRejected { + error: "CommandQueueClosed".to_owned(), + detail: "worker main loop is not accepting debug join commands".to_owned(), + } + } else { + response_rx + .await + .unwrap_or_else(|error| DebugJoinResponseWire::JoinRejected { + error: "CommandCancelled".to_owned(), + detail: error.to_string(), + }) + } + } + Err(response) => response, + }, + Err(error) => DebugJoinResponseWire::JoinRejected { + error: "MalformedCommand".to_owned(), + detail: format!("read request: {error}"), + }, + }; + let mut stream = reader.into_inner(); + if let Ok(line) = serde_json::to_string(&response) { + let _ = stream.write_all(line.as_bytes()).await; + let _ = stream.write_all(b"\n").await; + let _ = stream.flush().await; + } +} + +fn parse_debug_join_request(raw: &str) -> Result { + let value = + serde_json::from_str::(raw).map_err(|e| DebugJoinResponseWire::JoinRejected { + error: "MalformedCommand".to_owned(), + detail: e.to_string(), + })?; + let endpoint_decode_error = value.get("type").and_then(Value::as_str) == Some("JoinEndpoint") + && value.get("endpoint").is_some(); + serde_json::from_value::(value).map_err(|e| { + DebugJoinResponseWire::JoinRejected { + error: if endpoint_decode_error { + "MalformedEndpoint" + } else { + "MalformedCommand" + } + .to_owned(), + detail: e.to_string(), + } + }) +} + +fn drain_debug_join_commands( + debug_join_rx: &mut Option>, + driver: &mut IrohDriver, + config: &DeploymentConfig, + datastream: &mut DatastreamEmitter, +) { + let Some(rx) = debug_join_rx else { + return; + }; + while let Ok(command) = rx.try_recv() { + match command { + DebugJoinCommand::JoinEndpoint { endpoint, reply } => { + let peer_node_id = endpoint.id.to_string(); + let has_relay = endpoint.relay_urls().next().is_some(); + let direct_addr_count = endpoint.ip_addrs().count(); + driver.join(std::slice::from_ref(&endpoint)); + emit_node_event( + datastream, + config, + NODE_RUNTIME_CHANNEL, + "debug_join", + "queued", + json!({ + "peer_node_id":peer_node_id, + "has_relay":has_relay, + "direct_addr_count":direct_addr_count, + }), + ); + let _ = reply.send(DebugJoinResponseWire::JoinQueued { + peer_node_id, + has_relay, + direct_addr_count, + }); + } + } + } +} + fn spawn_host_gpu_sampler(handle: tokio::runtime::Handle, sink: DatastreamEventSink) { handle.spawn(async move { let mut seq = 0_u64; @@ -192,6 +489,11 @@ fn spawn_arena_sampler( } fn main() -> ExitCode { + let mut args = std::env::args().skip(1).collect::>(); + if args.first().map(String::as_str) == Some("debug-join") { + args.remove(0); + return debug_join_client_main(args); + } match run() { Ok(()) => ExitCode::SUCCESS, Err(error) => { @@ -218,6 +520,7 @@ fn run() -> Result<(), String> { "self_test_enabled":config.self_test_prompt.is_some(), "arena_bytes":config.arena_bytes, "arena_alignment":config.arena_alignment, + "debug_join_socket":config.debug_join_socket.as_deref().unwrap_or("disabled"), }), )?; emit_stdio_node_event( @@ -383,6 +686,45 @@ fn run() -> Result<(), String> { "ready", config.datastream_sink_detail(), )?; + let mut debug_join_rx = match &config.debug_join_socket { + Some(path) => { + match spawn_debug_join_listener(tokio.handle().clone(), PathBuf::from(path)) { + Ok(rx) => { + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "debug_join_socket", + "ready", + json!({"socket":path}), + ); + Some(rx) + } + Err(error) => { + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "debug_join_socket", + "failed", + json!({"socket":path,"error":error}), + ); + return Err(format!("bind debug join socket {}: {error}", path)); + } + } + } + None => { + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "debug_join_socket", + "skipped", + json!({"reason":"MVP_DEBUG_JOIN_SOCKET=disabled"}), + ); + None + } + }; let reports = match stack.runtime.new_inbox::() { Ok(inbox) => { @@ -512,34 +854,8 @@ fn run() -> Result<(), String> { return Err(error); } } - match stack.runtime.send_to( - node_actor, - NodeAgentMsg::RuntimeLoaded { - run_id: config.run_id, - node_id: config.logical_node_id, - stage_index: config.stage_index, - endpoint: driver.endpoint_addr(), - node_actor, - }, - ) { - Ok(()) => emit_stdio_node_event( - &config, - NODE_RUNTIME_CHANNEL, - "runtime_loaded", - "ready", - json!({"sent":"NodeAgentMsg::RuntimeLoaded","node_actor":node_actor}), - )?, - Err(error) => { - emit_stdio_node_event( - &config, - NODE_RUNTIME_CHANNEL, - "runtime_loaded", - "failed", - json!({"error":error.to_string()}), - )?; - return Err(format!("signal runtime loaded: {error}")); - } - } + let mut pending_runtime_ready = + PendingRuntimeReady::new(&config, driver.endpoint_addr(), node_actor); let ready = json!({ "type":"ready", @@ -552,23 +868,16 @@ fn run() -> Result<(), String> { emit_stdio_node_event( &config, NODE_BOOTSTRAP_CHANNEL, - "runtime_ready", + "runtime_ready_local", "ready", json!({ "endpoint":driver.endpoint_addr(), "node_actor":node_actor, "logical_node_id":config.logical_node_id, "stage_index":config.stage_index, + "readiness_id":pending_runtime_ready.readiness_id, }), )?; - datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string()); - emit_stdio_node_event( - &config, - NODE_BOOTSTRAP_CHANNEL, - "datastream_handoff", - "ready", - json!({"from":"stdio_envelope","to":"cluster_datastream","channel":NODE_BOOTSTRAP_CHANNEL}), - )?; if let Some(prompt) = &config.self_test_prompt { run_self_test( @@ -603,10 +912,11 @@ fn run() -> Result<(), String> { ); loop { pump_network(&mut driver, &stack); + drain_debug_join_commands(&mut debug_join_rx, &mut driver, &config, &mut datastream); datastream.tick(); worker.drain_stderr(&config, &mut datastream); while let Some(report) = reports.try_recv() { - handle_node_report( + match handle_node_report( report, &config, &stack, @@ -614,7 +924,72 @@ fn run() -> Result<(), String> { node_actor, &mut worker, &mut datastream, - )?; + )? { + NodeReportOutcome::None => {} + NodeReportOutcome::RuntimeReadyAck { + run_id, + node_id, + stage_index, + readiness_id, + } => { + if pending_runtime_ready.observe_ack(run_id, node_id, stage_index, readiness_id) + { + emit_node_event( + &mut datastream, + &config, + NODE_BOOTSTRAP_CHANNEL, + "runtime_ready_ack", + "ready", + json!({ + "readiness_id":readiness_id, + "attempts":pending_runtime_ready.attempts, + "endpoint":&pending_runtime_ready.endpoint, + "node_actor":pending_runtime_ready.node_actor, + }), + ); + datastream.submit_text(ChannelId::new("mvp.node.ready"), ready.to_string()); + emit_node_event( + &mut datastream, + &config, + NODE_BOOTSTRAP_CHANNEL, + "datastream_handoff", + "ready", + json!({"from":"runtime_ready_ack","to":"cluster_datastream","channel":"mvp.node.ready"}), + ); + } + } + } + } + if !pending_runtime_ready.swim_logged && pending_runtime_ready.swim_ready(&stack) { + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "coordinator_swim", + "ready", + json!({ + "coordinator":pending_runtime_ready + .coordinator + .map(|node| format!("{node:?}")) + .unwrap_or_else(|| "standalone".to_owned()), + "readiness_id":pending_runtime_ready.readiness_id, + }), + ); + pending_runtime_ready.swim_logged = true; + } + if !pending_runtime_ready.acked && pending_runtime_ready.maybe_send(&stack, node_actor)? { + emit_node_event( + &mut datastream, + &config, + NODE_RUNTIME_CHANNEL, + "runtime_ready_signal", + "sent", + json!({ + "readiness_id":pending_runtime_ready.readiness_id, + "attempts":pending_runtime_ready.attempts, + "next_backoff_ms":pending_runtime_ready.backoff.as_millis(), + }), + ); } if shutdown_rx.try_recv().is_ok() { emit_node_event( @@ -758,6 +1133,115 @@ impl FrameSink for JsonlFrameSink { } } +enum NodeReportOutcome { + None, + RuntimeReadyAck { + run_id: u64, + node_id: u64, + stage_index: u32, + readiness_id: u64, + }, +} + +struct PendingRuntimeReady { + run_id: u64, + node_id: u64, + stage_index: u32, + endpoint: EndpointAddr, + node_actor: ActorAddress, + coordinator: Option, + readiness_id: u64, + attempts: u32, + next_attempt_at: Instant, + backoff: Duration, + acked: bool, + swim_logged: bool, +} + +impl PendingRuntimeReady { + fn new(config: &DeploymentConfig, endpoint: EndpointAddr, node_actor: ActorAddress) -> Self { + Self { + run_id: config.run_id, + node_id: config.logical_node_id, + stage_index: config.stage_index, + endpoint, + node_actor, + coordinator: config + .coordinator_endpoint + .as_ref() + .map(|endpoint| DistNodeId(*endpoint.id.as_bytes())), + readiness_id: 1, + attempts: 0, + next_attempt_at: Instant::now(), + backoff: RUNTIME_READY_RETRY_INITIAL, + acked: false, + swim_logged: false, + } + } + + fn swim_ready(&self, stack: &DistributionRuntimeStack) -> bool { + let Some(coordinator) = self.coordinator else { + return true; + }; + stack.member_state(coordinator) == Some(MemberState::Alive) + } + + fn observe_ack( + &mut self, + run_id: u64, + node_id: u64, + stage_index: u32, + readiness_id: u64, + ) -> bool { + if self.acked + || self.run_id != run_id + || self.node_id != node_id + || self.stage_index != stage_index + || self.readiness_id != readiness_id + { + return false; + } + self.acked = true; + true + } + + fn maybe_send( + &mut self, + stack: &DistributionRuntimeStack, + node_actor: ActorAddress, + ) -> Result { + if self.acked || !self.swim_ready(stack) { + return Ok(false); + } + let now = Instant::now(); + if now < self.next_attempt_at { + return Ok(false); + } + stack + .runtime + .send_to( + node_actor, + NodeAgentMsg::RuntimeLoaded { + run_id: self.run_id, + node_id: self.node_id, + stage_index: self.stage_index, + endpoint: self.endpoint.clone(), + node_actor: self.node_actor, + readiness_id: self.readiness_id, + }, + ) + .map_err(|error| format!("signal runtime loaded: {error}"))?; + self.attempts = self.attempts.saturating_add(1); + self.next_attempt_at = now + self.backoff; + self.backoff = self + .backoff + .checked_mul(2) + .unwrap_or(RUNTIME_READY_RETRY_MAX) + .min(RUNTIME_READY_RETRY_MAX); + Ok(true) + } +} + fn handle_node_report( report: NodeAgentReport, config: &DeploymentConfig, @@ -766,11 +1250,12 @@ fn handle_node_report( node_actor: ActorAddress, worker: &mut TinygradWorker, datastream: &mut DatastreamEmitter, -) -> Result<(), String> { +) -> Result { let kind = match &report { NodeAgentReport::Command(_) => "Command", NodeAgentReport::Lifecycle(_) => "Lifecycle", NodeAgentReport::PromptRequested { .. } => "PromptRequested", + NodeAgentReport::RuntimeReadyAck { .. } => "RuntimeReadyAck", NodeAgentReport::Snapshot { .. } => "Snapshot", }; emit_node_event( @@ -782,9 +1267,12 @@ fn handle_node_report( json!({"kind":kind}), ); match report { - NodeAgentReport::Command(command) => handle_stage_command( - command, config, stack, driver, node_actor, worker, datastream, - ), + NodeAgentReport::Command(command) => { + handle_stage_command( + command, config, stack, driver, node_actor, worker, datastream, + )?; + Ok(NodeReportOutcome::None) + } NodeAgentReport::Lifecycle(event) => { let event = format!("{event:?}"); datastream.submit_text( @@ -799,17 +1287,31 @@ fn handle_node_report( "observed", json!({"event":event}), ); - Ok(()) + Ok(NodeReportOutcome::None) } NodeAgentReport::PromptRequested { request_id, prompt, max_tokens, reply_to, - } => handle_prompt_request( - request_id, prompt, max_tokens, reply_to, config, stack, driver, worker, datastream, - ), - NodeAgentReport::Snapshot { .. } => Ok(()), + } => { + handle_prompt_request( + request_id, prompt, max_tokens, reply_to, config, stack, driver, worker, datastream, + )?; + Ok(NodeReportOutcome::None) + } + NodeAgentReport::RuntimeReadyAck { + run_id, + node_id, + stage_index, + readiness_id, + } => Ok(NodeReportOutcome::RuntimeReadyAck { + run_id, + node_id, + stage_index, + readiness_id, + }), + NodeAgentReport::Snapshot { .. } => Ok(NodeReportOutcome::None), } } @@ -1267,6 +1769,7 @@ struct DeploymentConfig { orchestrator_actor: Option, datastream_sink_actor: Option, datastream_frame_log: Option, + debug_join_socket: Option, relay_mode: iroh::RelayMode, worker_script: String, device: String, @@ -1283,15 +1786,29 @@ struct DeploymentConfig { impl DeploymentConfig { fn from_env() -> Result { let run_id = env_u64("MVP_RUN_ID", 1)?; + let logical_node_id = env_u64("MVP_LOGICAL_NODE_ID", 1)?; let relay = relay_runtime_config_from_env(run_id)?; + let debug_join_socket = match env_optional("MVP_DEBUG_JOIN_SOCKET").as_deref() { + Some("disabled") => None, + Some(path) => Some(path.to_owned()), + None => Some( + std::env::temp_dir() + .join(format!( + "mvp-node-debug-join-{run_id}-{logical_node_id}.sock" + )) + .to_string_lossy() + .into_owned(), + ), + }; Ok(Self { run_id, - logical_node_id: env_u64("MVP_LOGICAL_NODE_ID", 1)?, + logical_node_id, stage_index: env_u32("MVP_STAGE_INDEX", 0)?, coordinator_endpoint: env_json("MVP_COORDINATOR_ENDPOINT")?, orchestrator_actor: env_json("MVP_ORCHESTRATOR_ACTOR")?, datastream_sink_actor: env_json("MVP_DATASTREAM_SINK_ACTOR")?, datastream_frame_log: env_optional("MVP_DATASTREAM_FRAME_LOG"), + debug_join_socket, relay_mode: relay.mode, worker_script: env_string("MVP_TINYGRAD_WORKER", DEFAULT_WORKER_SCRIPT), device: env_string("DEV", DEFAULT_DEVICE), @@ -1743,3 +2260,212 @@ fn tokenizer_from_env() -> TokenizerSource { .map(TokenizerSource::LocalPath) .unwrap_or(TokenizerSource::EmbeddedGguf) } + +#[cfg(test)] +mod tests { + use super::*; + use distribution::swim::actor::MembershipChanged; + use mvp_system::actors::orchestrator::OrchestratorMsg; + + fn endpoint(seed: u8) -> EndpointAddr { + EndpointAddr::new(iroh::SecretKey::from_bytes(&[seed; 32]).public()) + } + + fn test_config(coordinator_endpoint: Option) -> DeploymentConfig { + DeploymentConfig { + run_id: 7, + logical_node_id: 11, + stage_index: 3, + coordinator_endpoint, + orchestrator_actor: Some(ActorAddress::new_random()), + datastream_sink_actor: None, + datastream_frame_log: None, + debug_join_socket: None, + relay_mode: iroh::RelayMode::Disabled, + worker_script: DEFAULT_WORKER_SCRIPT.to_owned(), + device: DEFAULT_DEVICE.to_owned(), + model_id: DEFAULT_MODEL_ID.to_owned(), + gguf_source: GgufSource::LocalPath("/tmp/model.gguf".to_owned()), + tokenizer: TokenizerSource::EmbeddedGguf, + self_test_prompt: None, + self_test_layer_end: 16, + self_test_max_tokens: 1, + arena_bytes: DEFAULT_ARENA_BYTES, + arena_alignment: DEFAULT_ARENA_ALIGNMENT, + } + } + + fn test_stack() -> DistributionRuntimeStack { + DistributionRuntimeStack::new(DistNodeId([1; 32]), DistributedNodeConfig::default()) + } + + #[test] + fn debug_join_client_serializes_endpoint_from_stdin() { + let secret = iroh::SecretKey::from_bytes(&[7; 32]); + let endpoint = EndpointAddr::new(secret.public()).with_relay_url( + "http://relay.example.com" + .parse::() + .unwrap(), + ); + + let line = debug_join_request_line(endpoint).expect("serialize debug join request"); + let request: DebugJoinRequestWire = + serde_json::from_str(&line).expect("deserialize debug join request"); + + match request { + DebugJoinRequestWire::JoinEndpoint { endpoint } => { + assert_eq!( + endpoint.relay_urls().next().map(ToString::to_string), + Some("http://relay.example.com/".to_owned()) + ); + } + } + } + + #[test] + fn debug_join_listener_queues_join_endpoint() { + let root = std::env::temp_dir().join(format!( + "mvp-worker-debug-join-test-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system time after epoch") + .as_nanos() + )); + std::fs::create_dir(&root).expect("create debug join test temp dir"); + let socket_path = root.join("debug-join.sock"); + let runtime = tokio::runtime::Runtime::new().expect("create tokio runtime"); + let mut commands = spawn_debug_join_listener(runtime.handle().clone(), socket_path.clone()) + .expect("spawn debug join listener"); + + let secret = iroh::SecretKey::from_bytes(&[8; 32]); + let endpoint = EndpointAddr::new(secret.public()).with_relay_url( + "http://relay.example.com" + .parse::() + .unwrap(), + ); + let request = DebugJoinRequestWire::JoinEndpoint { endpoint }; + let mut request_line = + serde_json::to_string(&request).expect("serialize debug join request"); + request_line.push('\n'); + + runtime.block_on(async { + let mut stream = tokio::net::UnixStream::connect(&socket_path) + .await + .expect("connect to debug join listener"); + stream + .write_all(request_line.as_bytes()) + .await + .expect("write debug join request"); + stream.flush().await.expect("flush debug join request"); + + let DebugJoinCommand::JoinEndpoint { endpoint, reply } = + commands.recv().await.expect("receive debug join command"); + assert_eq!( + endpoint.relay_urls().next().map(ToString::to_string), + Some("http://relay.example.com/".to_owned()) + ); + let peer_node_id = endpoint.id.to_string(); + assert!( + reply + .send(DebugJoinResponseWire::JoinQueued { + peer_node_id, + has_relay: true, + direct_addr_count: 0, + }) + .is_ok() + ); + + let mut reader = tokio::io::BufReader::new(stream); + let mut response_line = String::new(); + reader + .read_line(&mut response_line) + .await + .expect("read debug join response"); + let response: DebugJoinResponseWire = + serde_json::from_str(&response_line).expect("deserialize debug join response"); + match response { + DebugJoinResponseWire::JoinQueued { has_relay, .. } => { + assert!(has_relay); + } + DebugJoinResponseWire::JoinRejected { error, detail } => { + panic!("debug join was rejected: {error}: {detail}"); + } + } + }); + + drop(commands); + let _ = std::fs::remove_file(&socket_path); + std::fs::remove_dir(&root).expect("remove debug join test temp dir"); + } + + #[test] + fn runtime_ready_retry_waits_for_swim() { + let stack = test_stack(); + let node_actor = ActorAddress::new_random(); + let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(3), node_actor); + pending.coordinator = Some(DistNodeId([2; 32])); + + assert_eq!(pending.maybe_send(&stack, node_actor), Ok(false)); + assert_eq!(pending.attempts, 0); + } + + #[test] + fn runtime_ready_retry_stops_after_matching_ack() { + let stack = test_stack(); + let node_actor = ActorAddress::new_random(); + let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(4), node_actor); + + assert!(pending.observe_ack( + pending.run_id, + pending.node_id, + pending.stage_index, + pending.readiness_id, + )); + assert!(pending.acked); + assert_eq!(pending.maybe_send(&stack, node_actor), Ok(false)); + } + + #[test] + fn runtime_ready_retry_backoff_caps() { + let stack = test_stack(); + let orchestrator_inbox = stack + .runtime + .new_inbox::() + .expect("orchestrator inbox"); + let node_actor = stack + .runtime + .spawn(NodeAgentActor::new( + stage::NodeId(11), + *orchestrator_inbox.addr(), + None, + )) + .expect("spawn node agent"); + let coordinator = DistNodeId([2; 32]); + stack + .runtime + .send_to( + stack.actors.membership_fanout, + MembershipChanged { + node_id: coordinator, + state: MemberState::Alive, + incarnation: 1, + }, + ) + .expect("send membership change"); + stack.pump_runtime_once(); + + let mut pending = PendingRuntimeReady::new(&test_config(None), endpoint(5), node_actor); + pending.coordinator = Some(coordinator); + for expected_attempts in 1..=4 { + pending.next_attempt_at = Instant::now(); + assert!( + pending + .maybe_send(&stack, node_actor) + .expect("runtime ready send") + ); + assert_eq!(pending.attempts, expected_attempts); + assert!(pending.backoff <= RUNTIME_READY_RETRY_MAX); + } + } +} diff --git a/crates/mvp-system/src/distribution_stack.rs b/crates/mvp-system/src/distribution_stack.rs index 7c7ea67..36ad73f 100644 --- a/crates/mvp-system/src/distribution_stack.rs +++ b/crates/mvp-system/src/distribution_stack.rs @@ -204,6 +204,18 @@ impl DistributionRuntimeStack { .filter(|entry| entry.state == MemberState::Alive) .count() } + + pub fn member_state(&self, node_id: NodeId) -> Option { + self.membership_mirror + .lock() + .ok()? + .get(&node_id) + .map(|entry| entry.state) + } + + pub fn route_owner(&self, actor: ActorAddress) -> Option { + self.route_view.read().ok()?.get(&actor).copied() + } } struct MembershipFanout { @@ -227,3 +239,41 @@ impl ActorInterface for MembershipFanout { let _ = ctx.send(self.directory, DirectoryIn::Membership(change)); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn distribution_stack_reports_member_state_and_route_owner() { + let stack = DistributionRuntimeStack::new( + distribution::types::NodeId([1; 32]), + DistributedNodeConfig::default(), + ); + let remote = distribution::types::NodeId([2; 32]); + + stack + .runtime + .send_to( + stack.actors.membership_fanout, + distribution::swim::actor::MembershipChanged { + node_id: remote, + state: MemberState::Alive, + incarnation: 1, + }, + ) + .expect("send membership change"); + stack.pump_runtime_once(); + assert_eq!(stack.member_state(remote), Some(MemberState::Alive)); + + let actor = ActorAddress::new_random(); + { + let mut route_view = match stack.route_view.write() { + Ok(route_view) => route_view, + Err(poisoned) => poisoned.into_inner(), + }; + route_view.insert(actor, remote); + } + assert_eq!(stack.route_owner(actor), Some(remote)); + } +} diff --git a/crates/mvp-system/src/vastai_provisioning.rs b/crates/mvp-system/src/vastai_provisioning.rs index d831a97..ab59408 100644 --- a/crates/mvp-system/src/vastai_provisioning.rs +++ b/crates/mvp-system/src/vastai_provisioning.rs @@ -387,7 +387,7 @@ pub trait VastAiBootstrapLauncher: Send { #[derive(Clone)] enum SshBootstrapMsg { - Stop { reason: BootstrapStopReason }, + Stop, } struct SshBootstrapActor { @@ -407,7 +407,7 @@ impl ActorInterface for SshBootstrapActor { fn handle(&mut self, _ctx: &Ctx, msg: Self::Incoming) { match msg { - SshBootstrapMsg::Stop { reason: _ } => { + SshBootstrapMsg::Stop => { self.stopping.store(true, Ordering::SeqCst); stop_ssh_child(&self.child); } @@ -481,10 +481,8 @@ impl VastAiBootstrapLauncher for SshCommandBootstrapLauncher { }) } - fn stop_bootstrap(&mut self, handle: &mut Self::Handle, reason: BootstrapStopReason) { - let _ = handle - .runtime - .send_to(handle.actor, SshBootstrapMsg::Stop { reason }); + fn stop_bootstrap(&mut self, handle: &mut Self::Handle, _reason: BootstrapStopReason) { + let _ = handle.runtime.send_to(handle.actor, SshBootstrapMsg::Stop); handle.runtime.tick(); } } @@ -1051,12 +1049,7 @@ mod tests { .expect("spawn ssh bootstrap actor"); runtime - .send_to( - actor, - SshBootstrapMsg::Stop { - reason: BootstrapStopReason::RuntimeReady, - }, - ) + .send_to(actor, SshBootstrapMsg::Stop) .expect("send stop"); runtime.tick(); diff --git a/crates/mvp-system/tests/local_e2e_cluster.rs b/crates/mvp-system/tests/local_e2e_cluster.rs index 206c513..b9694af 100644 --- a/crates/mvp-system/tests/local_e2e_cluster.rs +++ b/crates/mvp-system/tests/local_e2e_cluster.rs @@ -1,15 +1,20 @@ -use std::path::{Path, PathBuf}; +#![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") => return local_e2e_cluster::run_main(), + Some("cluster-supervisor" | "cluster-relay") => return local_e2e_cluster::run_main(), Some(role) => { eprintln!("unknown MVP_TEST_ROLE={role}"); return ExitCode::from(2); @@ -30,16 +35,20 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() { 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 output = Command::new(current_test_exe()) - .env("MVP_TEST_ROLE", "cluster-supervisor") - .arg("--prompt") - .arg("ping") - .env("MVP_LOCAL_E2E_CLUSTER_IMAGE", IMAGE) - .output() - .expect("run local e2e cluster supervisor"); + let docker = DockerRelayFixture::start(); + let output = docker.run_supervisor("ping"); assert!( output.status.success(), @@ -95,6 +104,17 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() { "{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() @@ -108,8 +128,166 @@ fn local_e2e_cluster_docker_cpu_pipeline_prompt() { "{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() @@ -176,7 +354,3 @@ fn copy_context_entry(source: &Path, dest: &Path) { fn phase(message: &str) { eprintln!("local-e2e-cluster: {message}"); } - -fn current_test_exe() -> PathBuf { - std::env::current_exe().expect("current test exe") -} diff --git a/crates/mvp-system/tests/local_e2e_cluster/Dockerfile b/crates/mvp-system/tests/local_e2e_cluster/Dockerfile index 10b4a19..ddaadfd 100644 --- a/crates/mvp-system/tests/local_e2e_cluster/Dockerfile +++ b/crates/mvp-system/tests/local_e2e_cluster/Dockerfile @@ -3,6 +3,7 @@ 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 && \ @@ -12,9 +13,12 @@ RUN apt-get update && \ ENV DEV=CPU ENV PYTHONDONTWRITEBYTECODE=1 -ENV CARGO_TARGET_DIR=/workspace/target +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 +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 ["sh", "-c", "test_bin=$(find /workspace/target/debug/deps -maxdepth 1 -type f -perm /111 \\( -name 'local_e2e_cluster-*' -o -name 'local-e2e-cluster-*' \\) | head -n 1); exec \"$test_bin\" \"$@\"", "local-e2e-cluster"] +ENTRYPOINT ["/usr/local/bin/local-e2e-cluster"] diff --git a/crates/mvp-system/tests/one_node_chat_e2e.rs b/crates/mvp-system/tests/one_node_chat_e2e.rs index 93cf15c..bca4170 100644 --- a/crates/mvp-system/tests/one_node_chat_e2e.rs +++ b/crates/mvp-system/tests/one_node_chat_e2e.rs @@ -14,12 +14,13 @@ 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 DEFAULT_CONTAINER: &str = "mvp-orchestrator-1-1"; +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 @@ -27,6 +28,9 @@ fn one_node_chat_docker_cuda_e2e() { .args(["mvp-chat"]) .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_PROMPT_MAX_TOKENS", "3") .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::piped()); @@ -60,7 +64,7 @@ fn one_node_chat_docker_cuda_e2e() { if result.is_ok() { result = assert_no_lower_layer_terminal_leaks(&stdout, &stderr); } - assert_container_removed(&root, DEFAULT_CONTAINER); + assert_containers_with_prefix_removed(&root, &container_prefix); if let Err(error) = result { panic!( @@ -84,9 +88,9 @@ fn run_full_flow( }) .map_err(|e| format!("provisioning frames not visible in dashboard: {e}"))?; wait_for_child_or(TEST_WATCHDOG, child, || { - dashboard_has_frame("mvp.worker.weights", "GgufDownloadProgress") + dashboard_has_frame("mvp.worker.weights", "LoadWeightsStarted") }) - .map_err(|e| format!("GGUF download progress not visible in dashboard: {e}"))?; + .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") }) @@ -111,10 +115,13 @@ fn run_full_flow( .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}"))?; - request_child_interrupt(child); + 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 Ctrl-C".to_owned())?; - if status.success() || status.code() == Some(130) || status.signal_name() == Some("SIGINT") { + .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}")) @@ -400,28 +407,6 @@ fn request_child_interrupt(child: &Child) { } } -trait ExitStatusSignalName { - fn signal_name(&self) -> Option<&'static str>; -} - -impl ExitStatusSignalName for std::process::ExitStatus { - fn signal_name(&self) -> Option<&'static str> { - #[cfg(target_os = "linux")] - { - use std::os::unix::process::ExitStatusExt; - match self.signal() { - Some(libc::SIGINT) => Some("SIGINT"), - Some(libc::SIGTERM) => Some("SIGTERM"), - _ => None, - } - } - #[cfg(not(target_os = "linux"))] - { - let _ = self; - None - } - } -} fn require_docker(root: &std::path::Path) { let version = Command::new("docker") @@ -437,29 +422,35 @@ fn require_docker(root: &std::path::Path) { ); } -fn assert_container_removed(root: &std::path::Path, container: &str) { - let output = Command::new("docker") - .current_dir(root) - .args([ - "ps", - "-a", - "--filter", - &format!("name=^{container}$"), - "--format", - "{{.Names}}", - ]) - .output() - .expect("run docker ps"); - assert!( - output.status.success(), - "docker ps failed\nstdout:\n{}\nstderr:\n{}", - String::from_utf8_lossy(&output.stdout), - String::from_utf8_lossy(&output.stderr) - ); - assert!( - String::from_utf8_lossy(&output.stdout).trim().is_empty(), - "container {container} still exists" - ); +fn 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 { diff --git a/crates/mvp-system/tests/support/local_e2e_cluster.rs b/crates/mvp-system/tests/support/local_e2e_cluster.rs index 43d4513..68311bb 100644 --- a/crates/mvp-system/tests/support/local_e2e_cluster.rs +++ b/crates/mvp-system/tests/support/local_e2e_cluster.rs @@ -1,5 +1,6 @@ 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::{ @@ -40,7 +41,8 @@ use mvp_system::observability_surface as obs; use mvp_system::orchestrator_run_fsm as fsm; use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream}; use mvp_system::relay_provisioning::{ - LocalShimRelayProvider, RelayProvider, RelayProvisionRequest, RelayPurpose, + 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; @@ -63,6 +65,11 @@ 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 { @@ -216,7 +223,9 @@ fn runtime_snapshot_interval_from_env() -> Result { pub fn run_main() -> ExitCode { install_signal_handlers(); let args = std::env::args().collect::>(); - let result = if args.iter().any(|arg| arg == "--role=node") { + 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() @@ -236,6 +245,53 @@ pub fn run_main() -> ExitCode { } } +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")] @@ -532,7 +588,8 @@ fn run_supervisor_once( .cloned() .ok_or_else(|| "engine builder did not assign stage 1".to_owned())?; - let self_endpoint_json = serde_json::to_string(&driver.endpoint_addr()) + 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}"))?; @@ -660,8 +717,22 @@ fn run_supervisor_once( 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() { @@ -827,7 +898,8 @@ fn run_supervisor_once( return Err(format!("run failed: {event:?}")); } }, - OrchestratorReport::Snapshot { .. } => {} + OrchestratorReport::NodeRuntimeReady { .. } + | OrchestratorReport::Snapshot { .. } => {} } } @@ -858,6 +930,10 @@ fn run_supervisor_once( }) .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", @@ -871,6 +947,15 @@ fn run_supervisor_once( }, "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, @@ -1132,7 +1217,9 @@ fn run_node_role(args: &[String]) -> Result<(), String> { .flush() .map_err(|e| format!("flush lifecycle stdout: {e}"))?; } - NodeAgentReport::PromptRequested { .. } | NodeAgentReport::Snapshot { .. } => {} + NodeAgentReport::PromptRequested { .. } + | NodeAgentReport::RuntimeReadyAck { .. } + | NodeAgentReport::Snapshot { .. } => {} } } @@ -1203,12 +1290,16 @@ fn run_node_role(args: &[String]) -> Result<(), String> { } fn new_driver(handle: tokio::runtime::Handle) -> Result { - let mut relay_provider = LocalShimRelayProvider; - let relay = relay_provider.provision_relay(RelayProvisionRequest { - run_id: RUN_ID, - purpose: RelayPurpose::Combined, - })?; - let relay_mode = relay_provider.relay_mode(&relay)?; + 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 { @@ -1234,8 +1325,14 @@ fn wait_for_routes( 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() @@ -1244,6 +1341,13 @@ fn wait_for_routes( 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)); } } @@ -2239,20 +2343,26 @@ fn local_docker_spec( 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: 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(), - ), - ], + env, args: vec![ "--role=node".to_owned(), "--logical-node-id".to_owned(), @@ -2308,14 +2418,15 @@ impl docker_provision::DockerCli for LocalE2eDockerCli { } let mut command = Command::new("docker"); - command - .arg("run") - .arg("--rm") - .arg("--add-host") - .arg("host.docker.internal:host-gateway") - .arg("--name") - .arg(&request.container_name) - .arg("-i"); + 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}")); } @@ -2939,6 +3050,29 @@ fn spawn_shutdown_listener() -> Receiver<()> { 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()