diff --git a/lib/llm/src/kv_router/prefill_router/admission.rs b/lib/llm/src/kv_router/prefill_router/admission.rs index d9bea2a96939..35cadd547250 100644 --- a/lib/llm/src/kv_router/prefill_router/admission.rs +++ b/lib/llm/src/kv_router/prefill_router/admission.rs @@ -6,7 +6,6 @@ use std::sync::Arc; use anyhow::Result; use futures::StreamExt; use tokio::sync::OwnedSemaphorePermit; -use tracing::Instrument; use dynamo_kv_router::selector::WorkerSelector; @@ -15,7 +14,7 @@ use dynamo_runtime::{ protocols::{annotated::Annotated, maybe_error::MaybeError}, }; -use super::{PrefillCompletion, PrefillError, PrefillRouter}; +use super::{PrefillCompletion, PrefillError, PrefillRouter, handoff::PrefillTask}; use crate::{ local_model::runtime_config::ModelRuntimeConfig, protocols::common::{ @@ -69,7 +68,7 @@ where && obj.contains_key("bootstrap_room") }); - if !is_bootstrap { + let completion = if !is_bootstrap { while let Some(next) = prefill_response.next().await { if let Some(error) = next.err() { let detail = format!("Prefill router returned error in output stream: {error}"); @@ -84,12 +83,16 @@ where .and_then(|usage| usage.prompt_tokens_details.clone()); } } + None } else { - tokio::spawn(async move { + Some(PrefillTask::spawn(async move { let _task_guard = task_guard; - while prefill_response.next().await.is_some() {} - }); - } + while let Some(output) = prefill_response.next().await { + PrefillTask::check_output(&output)?; + } + Ok(()) + })) + }; // A CTX request that reaches EOS/stop during its one-token prefill step // is already complete and does not establish a KV-cache handoff. The @@ -144,6 +147,7 @@ where prompt_tokens_details, }, worker_link: output.worker_trace_link.clone(), + completion, }) } @@ -152,19 +156,19 @@ where prefill_stream: ManyOut>, tracker: Option>, phase_transition_permit: OwnedSemaphorePermit, - ) { - let span = tracing::Span::current(); + ) -> PrefillTask { let task_guard = self.task_guard.clone(); - tokio::spawn( - async move { - drop(phase_transition_permit); - match Self::consume_prefill_stream(prefill_stream, tracker, task_guard).await { - Ok(_) => tracing::debug!("Prefill background task completed"), - Err(error) => tracing::warn!("Prefill background task error: {error:?}"), - } + PrefillTask::spawn(async move { + drop(phase_transition_permit); + match Self::consume_prefill_stream(prefill_stream, tracker, task_guard).await? { + PrefillCompletion::Handoff { + completion: Some(task), + .. + } => task.wait().await, + PrefillCompletion::Terminal { output } => PrefillTask::check_output(&output), + _ => Ok(()), } - .instrument(span), - ); + }) } } @@ -194,6 +198,59 @@ mod tests { }) } + #[tokio::test] + async fn bootstrap_tail_failures_reach_completion_observer() { + for failure in [ + Annotated::from_error("connection closed during prefill"), + Annotated::from_data(LLMEngineOutput::error("backend prefill failed".to_string())), + Annotated::from_data(LLMEngineOutput::cancelled()), + ] { + let first = Annotated::from_data(LLMEngineOutput { + disaggregated_params: Some(json!({ + "bootstrap_host": "prefill", "bootstrap_port": 1, "bootstrap_room": 7, + })), + ..Default::default() + }); + let (release_tx, release_rx) = tokio::sync::oneshot::channel(); + let stream = stream::iter([first]).chain(stream::once(async move { + release_rx.await.unwrap(); + failure + })); + let response = ResponseStream::new(Box::pin(stream), Arc::new(Controller::default())); + let result = PrefillRouter::::consume_prefill_stream( + response, None, None, + ) + .await + .unwrap(); + let PrefillCompletion::Handoff { + completion: Some(completion), + .. + } = result + else { + panic!("expected an early handoff with a completion observer"); + }; + release_tx.send(()).unwrap(); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = + ResponseStream::new(Box::pin(stream::pending()), context.clone()); + match completion + .forward_decode(async { Ok(decode) }, context) + .await + { + Err(error) => assert!(error.to_string().contains("Prefill")), + Ok(mut response) => { + let output = + tokio::time::timeout(std::time::Duration::from_secs(1), response.next()) + .await + .unwrap() + .unwrap(); + assert!(output.err().is_some()); + assert!(response.next().await.is_none()); + } + } + } + } + #[tokio::test] async fn bootstrap_drain_retains_teardown_guard_until_stream_finishes() { let first = Annotated::from_data(LLMEngineOutput { diff --git a/lib/llm/src/kv_router/prefill_router/handoff.rs b/lib/llm/src/kv_router/prefill_router/handoff.rs new file mode 100644 index 000000000000..02dd28d102aa --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/handoff.rs @@ -0,0 +1,344 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Keep consuming prefill while decode runs, and return failures to the caller. + +use std::{future::Future, sync::Arc}; + +use dynamo_runtime::{ + pipeline::{AsyncEngineContext, AsyncEngineContextProvider, ManyOut, ResponseStream}, + protocols::{annotated::Annotated, maybe_error::MaybeError}, +}; +use futures::StreamExt; +use tokio::task::JoinHandle; +use tracing::Instrument; + +use super::PrefillError; +use crate::protocols::common::llm_backend::{FinishReason, LLMEngineOutput}; + +/// Dropping the observer detaches the task: client disconnects must not abort a +/// prefill whose KV transfer may still be needed by the decode worker. +pub(super) struct PrefillTask(JoinHandle>); + +impl PrefillTask { + pub(super) fn spawn( + future: impl Future> + Send + 'static, + ) -> Self { + Self(tokio::spawn( + async move { + let result = future.await; + if let Err(error) = &result { + // Keep evidence even if the client has already disconnected. + tracing::warn!(%error, "Prefill background task failed"); + } + result + } + .instrument(tracing::Span::current()), + )) + } + + pub(super) async fn wait(self) -> Result<(), PrefillError> { + self.0.await.map_err(|error| { + PrefillError::PrefillError( + format!("Prefill background task did not complete: {error}"), + Some(Box::new(error)), + ) + })? + } + + pub(super) fn check_output(output: &Annotated) -> Result<(), PrefillError> { + if let Some(error) = output.err() { + return Err(PrefillError::PrefillError( + format!("Prefill router returned error in output stream: {error}"), + Some(Box::new(error)), + )); + } + match output + .data + .as_ref() + .and_then(|data| data.finish_reason.as_ref()) + { + Some(FinishReason::Error(message)) => Err(PrefillError::PrefillError( + format!("Prefill backend failed: {message}"), + None, + )), + Some(FinishReason::Cancelled) => Err(PrefillError::PrefillError( + "Prefill backend cancelled before completing the handoff".to_string(), + None, + )), + _ => Ok(()), + } + } + + /// Observe prefill failures both during decode dispatch and while streaming. + /// Successful prefill does not delay tokens; a successful decode terminal is + /// held until prefill finishes so a late failure cannot follow success. + pub(super) async fn forward_decode( + self, + decode: impl Future>>>, + request_context: Arc, + ) -> anyhow::Result>> { + let mut completion = Box::pin(self.wait()); + tokio::pin!(decode); + let mut response = tokio::select! { + biased; + result = &mut completion => { + if let Err(error) = result { + request_context.stop_generating(); + return Err(error.into()); + } + return decode.await; + } + result = &mut decode => result?, + }; + // Move the completion future into the response stream. The task itself + // keeps running if this stream is dropped during client cancellation. + let context = response.context(); + let stream_context = context.clone(); + let stream = async_stream::stream! { + loop { + tokio::select! { + biased; + result = &mut completion => { + if let Err(error) = result { + request_context.stop_generating(); + stream_context.stop_generating(); + yield Annotated::from_err(error); + return; + } + while let Some(output) = response.next().await { + yield output; + } + return; + } + output = response.next() => { + if output.as_ref().is_some_and(|output| { + output.err().is_some() || matches!( + output.data.as_ref().and_then(|data| data.finish_reason.as_ref()), + Some(FinishReason::Error(_) | FinishReason::Cancelled) + ) + }) { + yield output.unwrap(); + return; + } + let terminal = output.as_ref().is_none_or(|output| { + output.data.as_ref().is_some_and(|data| data.finish_reason.is_some()) + }); + if terminal + && let Err(error) = completion.as_mut().await + { + request_context.stop_generating(); + stream_context.stop_generating(); + yield Annotated::from_err(error); + return; + } + match output { + Some(output) => yield output, + None => return, + } + if terminal { + return; + } + } + } + } + }; + Ok(ResponseStream::new(Box::pin(stream), context)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_runtime::pipeline::context::Controller; + use futures::{FutureExt, stream}; + use tokio::sync::oneshot; + + fn pending_prefill() -> (oneshot::Sender>, PrefillTask) { + let (tx, rx) = oneshot::channel(); + let task = PrefillTask::spawn(async move { rx.await.unwrap() }); + (tx, task) + } + + fn failure() -> Result<(), PrefillError> { + Err(PrefillError::PrefillError( + "prefill connection lost".to_string(), + None, + )) + } + + #[tokio::test] + async fn failure_interrupts_pending_decode_dispatch() { + let (tx, task) = pending_prefill(); + let context = Arc::new(Controller::default()); + let mut response = Box::pin(task.forward_decode(std::future::pending(), context.clone())); + assert!(response.as_mut().now_or_never().is_none()); + tx.send(failure()).unwrap(); + let error = tokio::time::timeout(std::time::Duration::from_secs(1), response) + .await + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("prefill connection lost")); + assert!(context.is_stopped()); + } + + #[tokio::test] + async fn failure_interrupts_wait_for_decode_tokens() { + let (tx, task) = pending_prefill(); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = + ResponseStream::new(Box::pin(stream::pending()), context.clone()); + let mut response = task + .forward_decode(async { Ok(decode) }, context.clone()) + .await + .unwrap(); + tx.send(failure()).unwrap(); + let output = tokio::time::timeout(std::time::Duration::from_secs(1), response.next()) + .await + .unwrap() + .unwrap(); + assert!( + output + .err() + .unwrap() + .to_string() + .contains("prefill connection lost") + ); + assert!(context.is_stopped()); + assert!(response.next().await.is_none()); + } + + #[tokio::test] + async fn tokens_flow_while_prefill_is_pending_but_success_waits() { + let (tx, task) = pending_prefill(); + let context = Arc::new(Controller::default()); + let token = LLMEngineOutput { + token_ids: vec![42], + ..Default::default() + }; + let decode: ManyOut> = ResponseStream::new( + Box::pin(stream::iter([ + Annotated::from_data(token), + Annotated::from_data(LLMEngineOutput::length()), + ])), + context.clone(), + ); + let mut response = task + .forward_decode(async { Ok(decode) }, context.clone()) + .await + .unwrap(); + assert_eq!(response.next().await.unwrap().data.unwrap().token_ids, [42]); + assert!(response.next().now_or_never().is_none()); + tx.send(Ok(())).unwrap(); + assert_eq!( + response.next().await.unwrap().data.unwrap().finish_reason, + Some(FinishReason::Length) + ); + assert!(response.next().await.is_none()); + assert!(!context.is_stopped()); + } + + #[tokio::test] + async fn late_prefill_failure_replaces_decode_success() { + let (tx, task) = pending_prefill(); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = ResponseStream::new( + Box::pin(stream::iter([Annotated::from_data( + LLMEngineOutput::length(), + )])), + context.clone(), + ); + let mut response = task + .forward_decode(async { Ok(decode) }, context) + .await + .unwrap(); + assert!(response.next().now_or_never().is_none()); + tx.send(failure()).unwrap(); + assert!( + response + .next() + .await + .unwrap() + .err() + .unwrap() + .to_string() + .contains("prefill connection lost") + ); + assert!(response.next().await.is_none()); + } + + #[tokio::test] + async fn decode_failure_does_not_wait_for_prefill() { + let (tx, task) = pending_prefill(); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = ResponseStream::new( + Box::pin(stream::iter([Annotated::from_error("decode failed")])), + context.clone(), + ); + let mut response = task + .forward_decode(async { Ok(decode) }, context) + .await + .unwrap(); + assert!( + response + .next() + .now_or_never() + .unwrap() + .unwrap() + .err() + .unwrap() + .to_string() + .contains("decode failed") + ); + tx.send(Ok(())).unwrap(); + } + + #[tokio::test] + async fn completed_prefill_leaves_decode_stream_unchanged() { + let (tx, task) = pending_prefill(); + tx.send(Ok(())).unwrap(); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = ResponseStream::new( + Box::pin(stream::iter([Annotated::from_data( + LLMEngineOutput::length(), + )])), + context.clone(), + ); + let mut response = task + .forward_decode(async { Ok(decode) }, context.clone()) + .await + .unwrap(); + assert_eq!( + response.next().await.unwrap().data.unwrap().finish_reason, + Some(FinishReason::Length) + ); + assert!(response.next().await.is_none()); + assert!(!context.is_stopped()); + } + + #[tokio::test] + async fn disconnected_client_does_not_drop_prefill_transfer() { + let (release_tx, release_rx) = oneshot::channel(); + let (finished_tx, finished_rx) = oneshot::channel(); + let task = PrefillTask::spawn(async move { + release_rx.await.unwrap(); + finished_tx.send(()).unwrap(); + Ok(()) + }); + let context = Arc::new(Controller::default()); + let decode: ManyOut> = + ResponseStream::new(Box::pin(stream::pending()), context.clone()); + let response = task + .forward_decode(async { Ok(decode) }, context.clone()) + .await + .unwrap(); + context.stop_generating(); + drop(response); + // Even the dropped response observer must leave prefill able to finish. + release_tx.send(()).unwrap(); + tokio::time::timeout(std::time::Duration::from_secs(1), finished_rx) + .await + .unwrap() + .unwrap(); + } +} diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs index 416148eb4887..86951edf06dd 100644 --- a/lib/llm/src/kv_router/prefill_router/mod.rs +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -45,7 +45,9 @@ use crate::{ mod activation; mod admission; mod conditional_bypass; +mod handoff; mod query; +use handoff::PrefillTask; pub use query::PrefillReservation; #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -177,6 +179,7 @@ enum PrefillCompletion { Handoff { result: PrefillResult, worker_link: Option, + completion: Option, }, Terminal { output: Box>, @@ -326,7 +329,6 @@ where // Extract request data while preserving context let (mut req, mut context) = request.into_parts(); let request_id = context.id().to_string(); - let metadata = context.metadata().clone(); let engine_ctx = context.context(); // Conditional-disagg bypass is a router-owned decision. Drop any @@ -442,9 +444,7 @@ where .and_then(|r| r.prefill_worker_id); let tracker = prefill_req.tracker.clone(); - let mut prefill_context = - Context::with_id_and_metadata(prefill_req, request_id.clone(), metadata.clone()); - propagate_first_response_guard(&context, &mut prefill_context)?; + let mut prefill_context = independent_prefill_context(prefill_req, &context)?; if let Some(session_affinity) = session_affinity { prefill_context.insert( SESSION_AFFINITY_CONTEXT_KEY, @@ -466,6 +466,7 @@ where let router = &binding.router; let endpoint_id = &binding.endpoint_id; + let mut prefill_completion = None; let prefill_result: Result<(PrefillOutcome, Option)> = async { let (prepared, prefill_stream) = router .select_and_dispatch_prefill(prefill_context, |request, target| { @@ -474,7 +475,8 @@ where .await?; let topology_constraints = prepared.topology_constraints; let outcome = if let Some(bootstrap_info) = prepared.bootstrap_info { - self.spawn_prefill_task(prefill_stream, tracker, prefill_phase_barrier); + prefill_completion = + Some(self.spawn_prefill_task(prefill_stream, tracker, prefill_phase_barrier)); PrefillOutcome::Bootstrap { bootstrap_info, worker_id: prepared.worker_id, @@ -489,7 +491,9 @@ where PrefillCompletion::Handoff { result, worker_link, + completion, } => { + prefill_completion = completion; if let Some(bootstrap_info) = extract_bootstrap_info(&result.disaggregated_params) { @@ -584,10 +588,26 @@ where self.conditional_disagg_policy.is_enabled(), )); - next.generate(context.map(|_| decode_req)).await + let decode = next.generate(context.map(|_| decode_req)); + match prefill_completion { + Some(completion) => completion.forward_decode(decode, engine_ctx).await, + None => decode.await, + } } } +/// Prefill must survive client cancellation once decode needs its KV transfer. +/// Copy request identity and metadata, but keep a separate cancellation controller. +fn independent_prefill_context( + request: PreprocessedRequest, + source: &Context<()>, +) -> Result> { + let mut prefill = + Context::with_id_and_metadata(request, source.id().to_string(), source.metadata().clone()); + propagate_first_response_guard(source, &mut prefill)?; + Ok(prefill) +} + impl PrefillRouter where Sel: WorkerSelector + Send + 'static, @@ -813,6 +833,59 @@ mod tests { .unwrap() } + #[test] + fn client_cancellation_does_not_cancel_prefill_context() { + for kill in [false, true] { + let request = query_only_request(); + let mut source = Context::new(()); + source.insert_metadata("traceparent", "test-trace"); + let prefill = independent_prefill_context(request, &source).unwrap(); + assert_eq!(prefill.id(), source.id()); + assert_eq!(prefill.metadata(), source.metadata()); + + if kill { + source.context().kill(); + } else { + source.context().stop_generating(); + } + assert!(source.context().is_stopped()); + assert!(!prefill.context().is_stopped()); + + // Prefill still honors its own worker cancellation/shutdown. + prefill.context().kill(); + assert!(prefill.context().is_killed()); + } + } + + #[tokio::test] + async fn precomputed_bootstrap_observes_initial_and_late_prefill_failures() { + for after_handoff in [false, true] { + let router = active_conditional_router(); + let tracker = Arc::new(RequestTracker::new()); + let permit = tracker.set_phase(RequestPhase::Prefill).await; + let mut outputs = Vec::new(); + if after_handoff { + outputs.push(Annotated::from_data(LLMEngineOutput { + disaggregated_params: Some(serde_json::json!({ + "bootstrap_host": "prefill", "bootstrap_port": 1, "bootstrap_room": 7, + })), + ..Default::default() + })); + } + outputs.push(Annotated::from_error("prefill connection lost")); + let prefill = + ResponseStream::new(Box::pin(stream::iter(outputs)), Context::new(()).context()); + let task = router.spawn_prefill_task(prefill, Some(tracker.clone()), permit); + // Decode can enter its phase even though prefill is still observed. + let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; + let error = task + .forward_decode(std::future::pending(), Context::new(()).context()) + .await + .unwrap_err(); + assert!(error.to_string().contains("prefill connection lost")); + } + } + /// The decode leg must carry the cleanup obligation, on both outcomes that /// reach decode routing. Without it the routing host treats a disconnected /// decode request as ordinary traffic and drops it, which is the leak this diff --git a/lib/mocker/servers/sglang/tests/sidecar.rs b/lib/mocker/servers/sglang/tests/sidecar.rs index 66fbcca5c608..72453d6265c6 100644 --- a/lib/mocker/servers/sglang/tests/sidecar.rs +++ b/lib/mocker/servers/sglang/tests/sidecar.rs @@ -201,8 +201,9 @@ async fn prefill_handoff_round_trips_through_a_decode_server() { decode.start(1).await.unwrap(); let prefill_outputs = collect(&prefill, request(3)).await; - assert_eq!(prefill_outputs.len(), 1); + assert_eq!(prefill_outputs.len(), 2); assert!(prefill_outputs[0].token_ids.is_empty()); + assert!(prefill_outputs[0].finish_reason.is_none()); let handoff = prefill_outputs[0] .disaggregated_params .clone() @@ -210,6 +211,8 @@ async fn prefill_handoff_round_trips_through_a_decode_server() { assert_eq!(handoff["bootstrap_host"], "127.0.0.1"); assert_eq!(handoff["bootstrap_port"], 8_998); assert!(handoff["bootstrap_room"].is_number()); + assert_eq!(prefill_outputs[1].finish_reason, Some(FinishReason::Length)); + assert!(prefill_outputs[1].disaggregated_params.is_none()); let mut decode_request = request(3); decode_request.prefill_result = Some(PrefillResult { diff --git a/lib/sidecar/sglang/src/engine.rs b/lib/sidecar/sglang/src/engine.rs index 5357e6c33ff2..90d41623c8df 100644 --- a/lib/sidecar/sglang/src/engine.rs +++ b/lib/sidecar/sglang/src/engine.rs @@ -277,7 +277,7 @@ impl LLMEngine for SglangSidecarEngine { self.bootstrap_host.as_deref(), self.bootstrap_port, )?; - let prefill_handoff = if self.disaggregation_mode.is_prefill() { + let mut prefill_handoff = if self.disaggregation_mode.is_prefill() { grpc_request .disaggregated_params .as_ref() @@ -312,6 +312,20 @@ impl LLMEngine for SglangSidecarEngine { return; } }; + if is_prefill { + let Some(handoff) = prefill_handoff.take() else { + yield Err(client::protocol_error( + "SGLang gRPC prefill request is missing disaggregated params", + )); + return; + }; + // Publish the handoff only after the gRPC transport opens the + // response stream so decode can rendezvous while prefill runs. + yield Ok(LLMEngineOutput { + disaggregated_params: Some(handoff), + ..Default::default() + }); + } let mut generated = 0_u32; let mut observed_prompt_tokens = prompt_tokens; @@ -364,7 +378,7 @@ impl LLMEngine for SglangSidecarEngine { if is_prefill { if response.finished { - let mut terminal = match terminal_from_meta( + let terminal = match terminal_from_meta( &response.meta_info, observed_prompt_tokens, 0, @@ -375,7 +389,6 @@ impl LLMEngine for SglangSidecarEngine { break; } }; - terminal.disaggregated_params = prefill_handoff.clone(); yield Ok(terminal); break; } diff --git a/lib/sidecar/sglang/src/native_http.rs b/lib/sidecar/sglang/src/native_http.rs index 5c1000fb9f63..e6fe02a6cb1e 100644 --- a/lib/sidecar/sglang/src/native_http.rs +++ b/lib/sidecar/sglang/src/native_http.rs @@ -245,6 +245,7 @@ impl NativeHttp { ) -> BoxStream<'static, Result> { Box::pin(async_stream::stream! { let is_prefill = request.is_prefill; + let mut prefill_handoff = request.prefill_handoff; tracing::debug!(request_id = %ctx.id(), endpoint = %self.endpoint.with_path("/generate"), "sending native request to SGLang HTTP"); let opened = tokio::select! { biased; @@ -266,11 +267,25 @@ impl NativeHttp { return; } }; + if is_prefill { + let Some(handoff) = prefill_handoff.take() else { + yield Err(client::protocol_error( + "SGLang native prefill request is missing disaggregated params", + )); + return; + }; + // Publish the handoff only after the HTTP transport returns + // successful response headers so decode can rendezvous while + // prefill runs. + yield Ok(LLMEngineOutput { + disaggregated_params: Some(handoff), + ..Default::default() + }); + } let bytes = response.bytes_stream().map_err(io::Error::other); let reader = StreamReader::new(bytes); let mut lines = FramedRead::new(reader, LinesCodec::new_with_max_length(MAX_EVENT_BYTES)); - let mut prefill_handoff = request.prefill_handoff; let mut first_output_seen = false; loop { let selected = tokio::select! { @@ -319,7 +334,7 @@ impl NativeHttp { )); return; } - let response = match serde_json::from_str(data) { + let response: Value = match serde_json::from_str(data) { Ok(response) => response, Err(error) => { yield Err(client::protocol_error(format!( @@ -328,6 +343,16 @@ impl NativeHttp { return; } }; + // Prefill's raw HTTP payload is not forwarded to the caller. + // Surface backend failures before that payload is discarded. + if is_prefill + && let Some(finish) = response.pointer("/meta_info/finish_reason") + && let Some(kind @ ("abort" | "error" | "cancelled")) = + finish.get("type").and_then(Value::as_str) + { + yield Err(protocol::terminal_failure(kind, finish)); + return; + } let has_output = response_has_output(&response); let (mut output, terminal) = output(response, &mut prefill_handoff); if !first_output_seen && has_output && (!is_prefill || terminal) { @@ -420,7 +445,7 @@ mod tests { use dynamo_backend_common::engine::RoutingHints; use dynamo_backend_common::{ - BackendError, DisaggregationMode, ErrorType, GenerateContext, OutputOptions, + BackendError, DisaggregationMode, ErrorType, FinishReason, GenerateContext, OutputOptions, PreprocessedRequest, SamplingOptions, StopConditions, }; use dynamo_sidecar_common::{GrpcEndpoint, HttpEndpoint}; @@ -591,6 +616,133 @@ mod tests { server.await.unwrap(); } + #[tokio::test] + async fn prefill_stream_rejects_backend_error_before_bootstrap() { + let (port, server) = serve_once("rejected".to_string(), "500 Internal Server Error").await; + let ctx = GenerateContext::new(dynamo_backend_common::testing::mock_context(), None); + let mut stream = native_http(port).generate( + NativeRequest { + body: json!({"input_ids": [1], "stream": true}), + is_prefill: true, + prefill_handoff: Some(json!({ + "bootstrap_host": "prefill", + "bootstrap_port": 5000, + "bootstrap_room": 7 + })), + }, + ctx, + CancellationToken::new(), + ); + + let error = stream.next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains("HTTP 500")); + assert!(stream.next().await.is_none()); + server.await.unwrap(); + } + + #[tokio::test] + async fn prefill_stream_reports_failures_after_success_headers() { + for (body, expected) in [ + ("", "closed before a terminal response"), + ("data: broken-json\n\n", "invalid JSON"), + ( + "data: {\"meta_info\":{\"finish_reason\":{\"type\":\"abort\",\"message\":\"prefill rejected\",\"status_code\":400}}}\n\n", + "prefill rejected", + ), + ] { + let (port, server) = serve_once(body.to_string(), "200 OK").await; + let ctx = GenerateContext::new(dynamo_backend_common::testing::mock_context(), None); + let mut stream = native_http(port).generate( + NativeRequest { + body: json!({"input_ids": [1], "stream": true}), + is_prefill: true, + prefill_handoff: Some(json!({ + "bootstrap_host": "prefill", "bootstrap_port": 1, "bootstrap_room": 7, + })), + }, + ctx, + CancellationToken::new(), + ); + let handoff = stream.next().await.unwrap().unwrap(); + assert!(handoff.disaggregated_params.is_some()); + assert!(handoff.finish_reason.is_none()); + let error = stream.next().await.unwrap().unwrap_err(); + assert!(error.to_string().contains(expected), "{error}"); + assert!(stream.next().await.is_none()); + server.await.unwrap(); + } + } + + #[tokio::test] + async fn prefill_handoff_waits_for_backend_response_headers() { + let body = + "data: {\"output_ids\":[101],\"meta_info\":{\"finish_reason\":{\"type\":\"length\"}}}\n\n" + .to_string(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let (request_seen_tx, request_seen_rx) = tokio::sync::oneshot::channel(); + let (release_headers_tx, release_headers_rx) = tokio::sync::oneshot::channel(); + let (release_body_tx, release_body_rx) = tokio::sync::oneshot::channel(); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 4096]; + let _ = socket.read(&mut request).await.unwrap(); + request_seen_tx.send(()).unwrap(); + release_headers_rx.await.unwrap(); + let headers = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + socket.write_all(headers.as_bytes()).await.unwrap(); + release_body_rx.await.unwrap(); + socket.write_all(body.as_bytes()).await.unwrap(); + }); + + let ctx = GenerateContext::new(dynamo_backend_common::testing::mock_context(), None); + let mut stream = native_http(port).generate( + NativeRequest { + body: json!({"input_ids": [1], "stream": true}), + is_prefill: true, + prefill_handoff: Some(json!({ + "bootstrap_host": "prefill", + "bootstrap_port": 5000, + "bootstrap_room": 7 + })), + }, + ctx, + CancellationToken::new(), + ); + + let mut first_output = Box::pin(stream.next()); + tokio::select! { + result = request_seen_rx => result.unwrap(), + output = &mut first_output => panic!( + "prefill handoff arrived before the backend saw the request: {output:?}" + ), + } + assert!( + tokio::time::timeout(Duration::from_millis(25), &mut first_output) + .await + .is_err(), + "sending the request alone must not release the prefill handoff" + ); + + release_headers_tx.send(()).unwrap(); + let handoff = tokio::time::timeout(Duration::from_secs(1), first_output) + .await + .expect("handoff did not arrive after successful response headers") + .unwrap() + .unwrap(); + assert_eq!(handoff.disaggregated_params.unwrap()["bootstrap_room"], 7); + + release_body_tx.send(()).unwrap(); + let terminal = stream.next().await.unwrap().unwrap(); + assert_eq!(terminal.finish_reason, Some(FinishReason::Stop)); + assert!(terminal.disaggregated_params.is_none()); + assert!(stream.next().await.is_none()); + server.await.unwrap(); + } + #[tokio::test] async fn cancellation_is_a_typed_stream_error() { let context = dynamo_backend_common::testing::mock_context(); diff --git a/lib/sidecar/sglang/src/protocol.rs b/lib/sidecar/sglang/src/protocol.rs index b9afe744c5fb..84ad07ef56c4 100644 --- a/lib/sidecar/sglang/src/protocol.rs +++ b/lib/sidecar/sglang/src/protocol.rs @@ -390,7 +390,7 @@ pub(crate) fn terminal_from_meta( Ok(output) } -fn terminal_failure(finish_type: &str, finish: &Value) -> DynamoError { +pub(crate) fn terminal_failure(finish_type: &str, finish: &Value) -> DynamoError { let message = finish .get("message") .and_then(Value::as_str)