diff --git a/rust/Cargo.lock b/rust/Cargo.lock index c7b3d66e4c..4193655805 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5997,7 +5997,6 @@ dependencies = [ "tls-listener", "tokio", "tokio-openssl", - "tokio-stream", "tokio-util", "tonic", "tonic-health", diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 8b22b6344a..b7b00398ce 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -35,7 +35,6 @@ thiserror-ext.workspace = true tls-listener.workspace = true tokio.workspace = true tokio-openssl.workspace = true -tokio-stream.workspace = true tokio-util.workspace = true tonic.workspace = true tonic-health.workspace = true diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index a1e631b2e4..1484a9f9e3 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -5,16 +5,17 @@ use std::pin::Pin; use std::sync::Arc; use std::time::Instant; -use futures::{Stream, StreamExt as _}; +use asynk_strim_attr::{TryYielder, try_stream}; +use futures::{Stream, StreamExt as _, pin_mut}; use thiserror_ext::AsReport as _; -use tokio::sync::mpsc; -use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status}; use tracing::{Span, info, info_span, warn}; use tracing_futures::Instrument as _; use uuid::Uuid; use vllm_llm::current_unix_timestamp_secs; -use vllm_text::{DecodedTextEvent, Prompt, SampledDelta, TextOutputStreamExt as _, TextRequest}; +use vllm_text::{ + DecodedTextEvent, Prompt, SampledDelta, TextOutputStream, TextOutputStreamExt as _, TextRequest, +}; use super::convert::{self, ResponseOpts}; use super::{InferenceServer, pb}; @@ -304,69 +305,74 @@ impl pb::inference_server::Inference for InferenceServiceImpl { .await .map_err(|error| log_text_error(&request_span, started_at, "submission", error))?; - let (tx, rx) = mpsc::channel(32); - - let task_span = request_span.clone(); - tokio::spawn( - async move { - futures::pin_mut!(stream); - while let Some(event) = stream.next().await { - let response = match event { - Err(error) => Err(log_text_error(&task_span, started_at, "stream", error)), - Ok(DecodedTextEvent::Start { - prompt_token_ids, - prompt_logprobs, - }) => { - let prompt_info = convert::to_prompt_info( - &prompt_token_ids, - prompt_logprobs.as_ref(), - &response_opts, - ); - Ok(pb::GenerateResponse { - prompt_info: Some(prompt_info), - outputs: None, - }) - } - Ok(DecodedTextEvent::TextDelta { - decoded, - sampled: - SampledDelta { - token_ids, - logprobs, - }, - finished, - }) => convert::to_sequence_output( - &decoded.text, - &token_ids, - logprobs.as_ref(), - finished.as_deref(), - &response_opts, - ) - .map(|outputs| pb::GenerateResponse { - prompt_info: None, - outputs: Some(outputs), - }), - }; - - let failed = response.is_err(); - if tx.send(response).await.is_err() || failed { - break; - } - } - info!( - parent: &task_span, - elapsed_ms = started_at.elapsed().as_millis() as u64, - "gRPC inference stream closed" - ); - } - .instrument(request_span), - ); - - let response_stream = ReceiverStream::new(rx); + let response_stream = + generate_stream_responses(stream, response_opts, request_span.clone(), started_at) + .instrument(request_span); Ok(Response::new(Box::pin(response_stream))) } } +/// Convert one decoded text stream into gRPC streaming responses. +/// +/// The response body owns the generation stream, so a client disconnect drops +/// it and aborts the request in the engine, even before any output arrives. +#[try_stream] +async fn generate_stream_responses( + stream: impl TextOutputStream, + response_opts: ResponseOpts, + request_span: Span, + started_at: Instant, + mut y: TryYielder, +) -> Result<(), Status> { + pin_mut!(stream); + + while let Some(event) = stream.next().await { + let response = match event { + Err(error) => Err(log_text_error(&request_span, started_at, "stream", error)), + Ok(DecodedTextEvent::Start { + prompt_token_ids, + prompt_logprobs, + }) => { + let prompt_info = convert::to_prompt_info( + &prompt_token_ids, + prompt_logprobs.as_ref(), + &response_opts, + ); + Ok(pb::GenerateResponse { + prompt_info: Some(prompt_info), + outputs: None, + }) + } + Ok(DecodedTextEvent::TextDelta { + decoded, + sampled: + SampledDelta { + token_ids, + logprobs, + }, + finished, + }) => convert::to_sequence_output( + &decoded.text, + &token_ids, + logprobs.as_ref(), + finished.as_deref(), + &response_opts, + ) + .map(|outputs| pb::GenerateResponse { + prompt_info: None, + outputs: Some(outputs), + }), + }; + y.yield_ok(response?).await; + } + info!( + parent: &request_span, + elapsed_ms = started_at.elapsed().as_millis() as u64, + "gRPC inference stream closed" + ); + Ok(()) +} + fn text_error_to_status(error: vllm_text::Error) -> Status { let message = error.to_report_string(); if error.is_request_validation_error() { diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index 0ddbfb096e..f46c658b01 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -1827,6 +1827,70 @@ async fn grpc_without_keepalive_keeps_unresponsive_connection_open() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn dropping_generate_stream_aborts_before_engine_output() { + let (submitted_tx, submitted_rx) = tokio::sync::oneshot::channel(); + let (inference_service, control_service, engine_health, engine_task) = + setup_grpc_service_with_engine_script( + b"engine-grpc-cancel-pending-output".to_vec(), + default_ready_response(), + Arc::new(FakeTextBackend), + None, + move |dealer, _push| { + boxed_test_future(async move { + let add = recv_engine_message(dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = + rmp_serde::from_slice(&add[1]).expect("decode request"); + submitted_tx.send(()).expect("signal submitted request"); + + // A queued request produces no output before cancellation. + let abort = recv_engine_message(dealer).await; + assert_eq!(abort[0].as_ref(), &[0x01]); + let request_ids: Vec = + rmp_serde::from_slice(&abort[1]).expect("decode abort request IDs"); + assert_eq!(request_ids, vec![request.request_id]); + }) + }, + ) + .await; + let (channel, server_task) = start_grpc_test_server( + inference_service, + control_service, + engine_health, + tokio_util::sync::CancellationToken::new(), + ) + .await; + let mut client = InferenceClient::new(channel); + let stream = client + .generate_stream(pb::GenerateRequest { + request_id: "test-cancel-pending-output".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + stopping: Some(pb::StoppingCriteria { + max_new_tokens: 10, + ..Default::default() + }), + ..Default::default() + }) + .await + .expect("start generation") + .into_inner(); + tokio::time::timeout(Duration::from_secs(5), submitted_rx) + .await + .expect("timed out waiting for engine submission") + .expect("engine submission signal"); + + drop(stream); + + let result = tokio::time::timeout(Duration::from_secs(5), engine_task).await; + server_task.abort(); + result + .expect("disconnect must abort without waiting for engine output") + .expect("mock engine task"); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn control_abort_resolves_external_id_and_empty_is_noop() {