From 4851070bab7b59f7ed79de1e5b40d8b02b151925 Mon Sep 17 00:00:00 2001 From: sparkzky Date: Wed, 22 Jul 2026 01:46:42 +0800 Subject: [PATCH] fix(qwen3): proactive cancellation sweep for queued requests (#642) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cancelled requests in deferred/loading/prefilling/active queues were only retired reactively at token emission — by then the full prefill had already burned GPU. Under a disconnect storm, cancelled work consumed scheduler capacity for tens of seconds. Add sweep_cancelled_requests() at the top of each scheduler iteration (after drain, before admission) that checks token_tx.is_closed() across all queues and drops cancelled requests via executor.drop_request(). The 'loading' queue is excluded: drop_request blocks on prefetch DMA wait, which would stall the scheduler thread. Cancelled requests in 'loading' settle naturally via reclaim_ready_prefetch and are swept when they reach 'deferred' on the next iteration. publish_load() moved to after the sweep so load metrics reflect post-sweep reality before the scheduler parks idle. KV remove events are re-drained after the sweep drops requests so a KV-aware router sees freed blocks immediately. TOCTOU race fixed: retain is driven by collected IDs, not by re-checking is_closed() (Codex review). Tests: - 3 CPU-only scheduler tests (deferred, prefilling, active cancel) - GPU integration test (EngineHandle): burst + disconnect → 0/0 in 250ms - HTTP integration test (real server): 8 streaming clients → disconnect → 0/0 in 59ms, follow-up served Closes #642 Signed-off-by: sparkzky --- Cargo.lock | 3 + openinfer-qwen3/Cargo.toml | 2 + openinfer-qwen3/src/scheduler.rs | 185 ++++++++-- openinfer-qwen3/src/scheduler/tests.rs | 203 ++++++++++- openinfer-qwen3/tests/cancellation_sweep.rs | 219 ++++++++++++ openinfer-qwen3/tests/http_cancellation.rs | 372 ++++++++++++++++++++ 6 files changed, 944 insertions(+), 40 deletions(-) create mode 100644 openinfer-qwen3/tests/cancellation_sweep.rs create mode 100644 openinfer-qwen3/tests/http_cancellation.rs diff --git a/Cargo.lock b/Cargo.lock index 20d08baa1..324ef129d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3295,7 +3295,9 @@ dependencies = [ "openinfer-qwen3", "openinfer-sample", "openinfer-vllm-support", + "parking_lot", "rand 0.10.1", + "reqwest", "safetensors", "serde", "serde_json", @@ -4370,6 +4372,7 @@ dependencies = [ "base64 0.22.1", "bytes", "encoding_rs", + "futures-channel", "futures-core", "futures-util", "h2", diff --git a/openinfer-qwen3/Cargo.toml b/openinfer-qwen3/Cargo.toml index 66c7c6bed..dbb0d0e2a 100644 --- a/openinfer-qwen3/Cargo.toml +++ b/openinfer-qwen3/Cargo.toml @@ -47,6 +47,8 @@ test-fixtures = [] [dev-dependencies] openinfer-qwen3 = { workspace = true, features = ["test-fixtures"] } openinfer-vllm-support = { workspace = true } +parking_lot = { workspace = true } +reqwest = { workspace = true, features = ["blocking", "json"] } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } vllm-text = { workspace = true } diff --git a/openinfer-qwen3/src/scheduler.rs b/openinfer-qwen3/src/scheduler.rs index 56df3fac6..c303849e6 100644 --- a/openinfer-qwen3/src/scheduler.rs +++ b/openinfer-qwen3/src/scheduler.rs @@ -477,15 +477,18 @@ fn release_rejected(executor: &mut E, req: &PendingRequest) { // ── Main loop ─────────────────────────────────────────────────────────── -/// Republish live KV occupancy to the load-watch feed. Called once at the top of -/// every loop iteration (before this step admits/allocates), so it reports the -/// resident occupancy *between* steps — the steady-state load a router wants, -/// not a transient in-step peak. Top-of-loop placement guarantees exactly one -/// publish per iteration regardless of which `continue` the step takes, and the -/// post-completion free shows up at the next iteration's top before the loop -/// parks idle. `watch` coalesces (a consumer wakes at most once per step and -/// reads the latest); `send_replace` ignores a dropped receiver, so the -/// scheduler runs whether or not anyone is watching. +/// Republish live KV occupancy to the load-watch feed. Called once per loop +/// iteration, right after the proactive cancellation sweep and before this step +/// admits/allocates, so it reports post-sweep occupancy *between* steps — the +/// steady-state load a router wants, not a transient in-step peak. Placing it +/// after the sweep (not at the very top) guarantees exactly one publish per +/// iteration regardless of which `continue` the step takes, AND that a +/// just-disconnected request batch reads 0/0 before the loop parks idle on the +/// next `blocking_recv` — otherwise the watch would freeze on the stale +/// pre-sweep counts until a fresh request woke the scheduler. `watch` coalesces +/// (a consumer wakes at most once per step and reads the latest); `send_replace` +/// ignores a dropped receiver, so the scheduler runs whether or not anyone is +/// watching. /// `num_waiting_reqs` folds every not-yet-running queue (KV-deferred, /// prefetch-loading, post-control) into one number; the vLLM frontend exports /// it as `num_requests_waiting` and attributes all of it to @@ -507,6 +510,96 @@ fn publish_load( }); } +// ── Proactive cancellation sweep ───────────────────────────────────────── + +/// Remove every request whose [`TokenSink`] reports closed — the frontend +/// cancelled it, the client disconnected, or the whole engine demux is gone +/// — from the per-iteration scheduler queues, releasing its executor-side +/// state (KV blocks, parked prefetch DMA, saved cursor) via +/// [`ModelExecutor::drop_request`]. +/// +/// Called once at the top of each scheduler iteration, after new submissions +/// are drained and before admission, so a cancelled request never burns a +/// prefill (or decode) step. Between iterations no request's KV is being +/// written, so dropping lands at a safe boundary — the same one the reactive +/// retirement paths (`token_tx.send()` failure, the `ContinuePrefill` check) +/// already use, just reached one step earlier. +/// +/// Returns the number of requests dropped (useful for diagnostics). +/// +/// # What is *not* swept +/// +/// `loading` (requests waiting for async KV prefetch) is excluded because +/// `drop_request` blocks on `handle.wait()` until an in-flight H2D/RDMA/SSD +/// prefetch DMA completes, which would stall the scheduler thread and delay +/// unrelated decode steps. A cancelled request in `loading` settles naturally +/// via `reclaim_ready_prefetch`, enters `deferred`, and is swept on the next +/// iteration — no blocking, one-step delay at most. +/// +/// `inflight_prefill_pending` (decode-overlap mode) is deliberately +/// excluded. While a prefill runs on the overlap stream its KV is mid-write +/// and its queue entry is positionally zipped with the in-flight result, so +/// freeing it mid-compute would both corrupt KV and misalign the +/// request↔result pairing. Those requests live at most one step and are +/// retired reactively by the existing checks once the in-flight prefill +/// lands, so omitting them costs at most one extra step. +fn sweep_cancelled_requests( + executor: &mut E, + active: &mut Vec, + deferred: &mut Vec, + prefilling: &mut Vec, + post_control_deferred: Option<&mut Vec>, +) -> usize { + let mut dropped = 0; + dropped += sweep_pending(&mut *executor, prefilling); + dropped += sweep_pending(&mut *executor, deferred); + if let Some(post_control) = post_control_deferred { + dropped += sweep_pending(&mut *executor, post_control); + } + dropped += sweep_active(&mut *executor, active); + if dropped > 0 { + debug!("sweep_cancelled_requests: retired {dropped} cancelled request(s)"); + } + dropped +} + +/// Drop closed requests from a pending queue, releasing their state. Order- +/// preserving (`deferred`/`prefilling` are FIFO; a `swap_remove` sweep would +/// break that invariant). The ids are collected up front so the mutable +/// executor borrow needed for `drop_request` never overlaps the queue borrow. +/// +/// Retain is driven by the collected IDs, NOT by re-checking `is_closed()`. +/// A request can become closed between the two checks (frontend aborts it +/// concurrently); re-checking would remove it from the queue without calling +/// `drop_request`, leaking KV blocks and parked prefetch state. +fn sweep_pending(executor: &mut E, queue: &mut Vec) -> usize { + let dropped: Vec = queue + .iter() + .filter(|req| req.token_tx.is_closed()) + .map(|req| req.request_id) + .collect(); + for id in &dropped { + let _ = executor.drop_request(*id); + } + queue.retain(|req| !dropped.contains(&req.request_id)); + dropped.len() +} + +/// Drop closed requests from the active (decoding) set, releasing state. +/// See [`sweep_pending`] for the borrow/ordering/TOCTOU rationale. +fn sweep_active(executor: &mut E, queue: &mut Vec) -> usize { + let dropped: Vec = queue + .iter() + .filter(|req| req.token_tx.is_closed()) + .map(|req| req.request_id) + .collect(); + for id in &dropped { + let _ = executor.drop_request(*id); + } + queue.retain(|req| !dropped.contains(&req.request_id)); + dropped.len() +} + fn scheduler_loop( mut executor: E, mut submit_rx: mpsc::UnboundedReceiver, @@ -536,20 +629,11 @@ fn scheduler_loop( info!("Scheduler ready"); loop { - publish_load( - load_tx, - kv_total, - &executor, - (active.len() - + prefilling.len() - + inflight_prefill_pending.as_ref().map_or(0, Vec::len)) as u64, - (deferred.len() + loading.len()) as u64, - ); // Flush the prior step's cache changes to a router (no-op unless the - // event feed is on). Top-of-loop, like `publish_load`: one pass per - // iteration regardless of which branch the step takes, at the cost of a - // one-iteration announcement lag the router tolerates. Stores first so a - // block evicted the same step it registered is announced before removed. + // event feed is on). One pass per iteration regardless of which branch + // the step takes, at the cost of a one-iteration announcement lag the + // router tolerates. Stores first so a block evicted the same step it + // registered is announced before removed. if let Some(producer) = kv_producer.as_mut() { producer.emit_stores(executor.take_kv_store_events()); producer.drain_removes(); @@ -582,6 +666,37 @@ fn scheduler_loop( )); next_request_id += 1; } + // Proactive cancellation sweep: retire requests whose consumer is + // already gone (frontend cancel/disconnect) before they burn a + // prefill step. `inflight_prefill_pending` is excluded — see + // `sweep_cancelled_requests`. + let cancelled = sweep_cancelled_requests( + &mut executor, + &mut active, + &mut deferred, + &mut prefilling, + None, + ); + // The top-of-loop `drain_removes()` ran before the sweep, so block + // evictions from just-dropped requests are unpublished. Flush them + // now so a KV-aware router sees freed blocks before the next step. + if cancelled > 0 { + if let Some(producer) = kv_producer.as_mut() { + producer.drain_removes(); + } + } + // Publish live load AFTER the cancellation sweep so the metrics reflect + // post-sweep reality (a just-disconnected burst reads 0/0, not the stale + // pre-sweep counts) before the loop parks idle on the next `blocking_recv`. + publish_load( + load_tx, + kv_total, + &executor, + (active.len() + + prefilling.len() + + inflight_prefill_pending.as_ref().map_or(0, Vec::len)) as u64, + (deferred.len() + loading.len()) as u64, + ); // 2. Reclaim settled prefetches, then offer fresh requests to prefetch. let reserve_floor = admitted_future_blocks(&executor, &active, &prefilling); @@ -747,14 +862,6 @@ fn scheduler_loop_with_lora_control( info!("Scheduler ready with LoRA control"); loop { - publish_load( - load_tx, - kv_total, - &executor, - (active.len() + prefilling.len()) as u64, - (deferred.len() + loading.len() + post_control_deferred.len()) as u64, - ); - // 1. Drain incoming commands. Generation submitted after a pending // control command waits until that control command is handled at idle. while let Ok(command) = command_rx.try_recv() { @@ -766,6 +873,24 @@ fn scheduler_loop_with_lora_control( &mut next_request_id, ); } + // Proactive cancellation sweep: retire cancelled requests before + // admission so they never reach prefill. See `sweep_cancelled_requests`. + sweep_cancelled_requests( + &mut executor, + &mut active, + &mut deferred, + &mut prefilling, + Some(&mut post_control_deferred), + ); + // Publish live load AFTER the cancellation sweep so the metrics reflect + // post-sweep reality before the loop parks idle on the next `blocking_recv`. + publish_load( + load_tx, + kv_total, + &executor, + (active.len() + prefilling.len()) as u64, + (deferred.len() + loading.len() + post_control_deferred.len()) as u64, + ); // 1b. Reclaim settled prefetches and offer fresh requests. Control // commands gate generation, so only offer once no control is pending diff --git a/openinfer-qwen3/src/scheduler/tests.rs b/openinfer-qwen3/src/scheduler/tests.rs index 246549938..a6c209199 100644 --- a/openinfer-qwen3/src/scheduler/tests.rs +++ b/openinfer-qwen3/src/scheduler/tests.rs @@ -1,14 +1,18 @@ use std::collections::HashMap; use std::sync::Arc; -use std::sync::Mutex; +use std::sync::atomic::AtomicU8; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use anyhow::Result; use openinfer_core::engine::EngineControlError; use openinfer_core::engine::LoadLoraAdapterRequest; +use openinfer_core::engine::RequestAbortReason; +use openinfer_core::engine::TokenStreamReceiver; use openinfer_core::engine::UnloadLoraAdapterRequest; use openinfer_kv_cache::BlockPool; +use parking_lot::Mutex; use super::*; use crate::executor::DecodePlan; @@ -147,7 +151,7 @@ impl ModelExecutor for FakeExecutor { self.available_blocks += blocks_needed(tokens, self.block_size); } self.prefill_positions.remove(&request_id); - self.dropped.lock().unwrap().push(request_id.get()); + self.dropped.lock().push(request_id.get()); Ok(()) } @@ -158,7 +162,7 @@ impl ModelExecutor for FakeExecutor { _lora_adapter: Option<&str>, _reserve_floor: usize, ) -> bool { - self.prefetch_offers.lock().unwrap().push(request_id.get()); + self.prefetch_offers.lock().push(request_id.get()); false } @@ -780,12 +784,12 @@ fn echo_requests_are_never_offered_to_prefetch() { let offers = Arc::clone(&executor.prefetch_offers); let mut deferred = vec![pending(1, true), pending(2, false)]; - let mut loading = Vec::new(); + let mut loading: Vec = Vec::new(); offer_prefetch(&mut executor, &mut deferred, &mut loading, 0); // The plain request is probed; the echo request is skipped entirely, so // its prefill forwards the whole prompt without parking unspendable KV. - assert_eq!(*offers.lock().unwrap(), vec![2]); + assert_eq!(*offers.lock(), vec![2]); let echo = deferred.iter().find(|r| r.request_id.get() == 1).unwrap(); assert!(!echo.prefetch_offered, "echo request must stay un-probed"); let plain = deferred.iter().find(|r| r.request_id.get() == 2).unwrap(); @@ -909,10 +913,7 @@ fn decode_error_drops_request_state_and_scheduler_recovers() { _ => panic!("decode failure should surface as TokenEvent::Error"), } assert!( - wait_until(Duration::from_secs(1), || dropped - .lock() - .unwrap() - .contains(&0)), + wait_until(Duration::from_secs(1), || dropped.lock().contains(&0)), "failed request state should be dropped" ); @@ -996,7 +997,7 @@ fn retiring_multiple_active_requests_tolerates_unsorted_indices() { active.is_empty(), "all finished requests should retire without index drift" ); - let mut dropped = dropped.lock().unwrap().clone(); + let mut dropped = dropped.lock().clone(); dropped.sort_unstable(); assert_eq!(dropped, vec![1, 7, 10]); } @@ -1267,3 +1268,185 @@ fn speculative_resolves_each_request_independently() { if *request_id == RequestId(2) )); } + +// ── Proactive cancellation sweep (#642) ─────────────────────────────────── +// +// `TokenSink::standalone()` pins `abort_reason` to `None` forever, so it +// can't model a frontend-driven cancel. These helpers build a sink backed by +// a real shared `Arc` the test trips to an abort reason, then drive +// `sweep_cancelled_requests` directly with the FakeExecutor (CPU-only). + +fn cancellable_sink() -> (TokenSink, Arc, TokenStreamReceiver) { + let abort_reason = Arc::new(AtomicU8::new(RequestAbortReason::None as u8)); + let (tx, rx) = mpsc::unbounded_channel(); + let sink = TokenSink::new(Arc::from("test"), tx, Arc::clone(&abort_reason)); + (sink, abort_reason, rx) +} + +fn cancellable_pending(request_id: u64) -> (PendingRequest, Arc, TokenStreamReceiver) { + let (token_tx, abort_reason, rx) = cancellable_sink(); + ( + PendingRequest { + request_id: RequestId::new(request_id), + lora_adapter: None, + prompt_tokens: vec![1; 32], + params: SamplingParams::default(), + max_tokens: 1, + token_tx, + logprobs: 0, + echo: false, + queued_at_unix_s: None, + prefetch_offered: false, + prefill_pos: 0, + step_chunk: 0, + cached_tokens: 0, + }, + abort_reason, + rx, + ) +} + +fn cancellable_active(request_id: u64) -> (ActiveRequestState, Arc, TokenStreamReceiver) { + let (token_tx, abort_reason, rx) = cancellable_sink(); + ( + ActiveRequestState { + request_id: RequestId::new(request_id), + lora_adapter: None, + token_tx, + last_token: 100, + generated_count: 1, + max_tokens: 4, + prompt_len: 32, + params: SamplingParams::default(), + logprobs: 0, + }, + abort_reason, + rx, + ) +} + +/// A cancelled deferred request is dropped before it ever reaches prefill, +/// and the LoRA-control-only `post_control_deferred` queue (loop2 path) is +/// swept too. Live siblings are left untouched. +#[test] +fn sweep_drops_cancelled_deferred_request_before_prefill() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let mut executor = FakeExecutor::new(64, Arc::clone(&dropped)); + let mut active = Vec::new(); + let mut deferred = Vec::new(); + let _loading: Vec = Vec::new(); + let mut prefilling = Vec::new(); + let mut post_control_deferred = Vec::new(); + + let (cancelled, abort, _cancelled_rx) = cancellable_pending(1); + let (live, _live_abort, _live_rx) = cancellable_pending(2); + // Park a cancelled request in the LoRA-control-only queue as well, to + // cover the loop2 sweep path (the `Some(&mut ..)` branch). + let (post_cancelled, post_abort, _post_rx) = cancellable_pending(3); + deferred.push(cancelled); + deferred.push(live); + post_control_deferred.push(post_cancelled); + + abort.store(RequestAbortReason::Cancelled as u8, Ordering::Release); + post_abort.store(RequestAbortReason::Disconnected as u8, Ordering::Release); + + let n = sweep_cancelled_requests( + &mut executor, + &mut active, + &mut deferred, + &mut prefilling, + Some(&mut post_control_deferred), + ); + + assert_eq!( + n, 2, + "both cancelled requests (deferred + post-control) drop" + ); + let mut got = dropped.lock().clone(); + got.sort_unstable(); + assert_eq!(got, vec![1, 3]); + assert_eq!(deferred.len(), 1); + assert_eq!(deferred[0].request_id.get(), 2, "live request survives"); + assert!(post_control_deferred.is_empty()); + assert!( + prefilling.is_empty(), + "cancelled request must never reach prefilling" + ); +} + +/// A cancelled request mid-chunked-prefill is retired and its executor-side +/// progress state released; a live sibling keeps its place and its state. +#[test] +fn sweep_drops_cancelled_prefilling_request_and_releases_state() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let mut executor = FakeExecutor::new(64, Arc::clone(&dropped)); + let mut active = Vec::new(); + let mut deferred = Vec::new(); + let _loading: Vec = Vec::new(); + let mut prefilling = Vec::new(); + + let (cancelled, abort, _cancelled_rx) = cancellable_pending(1); + let (live, _live_abort, _live_rx) = cancellable_pending(2); + // Both already admitted into prefilling, holding prompt-progress state. + executor.prefill_positions.insert(RequestId::new(1), 16); + executor.prefill_positions.insert(RequestId::new(2), 16); + prefilling.push(cancelled); + prefilling.push(live); + + abort.store(RequestAbortReason::Cancelled as u8, Ordering::Release); + + let n = sweep_cancelled_requests( + &mut executor, + &mut active, + &mut deferred, + &mut prefilling, + None, + ); + + assert_eq!(n, 1); + assert_eq!(*dropped.lock(), vec![1]); + assert_eq!(prefilling.len(), 1); + assert_eq!(prefilling[0].request_id.get(), 2); + assert!(!executor.prefill_positions.contains_key(&RequestId::new(1))); + assert!(executor.prefill_positions.contains_key(&RequestId::new(2))); +} + +/// A cancelled active (decoding) request leaves the running set and its KV +/// reservation is returned; a live sibling keeps decoding. +#[test] +fn sweep_drops_cancelled_active_request_from_running_set() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let mut executor = FakeExecutor::new(64, Arc::clone(&dropped)); + let mut active = Vec::new(); + let mut deferred = Vec::new(); + let _loading: Vec = Vec::new(); + let mut prefilling = Vec::new(); + + let (cancelled, abort, _cancelled_rx) = cancellable_active(1); + let (live, _live_abort, _live_rx) = cancellable_active(2); + executor + .ensure_request_tokens(RequestId::new(1), 32) + .unwrap(); + executor + .ensure_request_tokens(RequestId::new(2), 32) + .unwrap(); + active.push(cancelled); + active.push(live); + + abort.store(RequestAbortReason::Cancelled as u8, Ordering::Release); + + let n = sweep_cancelled_requests( + &mut executor, + &mut active, + &mut deferred, + &mut prefilling, + None, + ); + + assert_eq!(n, 1); + assert_eq!(*dropped.lock(), vec![1]); + assert_eq!(active.len(), 1); + assert_eq!(active[0].request_id.get(), 2); + assert!(!executor.held_tokens.contains_key(&RequestId::new(1))); + assert!(executor.held_tokens.contains_key(&RequestId::new(2))); +} diff --git a/openinfer-qwen3/tests/cancellation_sweep.rs b/openinfer-qwen3/tests/cancellation_sweep.rs new file mode 100644 index 000000000..46494967b --- /dev/null +++ b/openinfer-qwen3/tests/cancellation_sweep.rs @@ -0,0 +1,219 @@ +//! Cancellation-sweep integration test for Qwen3-4B (issue #642). +//! +//! Issue #642 added a proactive cancellation sweep to the Qwen3 scheduler that +//! drops cancelled requests (their `token_tx` closed) before they ever reach +//! prefill. The scheduler also republishes load metrics *after* that sweep, so +//! a disconnected request batch is visible as a collapse to zero on the load +//! watch before the loop parks idle. This test proves both: when a whole batch +//! of clients disconnects mid-flight, the engine's load metrics +//! (`num_running_reqs`, `num_waiting_reqs`) collapse to zero promptly, and the +//! engine still serves a fresh request immediately after. +//! +//! It drives the real engine + `submit` rather than a mocked scheduler, so it +//! exercises the actual send-failure retirement and pre-schedule sweep paths. +//! +//! Requires a CUDA GPU and Qwen3-4B weights; skips cleanly when the model is +//! absent (point `OPENINFER_TEST_MODEL_PATH` at the weights to run it). + +use std::path::Path; +use std::time::Duration; +use std::time::Instant; + +use openinfer_core::engine::EngineHandle; +use openinfer_core::engine::EngineLoadOptions; +use openinfer_core::engine::GenerateRequest; +use openinfer_core::engine::TokenEvent; +use openinfer_core::engine::TokenSink; +use openinfer_core::sampler::SamplingParams; +use vllm_text::tokenizer::DynTokenizer; + +mod common; + +const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3-4B"); + +/// Number of concurrent requests in the disconnect burst. +const BURST_SIZE: usize = 10; +/// Prompt length (tokens) per burst request — long enough that only ~1 prefills +/// per scheduler step (`DEFAULT_MAX_PREFILL_TOKENS` = 1024), so a mix of running +/// and waiting requests is in flight when the disconnect happens. +const PROMPT_TOKENS: usize = 600; +/// Decode budget per burst request — keeps them active long enough to be in +/// flight at disconnect time. +const BURST_MAX_TOKENS: usize = 128; + +fn model_path_or_skip() -> Option { + match std::env::var("OPENINFER_TEST_MODEL_PATH") { + Ok(path) => Some(path), + Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { + Some(MODEL_PATH.to_string()) + } + Err(_) => { + eprintln!( + "skipping qwen3 cancellation_sweep: {MODEL_PATH}/config.json is missing; set OPENINFER_TEST_MODEL_PATH to run it" + ); + None + } + } +} + +/// Submit `prompt` and block until the request finishes; returns the decoded text. +fn generate_text( + handle: &EngineHandle, + tokenizer: &DynTokenizer, + prompt: &str, + max_tokens: usize, +) -> String { + let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); + let (token_tx, mut rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams::default(), + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + data_parallel_rank: None, + }) + .expect("submit failed"); + + let mut tokens = Vec::new(); + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, .. }) => tokens.push(id), + Some(TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. }) => {} + Some(TokenEvent::Finished { .. }) => break, + Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), + Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), + None => panic!("scheduler channel closed without Finished"), + } + } + tokenizer.decode(&tokens, true).expect("decode failed") +} + +/// A mass client disconnect must not leave cancelled requests occupying running +/// or waiting slots: the prefill cancellation sweep (issue #642) drops them +/// before prefill, and the scheduler retires in-flight ones when their sends +/// fail. We assert the live load metrics collapse to zero within a tight bound, +/// then prove the engine still serves a fresh request. +#[test] +fn cancelled_burst_metrics_collapse_promptly() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + + let handle = openinfer_qwen3::start_engine_with_offload( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0], + seed: 42, + ..EngineLoadOptions::default() + }, + openinfer_qwen3::Qwen3OffloadOptions::disabled(), + true, + openinfer_qwen3::DEFAULT_MAX_PREFILL_TOKENS, + openinfer_qwen3::Qwen3MemoryOptions::default(), + openinfer_qwen3::DecodeOverlap::Off, + true, + None, + false, + ) + .expect("failed to start engine"); + let tokenizer = common::load_tokenizer(&model_path); + + // Build a long prompt and trim it to a known token length so each request + // lands in prefill/waiting rather than retiring instantly. + let long_prompt = "The quick brown fox jumps over the lazy dog while a gentle breeze rustles the autumn leaves. ".repeat(64); + let prompt_tokens = { + let mut t = tokenizer + .encode(&long_prompt, false) + .expect("encode failed"); + assert!( + t.len() >= PROMPT_TOKENS, + "base prompt only tokenized to {} tokens; need >= {PROMPT_TOKENS}", + t.len() + ); + t.truncate(PROMPT_TOKENS); + t + }; + + // Submit a burst of concurrent requests, keeping every receiver alive so the + // scheduler admits and starts processing them. + let mut receivers = Vec::with_capacity(BURST_SIZE); + for _ in 0..BURST_SIZE { + let (token_tx, rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens: prompt_tokens.clone(), + params: SamplingParams::default(), + max_tokens: BURST_MAX_TOKENS, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + data_parallel_rank: None, + }) + .expect("submit failed"); + receivers.push(rx); + } + + // Let some requests enter prefilling/active so the disconnect hits real + // in-flight work, not a still-empty queue. + std::thread::sleep(Duration::from_millis(200)); + + let load_rx = handle + .load_watch() + .expect("engine did not wire a load feed"); + let pre_drop = *load_rx.borrow(); + eprintln!( + "[cancellation-sweep] pre-drop: running={} waiting={} kv_used={}", + pre_drop.num_running_reqs, pre_drop.num_waiting_reqs, pre_drop.kv_used_blocks + ); + // The burst is large and slow enough that something must be in flight at + // disconnect time; if not, the test would exercise nothing. + assert!( + pre_drop.num_running_reqs + pre_drop.num_waiting_reqs > 0, + "no burst requests were in flight at disconnect time; test exercised nothing" + ); + + // Mass disconnect: dropping every receiver closes every `token_tx`, so + // `token_tx.is_closed()` becomes true for all requests. + drop(receivers); + + // Poll the live load metrics until both running and waiting collapse to + // zero. The sweep + send-failure retirement (and the post-sweep publish) + // should make this fast. + let collapse_start = Instant::now(); + let deadline = collapse_start + Duration::from_secs(10); + let mut last = *load_rx.borrow(); + while last.num_running_reqs != 0 || last.num_waiting_reqs != 0 { + assert!( + Instant::now() < deadline, + "cancelled-request metrics did not collapse within 10s: running={} waiting={}", + last.num_running_reqs, + last.num_waiting_reqs + ); + std::thread::sleep(Duration::from_millis(50)); + last = *load_rx.borrow(); + } + eprintln!( + "[cancellation-sweep] collapsed to zero in {:.0}ms (kv_used={})", + collapse_start.elapsed().as_secs_f64() * 1000.0, + last.kv_used_blocks + ); + + // The engine must still serve a fresh, live request immediately after the + // sweep cleared the cancelled burst. + let follow_up = generate_text(&handle, &tokenizer, "Hello, how are you today?", 128); + assert!( + !follow_up.is_empty(), + "engine did not serve a follow-up request after cancellation sweep" + ); + eprintln!("[cancellation-sweep] follow-up reply: {follow_up:?}"); +} diff --git a/openinfer-qwen3/tests/http_cancellation.rs b/openinfer-qwen3/tests/http_cancellation.rs new file mode 100644 index 000000000..473048f78 --- /dev/null +++ b/openinfer-qwen3/tests/http_cancellation.rs @@ -0,0 +1,372 @@ +//! HTTP-level integration test for issue #642. +//! +//! `scheduler_robustness.rs` proves the scheduler retires an orphaned request +//! when its *engine* receiver is dropped, but it drives `EngineHandle` directly. +//! @xiaguan asked for the same guarantee through the **real HTTP disconnect +//! path**: a client hangs up on a streaming `/v1/completions` response and the +//! running server must tear that request down. +//! +//! This test starts the actual `openinfer` server binary as a subprocess, +//! opens a burst of concurrent streaming completions, verifies they show up as +//! in-flight in `/metrics`, disconnects every client at once, and polls +//! `/metrics` until running+waiting collapse to zero — then confirms a clean +//! follow-up request is still served. +//! +//! Requires a CUDA GPU, Qwen3-4B weights, and a built release binary at +//! `target/release/openinfer`. It skips cleanly when the model or binary is +//! absent; point `OPENINFER_TEST_MODEL_PATH` at the weights to run it. + +use std::net::TcpListener; +use std::path::Path; +use std::process::Child; +use std::process::Command; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::sync::mpsc; +use std::thread; +use std::time::Duration; +use std::time::Instant; + +use reqwest::blocking::Client; +use serde_json::Value; +use serde_json::json; + +/// Path to the release server binary, relative to this crate's manifest dir. +const SERVER_BIN: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../target/release/openinfer"); + +/// Default model weights location when `OPENINFER_TEST_MODEL_PATH` is unset. +const DEFAULT_MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3-4B"); + +/// How long to wait for the server to finish loading the model and answer +/// `/metrics`. Loading Qwen3-4B on a consumer GPU takes ~30s; leave headroom. +const SERVER_READY_TIMEOUT: Duration = Duration::from_secs(180); + +/// Number of concurrent streaming clients in the disconnect burst. +const NUM_CLIENTS: usize = 8; + +/// Prompt length in repeated words. Long enough that prefill keeps every +/// request firmly in-flight at the moment we disconnect. +const PROMPT_WORDS: usize = 800; + +/// How long to wait for in-flight + queued requests to drain to zero after the +/// mass disconnect. +const COLLAPSE_TIMEOUT: Duration = Duration::from_secs(15); + +/// How long to wait for at least one request to enter the running set. +const INFLIGHT_TIMEOUT: Duration = Duration::from_secs(15); + +/// `/metrics` polling cadence. +const POLL_INTERVAL: Duration = Duration::from_millis(50); + +// --------------------------------------------------------------------------- +// Skip guards +// --------------------------------------------------------------------------- + +fn model_path_or_skip() -> Option { + if let Ok(path) = std::env::var("OPENINFER_TEST_MODEL_PATH") { + return Some(path); + } + if Path::new(DEFAULT_MODEL_PATH).join("config.json").exists() { + return Some(DEFAULT_MODEL_PATH.to_string()); + } + eprintln!( + "skipping http_cancellation: model weights not found at {DEFAULT_MODEL_PATH}; \ + set OPENINFER_TEST_MODEL_PATH to run it" + ); + None +} + +fn server_bin_or_skip() -> Option<&'static str> { + if Path::new(SERVER_BIN).exists() { + Some(SERVER_BIN) + } else { + eprintln!( + "skipping http_cancellation: server binary not found at {SERVER_BIN}; \ + run `cargo build --release -p openinfer` first" + ); + None + } +} + +// --------------------------------------------------------------------------- +// Server lifecycle +// --------------------------------------------------------------------------- + +/// Kills the child on drop so a failing assertion cannot orphan the server. +struct ServerGuard { + child: Child, +} + +impl Drop for ServerGuard { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +/// Bind a loopback socket to grab a free ephemeral port, then release it so the +/// server can bind it. The inherent TOCTOU race is negligible on a test box and +/// this matches the rest of the suite. +fn reserve_loopback_port() -> u16 { + TcpListener::bind(("127.0.0.1", 0)) + .expect("failed to reserve loopback port") + .local_addr() + .expect("failed to read reserved port") + .port() +} + +fn wait_for_server(client: &Client, base_url: &str) -> bool { + let metrics_url = format!("{base_url}/metrics"); + let deadline = Instant::now() + SERVER_READY_TIMEOUT; + while Instant::now() < deadline { + if let Ok(resp) = client + .get(&metrics_url) + .timeout(Duration::from_secs(2)) + .send() + { + if resp.status().is_success() { + return true; + } + } + thread::sleep(Duration::from_millis(500)); + } + false +} + +// --------------------------------------------------------------------------- +// Metrics parsing +// --------------------------------------------------------------------------- + +/// Sum the trailing numeric value of every non-comment metrics line that +/// contains `name` (across all labelled series / engines). Prometheus emits one +/// line per series, so summing is correct for a single-engine server and robust +/// to the `engine="0"` label format. +fn metric_total(metrics_text: &str, name: &str) -> i64 { + metrics_text + .lines() + .filter(|line| !line.starts_with('#') && line.contains(name)) + .filter_map(|line| { + line.rsplit_once(' ') + .and_then(|(_, value)| value.trim().parse::().ok()) + .map(|value| value as i64) + }) + .sum() +} + +/// Returns `(running, waiting)` summed across engines, or `(-1, -1)` if the +/// metrics endpoint is temporarily unreachable. +fn metrics_counts(client: &Client, base_url: &str) -> (i64, i64) { + let url = format!("{base_url}/metrics"); + let Ok(resp) = client.get(&url).timeout(Duration::from_secs(5)).send() else { + return (-1, -1); + }; + let Ok(text) = resp.text() else { + return (-1, -1); + }; + ( + metric_total(&text, "num_requests_running"), + metric_total(&text, "num_requests_waiting"), + ) +} + +// --------------------------------------------------------------------------- +// The test +// --------------------------------------------------------------------------- + +/// Burst of streaming completions → mass disconnect → metrics collapse → clean +/// follow-up. See the module docs for the full rationale. +#[test] +fn http_disconnect_metrics_collapse() { + let Some(model_path) = model_path_or_skip() else { + return; + }; + let Some(server_bin) = server_bin_or_skip() else { + return; + }; + + let port = reserve_loopback_port(); + let base_url = format!("http://127.0.0.1:{port}"); + + eprintln!("=== starting {server_bin} --model-path {model_path} --port {port} ==="); + let child = Command::new(server_bin) + .arg("--model-path") + .arg(&model_path) + .arg("--port") + .arg(port.to_string()) + .spawn() + .unwrap_or_else(|e| panic!("failed to spawn server binary: {e}")); + // Killing `server` on drop guarantees cleanup even on assertion failure. + let server = ServerGuard { child }; + + let client = Client::builder() + .connect_timeout(Duration::from_secs(30)) + .timeout(Duration::from_secs(120)) + .build() + .expect("failed to build reqwest client"); + + // 1. Wait for the server to finish loading and answer /metrics. + eprintln!( + "=== waiting for server readiness (up to {}s) ===", + SERVER_READY_TIMEOUT.as_secs() + ); + assert!( + wait_for_server(&client, &base_url), + "server did not become ready within {}s", + SERVER_READY_TIMEOUT.as_secs() + ); + eprintln!("=== server ready ==="); + + // 2. Discover the model id from /v1/models (defaults to the model path). + let model_id: String = { + let resp = client + .get(format!("{base_url}/v1/models")) + .timeout(Duration::from_secs(10)) + .send() + .expect("GET /v1/models failed") + .error_for_status() + .expect("/v1/models returned non-2xx"); + let body: Value = resp.json().expect("/v1/models did not return JSON"); + body["data"][0]["id"] + .as_str() + .expect("/v1/models response missing data[0].id") + .to_string() + }; + eprintln!("=== model id: {model_id} ==="); + + // 3. Build the streaming request body shared by every burst client. + let burst_body = json!({ + "model": model_id.as_str(), + "prompt": "Hello ".repeat(PROMPT_WORDS), + "max_tokens": 128, + "stream": true, + "temperature": 0.7, + }); + + // 4. Launch the burst. Scoped threads borrow `client`/`base_url`/`burst_body` + // so only the per-client `Sender` and the shared stop flag need cloning. + let disconnect_start = thread::scope(|s| -> Instant { + let (tx, rx) = mpsc::channel::>(); + let stop = Arc::new(AtomicBool::new(false)); + + for id in 0..NUM_CLIENTS { + let tx = tx.clone(); + let stop = stop.clone(); + let client = &client; + let base_url = &base_url; + let body = &burst_body; + s.spawn(move || { + let url = format!("{base_url}/v1/completions"); + let resp = match client.post(&url).json(body).send() { + Ok(r) => r, + Err(e) => { + let _ = tx.send(Err(format!("client {id}: send failed: {e}"))); + return; + } + }; + // Headers received → the request is in-flight on the server. + let _ = tx.send(Ok(())); + // Hold the streaming response open until told to disconnect. + while !stop.load(Ordering::Relaxed) { + thread::sleep(Duration::from_millis(10)); + } + // Dropping `resp` closes the TCP connection: the server's next + // stream write fails and it retires the request. + drop(resp); + }); + } + drop(tx); // workers still hold their clones; rx sees all results. + + // 4a. Wait for every client to confirm it has an open streaming response. + let mut connected = 0usize; + while connected < NUM_CLIENTS { + #[allow(clippy::match_wild_err_arm)] + match rx.recv_timeout(Duration::from_secs(30)) { + Ok(Ok(())) => connected += 1, + Ok(Err(e)) => panic!("streaming client failed to connect: {e}"), + Err(_) => panic!( + "timed out waiting for {NUM_CLIENTS} streaming clients to connect \ + (only {connected} connected)" + ), + } + } + eprintln!("=== all {NUM_CLIENTS} streaming clients connected ==="); + + // 4b. Confirm the burst actually entered the running set. + let inflight_deadline = Instant::now() + INFLIGHT_TIMEOUT; + loop { + let (running, waiting) = metrics_counts(&client, &base_url); + if running > 0 { + eprintln!("=== in-flight: running={running} waiting={waiting} ==="); + break; + } + assert!( + Instant::now() < inflight_deadline, + "no requests entered the running set within {}s \ + (last running={running} waiting={waiting})", + INFLIGHT_TIMEOUT.as_secs() + ); + thread::sleep(POLL_INTERVAL); + } + + // 4c. Mass disconnect. + eprintln!("=== disconnecting all clients ==="); + let disconnect_start = Instant::now(); + stop.store(true, Ordering::Relaxed); + // Scope exit joins the workers, i.e. waits until every `resp` is dropped. + disconnect_start + }); + + // 5. Poll /metrics until running+waiting collapse to zero. + eprintln!("=== polling for metrics collapse ==="); + let collapse_deadline = disconnect_start + COLLAPSE_TIMEOUT; + loop { + let (running, waiting) = metrics_counts(&client, &base_url); + if running == 0 && waiting == 0 { + eprintln!( + "=== collapsed to running=0 waiting=0 in {}ms ===", + disconnect_start.elapsed().as_millis() + ); + break; + } + assert!( + Instant::now() < collapse_deadline, + "metrics did not collapse within {}s after disconnect \ + (running={running} waiting={waiting})", + COLLAPSE_TIMEOUT.as_secs() + ); + thread::sleep(POLL_INTERVAL); + } + + // 6. Clean follow-up request: the engine must still serve new work. + eprintln!("=== follow-up non-streaming request ==="); + let followup = json!({ + "model": model_id.as_str(), + "prompt": "Say hello", + "max_tokens": 5, + "stream": false, + }); + let resp = client + .post(format!("{base_url}/v1/completions")) + .json(&followup) + .timeout(Duration::from_secs(30)) + .send() + .expect("follow-up request failed"); + let status = resp.status(); + let body: Value = resp + .json() + .unwrap_or_else(|e| panic!("follow-up response is not JSON ({e})")); + assert!( + status.is_success(), + "follow-up request returned {status}: {body}" + ); + let text = body["choices"][0]["text"] + .as_str() + .expect("follow-up response missing choices[0].text"); + assert!(!text.trim().is_empty(), "follow-up returned empty text"); + eprintln!("=== follow-up ok: status={status} text={text:?} ==="); + + eprintln!("=== PASS: metrics collapsed after disconnect, follow-up served ==="); + // `server` dropped here → child killed. + drop(server); +}