//! T-integration: in-process end-to-end pipeline-parallel inference. //! //! TEST_SPEC §11. One `DistributedNode` per orchestrator + one per pipeline //! stage, all in the same test process, connected over real iroh QUIC in //! `RelayMode::Disabled`. Each stage runs the stub-mode `pp_tinygrad_worker.py` //! (no GPU, no GGUF). The pipeline runs the autoregressive loop end-to-end: //! orchestrator submits an `InferenceRequest`, stage 0 produces a //! `StageActivation` for prefill, each middle stage echoes activation control //! fields forward, the last stage samples a token and sends `NextToken` back //! to stage 0, and so on until EOS or `max_tokens`. The last stage emits the //! final `InferenceResponse` back to the orchestrator. //! //! The same shape the binary uses (see `pp_orchestrator.rs`), minus the child //! subprocesses — actors live in this test process and addresses are wired up //! directly without SWIM resolution. use std::collections::HashMap; use std::io::{BufRead, BufReader, Write}; use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; use std::sync::{Arc, Mutex, OnceLock}; use std::time::{Duration, Instant}; /// Process-wide lock that serialises pipeline lifetimes inside this /// test binary. The cargo default test-thread pool is `num_cpus`, and /// with each pipeline owning N+1 iroh drivers probing at /// `probe_interval=1` tick, the resulting fan-out (up to ~80 /// driver-pump loops at the default parallelism) reliably starved /// SWIM convergence at N=4/5 well past any reasonable cap. A /// concurrency limit of 1 (full serialisation of pipeline lifetimes) /// is the only setting that converges deterministically on this /// box across back-to-back full-suite runs — earlier values (3, 2) /// still flaked on N=4 cases under load. `Pipeline` holds the guard /// for its whole lifetime so cluster build *and* post-build /// inference work both run with exclusive access. static PIPELINE_LOCK: OnceLock> = OnceLock::new(); fn pipeline_lock() -> std::sync::MutexGuard<'static, ()> { PIPELINE_LOCK .get_or_init(|| Mutex::new(())) .lock() .unwrap_or_else(|poisoned| poisoned.into_inner()) } use distribution::iroh_driver::IrohDriverConfig; use distribution::node::DistributedNodeConfig; use distribution::registry::RegistryConfig; use distribution::swim::probe::SwimConfig; use iroh::{PublicKey, RelayMode}; use swactor::actor::ActorAddress; use swactor::runtime::{Inbox, Runtime}; use swactor_transport::CodecRegistry; use pipeline_parallel_inference::cluster::ClusterNode; use pipeline_parallel_inference::iroh_transport::{ ActorMessagePump, IrohActorTransport, ACTOR_ALPN, }; use pipeline_parallel_inference::messages::{ inference_codec_registry, InferenceRequest, InferenceResponse, NextToken, }; use pipeline_parallel_inference::stage_actor::{ ActivationBridge, NextTokenBridge, RequestBridge, StageActor, StageActorStatus, StageMsg, StageRole, }; use swactor_process::{ProcessMode, ProcessSpec}; // ─── Node / cluster helpers ───────────────────────────────────────────── fn test_node_config() -> DistributedNodeConfig { DistributedNodeConfig { // SWIM timing matches t_cluster's proven values: the tokio actor-bridge // engine drives probes via the test's pump loop, so the original 30ms // probe-timeout was too tight for N>=4 localhost convergence under load // (probes timed out before the ack was pumped). These are test-only knobs // — production uses SwimConfig::default(). swim: SwimConfig { probe_interval: Duration::from_millis(20), probe_timeout: Duration::from_millis(200), indirect_probes: 1, suspicion_timeout: Duration::from_millis(500), dead_reprobe_interval: Duration::ZERO, ..SwimConfig::default() }, cache_capacity: 100, registry: RegistryConfig::default(), metadata_lambda: 3, } } fn make_node() -> ClusterNode { ClusterNode::new( IrohDriverConfig { secret_key: None, relay_mode: RelayMode::Disabled, node: test_node_config(), peer_auth: None, additional_alpns: vec![ACTOR_ALPN.to_vec()], }, test_node_config(), inference_codec_registry(), |_| {}, ) .expect("failed to create cluster node") } fn pubkey_of(node: &ClusterNode) -> PublicKey { PublicKey::from_bytes(&node.node_id().0).unwrap() } fn sees_alive(node: &ClusterNode, peer_key: &PublicKey) -> bool { let peer = distribution::types::NodeId(*peer_key.as_bytes()); node.sees_alive(&peer) } /// Build a `num_stages + 1`-node cluster: node `[0]` is the orchestrator, /// nodes `[1..=num_stages]` are pipeline stages `0..num_stages-1`. Every /// non-orchestrator node joins via the orchestrator's seed address. Returns /// once every node sees every other node alive, or panics on timeout. /// /// Callers are expected to already hold the `PIPELINE_LOCK` (acquired /// by `build_pipeline`/`build_real_pipeline`), so SWIM convergence /// runs with exclusive access to the localhost network stack. fn make_cluster(num_stages: u32) -> Vec { assert!(num_stages >= 2, "pipeline tests require num_stages >= 2"); let total = num_stages as usize + 1; let mut nodes: Vec = (0..total).map(|_| make_node()).collect(); let seed = nodes[0].endpoint_addr(); for d in nodes.iter_mut().skip(1) { d.join(&[seed.clone()]); } let keys: Vec = nodes.iter().map(pubkey_of).collect(); let timeout = Duration::from_secs(60); let start = Instant::now(); let mut converged = false; while start.elapsed() < timeout { for n in nodes.iter_mut() { n.pump_once(); } let all_see_all = nodes.iter().enumerate().all(|(i, n)| { keys.iter() .enumerate() .all(|(j, k)| i == j || sees_alive(n, k)) }); if all_see_all { converged = true; break; } std::thread::sleep(Duration::from_millis(20)); } assert!( converged, "N={num_stages} cluster did not converge within {}s", timeout.as_secs(), ); nodes } // ─── Worker spec (stub mode) ────────────────────────────────────────────── fn worker_spec(stage: u32, num_stages: u32) -> ProcessSpec { let mut env = HashMap::new(); env.insert("STAGE".into(), stage.to_string()); env.insert("NUM_STAGES".into(), num_stages.to_string()); env.insert("PP_WORKER_STUB".into(), "1".into()); ProcessSpec { command: "python3".into(), args: vec![format!("{}/pp_tinygrad_worker.py", env!("CARGO_MANIFEST_DIR"))], env, working_dir: None, mode: ProcessMode::Automated, initial_pty_size: None, kill_timeout: Some(Duration::from_secs(2)), stdin_buffer_limit: None, } } fn is_process_alive(pid: u32) -> bool { std::fs::metadata(format!("/proc/{pid}")).is_ok() } fn kill_pid(pid: u32) { let _ = Command::new("kill").arg("-9").arg(pid.to_string()).status(); } // ─── Pipeline harness (N-stage) ─────────────────────────────────────────── /// Fully wired in-process pipeline of `num_stages` stages plus an /// orchestrator. Indexing: `drivers[0]` and `runtimes[0]` belong to the /// orchestrator; `drivers[s+1]` / `runtimes[s+1]` belong to stage `s`. /// `stage_actor_addrs[s]` / `stage_pids[s]` index stages directly. struct Pipeline { num_stages: u32, nodes: Vec, pumps: Vec, codecs: Arc, response_inbox: Inbox, /// Receives a copy of every `NextToken` the Last stage emits, when the /// pipeline was built with `observe_tokens=true`. Lives on the Last /// stage's runtime so the token-observer route is purely local. token_observer_inbox: Option>, inbox_addr: ActorAddress, request_bridge_addr: ActorAddress, stage_actor_addrs: Vec, stage_pids: Vec, /// Pipeline-lock guard held for the pipeline's whole lifetime then /// released on drop. See `PIPELINE_LOCK`. _pipeline_lock: std::sync::MutexGuard<'static, ()>, } /// Build an N-stage stub-worker pipeline. `max_tokens` caps the decode loop; /// `eos` optionally pins an EOS token id on the Last stage; `observe_tokens` /// installs an extra observer inbox on the Last stage's runtime that receives /// a clone of every `NextToken` it emits. fn build_pipeline(num_stages: u32, max_tokens: u32, eos: Option, observe_tokens: bool) -> Pipeline { assert!(num_stages >= 2, "pipeline requires num_stages >= 2"); let last_stage = num_stages - 1; let pipeline_guard = pipeline_lock(); let nodes = make_cluster(num_stages); let codecs = Arc::new(inference_codec_registry()); let pumps: Vec = (0..=num_stages).map(|_| ActorMessagePump::new()).collect(); // Orchestrator-side inbox for the final response. let response_inbox = nodes[0].rt.new_inbox::().unwrap(); let inbox_addr = *response_inbox.addr(); // Per-stage status inbox — `drain_until_ready` blocks until each stage // reports its worker pid. let status_inboxes: Vec> = (0..num_stages) .map(|s| nodes[s as usize + 1].rt.new_inbox::().unwrap()) .collect(); // Optional NextToken observer on the Last stage's runtime. let (token_observer_inbox, observer_addr): (Option>, Option) = if observe_tokens { let inbox = nodes[last_stage as usize + 1].rt.new_inbox::().unwrap(); let addr = *inbox.addr(); (Some(inbox), Some(addr)) } else { (None, None) }; let placeholder = ActorAddress([0; 32]); // Spawn each stage actor on its node's runtime. The First's `next_stage` // address and the Last's `prev_stage` address need addresses from other // runtimes that don't exist yet, so we spawn local placeholders and // patch the addresses via `SetNeighbors` after every actor exists. let mut stage_actor_addrs: Vec = Vec::with_capacity(num_stages as usize); for s in 0..num_stages { let rt = &nodes[s as usize + 1].rt; let sender = rt.create_sender(); let status_addr = *status_inboxes[s as usize].addr(); let actor_addr = if s == 0 { let actor = StageActor::first(worker_spec(s, num_stages), sender, placeholder) .with_status_addr(status_addr); rt.spawn(actor).unwrap() } else if s == last_stage { let mut actor = StageActor::last( worker_spec(s, num_stages), sender, placeholder, inbox_addr, max_tokens, ) .with_status_addr(status_addr); if let Some(eos_id) = eos { actor = actor.with_eos_token_id(eos_id); } if let Some(addr) = observer_addr { actor = actor.with_token_observer(addr); } rt.spawn(actor).unwrap() } else { let actor = StageActor::middle(worker_spec(s, num_stages), sender, placeholder) .with_status_addr(status_addr); rt.spawn(actor).unwrap() }; stage_actor_addrs.push(actor_addr); } // Per-stage bridges. Stage 0 carries Request + NextToken bridges (its // inbound network types); stages 1..N-1 (Middle) and stage N-1 (Last) // each carry an Activation bridge. let request_bridge_addr = nodes[1] .rt .spawn(RequestBridge { target: stage_actor_addrs[0] }) .unwrap(); let nt_bridge_addr = nodes[1] .rt .spawn(NextTokenBridge { target: stage_actor_addrs[0] }) .unwrap(); let mut activation_bridge_addrs: Vec> = vec![None; num_stages as usize]; for s in 1..num_stages { let rt = &nodes[s as usize + 1].rt; let addr = rt .spawn(ActivationBridge { target: stage_actor_addrs[s as usize] }) .unwrap(); activation_bridge_addrs[s as usize] = Some(addr); } // Wire neighbours on every stage actor. for s in 0..num_stages { let rt = &nodes[s as usize + 1].rt; let (next, prev) = if s == last_stage { (None, Some(nt_bridge_addr)) } else { ( activation_bridge_addrs[(s + 1) as usize], None, ) }; rt.send_to( stage_actor_addrs[s as usize], StageMsg::SetNeighbors { prev_stage: prev, next_stage: next, reply_to: None, }, ) .unwrap(); } // Wait for every worker to come up. The ready line for each stage is // delivered locally on that stage's runtime, so per-stage drain is // independent. let mut stage_pids: Vec = Vec::with_capacity(num_stages as usize); for s in 0..num_stages { let rt = &nodes[s as usize + 1].rt; let pid = drain_until_ready(rt, &status_inboxes[s as usize], Duration::from_secs(15)); stage_pids.push(pid); } // Transport routes. The ClusterNode already wired a `TransportRouter` // shared with the protocol actors and pre-installed it on `rt`; we // only need to add per-stage app routes to it. // // orch → stage 0's request bridge: let orch_to_first = Arc::new(IrohActorTransport::new( nodes[0].driver.endpoint().clone(), nodes[1].endpoint_addr(), nodes[0].driver.tokio_handle(), )); nodes[0].transport_router.add_route(request_bridge_addr, orch_to_first); // Stage s (0..N-1) → stage s+1's activation bridge. for s in 0..num_stages - 1 { let rt_idx = s as usize + 1; let next_idx = (s + 1) as usize + 1; let transport = Arc::new(IrohActorTransport::new( nodes[rt_idx].driver.endpoint().clone(), nodes[next_idx].endpoint_addr(), nodes[rt_idx].driver.tokio_handle(), )); nodes[rt_idx].transport_router.add_route( activation_bridge_addrs[(s + 1) as usize].expect("next stage has activation bridge"), transport, ); } // Last stage → stage 0's NextToken bridge AND → orchestrator's response // inbox. let last_idx = last_stage as usize + 1; let last_to_first = Arc::new(IrohActorTransport::new( nodes[last_idx].driver.endpoint().clone(), nodes[1].endpoint_addr(), nodes[last_idx].driver.tokio_handle(), )); let last_to_orch = Arc::new(IrohActorTransport::new( nodes[last_idx].driver.endpoint().clone(), nodes[0].endpoint_addr(), nodes[last_idx].driver.tokio_handle(), )); nodes[last_idx].transport_router.add_route(nt_bridge_addr, last_to_first); nodes[last_idx].transport_router.add_route(inbox_addr, last_to_orch); Pipeline { num_stages, nodes, pumps, codecs, response_inbox, token_observer_inbox, inbox_addr, request_bridge_addr, stage_actor_addrs, stage_pids, _pipeline_lock: pipeline_guard, } } fn drain_until_ready( rt: &Runtime, status_inbox: &Inbox, timeout: Duration, ) -> u32 { let start = Instant::now(); while start.elapsed() < timeout { rt.tick(); if let Some(status) = status_inbox.try_recv() { if let StageActorStatus::WorkerReady { pid } = status { return pid.expect("worker reports a pid"); } } std::thread::sleep(Duration::from_millis(10)); } panic!("worker did not become ready within {:?}", timeout); } impl Pipeline { /// One pass through every node: drain QUIC, pump actor messages into each /// runtime, tick every runtime. fn pump(&mut self) { for n in &mut self.nodes { n.pump_once(); } for (i, pump) in self.pumps.iter().enumerate() { pump.pump(&self.nodes[i].driver, &self.codecs, &self.nodes[i].rt); } } /// Pump only the orchestrator side. Used after one of the stages has been /// shut down; pumping a dead node's runtime contributes nothing and just /// clutters the loop. fn pump_orch_only(&mut self) { self.nodes[0].pump_once(); self.pumps[0].pump(&self.nodes[0].driver, &self.codecs, &self.nodes[0].rt); } fn submit(&self, prompt: &str, max_tokens: u32) { let req = InferenceRequest { reply_to: self.inbox_addr, prompt: prompt.into(), max_tokens, }; self.nodes[0] .rt .send_to(self.request_bridge_addr, req) .unwrap(); } fn await_response(&mut self, timeout: Duration) -> Result { let start = Instant::now(); while start.elapsed() < timeout { self.pump(); if let Some(r) = self.response_inbox.try_recv() { return Ok(r); } std::thread::sleep(Duration::from_millis(20)); } Err(format!( "no InferenceResponse within {:.0}s", timeout.as_secs_f32() )) } /// Await a response OR detect that a stage's node has gone non-alive in the /// orchestrator's SWIM view. With N stages plus self, the orchestrator's /// snapshot (which omits self) lists `num_stages` peers when healthy; /// `alive_count < num_stages` means at least one stage is no longer /// reachable. fn await_response_or_stage_failure( &mut self, timeout: Duration, ) -> Result { let start = Instant::now(); let expected = self.num_stages; while start.elapsed() < timeout { self.pump_orch_only(); if let Some(r) = self.response_inbox.try_recv() { return Ok(r); } let snap = self.nodes[0].snapshot(); if (snap.alive_count as u32) < expected { let dead: Vec<_> = snap .members .iter() .filter(|m| m.state != "alive") .map(|m| format!("{}={}", &m.node_id[..8.min(m.node_id.len())], m.state)) .collect(); return Err(format!( "stage failure detected via SWIM after {:.1}s ({:?})", start.elapsed().as_secs_f32(), dead )); } std::thread::sleep(Duration::from_millis(50)); } Err(format!( "no response and no SWIM failure within {:.0}s", timeout.as_secs_f32() )) } fn shutdown(mut self) { for (s, addr) in self.stage_actor_addrs.iter().enumerate() { let _ = self.nodes[s + 1].rt.stop_actor(*addr); } let deadline = Instant::now() + Duration::from_secs(3); while Instant::now() < deadline { for s in 0..self.num_stages as usize { self.nodes[s + 1].rt.tick(); } if self.stage_pids.iter().all(|p| !is_process_alive(*p)) { break; } std::thread::sleep(Duration::from_millis(20)); } for n in self.nodes.iter_mut() { n.driver.shutdown(); } } fn stage_pid(&self, stage: u32) -> u32 { self.stage_pids[stage as usize] } } /// Parse the stub `detokenize_stub` output `"tokens: [a b c]"` into a vector /// of token ids. The format is owned by `StageActor::detokenize_stub`. fn parse_stub_tokens(text: &str) -> Vec { let inside = text .trim_start_matches("tokens: [") .trim_end_matches(']') .trim(); if inside.is_empty() { return vec![]; } inside .split_whitespace() .map(|s| s.parse::().expect("stub token id parses as u32")) .collect() } // ═══════════════════════════════════════════════════════════════════════ // §11 — In-process N-stage integration tests // ═══════════════════════════════════════════════════════════════════════ /// Full pipeline at N stages: orchestrator submits one `InferenceRequest`, /// the chain drives the autoregressive loop with stub workers, the Last /// stage emits `InferenceResponse` back to the orchestrator within the /// test timeout, and the response text is non-empty. fn n_stage_stub_pipeline_returns_response_case(num_stages: u32) { let mut pipeline = build_pipeline(num_stages, 4, None, false); pipeline.submit("Say hello", 4); let response = pipeline .await_response(Duration::from_secs(60)) .expect("pipeline must produce an InferenceResponse"); assert!( !response.text.is_empty(), "N={num_stages}: response text must be non-empty, got: {:?}", response.text ); let tokens = parse_stub_tokens(&response.text); assert_eq!( tokens.len(), 4, "N={num_stages}: stub response must contain max_tokens=4 token ids: {:?}", response.text ); pipeline.shutdown(); } #[test] fn n_stage_stub_pipeline_returns_response_n_2() { n_stage_stub_pipeline_returns_response_case(2); } #[test] fn n_stage_stub_pipeline_returns_response_n_3() { n_stage_stub_pipeline_returns_response_case(3); } #[test] fn n_stage_stub_pipeline_returns_response_n_4() { n_stage_stub_pipeline_returns_response_case(4); } #[test] fn n_stage_stub_pipeline_returns_response_n_5() { n_stage_stub_pipeline_returns_response_case(5); } /// Configure the Last stage with an EOS token id and verify the decode loop /// terminates with exactly that many accumulated tokens. Done by running the /// pipeline twice with the same prompt: the first run discovers the /// deterministic stub token sequence at this N, the second pins EOS to one /// of those tokens and asserts the response truncates at the matching index. fn decode_loop_terminates_on_stub_eos_case(num_stages: u32) { const PROMPT: &str = "Say hello once"; const PROBE_TOKENS: u32 = 5; const TARGET_N: usize = 3; let probe_tokens = { let mut probe = build_pipeline(num_stages, PROBE_TOKENS, None, false); probe.submit(PROMPT, PROBE_TOKENS); let resp = probe .await_response(Duration::from_secs(60)) .expect("probe pipeline must produce a response"); let tokens = parse_stub_tokens(&resp.text); assert_eq!( tokens.len(), PROBE_TOKENS as usize, "N={num_stages}: probe (no EOS, max_tokens={PROBE_TOKENS}) should yield {PROBE_TOKENS} tokens" ); probe.shutdown(); tokens }; let eos = probe_tokens[TARGET_N - 1]; let mut real = build_pipeline(num_stages, PROBE_TOKENS, Some(eos), false); real.submit(PROMPT, PROBE_TOKENS); let resp = real .await_response(Duration::from_secs(60)) .expect("EOS pipeline must produce a response"); let tokens = parse_stub_tokens(&resp.text); assert_eq!( tokens.len(), TARGET_N, "N={num_stages}: EOS at index {} should truncate response to {TARGET_N} tokens, got {tokens:?}", TARGET_N - 1 ); assert_eq!( *tokens.last().unwrap(), eos, "N={num_stages}: last accumulated token must be the EOS id" ); real.shutdown(); } #[test] fn decode_loop_terminates_on_stub_eos_n_2() { decode_loop_terminates_on_stub_eos_case(2); } #[test] fn decode_loop_terminates_on_stub_eos_n_3() { decode_loop_terminates_on_stub_eos_case(3); } #[test] fn decode_loop_terminates_on_stub_eos_n_4() { decode_loop_terminates_on_stub_eos_case(4); } /// With no EOS configured, the decode loop runs until `max_tokens` and the /// Last stage emits `InferenceResponse` containing exactly `max_tokens` /// tokens. fn decode_loop_terminates_on_max_tokens_case(num_stages: u32) { const MAX: u32 = 4; let mut pipeline = build_pipeline(num_stages, MAX, None, false); pipeline.submit("max tokens stop", MAX); let response = pipeline .await_response(Duration::from_secs(60)) .expect("max_tokens pipeline must produce a response"); let tokens = parse_stub_tokens(&response.text); assert_eq!( tokens.len() as u32, MAX, "N={num_stages}: with no EOS, response must contain exactly max_tokens={MAX} tokens, got {tokens:?}" ); pipeline.shutdown(); } #[test] fn decode_loop_terminates_on_max_tokens_n_2() { decode_loop_terminates_on_max_tokens_case(2); } #[test] fn decode_loop_terminates_on_max_tokens_n_3() { decode_loop_terminates_on_max_tokens_case(3); } #[test] fn decode_loop_terminates_on_max_tokens_n_4() { decode_loop_terminates_on_max_tokens_case(4); } /// Mid-decode failure for the stage at index `victim_stage` in an N=4 /// pipeline. Kills the worker process AND shuts down the stage's iroh /// driver. The orchestrator's await loop must surface a stage failure /// (via SWIM) within the detection window — not hang. fn stage_failure_mid_decode_case(victim_stage: u32) { let num_stages: u32 = 4; let mut pipeline = build_pipeline(num_stages, 64, None, false); pipeline.submit("never finishes naturally", 64); // Let a few rounds happen so we are genuinely mid-decode. let pump_start = Instant::now(); while pump_start.elapsed() < Duration::from_millis(400) { pipeline.pump(); if pipeline.response_inbox.try_recv().is_some() { panic!( "N={num_stages}: pipeline returned a response too quickly to test mid-decode failure" ); } std::thread::sleep(Duration::from_millis(20)); } // Kill the victim stage's worker + driver. Worker death alone leaves // the SWIM node alive (it just stops doing pipeline work); killing the // driver too is what flips it to dead in the orchestrator's view. let victim_pid = pipeline.stage_pid(victim_stage); kill_pid(victim_pid); pipeline.nodes[victim_stage as usize + 1].driver.shutdown(); // SwimConfig probe/suspicion windows are tight in tests → dead within // ~8-10s. Give 20s for indirect probes + gossip propagation. let result = pipeline.await_response_or_stage_failure(Duration::from_secs(20)); assert!( result.is_err(), "orchestrator should NOT receive an InferenceResponse after stage {victim_stage} dies; got {result:?}", ); // Tear down survivors. for s in 0..num_stages { if s != victim_stage { let _ = pipeline.nodes[s as usize + 1].rt.stop_actor(pipeline.stage_actor_addrs[s as usize]); } } let deadline = Instant::now() + Duration::from_secs(2); while Instant::now() < deadline { for s in 0..num_stages { if s != victim_stage { pipeline.nodes[s as usize + 1].rt.tick(); } } let all_dead = (0..num_stages) .filter(|s| *s != victim_stage) .all(|s| !is_process_alive(pipeline.stage_pid(s))); if all_dead { break; } std::thread::sleep(Duration::from_millis(20)); } for (i, n) in pipeline.nodes.iter_mut().enumerate() { // Node at `victim_stage + 1` is already shut down. if i != victim_stage as usize + 1 { n.driver.shutdown(); } } } #[test] fn first_stage_failure_mid_decode_surfaces_as_error_to_orchestrator() { stage_failure_mid_decode_case(0); } #[test] fn middle_stage_failure_mid_decode_surfaces_as_error_to_orchestrator() { // N=4 has middle stages 1 and 2; pick 1 (closest to First). stage_failure_mid_decode_case(1); } #[test] fn last_stage_failure_mid_decode_surfaces_as_error_to_orchestrator() { stage_failure_mid_decode_case(3); } /// Every `NextToken` returning to stage 0 carries the same `request_id` as /// the `StageActivation` that traversed first → middle → middle → last. /// First's `next_request_id` starts at 1 and increments per outbound /// activation, so for a single `InferenceRequest` the round-trip rid is /// stable across every decode step. This test asserts that every observed /// `NextToken` carries the same rid and that the same rid produces a /// non-empty response. #[test] fn activation_request_id_round_trip_through_chain() { const MAX: u32 = 4; const NUM_STAGES: u32 = 4; let mut pipeline = build_pipeline(NUM_STAGES, MAX, None, true); pipeline.submit("rid round trip", MAX); // Drive the pipeline until the final response arrives, draining the // observer inbox on the way. Two observed tokens are enough to prove // "stable across hops"; collect all of them up to max_tokens for a // sharper assertion. let mut observed: Vec = Vec::with_capacity(MAX as usize); let start = Instant::now(); let mut response: Option = None; while response.is_none() && start.elapsed() < Duration::from_secs(60) { pipeline.pump(); if let Some(observer) = pipeline.token_observer_inbox.as_ref() { while let Some(nt) = observer.try_recv() { observed.push(nt); } } if let Some(r) = pipeline.response_inbox.try_recv() { response = Some(r); } std::thread::sleep(Duration::from_millis(20)); } let response = response.expect("N=4: pipeline must produce a response"); assert!(!response.text.is_empty(), "N=4: response text must be non-empty"); // Drain anything that arrived after the response. if let Some(observer) = pipeline.token_observer_inbox.as_ref() { while let Some(nt) = observer.try_recv() { observed.push(nt); } } assert!( !observed.is_empty(), "N=4: must have observed at least one NextToken from Last" ); let first_rid = observed[0].request_id; assert_ne!(first_rid, 0, "N=4: round-trip request_id must be non-zero"); for (i, nt) in observed.iter().enumerate() { assert_eq!( nt.request_id, first_rid, "N=4: NextToken #{i} has request_id {} (expected {first_rid}); rid must be stable across the full first->middle->middle->last->first round-trip", nt.request_id, ); } pipeline.shutdown(); } /// Each `StageActivation` after prefill has `position` exactly one greater /// than the previous, end-to-end through the chain. The Last stage's /// `NextToken.position` equals the inbound `activation.position + /// activation.seq_len`; for decode steps `seq_len == 1`, so successive /// `NextToken.position`s differ by 1 iff the activation positions did. /// Observing `NextToken`s is therefore a faithful proxy for the activation /// position stream. #[test] fn activation_position_advances_one_per_decode_step() { const MAX: u32 = 4; const NUM_STAGES: u32 = 3; const PROMPT: &str = "position advances"; let mut pipeline = build_pipeline(NUM_STAGES, MAX, None, true); pipeline.submit(PROMPT, MAX); let mut observed: Vec = Vec::with_capacity(MAX as usize); let start = Instant::now(); let mut response: Option = None; while response.is_none() && start.elapsed() < Duration::from_secs(60) { pipeline.pump(); if let Some(observer) = pipeline.token_observer_inbox.as_ref() { while let Some(nt) = observer.try_recv() { observed.push(nt); } } if let Some(r) = pipeline.response_inbox.try_recv() { response = Some(r); } std::thread::sleep(Duration::from_millis(20)); } response.expect("N=3: pipeline must produce a response"); if let Some(observer) = pipeline.token_observer_inbox.as_ref() { while let Some(nt) = observer.try_recv() { observed.push(nt); } } assert_eq!( observed.len(), MAX as usize, "N=3: expected {MAX} NextTokens for max_tokens={MAX} run, got {}: {observed:?}", observed.len(), ); // First NextToken's position == prompt_len (set by Last from the // prefill activation). Subsequent positions increment by 1 per decode // step. let prompt_len = PROMPT.split_whitespace().count() as u32; assert_eq!( observed[0].position, prompt_len, "N=3: first NextToken position must equal prompt_len ({prompt_len}); got {}", observed[0].position, ); for (i, win) in observed.windows(2).enumerate() { let prev = &win[0]; let next = &win[1]; assert_eq!( next.position, prev.position + 1, "N=3: position must advance by 1 between decode steps #{i} -> #{}; prev={prev:?}, next={next:?}", i + 1, ); } pipeline.shutdown(); } // ═══════════════════════════════════════════════════════════════════════ // §12 — Sliced-vs-full equivalence (gated) // ═══════════════════════════════════════════════════════════════════════ // // These tests prove that the 2-stage sliced pipeline produces the same // token sequence as a single-process run over the full model on the same // prompt, with argmax sampling. They: // // 1. Spawn a long-lived reference worker (single process, `NUM_STAGES=1`, // real GGUF). It serves a new `generate_full` op that runs the same // unjitted block iteration as the pipeline but over `model.blk[0..N]` // in one shot. // 2. Build a real-mode 3-node pipeline (workers without `--stub`, stage 0 // tokenises via the worker's `tokenize` op, stage 1 mirrors each // sampled `NextToken` to a test-owned observer inbox). // 3. For each prompt, submit through the pipeline, collect `max_tokens` // sampled token ids via the observer, ask the reference worker for // the same prompt, and assert equality. // // Both are `#[ignore]`: the GGUF fetch + per-token CPU forward dominate // wall-clock. Run with `cargo test -- --ignored`. const REAL_READY_TIMEOUT: Duration = Duration::from_secs(180); const REAL_PROMPT_TIMEOUT: Duration = Duration::from_secs(300); const EQUIVALENCE_MAX_TOKENS: u32 = 8; fn real_worker_spec(stage: u32, num_stages: u32) -> ProcessSpec { let mut env = HashMap::new(); env.insert("STAGE".into(), stage.to_string()); env.insert("NUM_STAGES".into(), num_stages.to_string()); env.insert("MODEL".into(), "llama3.2:1b".into()); env.insert("PP_WORKER_STUB".into(), "".into()); ProcessSpec { command: "python3".into(), args: vec![format!("{}/pp_tinygrad_worker.py", env!("CARGO_MANIFEST_DIR"))], env, working_dir: None, mode: ProcessMode::Automated, initial_pty_size: None, kill_timeout: Some(Duration::from_secs(5)), stdin_buffer_limit: None, } } /// Build a real-mode N-stage pipeline with a token observer on the Last /// stage. Same actor wiring as `build_pipeline` but with the /// real-tinygrad worker, `with_real_tokenization` on First (synthetic /// stub ids cannot be embedded against a real GGUF vocab), and a /// minute-scale ready timeout so the GGUF can load on a cold tinygrad /// cache. fn build_real_pipeline(num_stages: u32, max_tokens: u32) -> Pipeline { assert!(num_stages >= 2, "real pipeline requires num_stages >= 2"); let last_stage = num_stages - 1; let pipeline_guard = pipeline_lock(); let nodes = make_cluster(num_stages); let codecs = Arc::new(inference_codec_registry()); let pumps: Vec = (0..=num_stages).map(|_| ActorMessagePump::new()).collect(); let response_inbox = nodes[0].rt.new_inbox::().unwrap(); let inbox_addr = *response_inbox.addr(); let status_inboxes: Vec> = (0..num_stages) .map(|s| nodes[s as usize + 1].rt.new_inbox::().unwrap()) .collect(); let token_observer_inbox = nodes[last_stage as usize + 1] .rt .new_inbox::() .unwrap(); let observer_addr = *token_observer_inbox.addr(); let placeholder = ActorAddress([0; 32]); let mut stage_actor_addrs: Vec = Vec::with_capacity(num_stages as usize); for s in 0..num_stages { let rt = &nodes[s as usize + 1].rt; let sender = rt.create_sender(); let status_addr = *status_inboxes[s as usize].addr(); let actor_addr = if s == 0 { let actor = StageActor::first(real_worker_spec(s, num_stages), sender, placeholder) .with_status_addr(status_addr) .with_real_tokenization(); rt.spawn(actor).unwrap() } else if s == last_stage { let actor = StageActor::last( real_worker_spec(s, num_stages), sender, placeholder, inbox_addr, max_tokens, ) .with_status_addr(status_addr) .with_token_observer(observer_addr); rt.spawn(actor).unwrap() } else { let actor = StageActor::middle(real_worker_spec(s, num_stages), sender, placeholder) .with_status_addr(status_addr); rt.spawn(actor).unwrap() }; stage_actor_addrs.push(actor_addr); } // Bridges: Request + NextToken on stage 0; Activation on every stage // s >= 1. let request_bridge_addr = nodes[1] .rt .spawn(RequestBridge { target: stage_actor_addrs[0] }) .unwrap(); let nt_bridge_addr = nodes[1] .rt .spawn(NextTokenBridge { target: stage_actor_addrs[0] }) .unwrap(); let mut activation_bridge_addrs: Vec> = vec![None; num_stages as usize]; for s in 1..num_stages { let rt = &nodes[s as usize + 1].rt; let addr = rt .spawn(ActivationBridge { target: stage_actor_addrs[s as usize] }) .unwrap(); activation_bridge_addrs[s as usize] = Some(addr); } // Wire neighbours: First/Middle → next stage's activation bridge; // Last → first stage's NextToken bridge. for s in 0..num_stages { let rt = &nodes[s as usize + 1].rt; let (next, prev) = if s == last_stage { (None, Some(nt_bridge_addr)) } else { (activation_bridge_addrs[(s + 1) as usize], None) }; rt.send_to( stage_actor_addrs[s as usize], StageMsg::SetNeighbors { prev_stage: prev, next_stage: next, reply_to: None, }, ) .unwrap(); } // Real workers can take well over a minute to load the GGUF on a cold // tinygrad cache. Pumping every runtime in parallel keeps any one // stage from starving the others' process bridges. let rts_refs: Vec> = nodes.iter().skip(1).map(|n| Arc::clone(&n.rt)).collect(); let rts_borrows: Vec<&Runtime> = rts_refs.iter().map(|a| a.as_ref()).collect(); let mut stage_pids: Vec = Vec::with_capacity(num_stages as usize); for s in 0..num_stages { let label = match StageRole::for_stage(s, num_stages) { StageRole::First => format!("stage-{s} real worker (first)"), StageRole::Middle => format!("stage-{s} real worker (middle)"), StageRole::Last => format!("stage-{s} real worker (last)"), }; let pid = drain_until_ready_pumped( &rts_borrows, &status_inboxes[s as usize], REAL_READY_TIMEOUT, &label, ); stage_pids.push(pid); } drop(rts_borrows); drop(rts_refs); // Transport routes (mirror of build_pipeline). let orch_to_first = Arc::new(IrohActorTransport::new( nodes[0].driver.endpoint().clone(), nodes[1].endpoint_addr(), nodes[0].driver.tokio_handle(), )); nodes[0].transport_router.add_route(request_bridge_addr, orch_to_first); for s in 0..num_stages - 1 { let rt_idx = s as usize + 1; let next_idx = (s + 1) as usize + 1; let transport = Arc::new(IrohActorTransport::new( nodes[rt_idx].driver.endpoint().clone(), nodes[next_idx].endpoint_addr(), nodes[rt_idx].driver.tokio_handle(), )); nodes[rt_idx].transport_router.add_route( activation_bridge_addrs[(s + 1) as usize].expect("next stage has activation bridge"), transport, ); } let last_idx = last_stage as usize + 1; let last_to_first = Arc::new(IrohActorTransport::new( nodes[last_idx].driver.endpoint().clone(), nodes[1].endpoint_addr(), nodes[last_idx].driver.tokio_handle(), )); let last_to_orch = Arc::new(IrohActorTransport::new( nodes[last_idx].driver.endpoint().clone(), nodes[0].endpoint_addr(), nodes[last_idx].driver.tokio_handle(), )); nodes[last_idx].transport_router.add_route(nt_bridge_addr, last_to_first); nodes[last_idx].transport_router.add_route(inbox_addr, last_to_orch); Pipeline { num_stages, nodes, pumps, codecs, response_inbox, token_observer_inbox: Some(token_observer_inbox), inbox_addr, request_bridge_addr, stage_actor_addrs, stage_pids, _pipeline_lock: pipeline_guard, } } /// Tick every runtime in `rts` until `inbox` produces a `WorkerReady`, /// returning the reported pid. Used by `build_real_pipeline` where the /// model load can take much longer than `drain_until_ready`'s 10s budget /// — the swactor process bridge consumes stdout on its host runtime, so /// pumping is required for the ready line to surface. fn drain_until_ready_pumped( rts: &[&Runtime], inbox: &Inbox, timeout: Duration, label: &str, ) -> u32 { let start = Instant::now(); while start.elapsed() < timeout { for rt in rts { rt.tick(); } if let Some(status) = inbox.try_recv() { if let StageActorStatus::WorkerReady { pid } = status { return pid.unwrap_or_else(|| panic!("{label}: WorkerReady with no pid")); } } std::thread::sleep(Duration::from_millis(50)); } panic!("{label} did not become ready within {:?}", timeout); } impl Pipeline { /// Drive the pipeline for one prompt and return the `max_tokens` /// sampled token ids in order. Asserts that the run terminates on /// `done=true` (max_tokens reached) within `timeout`. fn run_real_prompt(&mut self, prompt: &str, max_tokens: u32, timeout: Duration) -> Vec { assert!( self.token_observer_inbox.is_some(), "pipeline has no token observer; built via build_pipeline instead of build_real_pipeline?" ); self.submit(prompt, max_tokens); let mut collected: Vec = Vec::with_capacity(max_tokens as usize); let start = Instant::now(); while collected.len() < max_tokens as usize && start.elapsed() < timeout { self.pump(); if let Some(observer) = self.token_observer_inbox.as_ref() { while let Some(nt) = observer.try_recv() { collected.push(nt.token_id); if nt.done { break; } } } std::thread::sleep(Duration::from_millis(20)); } // Drain the final InferenceResponse so a residual response in the // inbox doesn't leak across prompts when the harness is reused. let drain_until = Instant::now() + Duration::from_secs(10); while Instant::now() < drain_until { self.pump(); if self.response_inbox.try_recv().is_some() { break; } std::thread::sleep(Duration::from_millis(20)); } assert_eq!( collected.len(), max_tokens as usize, "pipeline produced {} tokens for prompt {prompt:?}, expected {max_tokens}", collected.len() ); collected } /// Reset every stage actor so a fresh prompt can be submitted. The /// worker's per-block KV cache is left alone — the next prefill at /// `position=0` rewrites the cache positions it needs. fn reset_for_next_prompt(&mut self) { for (s, addr) in self.stage_actor_addrs.iter().enumerate() { self.nodes[s + 1].rt.send_to(*addr, StageMsg::Reset).unwrap(); } for _ in 0..5 { self.pump(); std::thread::sleep(Duration::from_millis(10)); } } } // ─── Reference worker (single process, full-model generate) ───────────── /// A standalone real-mode worker driven directly over stdin/stdout. We do /// not need the actor/process-bridge infrastructure for the reference — /// it answers one op (`generate_full`) and is reused across prompts. struct ReferenceWorker { child: Child, stdout: BufReader, stdin: ChildStdin, next_rid: u64, } impl ReferenceWorker { fn spawn() -> Self { let mut cmd = Command::new("python3"); cmd.arg(format!( "{}/pp_tinygrad_worker.py", env!("CARGO_MANIFEST_DIR") )); cmd.env("STAGE", "0"); cmd.env("NUM_STAGES", "2"); cmd.env("MODEL", "llama3.2:1b"); cmd.env_remove("PP_WORKER_STUB"); cmd.stdin(Stdio::piped()); cmd.stdout(Stdio::piped()); cmd.stderr(Stdio::inherit()); let mut child = cmd.spawn().expect("spawn reference worker"); let stdout = BufReader::new(child.stdout.take().unwrap()); let stdin = child.stdin.take().unwrap(); let mut me = Self { child, stdout, stdin, next_rid: 1, }; let ready = me.read_json(); assert_eq!( ready.get("status").and_then(|v| v.as_str()), Some("ready"), "reference worker did not report ready: {ready}" ); me } fn read_json(&mut self) -> serde_json::Value { let mut line = String::new(); let n = self .stdout .read_line(&mut line) .expect("reference worker: read stdout"); if n == 0 { panic!("reference worker closed stdout before replying"); } serde_json::from_str(line.trim()) .unwrap_or_else(|e| panic!("reference worker: invalid JSON {line:?}: {e}")) } /// Ask the reference worker for the `max_tokens` argmax-sampled /// continuations of `prompt`. fn generate_full(&mut self, prompt: &str, max_tokens: u32) -> Vec { let rid = self.next_rid; self.next_rid += 1; let req = serde_json::json!({ "op": "generate_full", "request_id": rid, "prompt": prompt, "max_tokens": max_tokens, }); let line = serde_json::to_string(&req).unwrap(); self.stdin.write_all(line.as_bytes()).expect("write"); self.stdin.write_all(b"\n").expect("write newline"); self.stdin.flush().expect("flush"); let reply = self.read_json(); if let Some(err) = reply.get("error").and_then(|v| v.as_str()) { panic!("reference generate_full failed: {err}"); } let tokens = reply .get("tokens") .and_then(|v| v.as_array()) .unwrap_or_else(|| panic!("reference reply missing tokens: {reply}")); tokens .iter() .map(|t| { t.as_u64() .unwrap_or_else(|| panic!("reference token not u64: {t}")) as u32 }) .collect() } } impl Drop for ReferenceWorker { fn drop(&mut self) { let _ = self.child.kill(); let _ = self.child.wait(); } } /// The three "diverse" prompts §12 enumerates: a one-word prompt, a /// well-formed sentence, and a short string with punctuation. Catches /// per-prompt-shape bugs the single-prompt tests cannot — see TEST_SPEC /// §12 for why each shape matters. const DIVERSE_PROMPTS: &[&str] = &[ "Hello", "The quick brown fox jumps over the lazy dog.", "Wait... what?!", ]; /// Drive every prompt through both an `num_stages`-stage sliced pipeline /// and the single-process reference, asserting per-prompt token-id /// equality up to `EQUIVALENCE_MAX_TOKENS`. fn assert_pipeline_matches_reference(num_stages: u32, prompts: &[&str]) { assert!(!prompts.is_empty(), "need at least one prompt"); let mut reference = ReferenceWorker::spawn(); let mut pipeline = build_real_pipeline(num_stages, EQUIVALENCE_MAX_TOKENS); for (i, prompt) in prompts.iter().enumerate() { let ref_tokens = reference.generate_full(prompt, EQUIVALENCE_MAX_TOKENS); assert_eq!( ref_tokens.len() as u32, EQUIVALENCE_MAX_TOKENS, "reference returned wrong token count for prompt #{i} {prompt:?}: {ref_tokens:?}" ); let pipeline_tokens = pipeline.run_real_prompt(prompt, EQUIVALENCE_MAX_TOKENS, REAL_PROMPT_TIMEOUT); assert_eq!( pipeline_tokens, ref_tokens, "N={num_stages}: sliced pipeline diverged from full-model reference on prompt #{i} {prompt:?}\n pipeline: {pipeline_tokens:?}\n reference: {ref_tokens:?}" ); if i + 1 < prompts.len() { pipeline.reset_for_next_prompt(); } } pipeline.shutdown(); } #[ignore] #[test] fn sliced_two_stage_inference_matches_single_node_for_say_hello() { assert_pipeline_matches_reference(2, &["Say hello"]); } #[ignore] #[test] fn sliced_two_stage_inference_matches_single_node_for_three_diverse_prompts() { assert_pipeline_matches_reference(2, DIVERSE_PROMPTS); } #[ignore] #[test] fn sliced_three_stage_inference_matches_single_node_for_say_hello() { assert_pipeline_matches_reference(3, &["Say hello"]); } #[ignore] #[test] fn sliced_three_stage_inference_matches_single_node_for_three_diverse_prompts() { assert_pipeline_matches_reference(3, DIVERSE_PROMPTS); } #[ignore] #[test] fn sliced_four_stage_inference_matches_single_node_for_say_hello() { assert_pipeline_matches_reference(4, &["Say hello"]); } /// The N-invariance property: with argmax sampling the token sequence is /// independent of chain length. Runs the same prompt through pipelines at /// `N ∈ {2, 3, 4}` and asserts every run yields the same token ids. This /// is the test that catches silent state corruption a longer chain would /// introduce — e.g. a Middle stage that subtly mangles `hidden` would pass /// every individual N-vs-reference test if reference and pipeline mangled /// it the same way, but the N=2 run (which has no Middle) would diverge. #[ignore] #[test] fn sliced_n_stage_inference_is_invariant_to_n_for_argmax() { const PROMPT: &str = "Say hello"; let mut reference = ReferenceWorker::spawn(); let ref_tokens = reference.generate_full(PROMPT, EQUIVALENCE_MAX_TOKENS); let mut per_n_tokens: Vec<(u32, Vec)> = Vec::with_capacity(3); for &num_stages in &[2u32, 3, 4] { let mut pipeline = build_real_pipeline(num_stages, EQUIVALENCE_MAX_TOKENS); let tokens = pipeline.run_real_prompt(PROMPT, EQUIVALENCE_MAX_TOKENS, REAL_PROMPT_TIMEOUT); pipeline.shutdown(); per_n_tokens.push((num_stages, tokens)); } for (n, tokens) in &per_n_tokens { assert_eq!( tokens, &ref_tokens, "N={n}: pipeline disagrees with single-node reference\n pipeline: {tokens:?}\n reference: {ref_tokens:?}" ); } let (first_n, first_tokens) = &per_n_tokens[0]; for (n, tokens) in per_n_tokens.iter().skip(1) { assert_eq!( tokens, first_tokens, "N-invariance violated: N={first_n} produced {first_tokens:?} but N={n} produced {tokens:?}" ); } }