From a585b8ca7f3fff4eed0f008e4bbb76f514187f41 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Wed, 5 Aug 2026 17:19:39 -0700 Subject: [PATCH 1/4] fix(vllm): preserve decode handoff on cancellation Signed-off-by: Connor Carpenter --- lib/sidecar/vllm/README.md | 2 + lib/sidecar/vllm/src/engine.rs | 106 +++++++++++++++++++++++---------- lib/sidecar/vllm/src/tests.rs | 96 +++++++++++++++++++++++++++++ 3 files changed, 172 insertions(+), 32 deletions(-) diff --git a/lib/sidecar/vllm/README.md b/lib/sidecar/vllm/README.md index b25604b15c02..dc2426fcc901 100644 --- a/lib/sidecar/vllm/README.md +++ b/lib/sidecar/vllm/README.md @@ -67,6 +67,8 @@ pool. Override them with `--grpc-connect-attempt-timeout-secs`, `--grpc-retry-interval-secs`, and `--grpc-startup-deadline-secs`, or with the corresponding `DYN_SIDECAR_GRPC_*` environment variables. +Each request owns its response stream but borrows a channel from the shared pool. Aggregate and prefill cancellation drops only that request's stream. Decode cancellation first submits the decode request and retains its stream through the first output token so a NIXL receiver can complete and release the transferred KV; it then drops the stream. vLLM automatically aborts the corresponding engine request while the pooled HTTP/2 connection remains available to other requests. The sidecar does not call the Control `Abort` RPC. + ## Test without vLLM or a GPU Use the CPU-only `dynamo-vllm-mocker-server` to exercise the same Inference, Control, and health contracts: diff --git a/lib/sidecar/vllm/src/engine.rs b/lib/sidecar/vllm/src/engine.rs index 347e93e09db4..babf1787ef9f 100644 --- a/lib/sidecar/vllm/src/engine.rs +++ b/lib/sidecar/vllm/src/engine.rs @@ -191,18 +191,25 @@ impl LLMEngine for VllmSidecarEngine { let mut state = ResponseState::new(&request, self.mode); let mut proto_request = build_generate_request(request, request_id, self.mode)?; proto_request.model.clone_from(&self.model.served_name); + let defer_request_cancellation = self.mode.is_decode(); let stopped_ctx = ctx.inner_arc(); let shutdown = self.cancel.clone(); - let mut cancellation = Box::pin(async move { + let mut request_cancellation = Box::pin(async move { stopped_ctx.stopped().await }); + let mut shutdown_cancellation = Box::pin(async move { shutdown.cancelled().await }); + let stream = if defer_request_cancellation { + // Decode must reach vLLM so NIXL can release transferred KV. tokio::select! { - _ = stopped_ctx.stopped() => {} - _ = shutdown.cancelled() => {} + biased; + _ = shutdown_cancellation.as_mut() => None, + result = client.generate_stream(proto_request) => Some(result?), + } + } else { + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + _ = request_cancellation.as_mut() => None, + result = client.generate_stream(proto_request) => Some(result?), } - }); - let stream = tokio::select! { - biased; - _ = cancellation.as_mut() => None, - result = client.generate_stream(proto_request) => Some(result?), }; let Some(mut stream) = stream else { let output = cancelled(&state); @@ -210,41 +217,76 @@ impl LLMEngine for VllmSidecarEngine { }; Ok(Box::pin(async_stream::stream! { + let mut request_cancelled = false; + let mut first_token_observed = false; loop { - tokio::select! { - biased; - _ = cancellation.as_mut() => { - yield Ok(cancelled(&state)); - break; + let message = if request_cancelled { + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + message = stream.message() => Some(message), } - message = stream.message() => { - match message { - Ok(Some(response)) => match state.convert(response) { - Ok(Some(output)) => { - let terminal = output.finish_reason.is_some(); - yield Ok(output); - if terminal { - break; + } else { + tokio::select! { + biased; + _ = shutdown_cancellation.as_mut() => None, + _ = request_cancellation.as_mut() => { + if defer_request_cancellation && !first_token_observed { + request_cancelled = true; + continue; + } + None + } + message = stream.message() => Some(message), + } + }; + + let Some(message) = message else { + yield Ok(cancelled(&state)); + break; + }; + match message { + Ok(Some(response)) => { + let response_has_token = response + .outputs + .as_ref() + .is_some_and(|output| output.num_tokens > 0); + let transfer_completed = response.outputs.as_ref().is_some_and(|output| { + output.num_tokens > 0 || output.finish_info.is_some() + }); + match state.convert(response) { + Ok(Some(output)) => { + first_token_observed |= response_has_token; + if request_cancelled && transfer_completed { + if first_token_observed { + ctx.notify_first_token(); } + yield Ok(cancelled(&state)); + break; } - Ok(None) => {} - Err(error) => { - yield Err(error); + let terminal = output.finish_reason.is_some(); + yield Ok(output); + if terminal { break; } - }, - Ok(None) => { - yield Err(client::protocol_error( - "GenerateStream ended before a terminal response", - )); - break; } - Err(status) => { - yield Err(client::status_to_dynamo("GenerateStream", status)); + Ok(None) => {} + Err(error) => { + yield Err(error); break; } } } + Ok(None) => { + yield Err(client::protocol_error( + "GenerateStream ended before a terminal response", + )); + break; + } + Err(status) => { + yield Err(client::status_to_dynamo("GenerateStream", status)); + break; + } } } })) diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 05043e15cc10..4852ea40544c 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -38,6 +38,9 @@ struct FakeVllm { hang_before_headers: Arc, headers_pending: Arc, release_headers: Arc, + hold_before_first_token: Arc, + first_token_pending: Arc, + release_first_token: Arc, server_stream_dropped: Arc, } @@ -120,6 +123,9 @@ impl pb::inference_server::Inference for FakeVllm { "nested": {"flags": [true, null, "opaque"]}, }); let hang = self.hang.load(Ordering::SeqCst); + let hold_before_first_token = self.hold_before_first_token.load(Ordering::SeqCst); + let first_token_pending = self.first_token_pending.clone(); + let release_first_token = self.release_first_token.clone(); let dropped = self.server_stream_dropped.clone(); let stream = async_stream::try_stream! { @@ -147,6 +153,12 @@ impl pb::inference_server::Inference for FakeVllm { outputs: None, }; + if hold_before_first_token { + first_token_pending.store(true, Ordering::SeqCst); + release_first_token.notified().await; + first_token_pending.store(false, Ordering::SeqCst); + } + if hang { loop { yield sequence_response(false, wants_logprobs, None); @@ -815,6 +827,90 @@ async fn cancellation_interrupts_pending_response_headers() { server.service.release_headers.notify_waiters(); } +#[tokio::test] +async fn decode_cancellation_waits_for_submission_and_first_token() { + let service = FakeVllm::default(); + service.hang_before_headers.store(true, Ordering::SeqCst); + service + .hold_before_first_token + .store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); + engine.start(0).await.expect("start"); + + let mut decode_request = request(); + decode_request.prefill_result = Some(PrefillResult { + disaggregated_params: json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "remote_engine_id": "prefill-0", + "remote_host": "127.0.0.1", + "remote_port": 20097, + "remote_block_ids": [7, 8], + }), + prompt_tokens_details: None, + }); + let context = dynamo_backend_common::testing::mock_context(); + let generate = engine.generate(decode_request, GenerateContext::new(context.clone(), None)); + tokio::pin!(generate); + + tokio::select! { + _ = &mut generate => panic!("decode returned before response headers were gated"), + _ = async { + while !server.service.headers_pending.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } => {} + } + assert_eq!(server.service.requests.lock().await.len(), 1); + context.stop_generating(); + tokio::select! { + _ = &mut generate => panic!("decode cancellation returned before response headers"), + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {} + } + + server.service.release_headers.notify_one(); + let mut stream = tokio::time::timeout(std::time::Duration::from_secs(2), &mut generate) + .await + .expect("decode response headers") + .expect("decode stream"); + let next = stream.next(); + tokio::pin!(next); + tokio::select! { + _ = &mut next => panic!("decode returned before the first token was gated"), + _ = async { + while !server.service.first_token_pending.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + } => {} + } + assert!( + !server.service.server_stream_dropped.load(Ordering::SeqCst), + "decode stream dropped before the first token" + ); + tokio::select! { + _ = &mut next => panic!("decode cancellation completed before the first token"), + _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {} + } + + server.service.release_first_token.notify_one(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), &mut next) + .await + .expect("first token did not release decode cancellation") + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); + drop(stream); + + tokio::time::timeout(std::time::Duration::from_secs(2), async { + while !server.service.server_stream_dropped.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .expect("server stream dropped after first token"); +} + #[tokio::test] async fn unsupported_features_fail_before_rpc_submission() { let server = FakeServer::start(FakeVllm::default()).await; From 2a59ef66af094980b3492e390d33824ae0f543aa Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Thu, 6 Aug 2026 11:23:34 -0700 Subject: [PATCH 2/4] fix(vllm): preserve cancellation terminal semantics Signed-off-by: Connor Carpenter --- lib/sidecar/vllm/src/engine.rs | 23 ++++++++++++ lib/sidecar/vllm/src/tests.rs | 65 +++++++++++++++++++++++++++------- 2 files changed, 75 insertions(+), 13 deletions(-) diff --git a/lib/sidecar/vllm/src/engine.rs b/lib/sidecar/vllm/src/engine.rs index babf1787ef9f..5492ec888616 100644 --- a/lib/sidecar/vllm/src/engine.rs +++ b/lib/sidecar/vllm/src/engine.rs @@ -271,18 +271,41 @@ impl LLMEngine for VllmSidecarEngine { } } Ok(None) => {} + Err(error) if request_cancelled => { + tracing::warn!( + %error, + "vLLM response conversion failed after request cancellation" + ); + yield Ok(cancelled(&state)); + break; + } Err(error) => { yield Err(error); break; } } } + Ok(None) if request_cancelled => { + tracing::warn!( + "vLLM GenerateStream ended before transfer completion after request cancellation" + ); + yield Ok(cancelled(&state)); + break; + } Ok(None) => { yield Err(client::protocol_error( "GenerateStream ended before a terminal response", )); break; } + Err(status) if request_cancelled => { + tracing::warn!( + %status, + "vLLM GenerateStream failed before transfer completion after request cancellation" + ); + yield Ok(cancelled(&state)); + break; + } Err(status) => { yield Err(client::status_to_dynamo("GenerateStream", status)); break; diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 4852ea40544c..3cafdd850124 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -39,6 +39,7 @@ struct FakeVllm { headers_pending: Arc, release_headers: Arc, hold_before_first_token: Arc, + close_before_first_token: Arc, first_token_pending: Arc, release_first_token: Arc, server_stream_dropped: Arc, @@ -124,6 +125,7 @@ impl pb::inference_server::Inference for FakeVllm { }); let hang = self.hang.load(Ordering::SeqCst); let hold_before_first_token = self.hold_before_first_token.load(Ordering::SeqCst); + let close_before_first_token = self.close_before_first_token.load(Ordering::SeqCst); let first_token_pending = self.first_token_pending.clone(); let release_first_token = self.release_first_token.clone(); let dropped = self.server_stream_dropped.clone(); @@ -158,6 +160,9 @@ impl pb::inference_server::Inference for FakeVllm { release_first_token.notified().await; first_token_pending.store(false, Ordering::SeqCst); } + if close_before_first_token { + return; + } if hang { loop { @@ -494,6 +499,22 @@ fn request() -> PreprocessedRequest { .expect("request") } +fn decode_request() -> PreprocessedRequest { + let mut request = request(); + request.prefill_result = Some(PrefillResult { + disaggregated_params: json!({ + "do_remote_decode": false, + "do_remote_prefill": true, + "remote_engine_id": "prefill-0", + "remote_host": "127.0.0.1", + "remote_port": 20097, + "remote_block_ids": [7, 8], + }), + prompt_tokens_details: None, + }); + request +} + fn engine(endpoint: &str, mode: DisaggregationMode, connections: usize) -> VllmSidecarEngine { let transport = GrpcTransportConfig { connections: NonZeroUsize::new(connections).expect("non-zero connection count"), @@ -838,20 +859,11 @@ async fn decode_cancellation_waits_for_submission_and_first_token() { let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); engine.start(0).await.expect("start"); - let mut decode_request = request(); - decode_request.prefill_result = Some(PrefillResult { - disaggregated_params: json!({ - "do_remote_decode": false, - "do_remote_prefill": true, - "remote_engine_id": "prefill-0", - "remote_host": "127.0.0.1", - "remote_port": 20097, - "remote_block_ids": [7, 8], - }), - prompt_tokens_details: None, - }); let context = dynamo_backend_common::testing::mock_context(); - let generate = engine.generate(decode_request, GenerateContext::new(context.clone(), None)); + let generate = engine.generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ); tokio::pin!(generate); tokio::select! { @@ -911,6 +923,33 @@ async fn decode_cancellation_waits_for_submission_and_first_token() { .expect("server stream dropped after first token"); } +#[tokio::test] +async fn decode_cancellation_maps_premature_eof_to_cancelled() { + let service = FakeVllm::default(); + service + .close_before_first_token + .store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); + engine.start(0).await.expect("start"); + + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ) + .await + .expect("decode stream"); + context.stop_generating(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("premature EOF did not release decode cancellation") + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); +} + #[tokio::test] async fn unsupported_features_fail_before_rpc_submission() { let server = FakeServer::start(FakeVllm::default()).await; From 8744c0773f334bb2c3fb66aa9bbe636d29760bb8 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Mon, 10 Aug 2026 11:04:17 -0700 Subject: [PATCH 3/4] fix(vllm): clarify cancelled decode failures Signed-off-by: Connor Carpenter --- lib/sidecar/vllm/README.md | 2 +- lib/sidecar/vllm/src/tests.rs | 79 +++++++++++++++++++++++------------ 2 files changed, 53 insertions(+), 28 deletions(-) diff --git a/lib/sidecar/vllm/README.md b/lib/sidecar/vllm/README.md index dc2426fcc901..5dc340ce11ab 100644 --- a/lib/sidecar/vllm/README.md +++ b/lib/sidecar/vllm/README.md @@ -67,7 +67,7 @@ pool. Override them with `--grpc-connect-attempt-timeout-secs`, `--grpc-retry-interval-secs`, and `--grpc-startup-deadline-secs`, or with the corresponding `DYN_SIDECAR_GRPC_*` environment variables. -Each request owns its response stream but borrows a channel from the shared pool. Aggregate and prefill cancellation drops only that request's stream. Decode cancellation first submits the decode request and retains its stream through the first output token so a NIXL receiver can complete and release the transferred KV; it then drops the stream. vLLM automatically aborts the corresponding engine request while the pooled HTTP/2 connection remains available to other requests. The sidecar does not call the Control `Abort` RPC. +Each request owns its response stream but borrows a channel from the shared pool. Aggregate and prefill cancellation drops only that request's stream. Decode cancellation first submits the decode request and retains its stream until the first output token or a response containing `finish_info`, so a NIXL receiver can complete and release the transferred KV; it then drops the stream. If the stream ends early, returns a gRPC error, or produces an invalid response after cancellation, the sidecar logs the failure and reports the request as cancelled. vLLM automatically aborts the corresponding engine request while the pooled HTTP/2 connection remains available to other requests. The sidecar does not call the Control `Abort` RPC. ## Test without vLLM or a GPU diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 3cafdd850124..4c1ca64fa00f 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -28,6 +28,15 @@ use crate::json::{json_to_struct, struct_to_json}; use crate::model::DiscoveredModel; use crate::proto as pb; +#[derive(Clone, Copy, Default)] +enum BeforeFirstTokenFailure { + #[default] + None, + Eof, + GrpcStatus, + InvalidResponse, +} + #[derive(Clone, Default)] struct FakeVllm { requests: Arc>>, @@ -39,7 +48,7 @@ struct FakeVllm { headers_pending: Arc, release_headers: Arc, hold_before_first_token: Arc, - close_before_first_token: Arc, + before_first_token_failure: BeforeFirstTokenFailure, first_token_pending: Arc, release_first_token: Arc, server_stream_dropped: Arc, @@ -125,7 +134,7 @@ impl pb::inference_server::Inference for FakeVllm { }); let hang = self.hang.load(Ordering::SeqCst); let hold_before_first_token = self.hold_before_first_token.load(Ordering::SeqCst); - let close_before_first_token = self.close_before_first_token.load(Ordering::SeqCst); + let before_first_token_failure = self.before_first_token_failure; let first_token_pending = self.first_token_pending.clone(); let release_first_token = self.release_first_token.clone(); let dropped = self.server_stream_dropped.clone(); @@ -160,8 +169,18 @@ impl pb::inference_server::Inference for FakeVllm { release_first_token.notified().await; first_token_pending.store(false, Ordering::SeqCst); } - if close_before_first_token { - return; + match before_first_token_failure { + BeforeFirstTokenFailure::None => {} + BeforeFirstTokenFailure::Eof => return, + BeforeFirstTokenFailure::GrpcStatus => { + Err(Status::internal("failed before first token"))?; + } + BeforeFirstTokenFailure::InvalidResponse => { + let mut response = sequence_response(false, wants_logprobs, None); + response.outputs.as_mut().expect("sequence output").num_tokens = 2; + yield response; + return; + } } if hang { @@ -924,30 +943,36 @@ async fn decode_cancellation_waits_for_submission_and_first_token() { } #[tokio::test] -async fn decode_cancellation_maps_premature_eof_to_cancelled() { - let service = FakeVllm::default(); - service - .close_before_first_token - .store(true, Ordering::SeqCst); - let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); - engine.start(0).await.expect("start"); +async fn decode_cancellation_maps_stream_failures_to_cancelled() { + for (failure, case) in [ + (BeforeFirstTokenFailure::Eof, "premature EOF"), + (BeforeFirstTokenFailure::GrpcStatus, "gRPC failure"), + (BeforeFirstTokenFailure::InvalidResponse, "invalid response"), + ] { + let service = FakeVllm { + before_first_token_failure: failure, + ..Default::default() + }; + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); + engine.start(0).await.expect("start"); - let context = dynamo_backend_common::testing::mock_context(); - let mut stream = engine - .generate( - decode_request(), - GenerateContext::new(context.clone(), None), - ) - .await - .expect("decode stream"); - context.stop_generating(); - let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) - .await - .expect("premature EOF did not release decode cancellation") - .expect("cancelled terminal") - .expect("cancelled output"); - assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ) + .await + .expect("decode stream"); + context.stop_generating(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .unwrap_or_else(|_| panic!("{case} did not release decode cancellation")) + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); + } } #[tokio::test] From a67345b34949f0e0be066b4fbb7c15d240f4e640 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Mon, 10 Aug 2026 11:07:37 -0700 Subject: [PATCH 4/4] test(vllm): restore focused cancellation coverage Signed-off-by: Connor Carpenter --- lib/sidecar/vllm/src/tests.rs | 79 ++++++++++++----------------------- 1 file changed, 27 insertions(+), 52 deletions(-) diff --git a/lib/sidecar/vllm/src/tests.rs b/lib/sidecar/vllm/src/tests.rs index 4c1ca64fa00f..3cafdd850124 100644 --- a/lib/sidecar/vllm/src/tests.rs +++ b/lib/sidecar/vllm/src/tests.rs @@ -28,15 +28,6 @@ use crate::json::{json_to_struct, struct_to_json}; use crate::model::DiscoveredModel; use crate::proto as pb; -#[derive(Clone, Copy, Default)] -enum BeforeFirstTokenFailure { - #[default] - None, - Eof, - GrpcStatus, - InvalidResponse, -} - #[derive(Clone, Default)] struct FakeVllm { requests: Arc>>, @@ -48,7 +39,7 @@ struct FakeVllm { headers_pending: Arc, release_headers: Arc, hold_before_first_token: Arc, - before_first_token_failure: BeforeFirstTokenFailure, + close_before_first_token: Arc, first_token_pending: Arc, release_first_token: Arc, server_stream_dropped: Arc, @@ -134,7 +125,7 @@ impl pb::inference_server::Inference for FakeVllm { }); let hang = self.hang.load(Ordering::SeqCst); let hold_before_first_token = self.hold_before_first_token.load(Ordering::SeqCst); - let before_first_token_failure = self.before_first_token_failure; + let close_before_first_token = self.close_before_first_token.load(Ordering::SeqCst); let first_token_pending = self.first_token_pending.clone(); let release_first_token = self.release_first_token.clone(); let dropped = self.server_stream_dropped.clone(); @@ -169,18 +160,8 @@ impl pb::inference_server::Inference for FakeVllm { release_first_token.notified().await; first_token_pending.store(false, Ordering::SeqCst); } - match before_first_token_failure { - BeforeFirstTokenFailure::None => {} - BeforeFirstTokenFailure::Eof => return, - BeforeFirstTokenFailure::GrpcStatus => { - Err(Status::internal("failed before first token"))?; - } - BeforeFirstTokenFailure::InvalidResponse => { - let mut response = sequence_response(false, wants_logprobs, None); - response.outputs.as_mut().expect("sequence output").num_tokens = 2; - yield response; - return; - } + if close_before_first_token { + return; } if hang { @@ -943,36 +924,30 @@ async fn decode_cancellation_waits_for_submission_and_first_token() { } #[tokio::test] -async fn decode_cancellation_maps_stream_failures_to_cancelled() { - for (failure, case) in [ - (BeforeFirstTokenFailure::Eof, "premature EOF"), - (BeforeFirstTokenFailure::GrpcStatus, "gRPC failure"), - (BeforeFirstTokenFailure::InvalidResponse, "invalid response"), - ] { - let service = FakeVllm { - before_first_token_failure: failure, - ..Default::default() - }; - let server = FakeServer::start(service).await; - let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); - engine.start(0).await.expect("start"); +async fn decode_cancellation_maps_premature_eof_to_cancelled() { + let service = FakeVllm::default(); + service + .close_before_first_token + .store(true, Ordering::SeqCst); + let server = FakeServer::start(service).await; + let engine = engine(&server.endpoint, DisaggregationMode::Decode, 1); + engine.start(0).await.expect("start"); - let context = dynamo_backend_common::testing::mock_context(); - let mut stream = engine - .generate( - decode_request(), - GenerateContext::new(context.clone(), None), - ) - .await - .expect("decode stream"); - context.stop_generating(); - let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) - .await - .unwrap_or_else(|_| panic!("{case} did not release decode cancellation")) - .expect("cancelled terminal") - .expect("cancelled output"); - assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); - } + let context = dynamo_backend_common::testing::mock_context(); + let mut stream = engine + .generate( + decode_request(), + GenerateContext::new(context.clone(), None), + ) + .await + .expect("decode stream"); + context.stop_generating(); + let terminal = tokio::time::timeout(std::time::Duration::from_secs(2), stream.next()) + .await + .expect("premature EOF did not release decode cancellation") + .expect("cancelled terminal") + .expect("cancelled output"); + assert_eq!(terminal.finish_reason, Some(FinishReason::Cancelled)); } #[tokio::test]