Zachery Aaron Shores-Chmielewski
f4dee176c0
Move iroh_driver and the relay binary out of distribution into a dedicated crates/iroh-driver (lib re-exports IrohDriver; relay bin renamed). Remove the node crate and the single-gpu-inference example; drop the docker/datastream demo. Slim pipeline-parallel vastai. Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
1201 lines
44 KiB
Rust
1201 lines
44 KiB
Rust
//! T-actor: `StageActor` per-role + process-bridge component tests.
|
|
//!
|
|
//! All tests run against the stub-mode `pp_tinygrad_worker`. No GPU, no
|
|
//! real model. Test names match TEST_SPEC §6 verbatim.
|
|
//!
|
|
//! Stage 3 landed §6.1 (Common), §6.2 (First, minus defensive drops), and
|
|
//! §6.4 (Last, minus defensive drops). Stage 4 adds §6.3 (Middle), the
|
|
//! Middle property test, and the defensive-drop tests in §6.2 / §6.4.
|
|
//!
|
|
//! Worker stub constants we depend on (defined in `pp_tinygrad_worker.py`):
|
|
//! - hidden bytes per token = STUB_HIDDEN_DIM * BYTES_PER_ELEM = 16 * 2 = 32
|
|
//! - token_id range = [0, STUB_VOCAB_SIZE) = [0, 32)
|
|
//! Tests assert membership, never specific values, so the actor's contract
|
|
//! stays decoupled from the worker's pseudo-hash.
|
|
|
|
use std::collections::HashMap;
|
|
use std::process::Command;
|
|
use std::time::{Duration, Instant};
|
|
|
|
use swactor::actor::{ActorAddress, ActorInterface};
|
|
use swactor::runtime::{Inbox, Runtime, RuntimeConfig};
|
|
|
|
use pipeline_parallel_inference::messages::{
|
|
InferenceRequest, InferenceResponse, NextToken, StageActivation,
|
|
};
|
|
use pipeline_parallel_inference::stage_actor::{
|
|
StageActor, StageActorStatus, StageMsg, StageRole, stub_tokenize_prompt,
|
|
};
|
|
use swactor_process::{ProcessMode, ProcessSpec};
|
|
|
|
const STUB_BYTES_PER_TOKEN: usize = 16 * 2; // STUB_HIDDEN_DIM * BYTES_PER_ELEM
|
|
const STUB_VOCAB_SIZE: u32 = 32;
|
|
|
|
// ─── Helpers ──────────────────────────────────────────────────────────────
|
|
|
|
/// Build a stub-mode worker spec for `(stage, num_stages)`. Caller picks
|
|
/// values matching the role they want to test: `(0, 2)` for First,
|
|
/// `(1, 3)` for Middle, `(N-1, N)` for Last.
|
|
fn worker_spec(stage: u32, num_stages: u32) -> ProcessSpec {
|
|
let mut env = HashMap::new();
|
|
env.insert("STAGE".to_string(), stage.to_string());
|
|
env.insert("NUM_STAGES".to_string(), num_stages.to_string());
|
|
env.insert("PP_WORKER_STUB".to_string(), "1".to_string());
|
|
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,
|
|
}
|
|
}
|
|
|
|
/// Per-role concrete worker spec — keeps tests role-agnostic.
|
|
fn role_worker_spec(role: StageRole) -> ProcessSpec {
|
|
match role {
|
|
StageRole::First => worker_spec(0, 2),
|
|
// The smallest configuration that yields a Middle stage.
|
|
StageRole::Middle => worker_spec(1, 3),
|
|
StageRole::Last => worker_spec(1, 2),
|
|
}
|
|
}
|
|
|
|
/// Spawn a `StageActor` of the given role with sentinel routing
|
|
/// addresses and the supplied status inbox. The status inbox is required
|
|
/// — the `drain_until_ready` helper reads from it.
|
|
fn spawn_role_actor(
|
|
rt: &Runtime,
|
|
role: StageRole,
|
|
status_addr: ActorAddress,
|
|
max_tokens: u32,
|
|
) -> ActorAddress {
|
|
let placeholder = ActorAddress([0; 32]);
|
|
let sender = rt.create_sender();
|
|
let actor = match role {
|
|
StageRole::First => StageActor::first(role_worker_spec(role), sender, placeholder)
|
|
.with_status_addr(status_addr),
|
|
StageRole::Middle => StageActor::middle(role_worker_spec(role), sender, placeholder)
|
|
.with_status_addr(status_addr),
|
|
StageRole::Last => StageActor::last(
|
|
role_worker_spec(role),
|
|
sender,
|
|
placeholder,
|
|
placeholder,
|
|
max_tokens,
|
|
)
|
|
.with_status_addr(status_addr),
|
|
};
|
|
rt.spawn(actor).unwrap()
|
|
}
|
|
|
|
fn tick_until_recv<M: swactor::actor::Message>(
|
|
rt: &Runtime,
|
|
inbox: &Inbox<M>,
|
|
timeout: Duration,
|
|
) -> Option<M> {
|
|
let start = Instant::now();
|
|
while start.elapsed() < timeout {
|
|
rt.tick();
|
|
if let Some(msg) = inbox.try_recv() {
|
|
return Some(msg);
|
|
}
|
|
std::thread::sleep(Duration::from_millis(5));
|
|
}
|
|
None
|
|
}
|
|
|
|
fn drain_until_ready(
|
|
rt: &Runtime,
|
|
status_inbox: &Inbox<StageActorStatus>,
|
|
timeout: Duration,
|
|
) -> u32 {
|
|
let mut saw_started = false;
|
|
let mut pid: Option<u32> = None;
|
|
let start = Instant::now();
|
|
while start.elapsed() < timeout {
|
|
if let Some(status) = tick_until_recv(rt, status_inbox, Duration::from_millis(100)) {
|
|
match status {
|
|
StageActorStatus::ProcessStarted => saw_started = true,
|
|
StageActorStatus::WorkerReady { pid: p } => {
|
|
assert!(saw_started, "WorkerReady should come after ProcessStarted");
|
|
pid = p;
|
|
break;
|
|
}
|
|
other => panic!("unexpected status during startup: {other:?}"),
|
|
}
|
|
}
|
|
}
|
|
pid.expect("worker should report a pid within timeout")
|
|
}
|
|
|
|
fn is_process_alive(pid: u32) -> bool {
|
|
std::fs::metadata(format!("/proc/{pid}")).is_ok()
|
|
}
|
|
|
|
fn dummy_addr(rt: &Runtime) -> ActorAddress {
|
|
// Some message types require an `ActorAddress` field the actor under
|
|
// test doesn't read — point them at an unread inbox so they're real
|
|
// but inert.
|
|
*rt.new_inbox::<DummySink>().unwrap().addr()
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
struct DummySink;
|
|
|
|
// A trivial runtime-survival probe; mirrors the single-GPU pattern.
|
|
#[derive(Clone)]
|
|
struct Ping {
|
|
reply_to: ActorAddress,
|
|
}
|
|
#[derive(Clone, Debug, PartialEq)]
|
|
struct Pong;
|
|
|
|
struct PongActor;
|
|
impl ActorInterface for PongActor {
|
|
type Incoming = Ping;
|
|
type Response = Pong;
|
|
fn handle(&mut self, ctx: &swactor::runtime::Ctx, msg: Ping) {
|
|
let _ = ctx.send(msg.reply_to, Pong);
|
|
}
|
|
}
|
|
|
|
fn assert_runtime_alive(rt: &Runtime) {
|
|
let pong_inbox = rt.new_inbox::<Pong>().unwrap();
|
|
let pong_addr = rt.spawn(PongActor).unwrap();
|
|
rt.send_to(
|
|
pong_addr,
|
|
Ping {
|
|
reply_to: *pong_inbox.addr(),
|
|
},
|
|
)
|
|
.unwrap();
|
|
let pong = tick_until_recv(rt, &pong_inbox, Duration::from_secs(2));
|
|
assert_eq!(pong, Some(Pong), "runtime should still be functional");
|
|
}
|
|
|
|
fn kill_pid(pid: u32) {
|
|
let _ = Command::new("kill").arg("-9").arg(pid.to_string()).status();
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.1 Common — three lifecycle tests, one per role
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
//
|
|
// Stable Rust has no built-in `paste`-style identifier concatenation, so
|
|
// each per-role test is spelled out at the bottom of this section. The
|
|
// shared body lives in a `body_*` helper directly above each cluster of
|
|
// three `#[test]` entry points.
|
|
|
|
fn body_spawns_worker_and_reports_ready(role: StageRole) {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let _addr = spawn_role_actor(&rt, role, *status_inbox.addr(), 8);
|
|
let pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
assert!(
|
|
is_process_alive(pid),
|
|
"stage worker ({role:?}) should be running"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_spawns_worker_and_reports_ready_first() {
|
|
body_spawns_worker_and_reports_ready(StageRole::First);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_spawns_worker_and_reports_ready_middle() {
|
|
body_spawns_worker_and_reports_ready(StageRole::Middle);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_spawns_worker_and_reports_ready_last() {
|
|
body_spawns_worker_and_reports_ready(StageRole::Last);
|
|
}
|
|
|
|
fn body_worker_crash_emits_process_exited(role: StageRole) {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let _addr = spawn_role_actor(&rt, role, *status_inbox.addr(), 8);
|
|
let pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
kill_pid(pid);
|
|
|
|
let status = tick_until_recv(&rt, &status_inbox, Duration::from_secs(5))
|
|
.expect("actor should report ProcessExited after kill");
|
|
match status {
|
|
StageActorStatus::ProcessExited { .. } => {}
|
|
other => panic!("expected ProcessExited, got {other:?}"),
|
|
}
|
|
assert_runtime_alive(&rt);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_worker_crash_emits_process_exited_first() {
|
|
body_worker_crash_emits_process_exited(StageRole::First);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_worker_crash_emits_process_exited_middle() {
|
|
body_worker_crash_emits_process_exited(StageRole::Middle);
|
|
}
|
|
|
|
#[test]
|
|
fn stage_actor_worker_crash_emits_process_exited_last() {
|
|
body_worker_crash_emits_process_exited(StageRole::Last);
|
|
}
|
|
|
|
fn body_stopping_stage_actor_kills_child_worker(role: StageRole) {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let addr = spawn_role_actor(&rt, role, *status_inbox.addr(), 8);
|
|
let pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
assert!(is_process_alive(pid), "worker alive before stop");
|
|
|
|
rt.stop_actor(addr).unwrap();
|
|
|
|
let start = Instant::now();
|
|
let timeout = Duration::from_secs(5);
|
|
while start.elapsed() < timeout {
|
|
rt.tick();
|
|
std::thread::sleep(Duration::from_millis(10));
|
|
if !is_process_alive(pid) {
|
|
break;
|
|
}
|
|
}
|
|
assert!(
|
|
!is_process_alive(pid),
|
|
"worker ({role:?}, pid {pid}) should be dead after actor stop"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn stopping_stage_actor_kills_child_worker_first() {
|
|
body_stopping_stage_actor_kills_child_worker(StageRole::First);
|
|
}
|
|
|
|
#[test]
|
|
fn stopping_stage_actor_kills_child_worker_middle() {
|
|
body_stopping_stage_actor_kills_child_worker(StageRole::Middle);
|
|
}
|
|
|
|
#[test]
|
|
fn stopping_stage_actor_kills_child_worker_last() {
|
|
body_stopping_stage_actor_kills_child_worker(StageRole::Last);
|
|
}
|
|
|
|
/// `SetNeighbors` must redirect outbound traffic. Verified for First by
|
|
/// observing that an `InferenceRequest` lands at the original `next`,
|
|
/// then at a different `next` after a `SetNeighbors` send. We do not
|
|
/// re-verify on Last separately — the actor's SetNeighbors handler is
|
|
/// role-agnostic (every Some-valued field overwrites unconditionally),
|
|
/// and exercising the same field-update code on multiple roles adds no
|
|
/// signal.
|
|
#[test]
|
|
fn set_neighbors_updates_routing_addresses() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let first_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
let second_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *first_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
// First request lands at the original next address.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "first prompt".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let act_a = tick_until_recv(&rt, &first_inbox, Duration::from_secs(5))
|
|
.expect("first request must produce an activation at the original next");
|
|
assert!(act_a.is_prefill);
|
|
|
|
// Repoint next via SetNeighbors. prev_stage and reply_to are unused
|
|
// by First; passing None for them exercises the optional-field
|
|
// contract.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::SetNeighbors {
|
|
prev_stage: None,
|
|
next_stage: Some(*second_inbox.addr()),
|
|
reply_to: None,
|
|
},
|
|
)
|
|
.unwrap();
|
|
|
|
// Second request must land at the new next, not the old one.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "second prompt".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let act_b = tick_until_recv(&rt, &second_inbox, Duration::from_secs(5))
|
|
.expect("after SetNeighbors, request must produce activation at the new next");
|
|
assert!(act_b.is_prefill);
|
|
let leftover = tick_until_recv(&rt, &first_inbox, Duration::from_millis(400));
|
|
assert!(
|
|
leftover.is_none(),
|
|
"after SetNeighbors, no activation should leak back to the previous next address; got {leftover:?}"
|
|
);
|
|
}
|
|
|
|
/// `Reset` clears per-request state so the actor accepts a fresh request
|
|
/// without being respawned. Verified for First: send a request, observe
|
|
/// the activation, send `Reset`, then send a *new* request and observe
|
|
/// its activation. The two activations carry distinct request ids — the
|
|
/// second is produced after Reset cleared the first's pending map.
|
|
#[test]
|
|
fn reset_clears_per_request_state() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "alpha".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let first = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("activation for first request");
|
|
|
|
rt.send_to(addr, StageMsg::Reset).unwrap();
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "beta".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let second = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("activation for second request after Reset");
|
|
|
|
assert!(first.is_prefill && second.is_prefill);
|
|
assert_ne!(
|
|
first.request_id, second.request_id,
|
|
"request ids should differ across submissions"
|
|
);
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.2 First (Stage 3 — non-defensive-drop subset)
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
|
|
#[test]
|
|
fn first_role_inference_request_produces_stage_activation_to_next() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
let prompt = "say hello there";
|
|
let expected_token_count = stub_tokenize_prompt(prompt).len() as u32;
|
|
assert!(expected_token_count > 0);
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: prompt.into(),
|
|
max_tokens: 8,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let activation = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("first-role actor must emit StageActivation");
|
|
|
|
assert!(
|
|
activation.is_prefill,
|
|
"prefill flag must be set on first hop"
|
|
);
|
|
assert_eq!(activation.seq_len, expected_token_count);
|
|
assert_eq!(activation.position, 0);
|
|
assert_eq!(
|
|
activation.hidden.len(),
|
|
expected_token_count as usize * STUB_BYTES_PER_TOKEN,
|
|
"hidden byte length must match seq_len * hidden_dim * 2"
|
|
);
|
|
assert!(
|
|
activation.hidden.iter().any(|&b| b != 0),
|
|
"stub hidden state must be non-trivial"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn first_role_next_token_produces_stage_activation_for_decode_step() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::NextToken(NextToken {
|
|
request_id: 0xCAFE,
|
|
token_id: 7,
|
|
position: 11,
|
|
done: false,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let activation = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("decode-step activation");
|
|
|
|
assert!(!activation.is_prefill);
|
|
assert_eq!(activation.seq_len, 1);
|
|
assert_eq!(
|
|
activation.position, 11,
|
|
"decode position must match NextToken"
|
|
);
|
|
assert_eq!(activation.request_id, 0xCAFE);
|
|
assert_eq!(activation.hidden.len(), STUB_BYTES_PER_TOKEN);
|
|
}
|
|
|
|
#[test]
|
|
fn first_role_next_token_with_done_produces_no_further_activations() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::NextToken(NextToken {
|
|
request_id: 1,
|
|
token_id: 7,
|
|
position: 5,
|
|
done: true,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let observed = tick_until_recv(&rt, &activation_inbox, Duration::from_millis(800));
|
|
assert!(
|
|
observed.is_none(),
|
|
"no StageActivation after done=true; got {observed:?}"
|
|
);
|
|
}
|
|
|
|
/// An `InferenceRequest` arriving before the worker reports ready is
|
|
/// dropped silently; subsequent requests post-ready still succeed.
|
|
#[test]
|
|
fn first_role_request_before_ready_is_dropped_without_panic() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
|
|
// Submit before WorkerReady. The runtime is freshly created — the
|
|
// worker subprocess has not started, much less printed its ready
|
|
// line, so `ready` is false and the actor must drop the request.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "submitted too early".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
// After becoming ready, no activation from the early request should
|
|
// arrive — pump for long enough that any spuriously-buffered work
|
|
// would have completed.
|
|
let stale = tick_until_recv(&rt, &activation_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stale.is_none(),
|
|
"request submitted before ready must not produce an activation; got {stale:?}"
|
|
);
|
|
|
|
// A post-ready request still works.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "submitted after ready".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let activation = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("post-ready request must produce an activation");
|
|
assert!(activation.is_prefill);
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.4 Last (Stage 3 — non-defensive-drop subset)
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
|
|
/// One activation with `seq_len=1`; hidden bytes are deterministic from
|
|
/// `seed` so the stub worker's sampling is deterministic too. Bytes are
|
|
/// not all-zero (the stub hash folds them in but its determinism does
|
|
/// not require it).
|
|
fn make_activation(request_id: u64, position: u32, seed: u8) -> StageActivation {
|
|
let mut hidden = vec![0u8; STUB_BYTES_PER_TOKEN];
|
|
for (i, b) in hidden.iter_mut().enumerate() {
|
|
*b = seed.wrapping_add(i as u8);
|
|
}
|
|
StageActivation {
|
|
request_id,
|
|
position,
|
|
hidden,
|
|
seq_len: 1,
|
|
is_prefill: false,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn last_role_activation_produces_next_token_to_first() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
100,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
let activation = make_activation(0x2A, 5, 0x11);
|
|
rt.send_to(addr, StageMsg::Activation(activation)).unwrap();
|
|
|
|
let nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("last-role actor must emit NextToken");
|
|
|
|
assert_eq!(nt.request_id, 0x2A);
|
|
assert_eq!(
|
|
nt.position, 6,
|
|
"next position must be activation.position + seq_len"
|
|
);
|
|
assert!(
|
|
!nt.done,
|
|
"single activation under max_tokens=100 must not be done"
|
|
);
|
|
assert!(
|
|
nt.token_id < STUB_VOCAB_SIZE,
|
|
"token_id must lie in stub vocab range"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn last_role_accumulates_tokens_and_emits_response_on_eos() {
|
|
// Probe: discover what token the stub produces for a specific activation.
|
|
let probe_token = {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
100,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(addr, StageMsg::Activation(make_activation(1, 0, 0x42)))
|
|
.unwrap();
|
|
let nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("probe should observe a NextToken");
|
|
rt.stop_actor(addr).unwrap();
|
|
nt.token_id
|
|
};
|
|
|
|
// Real: configure EOS at the observed token, replay the activation,
|
|
// expect done + InferenceResponse.
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
100,
|
|
)
|
|
.with_eos_token_id(probe_token)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(addr, StageMsg::Activation(make_activation(1, 0, 0x42)))
|
|
.unwrap();
|
|
|
|
let nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("should receive NextToken on EOS step");
|
|
assert_eq!(nt.token_id, probe_token);
|
|
assert!(nt.done, "NextToken on EOS must have done=true");
|
|
|
|
let response = tick_until_recv(&rt, &response_inbox, Duration::from_secs(5))
|
|
.expect("InferenceResponse must arrive on EOS");
|
|
assert!(
|
|
!response.text.is_empty(),
|
|
"stub detokenization must produce non-empty text"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn last_role_emits_response_when_max_tokens_reached() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let cap: u32 = 3;
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
cap,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
// The pipeline protocol is strictly sequential — the first stage only
|
|
// emits a new activation after the previous NextToken arrives. Mirror
|
|
// that here: send one, await the reply, send the next.
|
|
let mut saw_done = false;
|
|
for i in 0..cap {
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Activation(make_activation(99, i, 0x10 + i as u8)),
|
|
)
|
|
.unwrap();
|
|
let nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("should receive NextToken per activation");
|
|
let expected_done = i + 1 == cap;
|
|
assert_eq!(
|
|
nt.done,
|
|
expected_done,
|
|
"done should be set exactly on the max_tokens-th NextToken (#{})",
|
|
i + 1
|
|
);
|
|
if nt.done {
|
|
saw_done = true;
|
|
}
|
|
}
|
|
assert!(saw_done);
|
|
|
|
let response = tick_until_recv(&rt, &response_inbox, Duration::from_secs(5))
|
|
.expect("InferenceResponse must arrive when max_tokens is hit");
|
|
assert!(!response.text.is_empty());
|
|
}
|
|
|
|
/// After the actor finalises a request (`done=true` + `InferenceResponse`),
|
|
/// any further `StageActivation` for the same `request_id` must not
|
|
/// duplicate the response. (The finished flag stays set; the trailing
|
|
/// activation is dropped silently.)
|
|
#[test]
|
|
fn last_role_terminate_clears_pending_state() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let cap: u32 = 3;
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
cap,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
// Drive to termination via max_tokens.
|
|
for i in 0..cap {
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Activation(make_activation(7, i, 0x55 + i as u8)),
|
|
)
|
|
.unwrap();
|
|
let _nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("NextToken per activation");
|
|
}
|
|
let first_response = tick_until_recv(&rt, &response_inbox, Duration::from_secs(5))
|
|
.expect("first InferenceResponse after termination");
|
|
assert!(!first_response.text.is_empty());
|
|
|
|
// Trailing activation with the same request id must not produce
|
|
// another NextToken or another InferenceResponse.
|
|
rt.send_to(addr, StageMsg::Activation(make_activation(7, cap, 0xFF)))
|
|
.unwrap();
|
|
let trailing_nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
trailing_nt.is_none(),
|
|
"no NextToken after termination; got {trailing_nt:?}"
|
|
);
|
|
let trailing_resp = tick_until_recv(&rt, &response_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
trailing_resp.is_none(),
|
|
"no duplicate InferenceResponse after termination; got {trailing_resp:?}"
|
|
);
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.2 First — defensive drop (Stage 4)
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
|
|
/// First never receives activations on the wire (they flow forward, not
|
|
/// back to the entry stage), but the swactor inbox is single-typed and the
|
|
/// network may technically deliver one. The actor must drop without
|
|
/// emitting an activation onward.
|
|
#[test]
|
|
fn first_role_drops_activation_messages_defensively() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::first(worker_spec(0, 2), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
// Construct a syntactically valid StageActivation. First should never
|
|
// forward it nor crash the worker.
|
|
rt.send_to(addr, StageMsg::Activation(make_activation(0xF1, 0, 0x33)))
|
|
.unwrap();
|
|
|
|
let stray_act = tick_until_recv(&rt, &activation_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stray_act.is_none(),
|
|
"First must not emit a StageActivation in response to an inbound activation; got {stray_act:?}"
|
|
);
|
|
let stray_nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_millis(100));
|
|
assert!(
|
|
stray_nt.is_none(),
|
|
"First emits no NextToken from a stray Activation"
|
|
);
|
|
|
|
// Subsequent real request still works — drop did not poison anything.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "still works".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
let act = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("post-drop InferenceRequest must still produce an activation");
|
|
assert!(act.is_prefill);
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.3 Middle (Stage 4)
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
|
|
/// Activation builder for Middle tests. `seq_len` is arbitrary; hidden
|
|
/// length is exactly `seq_len * STUB_BYTES_PER_TOKEN` so the worker's
|
|
/// length-check passes.
|
|
fn make_middle_activation(
|
|
request_id: u64,
|
|
position: u32,
|
|
seq_len: u32,
|
|
is_prefill: bool,
|
|
seed: u8,
|
|
) -> StageActivation {
|
|
let total = seq_len as usize * STUB_BYTES_PER_TOKEN;
|
|
let mut hidden = vec![0u8; total];
|
|
for (i, b) in hidden.iter_mut().enumerate() {
|
|
*b = seed.wrapping_add(i as u8);
|
|
}
|
|
StageActivation {
|
|
request_id,
|
|
position,
|
|
hidden,
|
|
seq_len,
|
|
is_prefill,
|
|
}
|
|
}
|
|
|
|
/// Spawn a `Middle` actor wired to a caller-supplied `next_stage` inbox
|
|
/// and return its address. Lets multiple Middle tests share boot.
|
|
fn spawn_middle_actor(
|
|
rt: &Runtime,
|
|
next_stage: ActorAddress,
|
|
status_addr: ActorAddress,
|
|
) -> ActorAddress {
|
|
let sender = rt.create_sender();
|
|
let actor =
|
|
StageActor::middle(worker_spec(1, 3), sender, next_stage).with_status_addr(status_addr);
|
|
rt.spawn(actor).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn middle_role_activation_produces_activation_to_next() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let addr = spawn_middle_actor(&rt, *activation_inbox.addr(), *status_inbox.addr());
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
let inbound = make_middle_activation(0xAB, 5, 3, true, 0x21);
|
|
rt.send_to(addr, StageMsg::Activation(inbound.clone()))
|
|
.unwrap();
|
|
|
|
let outbound = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("middle-role actor must emit an outbound StageActivation");
|
|
|
|
assert_eq!(outbound.request_id, inbound.request_id);
|
|
assert_eq!(outbound.position, inbound.position);
|
|
assert_eq!(outbound.seq_len, inbound.seq_len);
|
|
assert_eq!(outbound.is_prefill, inbound.is_prefill);
|
|
assert_eq!(
|
|
outbound.hidden.len(),
|
|
inbound.seq_len as usize * STUB_BYTES_PER_TOKEN,
|
|
"outbound hidden byte length must match seq_len * hidden_dim * 2"
|
|
);
|
|
}
|
|
|
|
/// Middle is a pure pass-through for the four control fields; only
|
|
/// `hidden` is allowed to change. Hand-rolled property fuzz: iterate
|
|
/// across a handful of distinct `(request_id, position, seq_len,
|
|
/// is_prefill)` tuples chosen to span the cases the orchestrator drives
|
|
/// (prefill vs decode, varied positions, varied seq lengths).
|
|
#[test]
|
|
fn middle_role_preserves_control_fields_property() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let addr = spawn_middle_actor(&rt, *activation_inbox.addr(), *status_inbox.addr());
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
// Distinct request ids so pending-fwd entries don't collide. The
|
|
// 4-tuple sweep covers prefill (seq_len > 1) and decode (seq_len = 1)
|
|
// at varied positions, with both is_prefill polarities.
|
|
let cases: &[(u64, u32, u32, bool, u8)] = &[
|
|
(1, 0, 1, true, 0x01),
|
|
(2, 0, 4, true, 0x02),
|
|
(3, 4, 1, false, 0x03),
|
|
(4, 11, 1, false, 0x04),
|
|
(5, 0, 8, true, 0x05),
|
|
(6, 32, 1, false, 0x06),
|
|
(7, 1, 2, true, 0x07),
|
|
(8, 99, 1, false, 0x08),
|
|
];
|
|
|
|
for (rid, position, seq_len, is_prefill, seed) in cases.iter().copied() {
|
|
let inbound = make_middle_activation(rid, position, seq_len, is_prefill, seed);
|
|
rt.send_to(addr, StageMsg::Activation(inbound.clone()))
|
|
.unwrap();
|
|
|
|
let outbound = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("middle must emit an outbound activation per inbound activation");
|
|
|
|
assert_eq!(
|
|
outbound.request_id, inbound.request_id,
|
|
"rid changed (in {inbound:?})"
|
|
);
|
|
assert_eq!(
|
|
outbound.position, inbound.position,
|
|
"position changed (in {inbound:?})"
|
|
);
|
|
assert_eq!(
|
|
outbound.seq_len, inbound.seq_len,
|
|
"seq_len changed (in {inbound:?})"
|
|
);
|
|
assert_eq!(
|
|
outbound.is_prefill, inbound.is_prefill,
|
|
"is_prefill changed (in {inbound:?})"
|
|
);
|
|
assert_eq!(
|
|
outbound.hidden.len(),
|
|
inbound.seq_len as usize * STUB_BYTES_PER_TOKEN,
|
|
);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn middle_role_drops_inference_requests_defensively() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let addr = spawn_middle_actor(&rt, *activation_inbox.addr(), *status_inbox.addr());
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: dummy_addr(&rt),
|
|
prompt: "should be dropped".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let stray = tick_until_recv(&rt, &activation_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stray.is_none(),
|
|
"Middle must not emit an activation in response to InferenceRequest; got {stray:?}"
|
|
);
|
|
|
|
// Verify the drop did not poison the actor: a real activation still
|
|
// produces a real outbound.
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Activation(make_middle_activation(0xCC, 0, 2, true, 0xAA)),
|
|
)
|
|
.unwrap();
|
|
let outbound = tick_until_recv(&rt, &activation_inbox, Duration::from_secs(5))
|
|
.expect("post-drop activation must still produce an outbound");
|
|
assert_eq!(outbound.request_id, 0xCC);
|
|
}
|
|
|
|
#[test]
|
|
fn middle_role_drops_next_tokens_defensively() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
|
|
let addr = spawn_middle_actor(&rt, *activation_inbox.addr(), *status_inbox.addr());
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::NextToken(NextToken {
|
|
request_id: 0x9F,
|
|
token_id: 3,
|
|
position: 7,
|
|
done: false,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let stray = tick_until_recv(&rt, &activation_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stray.is_none(),
|
|
"Middle must not emit an activation in response to NextToken; got {stray:?}"
|
|
);
|
|
}
|
|
|
|
/// Middle never reaches the sampler / detokenizer — it has no concept of
|
|
/// EOS or max_tokens. Even a stream of activations that would trip EOS on
|
|
/// Last produces zero `InferenceResponse`s here.
|
|
#[test]
|
|
fn middle_role_does_not_emit_inference_response() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let activation_inbox = rt.new_inbox::<StageActivation>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
// Set reply_to too — proves Middle won't write to it even when given
|
|
// an address. (Middle's constructor doesn't take reply_to, so we
|
|
// route through SetNeighbors to be thorough.)
|
|
let actor = StageActor::middle(worker_spec(1, 3), sender, *activation_inbox.addr())
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::SetNeighbors {
|
|
prev_stage: None,
|
|
next_stage: None,
|
|
reply_to: Some(*response_inbox.addr()),
|
|
},
|
|
)
|
|
.unwrap();
|
|
|
|
// Drive several activations and an InferenceRequest. None should
|
|
// produce an InferenceResponse.
|
|
for i in 0..4u32 {
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Activation(make_middle_activation(
|
|
0xD0 + i as u64,
|
|
i,
|
|
1,
|
|
false,
|
|
0x40 + i as u8,
|
|
)),
|
|
)
|
|
.unwrap();
|
|
}
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: *response_inbox.addr(),
|
|
prompt: "should be dropped".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
// Pump the runtime long enough to be sure any spuriously-buffered
|
|
// response would have arrived.
|
|
let resp = tick_until_recv(&rt, &response_inbox, Duration::from_millis(800));
|
|
assert!(
|
|
resp.is_none(),
|
|
"Middle must never emit InferenceResponse; got {resp:?}"
|
|
);
|
|
}
|
|
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
// §6.4 Last — defensive drops (Stage 4)
|
|
// ═══════════════════════════════════════════════════════════════════════
|
|
|
|
#[test]
|
|
fn last_role_drops_inference_requests_defensively() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
8,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::Inference(InferenceRequest {
|
|
reply_to: *response_inbox.addr(),
|
|
prompt: "should be dropped".into(),
|
|
max_tokens: 4,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let stray_nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stray_nt.is_none(),
|
|
"Last must not emit NextToken in response to InferenceRequest; got {stray_nt:?}"
|
|
);
|
|
let stray_resp = tick_until_recv(&rt, &response_inbox, Duration::from_millis(100));
|
|
assert!(
|
|
stray_resp.is_none(),
|
|
"Last must not emit InferenceResponse in response to InferenceRequest; got {stray_resp:?}"
|
|
);
|
|
|
|
// Drop did not poison the actor — a real activation still works.
|
|
rt.send_to(addr, StageMsg::Activation(make_activation(1, 0, 0x77)))
|
|
.unwrap();
|
|
let nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_secs(5))
|
|
.expect("post-drop activation must still produce a NextToken");
|
|
assert_eq!(nt.request_id, 1);
|
|
}
|
|
|
|
#[test]
|
|
fn last_role_drops_next_tokens_defensively() {
|
|
let rt = Runtime::new(RuntimeConfig::default());
|
|
let status_inbox = rt.new_inbox::<StageActorStatus>().unwrap();
|
|
let next_token_inbox = rt.new_inbox::<NextToken>().unwrap();
|
|
let response_inbox = rt.new_inbox::<InferenceResponse>().unwrap();
|
|
|
|
let sender = rt.create_sender();
|
|
let actor = StageActor::last(
|
|
worker_spec(1, 2),
|
|
sender,
|
|
*next_token_inbox.addr(),
|
|
*response_inbox.addr(),
|
|
8,
|
|
)
|
|
.with_status_addr(*status_inbox.addr());
|
|
let addr = rt.spawn(actor).unwrap();
|
|
let _pid = drain_until_ready(&rt, &status_inbox, Duration::from_secs(5));
|
|
|
|
rt.send_to(
|
|
addr,
|
|
StageMsg::NextToken(NextToken {
|
|
request_id: 0xA1,
|
|
token_id: 4,
|
|
position: 6,
|
|
done: false,
|
|
}),
|
|
)
|
|
.unwrap();
|
|
|
|
let stray_nt = tick_until_recv(&rt, &next_token_inbox, Duration::from_millis(500));
|
|
assert!(
|
|
stray_nt.is_none(),
|
|
"Last must not echo NextTokens it receives back onto the wire; got {stray_nt:?}"
|
|
);
|
|
let stray_resp = tick_until_recv(&rt, &response_inbox, Duration::from_millis(100));
|
|
assert!(
|
|
stray_resp.is_none(),
|
|
"Last must not emit InferenceResponse in response to NextToken; got {stray_resp:?}"
|
|
);
|
|
}
|