diff --git a/e2e_test/completions/test_basic.py b/e2e_test/completions/test_basic.py index 865bbcf87..490a2b67c 100644 --- a/e2e_test/completions/test_basic.py +++ b/e2e_test/completions/test_basic.py @@ -261,3 +261,126 @@ def test_streaming_echo_max_tokens_zero(self, model, api_client): assert full_text == prompt, f"Expected echoed prompt, got: {full_text!r}" assert len(finish_reasons) == 1 assert finish_reasons[0] in ("stop", "length") + + +@pytest.mark.engine("sglang", "vllm", "tokenspeed") +@pytest.mark.gpu(1) +@pytest.mark.model("meta-llama/Llama-3.1-8B-Instruct") +@pytest.mark.parametrize("setup_backend", ["grpc"], indirect=True) +class TestCompletionBatch: + """Tests for batched prompt arrays on /v1/completions.""" + + PROMPTS = ["The capital of France is", "The capital of Germany is"] + + @staticmethod + def _collect_stream_by_index(stream): + """Consume a stream, returning per-index text, finish reasons, usage, and usage-chunk count.""" + texts = {} + finish_reasons = {} + usage = None + usage_chunks = 0 + for chunk in stream: + assert chunk.object == "text_completion" + if chunk.usage is not None: + usage = chunk.usage + usage_chunks += 1 + for choice in chunk.choices: + if choice.text: + texts.setdefault(choice.index, []).append(choice.text) + if choice.finish_reason: + finish_reasons.setdefault(choice.index, []).append(choice.finish_reason) + return ( + {index: "".join(parts) for index, parts in texts.items()}, + finish_reasons, + usage, + usage_chunks, + ) + + def test_batch_non_streaming(self, model, api_client): + """Test that a prompt array returns one choice per prompt with global indices.""" + + response = api_client.completions.create( + model=model, + prompt=self.PROMPTS, + max_tokens=20, + temperature=0, + ) + + assert len(response.choices) == len(self.PROMPTS) + assert sorted(choice.index for choice in response.choices) == [0, 1] + for choice in response.choices: + assert isinstance(choice.text, str) + assert len(choice.text) > 0 + assert choice.finish_reason in ("stop", "length") + + assert response.usage is not None + assert response.usage.prompt_tokens > 0 + assert response.usage.completion_tokens > 0 + assert response.usage.total_tokens == ( + response.usage.prompt_tokens + response.usage.completion_tokens + ) + + def test_batch_non_streaming_with_n(self, model, api_client): + """Test prompt-major global indices with n > 1: index = prompt_index * n + i.""" + + n = 2 + response = api_client.completions.create( + model=model, + prompt=self.PROMPTS, + max_tokens=20, + temperature=0.7, + n=n, + echo=True, + ) + + assert len(response.choices) == len(self.PROMPTS) * n + assert [choice.index for choice in response.choices] == [0, 1, 2, 3] + for choice in response.choices: + assert choice.text.startswith(self.PROMPTS[choice.index // n]) + + def test_batch_echo_maps_prompts_to_choices(self, model, api_client): + """Test that echo=True prepends each prompt to its own choice.""" + + response = api_client.completions.create( + model=model, + prompt=self.PROMPTS, + max_tokens=10, + temperature=0, + echo=True, + ) + + choices = {choice.index: choice for choice in response.choices} + assert len(choices) == len(self.PROMPTS) + for prompt_index, prompt in enumerate(self.PROMPTS): + assert choices[prompt_index].text.startswith(prompt) + + def test_batch_streaming(self, model, api_client): + """Test streaming with a prompt array and n > 1: global per-choice deltas + and exactly one aggregated usage chunk.""" + + n = 2 + stream = api_client.completions.create( + model=model, + prompt=self.PROMPTS, + max_tokens=20, + temperature=0.7, + n=n, + stream=True, + stream_options={"include_usage": True}, + ) + + texts, finish_reasons, usage, usage_chunks = self._collect_stream_by_index(stream) + + expected_indices = list(range(len(self.PROMPTS) * n)) + assert sorted(texts) == expected_indices, ( + f"Expected deltas for all choices, got {sorted(texts)}" + ) + assert sorted(finish_reasons) == expected_indices + for index in expected_indices: + assert len(texts[index]) > 0 + assert finish_reasons[index] in (["stop"], ["length"]) + + assert usage_chunks == 1, f"Expected exactly one usage chunk, got {usage_chunks}" + assert usage is not None + assert usage.prompt_tokens > 0 + assert usage.completion_tokens > 0 diff --git a/model_gateway/src/routers/grpc/common/response_collection.rs b/model_gateway/src/routers/grpc/common/response_collection.rs index 30ee5b9df..b8bf71e70 100644 --- a/model_gateway/src/routers/grpc/common/response_collection.rs +++ b/model_gateway/src/routers/grpc/common/response_collection.rs @@ -67,6 +67,14 @@ pub(crate) async fn collect_responses( "Embedding result encountered in response collection", )); } + // Batches are split into per-prompt results by the completion processor + // before collection. + ExecutionResult::Batch { .. } => { + return Err(error::internal_error( + "invalid_execution_mode", + "Batch result encountered in response collection", + )); + } }; if all_responses.is_empty() { diff --git a/model_gateway/src/routers/grpc/common/stages/request_execution.rs b/model_gateway/src/routers/grpc/common/stages/request_execution.rs index e6b05d983..c0957be4e 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -4,7 +4,7 @@ use std::time::Instant; use async_trait::async_trait; use axum::response::Response; -use futures::future::join_all; +use futures::future::{join_all, try_join_all}; use tracing::{debug, error, info_span, Instrument}; use super::PipelineStage; @@ -15,8 +15,8 @@ use crate::{ grpc::{ common::stages::encode::EncodeDispatchPlan, context::{ - ClientSelection, ExecutionPlan, ExecutionResult, LoadGuards, PdTiming, - RequestContext, WorkerSelection, + ClientSelection, ExecutionPlan, ExecutionPlanKind, ExecutionResult, LoadGuards, + PdTiming, RequestContext, WorkerSelection, }, proto_wrapper::{ ProtoEmbedRequest, ProtoGenerateRequest, ProtoRequest, ProtoResponseVariant, @@ -167,7 +167,15 @@ impl PipelineStage for RequestExecutionStage { ) })?; - ctx.state.load_guards = Some(LoadGuards::new(workers, ctx.input.headers.as_ref())); + let sub_requests = match &execution_plan { + ExecutionPlan::Batch { requests, .. } => requests.len(), + _ => 1, + }; + ctx.state.load_guards = Some(LoadGuards::scaled( + workers, + ctx.input.headers.as_ref(), + sub_requests, + )); // Extract dispatch metadata for tracing span let dispatch = ctx.state.dispatch.as_ref(); @@ -204,6 +212,10 @@ impl PipelineStage for RequestExecutionStage { self.execute_epd_dispatch(request, clients, workers, model, encode_dispatch) .await } + ExecutionPlan::Batch { kind, requests, .. } => { + self.execute_batch_dispatch(kind, requests, clients, workers, model) + .await + } } } .instrument(span) @@ -329,6 +341,37 @@ impl RequestExecutionStage { }); } + /// Dispatch one backend request per batched prompt concurrently, preserving + /// prompt order. Fail-fast: the first failed dispatch fails the batch and + /// drops the remaining streams (abort-on-drop reclaims them backend-side). + async fn execute_batch_dispatch( + &self, + kind: ExecutionPlanKind, + requests: Vec, + clients: &ClientSelection, + workers: &WorkerSelection, + model: &str, + ) -> Result { + let dispatches = requests.into_iter().map(|request| { + let mut clients = clients.clone(); + async move { + match kind { + ExecutionPlanKind::Single => { + self.execute_single(request, &mut clients, workers).await + } + // Completion EPD carries no encode jobs; sub-requests dispatch as PD. + ExecutionPlanKind::PrefillDecode | ExecutionPlanKind::EncodePrefillDecode => { + self.execute_pd_dispatch(request, &mut clients, workers, model) + .await + } + } + } + }); + + let results = try_join_all(dispatches).await?; + Ok(ExecutionResult::Batch { results }) + } + async fn execute_single( &self, mut proto_request: ProtoGenerateRequest, diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index 77d457561..668a11021 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -163,7 +163,18 @@ pub(crate) struct EncodeOutputs { pub(crate) enum ExecutionPlan { Single(ProtoRequest), PrefillDecode(ProtoGenerateRequest), - EncodePrefillDecode { request: ProtoGenerateRequest }, + EncodePrefillDecode { + request: ProtoGenerateRequest, + }, + /// Batched completion fan-out: one backend request per prompt, all + /// dispatched with the disaggregation shape given by `kind`. Sub-request + /// ids are `{shared_request_id}-p{i}`; the client-visible response id is + /// `shared_request_id`. + Batch { + kind: ExecutionPlanKind, + shared_request_id: String, + requests: Vec, + }, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -192,6 +203,9 @@ impl ExecutionPlan { Self::PrefillDecode(request) | Self::EncodePrefillDecode { request, .. } => { request.request_id() } + Self::Batch { + shared_request_id, .. + } => shared_request_id, } } @@ -199,7 +213,8 @@ impl ExecutionPlan { match self { Self::Single(ProtoRequest::Generate(_)) | Self::PrefillDecode(_) - | Self::EncodePrefillDecode { .. } => "generate", + | Self::EncodePrefillDecode { .. } + | Self::Batch { .. } => "generate", Self::Single(ProtoRequest::Embed(_)) => "embed", } } @@ -209,6 +224,11 @@ impl ExecutionPlan { Self::Single(_) => "single", Self::PrefillDecode(_) => "prefill_decode", Self::EncodePrefillDecode { .. } => "encode_prefill_decode", + Self::Batch { kind, .. } => match kind { + ExecutionPlanKind::Single => "single", + ExecutionPlanKind::PrefillDecode => "prefill_decode", + ExecutionPlanKind::EncodePrefillDecode => "encode_prefill_decode", + }, } } } @@ -229,8 +249,11 @@ pub(crate) enum PreparationOutput { tool_constraints: Option<(String, String)>, }, Completion { - original_text: String, - token_ids: Vec, + /// One entry per prompt; scalar requests carry exactly one. + items: Vec, + /// `Some` iff multiple prompts: their texts joined for routing, + /// mirroring the HTTP router's `extract_text_for_routing`. + joined_routing_text: Option, }, Generate { original_text: Option, @@ -252,16 +275,25 @@ pub(crate) enum PreparationOutput { }, } +/// One tokenized completion prompt. +pub(crate) struct CompletionItem { + pub text: String, + pub token_ids: Vec, +} + impl PreparationOutput { - /// Token IDs (common to all variants) + /// Token IDs (common to all variants). Batched completions expose the + /// first prompt's tokens as the routing-affinity proxy. pub fn token_ids(&self) -> &[u32] { match self { Self::Chat { token_ids, .. } | Self::Messages { token_ids, .. } - | Self::Completion { token_ids, .. } | Self::Generate { token_ids, .. } | Self::Embedding { token_ids, .. } | Self::Harmony { token_ids, .. } => token_ids, + Self::Completion { items, .. } => { + items.first().map_or(&[], |item| item.token_ids.as_slice()) + } } } @@ -275,9 +307,13 @@ impl PreparationOutput { | Self::Messages { processed_messages, .. } => Some(&processed_messages.text), - Self::Completion { original_text, .. } | Self::Embedding { original_text, .. } => { - Some(original_text) - } + Self::Completion { + items, + joined_routing_text, + } => joined_routing_text + .as_deref() + .or_else(|| items.first().map(|item| item.text.as_str())), + Self::Embedding { original_text, .. } => Some(original_text), Self::Generate { original_text, .. } => original_text.as_deref(), Self::Harmony { selection_text, .. } => Some(selection_text), } @@ -306,6 +342,7 @@ pub(crate) enum WorkerSelection { } /// Client selection (Step 3) +#[derive(Clone)] pub(crate) enum ClientSelection { Single { client: GrpcClient, @@ -339,6 +376,11 @@ pub(crate) enum LoadGuards { _prefill: WorkerLoadGuard, _decode: WorkerLoadGuard, }, + /// Batched completion fan-out: one guard set per sub-request so load-aware + /// policies see the real backend concurrency. + Batch { + _guards: Vec, + }, } impl LoadGuards { @@ -355,6 +397,17 @@ impl LoadGuards { }, } } + + /// One guard set per concurrent sub-request. + pub fn scaled(selection: &WorkerSelection, headers: Option<&HeaderMap>, count: usize) -> Self { + if count <= 1 { + Self::new(selection, headers) + } else { + Self::Batch { + _guards: (0..count).map(|_| Self::new(selection, headers)).collect(), + } + } + } } /// Response processing state (Step 6) @@ -825,6 +878,10 @@ pub(crate) enum ExecutionResult { Embedding { response: ProtoEmbedComplete, }, + /// Batched completion fan-out: one result per prompt, in prompt order. + Batch { + results: Vec, + }, } /// Timing context threaded from PD execution into the streaming layer so the @@ -852,3 +909,45 @@ pub(crate) enum FinalResponse { /// Messages API response Messages(Message), } + +#[cfg(test)] +mod tests { + use super::*; + + fn completion_prep(texts: &[&str], joined: Option<&str>) -> PreparationOutput { + PreparationOutput::Completion { + items: texts + .iter() + .enumerate() + .map(|(i, text)| CompletionItem { + text: (*text).to_string(), + token_ids: vec![i as u32], + }) + .collect(), + joined_routing_text: joined.map(str::to_string), + } + } + + #[test] + fn completion_preparation_routes_by_first_item_or_joined_text() { + let scalar = completion_prep(&["hello"], None); + assert_eq!(scalar.routing_text(), Some("hello")); + assert_eq!(scalar.token_ids(), &[0]); + + let batch = completion_prep(&["a", "b"], Some("a b")); + assert_eq!(batch.routing_text(), Some("a b")); + assert_eq!(batch.token_ids(), &[0]); + } + + #[test] + fn batch_execution_plan_reports_shared_id_and_kind_label() { + let plan = ExecutionPlan::Batch { + kind: ExecutionPlanKind::PrefillDecode, + shared_request_id: "cmpl_shared".to_string(), + requests: vec![], + }; + assert_eq!(plan.request_id(), "cmpl_shared"); + assert_eq!(plan.request_type(), "generate"); + assert_eq!(plan.mode_label(), "prefill_decode"); + } +} diff --git a/model_gateway/src/routers/grpc/harmony/streaming.rs b/model_gateway/src/routers/grpc/harmony/streaming.rs index cb5076b5e..93e46b9d1 100644 --- a/model_gateway/src/routers/grpc/harmony/streaming.rs +++ b/model_gateway/src/routers/grpc/harmony/streaming.rs @@ -155,6 +155,16 @@ impl HarmonyStreamingProcessor { ); let _ = tx.send(Ok(SseEncoder::done())); } + // Batch results exist only on the completions pipeline. + context::ExecutionResult::Batch { .. } => { + error!("Harmony streaming not supported for batched results"); + utils::send_error_sse( + &tx, + "Batched results not supported in Harmony streaming", + "invalid_request_error", + ); + let _ = tx.send(Ok(SseEncoder::done())); + } } // Return SSE response @@ -549,6 +559,10 @@ impl HarmonyStreamingProcessor { context::ExecutionResult::Embedding { .. } => { Err("Embeddings not supported in Responses API streaming".to_string()) } + // Batch results exist only on the completions pipeline. + context::ExecutionResult::Batch { .. } => { + Err("Batched results not supported in Responses API streaming".to_string()) + } } } diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index fadf4f8b4..6d1d8d9e0 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -5,13 +5,16 @@ use std::{sync::Arc, time::Instant}; +use futures::future::try_join_all; use llm_tokenizer::{ stop::{SequenceDecoderOutput, StopSequenceDecoder}, traits::Tokenizer, }; use openai_protocol::{ chat::{ChatChoice, ChatCompletionMessage, ChatCompletionRequest, ChatCompletionResponse}, - common::{FunctionCallResponse, Tool, ToolCall, ToolChoice, ToolChoiceValue, Usage}, + common::{ + FunctionCallResponse, StringOrArray, Tool, ToolCall, ToolChoice, ToolChoiceValue, Usage, + }, completion::{CompletionChoice, CompletionRequest, CompletionResponse}, generate::{GenerateMetaInfo, GenerateRequest, GenerateResponse}, messages::{self, CreateMessageRequest, Message}, @@ -771,9 +774,10 @@ impl ResponseProcessor { /// Process non-streaming completion response /// - /// Collects all responses (supports n>1), decodes tokens through stop decoder, - /// applies `echo` and `suffix`, and builds `CompletionResponse` with legacy - /// `LogProbs` format. + /// Collects all responses (supports n>1 and batched prompts), decodes tokens + /// through the stop decoder, applies `echo` and `suffix`, and builds one + /// `CompletionResponse` with prompt-major global choice indices + /// (`prompt_index * n + choice_index`). pub async fn process_non_streaming_completion_response( &self, execution_result: ExecutionResult, @@ -781,90 +785,116 @@ impl ResponseProcessor { dispatch: DispatchMetadata, _tokenizer: Arc, stop_decoder: &mut StopSequenceDecoder, - prompt_text: &str, ) -> Result { let request_logprobs = completion_req.logprobs.is_some(); - let all_responses = - response_collection::collect_responses(execution_result, request_logprobs).await?; + let per_prompt_results = match execution_result { + ExecutionResult::Batch { results } => results, + other => vec![other], + }; + let choices_per_prompt = completion_req.n.unwrap_or(1).max(1); + let prompt_texts: Vec<&str> = match &completion_req.prompt { + StringOrArray::String(text) => vec![text.as_str()], + StringOrArray::Array(texts) => texts.iter().map(String::as_str).collect(), + }; + + // Drain all sub-streams concurrently; decoding below stays sequential + // (shared stop decoder). + let collected = try_join_all( + per_prompt_results + .into_iter() + .map(|result| response_collection::collect_responses(result, request_logprobs)), + ) + .await?; let mut total_prompt = 0u32; let mut total_completion = 0u32; let mut choices = Vec::new(); - for (i, complete) in all_responses.into_iter().enumerate() { - stop_decoder.reset(); + for (prompt_index, all_responses) in collected.into_iter().enumerate() { + let prompt_text = prompt_texts.get(prompt_index).copied().unwrap_or_default(); + let index_offset = prompt_index as u32 * choices_per_prompt; + // n>1 choices share one prompt: max within a prompt, summed across prompts. + let mut prompt_tokens = 0u32; - let outputs = match stop_decoder.process_tokens(complete.output_ids()) { - Ok(outputs) => outputs, - Err(e) => { - return Err(error::internal_error( - "process_tokens_failed", - format!("Failed to process tokens: {e}"), - )) - } - }; + // Arrival order, not `complete.index()`: SGLang non-streaming + // Complete frames carry index 0 for every choice. + for (i, complete) in all_responses.into_iter().enumerate() { + stop_decoder.reset(); - let mut decoded_text = String::new(); - for output in outputs { - match output { - SequenceDecoderOutput::Text(t) => decoded_text.push_str(&t), - SequenceDecoderOutput::StoppedWithText(t) => { - decoded_text.push_str(&t); - break; + let outputs = match stop_decoder.process_tokens(complete.output_ids()) { + Ok(outputs) => outputs, + Err(e) => { + return Err(error::internal_error( + "process_tokens_failed", + format!("Failed to process tokens: {e}"), + )) + } + }; + + let mut decoded_text = String::new(); + for output in outputs { + match output { + SequenceDecoderOutput::Text(t) => decoded_text.push_str(&t), + SequenceDecoderOutput::StoppedWithText(t) => { + decoded_text.push_str(&t); + break; + } + SequenceDecoderOutput::Stopped => break, + SequenceDecoderOutput::Held => {} } - SequenceDecoderOutput::Stopped => break, - SequenceDecoderOutput::Held => {} } - } - if let SequenceDecoderOutput::Text(t) = stop_decoder.flush() { - decoded_text.push_str(&t); - } + if let SequenceDecoderOutput::Text(t) = stop_decoder.flush() { + decoded_text.push_str(&t); + } - total_prompt = total_prompt.max(complete.prompt_tokens()); - total_completion += complete.completion_tokens(); - - let finish_reason = { - let reason = complete.finish_reason(); - if reason.is_empty() { - None - } else if reason == "stop" || reason == "length" { - Some(reason.to_string()) - } else if let Ok(json) = serde_json::from_str::(reason) { - json.get("type").and_then(|v| v.as_str()).map(|s| match s { - "length" => "length".to_string(), - "stop" => "stop".to_string(), - other => other.to_string(), - }) + prompt_tokens = prompt_tokens.max(complete.prompt_tokens()); + total_completion += complete.completion_tokens(); + + let finish_reason = { + let reason = complete.finish_reason(); + if reason.is_empty() { + None + } else if reason == "stop" || reason == "length" { + Some(reason.to_string()) + } else if let Ok(json) = serde_json::from_str::(reason) { + json.get("type").and_then(|v| v.as_str()).map(|s| match s { + "length" => "length".to_string(), + "stop" => "stop".to_string(), + other => other.to_string(), + }) + } else { + Some(reason.to_string()) + } + }; + + let matched_stop = complete.matched_stop_json(); + + let suffix_len = completion_req.suffix.as_ref().map_or(0, |s| s.len()); + let echo_len = if completion_req.echo { + prompt_text.len() } else { - Some(reason.to_string()) + 0 + }; + let mut text = String::with_capacity(echo_len + decoded_text.len() + suffix_len); + if completion_req.echo { + text.push_str(prompt_text); + } + text.push_str(&decoded_text); + if let Some(ref sfx) = completion_req.suffix { + text.push_str(sfx); } - }; - - let matched_stop = complete.matched_stop_json(); - let suffix_len = completion_req.suffix.as_ref().map_or(0, |s| s.len()); - let echo_len = if completion_req.echo { - prompt_text.len() - } else { - 0 - }; - let mut text = String::with_capacity(echo_len + decoded_text.len() + suffix_len); - if completion_req.echo { - text.push_str(prompt_text); - } - text.push_str(&decoded_text); - if let Some(ref sfx) = completion_req.suffix { - text.push_str(sfx); + choices.push(CompletionChoice { + text, + index: index_offset + i as u32, + logprobs: None, // TODO: wire legacy LogProbs from backend token_logprobs + finish_reason: finish_reason.or_else(|| Some("stop".to_string())), + matched_stop, + }); } - choices.push(CompletionChoice { - text, - index: i as u32, - logprobs: None, // TODO: wire legacy LogProbs from backend token_logprobs - finish_reason: finish_reason.or_else(|| Some("stop".to_string())), - matched_stop, - }); + total_prompt += prompt_tokens; } Ok(CompletionResponse { diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs b/model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs index 2cc1f331c..4060e27a5 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/preparation.rs @@ -1,4 +1,4 @@ -//! Completion preparation stage: resolve prompt, tokenize, create stop decoder. +//! Completion preparation stage: resolve prompt(s), tokenize, create stop decoder. //! //! This is the `/v1/completions` Stage 1 equivalent. It intentionally builds on top of //! the native completion pipeline typing introduced in PR #840. It keeps @@ -7,6 +7,7 @@ use async_trait::async_trait; use axum::response::Response; +use futures::future::try_join_all; use openai_protocol::common::StringOrArray; use tracing::error; @@ -14,7 +15,7 @@ use crate::routers::{ error, grpc::{ common::stages::PipelineStage, - context::{PreparationOutput, RequestContext}, + context::{CompletionItem, PreparationOutput, RequestContext}, utils, }, }; @@ -29,26 +30,43 @@ impl PipelineStage for CompletionPreparationStage { let tokenizer = utils::resolve_tokenizer(ctx, "CompletionPreparationStage::execute").map_err(|e| *e)?; - let prompt_text = match &request.prompt { - StringOrArray::String(text) => text.clone(), - StringOrArray::Array(_) => { - return Err(error::bad_request( - "batch_prompts_not_supported", - "Batched prompt arrays are not supported for gRPC /v1/completions yet", - )); - } + let prompts: Vec = match &request.prompt { + StringOrArray::String(text) => vec![text.clone()], + // Empty arrays are rejected at the boundary (`validate_completion_prompt`). + StringOrArray::Array(texts) => texts.clone(), }; - let encoding = utils::encode_blocking(tokenizer.clone(), prompt_text.clone(), false) - .await - .map_err(|e| { - error!( - function = "CompletionPreparationStage::execute", - error = %e, - "Tokenization failed" - ); - error::bad_request("tokenization_failed", format!("Tokenization failed: {e}")) - })?; + let encodings = try_join_all( + prompts + .iter() + .map(|prompt| utils::encode_blocking(tokenizer.clone(), prompt.clone(), false)), + ) + .await + .map_err(|e| { + error!( + function = "CompletionPreparationStage::execute", + error = %e, + "Tokenization failed" + ); + error::bad_request("tokenization_failed", format!("Tokenization failed: {e}")) + })?; + + let items: Vec = prompts + .into_iter() + .zip(encodings) + .map(|(text, encoding)| CompletionItem { + text, + token_ids: encoding.token_ids().to_vec(), + }) + .collect(); + + let joined_routing_text = (items.len() > 1).then(|| { + items + .iter() + .map(|item| item.text.as_str()) + .collect::>() + .join(" ") + }); let stop_decoder = utils::create_stop_decoder( &tokenizer, @@ -60,8 +78,8 @@ impl PipelineStage for CompletionPreparationStage { ); ctx.state.preparation = Some(PreparationOutput::Completion { - original_text: prompt_text, - token_ids: encoding.token_ids().to_vec(), + items, + joined_routing_text, }); ctx.state.response.stop_decoder = Some(stop_decoder); diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs index 17c046c42..df790ecde 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -1,8 +1,8 @@ -//! Completion request building stage: build proto GenerateRequest from CompletionRequest +//! Completion request building stage: build proto GenerateRequest(s) from CompletionRequest //! //! Stage 4 for the `/v1/completions` pipeline, parallel to `MessageRequestBuildingStage` -//! from the Messages rollout. Builds backend-specific proto `GenerateRequest` from -//! `PreparationOutput` + `CompletionRequest` sampling parameters. +//! from the Messages rollout. Builds backend-specific proto `GenerateRequest`s from +//! `PreparationOutput` + `CompletionRequest` sampling parameters — one per prompt. //! //! Completions has richer sampling knobs than Messages (frequency_penalty, presence_penalty, //! repetition_penalty, min_p, n, logprobs, structured output constraints) but no tools @@ -10,14 +10,20 @@ use async_trait::async_trait; use axum::response::Response; +use openai_protocol::completion::CompletionRequest; use tracing::error; use uuid::Uuid; use crate::routers::{ error, grpc::{ + client::GrpcClient, common::stages::{helpers, PipelineStage}, - context::{ClientSelection, ExecutionPlan, ExecutionPlanKind, RequestContext}, + context::{ + ClientSelection, CompletionItem, ExecutionPlan, ExecutionPlanKind, PreparationOutput, + RequestContext, RequestType, WorkerSelection, + }, + proto_wrapper::ProtoGenerateRequest, }, }; @@ -33,6 +39,63 @@ impl CompletionRequestBuildingStage { plan_kind, } } + + /// Build one backend request for one prompt. PD bootstrap rooms are minted + /// per call, so injection runs per sub-request rather than + /// build-once-then-clone. + #[expect( + clippy::result_large_err, + reason = "Response is the standard error type in the pipeline stage pattern" + )] + fn build_proto_request( + &self, + builder_client: &GrpcClient, + request_id: String, + item: &CompletionItem, + completion_request: &CompletionRequest, + request_type: &RequestType, + workers: Option<&WorkerSelection>, + ) -> Result { + let mut proto_request = builder_client + .build_completion_request( + request_id, + completion_request, + item.text.clone(), + item.token_ids.clone(), + ) + .map_err(|e| { + error!( + function = "CompletionRequestBuildingStage::execute", + error = %e, + "Failed to build generate request" + ); + error::bad_request( + "invalid_request_parameters", + format!("Invalid request parameters: {e}"), + ) + })?; + + helpers::apply_sampling_defaults_to_generate_request( + &mut proto_request, + request_type, + workers, + ); + + if self.inject_pd_metadata { + if let Some(workers) = workers { + helpers::maybe_inject_pd_metadata(&mut proto_request, workers); + } + } + + // EPD: inject the prefill->decode KV rendezvous. Completion EPD is + // text-only (no encode jobs), so this is the only EPD injection here. + // No-op unless the backend carries it in the request. + if let Some(workers) = workers { + helpers::maybe_inject_pd_rendezvous(&mut proto_request, workers); + } + + Ok(proto_request) + } } #[async_trait] @@ -46,6 +109,17 @@ impl PipelineStage for CompletionRequestBuildingStage { error::internal_error("preparation_not_completed", "Preparation not completed") })?; + let PreparationOutput::Completion { items, .. } = prep else { + error!( + function = "CompletionRequestBuildingStage::execute", + "Preparation output is not a completion" + ); + return Err(error::internal_error( + "unexpected_preparation_output", + "Preparation output is not a completion", + )); + }; + let clients = ctx.state.clients.as_ref().ok_or_else(|| { error!( function = "CompletionRequestBuildingStage::execute", @@ -64,47 +138,49 @@ impl PipelineStage for CompletionRequestBuildingStage { ClientSelection::Disaggregated { prefill, .. } => prefill, }; - let request_id = format!("cmpl_{}", Uuid::now_v7()); - - let mut proto_request = builder_client - .build_completion_request( - request_id, - &completion_request, - prep.routing_text().unwrap_or_default().to_string(), - prep.token_ids().to_vec(), - ) - .map_err(|e| { - error!( - function = "CompletionRequestBuildingStage::execute", - error = %e, - "Failed to build generate request" - ); - error::bad_request( - "invalid_request_parameters", - format!("Invalid request parameters: {e}"), - ) - })?; - - helpers::apply_sampling_defaults_to_generate_request( - &mut proto_request, - &ctx.input.request_type, - ctx.state.workers.as_ref(), - ); + let shared_request_id = format!("cmpl_{}", Uuid::now_v7()); + let request_type = &ctx.input.request_type; + let workers = ctx.state.workers.as_ref(); - if self.inject_pd_metadata { - if let Some(workers) = ctx.state.workers.as_ref() { - helpers::maybe_inject_pd_metadata(&mut proto_request, workers); + let plan = match items.as_slice() { + [] => { + return Err(error::internal_error( + "preparation_not_completed", + "No prompts prepared", + )) } - } - - // EPD: inject the prefill->decode KV rendezvous. Completion EPD is - // text-only (no encode jobs), so this is the only EPD injection here. - // No-op unless the backend carries it in the request. - if let Some(workers) = ctx.state.workers.as_ref() { - helpers::maybe_inject_pd_rendezvous(&mut proto_request, workers); - } + [item] => ExecutionPlan::generate( + self.plan_kind, + self.build_proto_request( + builder_client, + shared_request_id, + item, + &completion_request, + request_type, + workers, + )?, + ), + batch_items => { + let mut requests = Vec::with_capacity(batch_items.len()); + for (i, item) in batch_items.iter().enumerate() { + requests.push(self.build_proto_request( + builder_client, + format!("{shared_request_id}-p{i}"), + item, + &completion_request, + request_type, + workers, + )?); + } + ExecutionPlan::Batch { + kind: self.plan_kind, + shared_request_id, + requests, + } + } + }; - ctx.state.execution_plan = Some(ExecutionPlan::generate(self.plan_kind, proto_request)); + ctx.state.execution_plan = Some(plan); Ok(None) } diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs index 47af78c31..e373c3066 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/response_processing.rs @@ -112,13 +112,6 @@ impl PipelineStage for CompletionResponseProcessingStage { ) })?; - let prompt_text = ctx - .state - .preparation - .as_ref() - .and_then(|p| p.routing_text()) - .unwrap_or(""); - let response = self .processor .process_non_streaming_completion_response( @@ -127,7 +120,6 @@ impl PipelineStage for CompletionResponseProcessingStage { dispatch, tokenizer, stop_decoder, - prompt_text, ) .await?; diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 0cb61ab91..505a8dc3b 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -11,6 +11,7 @@ use std::{ use axum::response::Response; use bytes::Bytes; +use futures::future::try_join_all; use llm_tokenizer::{ stop::{SequenceDecoderOutput, StopSequenceDecoder}, traits::Tokenizer, @@ -47,6 +48,26 @@ use crate::{ }, }; +/// One backend stream of a `/v1/completions` request. Batched requests fan +/// out into several, each remapped by a prompt-major choice-index offset. +enum CompletionStreamUnit { + Single(ProtoStream), + PrefillDecode { + prefill: ProtoStream, + decode: Box, + }, +} + +/// Per-stream token/timing totals returned by the completion chunk loop; the +/// coordinator aggregates them into the single usage chunk and metrics record. +struct CompletionStreamOutcome { + prompt_tokens: u32, + cached_tokens: u32, + reasoning_tokens: u32, + completion_tokens: u32, + first_token_time: Option, +} + /// Shared streaming processor for both single and prefill/decode dispatch modes #[derive(Clone)] pub(crate) struct StreamingProcessor { @@ -183,6 +204,15 @@ impl StreamingProcessor { ); let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); } + // Batch results exist only on the completions pipeline. + context::ExecutionResult::Batch { .. } => { + utils::send_error_sse( + &tx, + "Batched results not supported in chat streaming", + "invalid_request_error", + ); + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + } } // Return SSE response @@ -784,6 +814,15 @@ impl StreamingProcessor { ); let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); } + // Batch results exist only on the completions pipeline. + context::ExecutionResult::Batch { .. } => { + utils::send_error_sse( + &tx, + "Batched results not supported in streaming generate", + "invalid_request_error", + ); + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + } } // Return SSE response @@ -1632,6 +1671,17 @@ impl StreamingProcessor { let mut buf = Vec::with_capacity(256); let _ = Self::send_messages_event(&tx, &mut buf, &error_event); } + // Batch results exist only on the completions pipeline. + context::ExecutionResult::Batch { .. } => { + let error_event = MessageStreamEvent::Error { + error: messages::ErrorResponse { + error_type: "invalid_request_error".to_string(), + message: "Batched results not supported for Messages API".to_string(), + }, + }; + let mut buf = Vec::with_capacity(256); + let _ = Self::send_messages_event(&tx, &mut buf, &error_event); + } } build_sse_response(rx) @@ -2311,6 +2361,11 @@ impl StreamingProcessor { // ========================================================================= /// Entry point for `/v1/completions` streaming. + /// + /// Batched requests fan out into one stream unit per prompt; all units + /// write typed SSE events (with prompt-major index offsets) into one + /// channel, and the coordinator emits the single aggregated usage chunk, + /// metrics record, and `[DONE]`. pub fn process_completion_streaming_response( self: Arc, execution_result: context::ExecutionResult, @@ -2318,91 +2373,196 @@ impl StreamingProcessor { dispatch: context::DispatchMetadata, tokenizer: Arc, ) -> Response { - let stop_params = ( - completion_request.stop.clone(), - completion_request.stop_token_ids.clone(), - completion_request.skip_special_tokens, - completion_request.no_stop_trim, - completion_request.ignore_eos, - ); - let (tx, rx) = mpsc::unbounded_channel::>(); - match execution_result { - context::ExecutionResult::Single { stream } => { - let processor = self.clone(); - #[expect( - clippy::disallowed_methods, - reason = "streaming task is fire-and-forget; client disconnect terminates it" - )] - tokio::spawn(async move { - let result = processor - .process_completion_streaming_chunks( - stream, - dispatch, - tokenizer, - stop_params, - completion_request, - &tx, - ) - .await; + let units = match Self::completion_stream_units(execution_result) { + Ok(units) => units, + Err(message) => { + utils::send_error_sse(&tx, message, "invalid_request_error"); + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + return build_sse_response(rx); + } + }; - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error"); + let processor = self; + #[expect( + clippy::disallowed_methods, + reason = "streaming task is fire-and-forget; client disconnect terminates it" + )] + tokio::spawn(async move { + let start_time = Instant::now(); + let choices_per_prompt = completion_request.n.unwrap_or(1).max(1); + let echo = completion_request.echo; + let include_usage = completion_request + .stream_options + .as_ref() + .and_then(|opts| opts.include_usage) + .unwrap_or(false); + let prompt_texts: Vec<&str> = match &completion_request.prompt { + StringOrArray::String(text) => vec![text.as_str()], + StringOrArray::Array(texts) => texts.iter().map(String::as_str).collect(), + }; + + // Fail-fast: the first stream error cancels the remaining units + // (their streams abort on drop) and fails the whole request. + let outcomes = + try_join_all(units.into_iter().enumerate().map(|(prompt_index, unit)| { + let stop_params = ( + completion_request.stop.clone(), + completion_request.stop_token_ids.clone(), + completion_request.skip_special_tokens, + completion_request.no_stop_trim, + completion_request.ignore_eos, + ); + let dispatch = dispatch.clone(); + let tokenizer = tokenizer.clone(); + let completion_request = completion_request.clone(); + let prompt_text = if echo { + prompt_texts.get(prompt_index).copied().unwrap_or_default() + } else { + "" + }; + let index_offset = prompt_index as u32 * choices_per_prompt; + let processor = &processor; + let tx = &tx; + async move { + match unit { + CompletionStreamUnit::Single(stream) => { + processor + .process_completion_streaming_chunks( + stream, + dispatch, + tokenizer, + stop_params, + completion_request, + prompt_text, + index_offset, + tx, + ) + .await + } + CompletionStreamUnit::PrefillDecode { prefill, decode } => { + processor + .process_prefill_decode_completion_streaming_chunks( + prefill, + *decode, + dispatch, + tokenizer, + stop_params, + completion_request, + prompt_text, + index_offset, + tx, + ) + .await + } + } } + })) + .await; - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); - }); + match outcomes { + Err(e) => { + utils::send_error_sse(&tx, &e, "internal_error"); + } + Ok(outcomes) => { + let mut total_prompt = 0u32; + let mut total_cached = 0u32; + let mut total_reasoning = 0u32; + let mut total_completion = 0u32; + let mut first_token_time: Option = None; + for outcome in outcomes { + total_prompt += outcome.prompt_tokens; + total_cached += outcome.cached_tokens; + total_reasoning += outcome.reasoning_tokens; + total_completion += outcome.completion_tokens; + first_token_time = match (first_token_time, outcome.first_token_time) { + (Some(current), Some(candidate)) => Some(current.min(candidate)), + (current, candidate) => current.or(candidate), + }; + } + + if include_usage { + let usage_chunk = CompletionStreamResponse { + id: dispatch.request_id.clone(), + object: "text_completion".to_string(), + created: dispatch.created, + choices: vec![], + model: dispatch.model.clone(), + system_fingerprint: dispatch.weight_version.clone(), + usage: Some(Self::build_completion_streaming_usage( + total_prompt, + total_completion, + total_cached, + total_reasoning, + )), + }; + let mut sse_buffer = Vec::with_capacity(256); + Self::format_completion_sse_into(&mut sse_buffer, &usage_chunk); + let _ = tx.send(Ok(Bytes::from(sse_buffer))); + } + Metrics::record_streaming_metrics(StreamingMetricsParams { + router_type: metrics_labels::ROUTER_GRPC, + backend_type: processor.backend_type, + model_id: &dispatch.model, + endpoint: metrics_labels::ENDPOINT_COMPLETIONS, + ttft: first_token_time.map(|t| t.duration_since(start_time)), + generation_duration: start_time.elapsed(), + input_tokens: Some(total_prompt as u64), + output_tokens: total_completion as u64, + }); + } + } + + let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + }); + + build_sse_response(rx) + } + + /// Split an execution result into per-prompt stream units, in prompt order. + fn completion_stream_units( + execution_result: context::ExecutionResult, + ) -> Result, &'static str> { + match execution_result { + context::ExecutionResult::Single { stream } => { + Ok(vec![CompletionStreamUnit::Single(stream)]) } context::ExecutionResult::PrefillDecode { // TODO(#1781 follow-up): thread pd_timing for honest PD TTFT prefill, decode, .. - } => { - let processor = self.clone(); - #[expect( - clippy::disallowed_methods, - reason = "streaming task is fire-and-forget; client disconnect terminates it" - )] - tokio::spawn(async move { - let result = processor - .process_prefill_decode_completion_streaming_chunks( - prefill, - *decode, - dispatch, - tokenizer, - stop_params, - completion_request, - &tx, - ) - .await; - - if let Err(e) = result { - utils::send_error_sse(&tx, &e, "internal_error"); + } => Ok(vec![CompletionStreamUnit::PrefillDecode { + prefill, + decode, + }]), + context::ExecutionResult::Batch { results } => results + .into_iter() + .map(|result| match result { + context::ExecutionResult::Single { stream } => { + Ok(CompletionStreamUnit::Single(stream)) } - - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); - }); - } + context::ExecutionResult::PrefillDecode { + prefill, decode, .. + } => Ok(CompletionStreamUnit::PrefillDecode { prefill, decode }), + _ => Err("Nested batch or embedding result in completion streaming"), + }) + .collect(), context::ExecutionResult::Embedding { .. } => { - utils::send_error_sse( - &tx, - "Embeddings not supported in streaming mode", - "invalid_request_error", - ); - let _ = tx.send(Ok(Bytes::from("data: [DONE]\n\n"))); + Err("Embeddings not supported in streaming mode") } } - - build_sse_response(rx) } /// Process completion streaming chunks from a single stream. /// /// Decodes tokens through stop decoder, handles `echo` (first chunk) and /// `suffix` (after final chunk), and emits `CompletionStreamResponse` SSE - /// events. Supports n>1 via per-index tracking. + /// events with `index_offset`-shifted choice indices. Supports n>1 via + /// per-index tracking. Returns per-stream totals for the coordinator's + /// usage chunk and metrics record. + #[expect(clippy::too_many_arguments)] async fn process_completion_streaming_chunks( &self, mut grpc_stream: ProtoStream, @@ -2410,9 +2570,10 @@ impl StreamingProcessor { tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), completion_request: Arc, + prompt_text: &str, + index_offset: u32, tx: &UnboundedSender>, - ) -> Result<(), String> { - let start_time = Instant::now(); + ) -> Result { let mut first_token_time: Option = None; let request_id = &dispatch.request_id; @@ -2421,28 +2582,9 @@ impl StreamingProcessor { let system_fingerprint = dispatch.weight_version.as_deref(); let echo = completion_request.echo; - let prompt_text: &str = if echo { - match &completion_request.prompt { - StringOrArray::String(s) => s.as_str(), - // Array prompts are rejected by CompletionPreparationStage (Stage 1). - // If this arm is ever reached, it means a new code path bypassed Stage 1. - StringOrArray::Array(_) => { - debug_assert!(false, "Array prompt reached streaming — CompletionPreparationStage should have rejected it"); - warn!("Array prompt reached completion streaming — CompletionPreparationStage should have rejected it"); - "" - } - } - } else { - "" - }; let suffix = completion_request.suffix.as_deref(); // TODO: wire per-token logprob streaming when backend support is available let _request_logprobs = completion_request.logprobs.is_some(); - let include_usage = completion_request - .stream_options - .as_ref() - .and_then(|opts| opts.include_usage) - .unwrap_or(false); let mut stop_decoders: HashMap = HashMap::new(); let mut is_firsts: HashMap = HashMap::new(); @@ -2465,7 +2607,7 @@ impl StreamingProcessor { first_token_time = Some(Instant::now()); } - let index = chunk.index(); + let index = index_offset + chunk.index(); if stopped_indices.contains(&index) { continue; @@ -2571,7 +2713,7 @@ impl StreamingProcessor { } } ProtoResponseVariant::Complete(complete) => { - let index = complete.index(); + let index = index_offset + complete.index(); total_prompt = total_prompt.max(complete.prompt_tokens()); total_cached = total_cached.max(complete.cached_tokens()); reasoning_tokens.insert(index, complete.reasoning_tokens()); @@ -2698,43 +2840,15 @@ impl StreamingProcessor { } } - // Mark stream as completed before usage emission — the backend stream - // is fully consumed at this point, so an abort would be a no-op. - // This ensures mark_completed runs even if the usage send fails. grpc_stream.mark_completed(); - if include_usage { - let total_reasoning: u32 = reasoning_tokens.values().sum(); - let usage_chunk = CompletionStreamResponse { - id: request_id.clone(), - object: "text_completion".to_string(), - created, - choices: vec![], - model: model.clone(), - system_fingerprint: system_fingerprint.map(String::from), - usage: Some(Self::build_completion_streaming_usage( - total_prompt, - total_completion.total(), - total_cached, - total_reasoning, - )), - }; - - Self::format_completion_sse_into(&mut sse_buffer, &usage_chunk); - let _ = tx.send(Ok(Bytes::from(sse_buffer.clone()))); - } - Metrics::record_streaming_metrics(StreamingMetricsParams { - router_type: metrics_labels::ROUTER_GRPC, - backend_type: self.backend_type, - model_id: model, - endpoint: metrics_labels::ENDPOINT_COMPLETIONS, - ttft: first_token_time.map(|t| t.duration_since(start_time)), - generation_duration: start_time.elapsed(), - input_tokens: Some(total_prompt as u64), - output_tokens: total_completion.total() as u64, - }); - - Ok(()) + Ok(CompletionStreamOutcome { + prompt_tokens: total_prompt, + cached_tokens: total_cached, + reasoning_tokens: reasoning_tokens.values().sum(), + completion_tokens: total_completion.total(), + first_token_time, + }) } /// PD prefill/decode variant: consume prefill stream, then delegate decode @@ -2748,8 +2862,10 @@ impl StreamingProcessor { tokenizer: Arc, stop_params: (Option, Option>, bool, bool, bool), original_request: Arc, + prompt_text: &str, + index_offset: u32, tx: &UnboundedSender>, - ) -> Result<(), String> { + ) -> Result { while let Some(response) = prefill_stream.next().await { let gen_response = response.map_err(|e| format!("Prefill stream error: {}", e.message()))?; @@ -2767,6 +2883,8 @@ impl StreamingProcessor { tokenizer, stop_params, original_request, + prompt_text, + index_offset, tx, ) .await; diff --git a/model_gateway/tests/routing/grpc_completion_batch_test.rs b/model_gateway/tests/routing/grpc_completion_batch_test.rs new file mode 100644 index 000000000..0ec80e0a8 --- /dev/null +++ b/model_gateway/tests/routing/grpc_completion_batch_test.rs @@ -0,0 +1,169 @@ +//! gRPC /v1/completions prompt-array acceptance (issue #1903). +//! +//! Prompt arrays previously short-circuited in `CompletionPreparationStage` +//! with `batch_prompts_not_supported` (400). They now tokenize per prompt and +//! flow into the shared pipeline in every gRPC mode, so with no workers +//! registered the request must die at worker selection instead. + +use std::sync::{Arc, OnceLock}; + +use axum::http::StatusCode; +use llm_tokenizer::registry::TokenizerRegistry; +use openai_protocol::completion::CompletionRequest; +use reasoning_parser::ParserFactory as ReasoningParserFactory; +use smg::{ + app_context::AppContext, + config::{PolicyConfig, RouterConfig, RoutingMode}, + middleware::TenantRequestMeta, + policies::PolicyRegistry, + routers::{error::extract_error_code_from_response, RouterFactory}, + tenant::TenantKey, + worker::WorkerRegistry, +}; +use smg_data_connector::{ + MemoryConversationItemStorage, MemoryConversationStorage, MemoryResponseStorage, +}; +use smg_mcp::{McpConfig, McpOrchestrator}; +use tool_parser::ParserFactory as ToolParserFactory; + +use crate::common::ensure_tokenizer_cached; + +const MODEL: &str = "test-model"; + +fn grpc_modes() -> Vec { + vec![ + RoutingMode::Regular { + worker_urls: vec![], + }, + RoutingMode::PrefillDecode { + prefill_urls: vec![], + decode_urls: vec![], + prefill_policy: None, + decode_policy: None, + }, + RoutingMode::EncodePrefillDecode { + encode_urls: vec![], + prefill_urls: vec![], + decode_urls: vec![], + encode_policy: None, + prefill_policy: None, + decode_policy: None, + }, + ] +} + +#[expect( + clippy::expect_used, + reason = "test setup helper; failures should panic" +)] +async fn grpc_ctx(mode: RoutingMode) -> Arc { + let config = RouterConfig::builder() + .mode(mode) + .grpc_connection() + .policy(PolicyConfig::Random) + .host("127.0.0.1") + .port(3001) + .max_payload_size(1024 * 1024) + .request_timeout_secs(60) + .worker_startup_timeout_secs(10) + .worker_startup_check_interval_secs(1) + .max_concurrent_requests(64) + .queue_timeout_secs(60) + .build_unchecked(); + + let tokenizer_registry = Arc::new(TokenizerRegistry::new()); + // Blocking download helper; must run off the async runtime. + let tokenizer_path = tokio::task::spawn_blocking(ensure_tokenizer_cached) + .await + .expect("tokenizer download"); + let tokenizer_source = tokenizer_path.to_string_lossy().to_string(); + tokenizer_registry + .load("test-tokenizer", MODEL, &tokenizer_source, || async { + llm_tokenizer::factory::create_tokenizer_from_file(&tokenizer_source) + .map_err(|e| e.to_string()) + }) + .await + .expect("load tokenizer"); + + let mcp_orchestrator = Arc::new(OnceLock::new()); + mcp_orchestrator + .set(Arc::new( + McpOrchestrator::new(McpConfig::default()) + .await + .expect("mcp orchestrator"), + )) + .ok(); + + Arc::new( + AppContext::builder() + .router_config(config.clone()) + .client(reqwest::Client::new()) + .tokenizer_registry(tokenizer_registry) + .reasoning_parser_factory(Some(ReasoningParserFactory::new())) + .tool_parser_factory(Some(ToolParserFactory::new())) + .worker_registry(Arc::new(WorkerRegistry::new())) + .policy_registry(Arc::new(PolicyRegistry::new(config.policy.clone()))) + .response_storage(Arc::new(MemoryResponseStorage::new())) + .conversation_storage(Arc::new(MemoryConversationStorage::new())) + .conversation_item_storage(Arc::new(MemoryConversationItemStorage::new())) + .worker_job_queue(Arc::new(OnceLock::new())) + .workflow_engines(Arc::new(OnceLock::new())) + .mcp_orchestrator(mcp_orchestrator) + .build() + .expect("app context"), + ) +} + +#[expect( + clippy::expect_used, + reason = "test setup helper; failures should panic" +)] +fn completion_request(prompt: serde_json::Value) -> CompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": MODEL, + "prompt": prompt, + "max_tokens": 8, + })) + .expect("completion request") +} + +/// Every gRPC mode must carry prompt arrays past preparation: no workers are +/// registered, so both scalar and array requests hit the same worker-selection +/// wall instead of the removed array gate. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn grpc_completion_accepts_prompt_arrays_in_every_mode() { + for mode in grpc_modes() { + let mode_label = format!("{mode:?}"); + let ctx = grpc_ctx(mode).await; + let router = RouterFactory::create_router(&ctx).await.expect("router"); + let tenant_meta = TenantRequestMeta::new(TenantKey::new("test-tenant")); + + let scalar = completion_request(serde_json::json!("Hello world")); + let scalar_response = router + .route_completion(None, &tenant_meta, &scalar, MODEL) + .await; + let scalar_status = scalar_response.status(); + let scalar_code = extract_error_code_from_response(&scalar_response).to_string(); + + let array = completion_request(serde_json::json!(["Hello world", "Hello test"])); + let array_response = router + .route_completion(None, &tenant_meta, &array, MODEL) + .await; + + assert_ne!( + array_response.status(), + StatusCode::BAD_REQUEST, + "array prompt must not be rejected at preparation ({mode_label})" + ); + let array_code = extract_error_code_from_response(&array_response).to_string(); + assert_ne!( + array_code, "batch_prompts_not_supported", + "array gate must be gone ({mode_label})" + ); + assert_eq!( + (array_response.status(), array_code), + (scalar_status, scalar_code), + "array and scalar prompts must hit the same worker-selection wall ({mode_label})" + ); + } +} diff --git a/model_gateway/tests/routing/mod.rs b/model_gateway/tests/routing/mod.rs index 25293eba4..f5bfbb5ba 100644 --- a/model_gateway/tests/routing/mod.rs +++ b/model_gateway/tests/routing/mod.rs @@ -1,6 +1,7 @@ //! Routing integration tests pub mod cache_aware_backward_compat_test; +pub mod grpc_completion_batch_test; pub mod header_forwarding_test; pub mod load_balancing_test; pub mod manual_routing_test;