From 3b8ac7658d47cf52ed5479d3c30d4b4b37129dee Mon Sep 17 00:00:00 2001 From: VS Chandra Mourya Date: Wed, 25 Mar 2026 22:37:09 -0700 Subject: [PATCH] feat(completions): add CompletionRequestBuildingStage and backend sampling params Signed-off-by: VS Chandra Mourya --- crates/grpc_client/src/sglang_scheduler.rs | 92 +++++++++++++ crates/grpc_client/src/trtllm_service.rs | 127 ++++++++++++++++++ crates/grpc_client/src/vllm_engine.rs | 91 +++++++++++++ model_gateway/src/routers/grpc/client.rs | 42 +++++- .../grpc/regular/stages/completion/mod.rs | 13 +- .../stages/completion/request_building.rs | 99 ++++++++++++++ 6 files changed, 455 insertions(+), 9 deletions(-) create mode 100644 model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs diff --git a/crates/grpc_client/src/sglang_scheduler.rs b/crates/grpc_client/src/sglang_scheduler.rs index d09f5c1bba..b49c6d0a93 100644 --- a/crates/grpc_client/src/sglang_scheduler.rs +++ b/crates/grpc_client/src/sglang_scheduler.rs @@ -11,6 +11,7 @@ use std::{ use openai_protocol::{ chat::ChatCompletionRequest, common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue}, + completion::CompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, responses::ResponsesRequest, @@ -669,6 +670,97 @@ impl SglangSchedulerClient { }) } + /// Build a GenerateRequest from CompletionRequest (`/v1/completions`) + #[expect( + clippy::unused_self, + reason = "method receiver kept for consistent public API" + )] + pub fn build_generate_request_from_completion( + &self, + request_id: String, + body: &CompletionRequest, + original_text: String, + token_ids: Vec, + ) -> Result { + let sampling_params = Self::build_grpc_sampling_params_from_completion(body)?; + + let grpc_request = proto::GenerateRequest { + request_id, + tokenized: Some(proto::TokenizedInput { + original_text, + input_ids: token_ids, + }), + mm_inputs: None, + sampling_params: Some(sampling_params), + return_logprob: body.logprobs.is_some(), + logprob_start_len: -1, + top_logprobs_num: body.logprobs.unwrap_or(0) as i32, + return_hidden_states: body.return_hidden_states, + stream: body.stream, + ..Default::default() + }; + + Ok(grpc_request) + } + + fn build_grpc_sampling_params_from_completion( + request: &CompletionRequest, + ) -> Result { + let stop_sequences = match &request.stop { + Some(StringOrArray::String(s)) => vec![s.clone()], + Some(StringOrArray::Array(arr)) => arr.clone(), + None => vec![], + }; + + let constraint = Self::build_single_constraint_from_completion(request)?; + + Ok(proto::SamplingParams { + temperature: request.temperature.unwrap_or(1.0), + top_p: request.top_p.unwrap_or(1.0), + top_k: request.top_k.unwrap_or(-1), + min_p: request.min_p.unwrap_or(0.0), + frequency_penalty: request.frequency_penalty.unwrap_or(0.0), + presence_penalty: request.presence_penalty.unwrap_or(0.0), + repetition_penalty: request.repetition_penalty.unwrap_or(1.0), + max_new_tokens: request.max_tokens, + min_new_tokens: request.min_tokens.unwrap_or(0), + stop: stop_sequences, + stop_token_ids: request.stop_token_ids.clone().unwrap_or_default(), + skip_special_tokens: request.skip_special_tokens, + spaces_between_special_tokens: true, + ignore_eos: request.ignore_eos, + no_stop_trim: request.no_stop_trim, + n: request.n.unwrap_or(1), + constraint, + ..Default::default() + }) + } + + fn build_single_constraint_from_completion( + request: &CompletionRequest, + ) -> Result, String> { + let mut constraints = Vec::new(); + if let Some(json_schema) = &request.json_schema { + constraints.push(proto::sampling_params::Constraint::JsonSchema( + json_schema.clone(), + )); + } + if let Some(regex) = &request.regex { + constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); + } + if let Some(ebnf) = &request.ebnf { + constraints.push(proto::sampling_params::Constraint::EbnfGrammar( + ebnf.clone(), + )); + } + + match constraints.len() { + 0 => Ok(None), + 1 => Ok(constraints.pop()), + _ => Err("Multiple structured constraints are not allowed".to_string()), + } + } + fn build_single_constraint_from_plain( params: &GenerateSamplingParams, ) -> Result, String> { diff --git a/crates/grpc_client/src/trtllm_service.rs b/crates/grpc_client/src/trtllm_service.rs index b1724a67e5..2f49bec70a 100644 --- a/crates/grpc_client/src/trtllm_service.rs +++ b/crates/grpc_client/src/trtllm_service.rs @@ -11,6 +11,7 @@ use std::{ use openai_protocol::{ chat::ChatCompletionRequest, common::{ResponseFormat, StringOrArray}, + completion::CompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, responses::ResponsesRequest, @@ -715,6 +716,132 @@ impl TrtllmServiceClient { } } + /// Build a GenerateRequest from CompletionRequest (`/v1/completions`) + #[expect( + clippy::unused_self, + reason = "method receiver kept for consistent public API" + )] + pub fn build_generate_request_from_completion( + &self, + request_id: String, + body: &CompletionRequest, + original_text: String, + token_ids: Vec, + ) -> Result { + let sampling_config = Self::build_sampling_config_from_completion(body); + let output_config = proto::OutputConfig { + logprobs: body.logprobs.map(|v| v.min(5) as i32), + prompt_logprobs: None, + return_context_logits: false, + return_generation_logits: false, + exclude_input_from_output: true, + return_encoder_output: false, + return_perf_metrics: false, + }; + + let guided_decoding = Self::build_guided_decoding_from_completion(body)?; + + let stop = match &body.stop { + Some(StringOrArray::String(s)) => vec![s.clone()], + Some(StringOrArray::Array(arr)) => arr.clone(), + None => vec![], + }; + + let grpc_request = proto::GenerateRequest { + request_id, + tokenized: Some(proto::TokenizedInput { + original_text, + input_token_ids: token_ids, + query_token_ids: vec![], + }), + sampling_config: Some(sampling_config), + output_config: Some(output_config), + max_tokens: body.max_tokens.unwrap_or(16), + streaming: body.stream, + stop, + stop_token_ids: body.stop_token_ids.clone().unwrap_or_default(), + ignore_eos: body.ignore_eos, + bad: vec![], + bad_token_ids: vec![], + guided_decoding, + embedding_bias: vec![], + lora_config: None, + prompt_tuning_config: None, + multimodal_input: None, + kv_cache_retention: None, + disaggregated_params: None, + lookahead_config: None, + cache_salt_id: None, + arrival_time: None, + include_stop_token_in_output: body.no_stop_trim, + }; + + Ok(grpc_request) + } + + fn build_sampling_config_from_completion(request: &CompletionRequest) -> proto::SamplingConfig { + proto::SamplingConfig { + beam_width: 1, + num_return_sequences: request.n.unwrap_or(1), + top_k: request.top_k.map(|v| v.max(0)), + top_p: Some(request.top_p.unwrap_or(1.0)), + top_p_min: None, + top_p_reset_ids: None, + top_p_decay: None, + seed: request.seed.map(|s| s as u64), + temperature: Some(request.temperature.unwrap_or(1.0)), + min_tokens: request.min_tokens, + beam_search_diversity_rate: None, + repetition_penalty: Some(request.repetition_penalty.unwrap_or(1.0)), + presence_penalty: request.presence_penalty, + frequency_penalty: request.frequency_penalty, + prompt_ignore_length: None, + length_penalty: None, + early_stopping: None, + no_repeat_ngram_size: None, + min_p: request.min_p, + beam_width_array: vec![], + } + } + + fn build_guided_decoding_from_completion( + request: &CompletionRequest, + ) -> Result, String> { + let mut guides = Vec::new(); + + if let Some(json_schema) = &request.json_schema { + guides.push(( + proto::guided_decoding_params::GuideType::JsonSchema, + json_schema.clone(), + )); + } + if let Some(regex) = &request.regex { + guides.push(( + proto::guided_decoding_params::GuideType::Regex, + regex.clone(), + )); + } + if let Some(ebnf) = &request.ebnf { + guides.push(( + proto::guided_decoding_params::GuideType::EbnfGrammar, + ebnf.clone(), + )); + } + + match guides.len() { + 0 => Ok(None), + 1 => { + #[expect(clippy::expect_used, reason = "INVARIANT: checked len == 1 above")] + let (guide_type, guide) = guides.pop().expect("checked len == 1"); + Ok(Some(proto::GuidedDecodingParams { + guide_type: guide_type as i32, + guide, + })) + } + _ => Err("Multiple structured constraints are not allowed".to_string()), + } + } + fn build_sampling_config_from_plain( params: Option<&GenerateSamplingParams>, ) -> proto::SamplingConfig { diff --git a/crates/grpc_client/src/vllm_engine.rs b/crates/grpc_client/src/vllm_engine.rs index 13993d6444..86417b4957 100644 --- a/crates/grpc_client/src/vllm_engine.rs +++ b/crates/grpc_client/src/vllm_engine.rs @@ -11,6 +11,7 @@ use std::{ use openai_protocol::{ chat::ChatCompletionRequest, common::{ResponseFormat, StringOrArray, ToolChoice, ToolChoiceValue}, + completion::CompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, responses::ResponsesRequest, @@ -600,6 +601,96 @@ impl VllmEngineClient { }) } + /// Build a GenerateRequest from CompletionRequest (`/v1/completions`) + #[expect( + clippy::unused_self, + reason = "method receiver kept for consistent public API" + )] + pub fn build_generate_request_from_completion( + &self, + request_id: String, + body: &CompletionRequest, + original_text: String, + token_ids: Vec, + ) -> Result { + let sampling_params = Self::build_grpc_sampling_params_from_completion(body)?; + + let grpc_request = proto::GenerateRequest { + request_id, + input: Some(proto::generate_request::Input::Tokenized( + proto::TokenizedInput { + original_text, + input_ids: token_ids, + }, + )), + sampling_params: Some(sampling_params), + stream: body.stream, + kv_transfer_params: None, + mm_inputs: None, + }; + + Ok(grpc_request) + } + + fn build_grpc_sampling_params_from_completion( + request: &CompletionRequest, + ) -> Result { + let stop_sequences = match &request.stop { + Some(StringOrArray::String(s)) => vec![s.clone()], + Some(StringOrArray::Array(arr)) => arr.clone(), + None => vec![], + }; + + let logprobs = request.logprobs.map(|v| v.min(5) as i32); + + let constraint = Self::build_single_constraint_from_completion(request)?; + + Ok(proto::SamplingParams { + temperature: request.temperature, + top_p: request.top_p.unwrap_or(1.0), + top_k: request.top_k.map(|v| v.max(0) as u32).unwrap_or(0), + min_p: request.min_p.unwrap_or(0.0), + frequency_penalty: request.frequency_penalty.unwrap_or(0.0), + presence_penalty: request.presence_penalty.unwrap_or(0.0), + repetition_penalty: request.repetition_penalty.unwrap_or(1.0), + max_tokens: request.max_tokens, + min_tokens: request.min_tokens.unwrap_or(0), + stop: stop_sequences, + stop_token_ids: request.stop_token_ids.clone().unwrap_or_default(), + skip_special_tokens: request.skip_special_tokens, + spaces_between_special_tokens: true, + ignore_eos: request.ignore_eos, + include_stop_str_in_output: request.no_stop_trim, + n: request.n.unwrap_or(1), + logprobs, + constraint, + ..Default::default() + }) + } + + fn build_single_constraint_from_completion( + request: &CompletionRequest, + ) -> Result, String> { + let mut constraints = Vec::new(); + if let Some(json_schema) = &request.json_schema { + constraints.push(proto::sampling_params::Constraint::JsonSchema( + json_schema.clone(), + )); + } + if let Some(regex) = &request.regex { + constraints.push(proto::sampling_params::Constraint::Regex(regex.clone())); + } + if let Some(ebnf) = &request.ebnf { + constraints.push(proto::sampling_params::Constraint::Grammar(ebnf.clone())); + } + + match constraints.len() { + 0 => Ok(None), + 1 => Ok(constraints.pop()), + _ => Err("Multiple structured constraints are not allowed".to_string()), + } + } + fn build_single_constraint_from_plain( params: &GenerateSamplingParams, ) -> Result, String> { diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index ecbcb41245..17c268a90f 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -3,8 +3,8 @@ use std::collections::HashMap; use openai_protocol::{ - chat::ChatCompletionRequest, generate::GenerateRequest, messages::CreateMessageRequest, - worker::WorkerLoadResponse, + chat::ChatCompletionRequest, completion::CompletionRequest, generate::GenerateRequest, + messages::CreateMessageRequest, worker::WorkerLoadResponse, }; use smg_grpc_client::{ tokenizer_bundle, tokenizer_bundle::StreamBundle, SglangSchedulerClient, TrtllmServiceClient, @@ -384,6 +384,44 @@ impl GrpcClient { } } + pub fn build_completion_request( + &self, + request_id: String, + body: &CompletionRequest, + original_text: String, + token_ids: Vec, + ) -> Result { + match self { + Self::Sglang(client) => { + let req = client.build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Sglang(Box::new(req))) + } + Self::Vllm(client) => { + let req = client.build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Vllm(Box::new(req))) + } + Self::Trtllm(client) => { + let req = client.build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Trtllm(Box::new(req))) + } + } + } + pub fn build_generate_request( &self, request_id: String, diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/mod.rs b/model_gateway/src/routers/grpc/regular/stages/completion/mod.rs index 04d8746132..523cda4e55 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/mod.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/mod.rs @@ -1,13 +1,12 @@ //! Completion API endpoint pipeline stages //! -//! This module is the next stacked step after the native completion typing work -//! landed in PR #840. That PR introduced `RequestType::Completion`, -//! `FinalResponse::Completion`, and `execute_completion()`. This branch begins -//! the endpoint-specific stage stack with `CompletionPreparationStage`. -//! -//! Later follow-up PRs can add completion-specific request building and -//! response processing here, similar to the Messages API pipeline. +//! This module continues the native `/v1/completions` stage stack after the +//! scaffolding in PR #840 and preparation in PR #907. It adds completion-specific +//! request building, with response processing deferred to a follow-up PR. mod preparation; +mod request_building; pub(crate) use preparation::CompletionPreparationStage; +#[expect(unused_imports, reason = "wired in pipeline factory follow-up PR")] +pub(crate) use request_building::CompletionRequestBuildingStage; 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 new file mode 100644 index 0000000000..ab7e4ae4a6 --- /dev/null +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -0,0 +1,99 @@ +#![allow(dead_code)] +//! Completion request building stage: build proto GenerateRequest 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. +//! +//! Completions has richer sampling knobs than Messages (frequency_penalty, presence_penalty, +//! repetition_penalty, min_p, n, logprobs, structured output constraints) but no tools +//! and no multimodal. + +use async_trait::async_trait; +use axum::response::Response; +use tracing::error; +use uuid::Uuid; + +use crate::routers::{ + error, + grpc::{ + common::stages::{helpers, PipelineStage}, + context::{ClientSelection, RequestContext}, + proto_wrapper::ProtoRequest, + }, +}; + +pub(crate) struct CompletionRequestBuildingStage { + inject_pd_metadata: bool, +} + +impl CompletionRequestBuildingStage { + pub fn new(inject_pd_metadata: bool) -> Self { + Self { inject_pd_metadata } + } +} + +#[async_trait] +impl PipelineStage for CompletionRequestBuildingStage { + async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { + let prep = ctx.state.preparation.take().ok_or_else(|| { + error!( + function = "CompletionRequestBuildingStage::execute", + "Preparation not completed" + ); + error::internal_error("preparation_not_completed", "Preparation not completed") + })?; + + let clients = ctx.state.clients.as_ref().ok_or_else(|| { + error!( + function = "CompletionRequestBuildingStage::execute", + "Client acquisition not completed" + ); + error::internal_error( + "client_acquisition_not_completed", + "Client acquisition not completed", + ) + })?; + + let completion_request = ctx.completion_request_arc(); + + let builder_client = match clients { + ClientSelection::Single { client } => client, + ClientSelection::Dual { prefill, .. } => prefill, + }; + + let request_id = format!("cmpl_{}", Uuid::now_v7()); + + let mut proto_request = builder_client + .build_completion_request( + request_id, + &completion_request, + prep.original_text.unwrap_or_default(), + prep.token_ids, + ) + .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}"), + ) + })?; + + if self.inject_pd_metadata { + if let Some(workers) = ctx.state.workers.as_ref() { + helpers::maybe_inject_pd_metadata(&mut proto_request, workers); + } + } + + ctx.state.proto_request = Some(ProtoRequest::Generate(proto_request)); + Ok(None) + } + + fn name(&self) -> &'static str { + "CompletionRequestBuilding" + } +}