diff --git a/model_gateway/Cargo.toml b/model_gateway/Cargo.toml index 909235103..a7b56e90b 100644 --- a/model_gateway/Cargo.toml +++ b/model_gateway/Cargo.toml @@ -94,7 +94,7 @@ rayon = "1.12" rustix = { version = "1", features = ["fs", "process"] } tower = { version = "0.5", features = ["full"] } tower-http = { version = "0.7", features = ["trace", "compression-gzip", "cors", "timeout", "limit", "request-id", "util"] } -serde_json = { version = "1.0", default-features = false, features = ["std", "preserve_order"] } +serde_json = { version = "1.0", default-features = false, features = ["std", "preserve_order", "raw_value"] } bytes = "1.12.0" http-body = "1.0" http-body-util = "0.1" diff --git a/model_gateway/src/routers/http/mod.rs b/model_gateway/src/routers/http/mod.rs index beeef2386..85333edbb 100644 --- a/model_gateway/src/routers/http/mod.rs +++ b/model_gateway/src/routers/http/mod.rs @@ -2,6 +2,7 @@ pub mod pd_router; pub mod pd_types; +pub(crate) mod request_body; pub mod router; use serde_json::Value; diff --git a/model_gateway/src/routers/http/request_body.rs b/model_gateway/src/routers/http/request_body.rs new file mode 100644 index 000000000..8afb3d846 --- /dev/null +++ b/model_gateway/src/routers/http/request_body.rs @@ -0,0 +1,367 @@ +//! Outbound proxy body construction for typed requests. +//! +//! Serializes the typed request straight to bytes and edits the top-level +//! object as borrowed [`RawValue`] slices, so token-heavy payloads +//! (`input_ids`, messages) are never materialized as a `serde_json::Value` +//! tree. Workers whose `prepare_request` rewrites the body still take the +//! `Value` path (that hook is defined on `Value`), as does any body the +//! raw editor cannot parse. + +use serde::{ + de::{MapAccess, Visitor}, + ser::SerializeMap, + Deserialize, Deserializer, Serialize, Serializer, +}; +use serde_json::value::{to_raw_value, RawValue}; + +use crate::{ + routers::openai::{is_stripped_sglang_default, strip_default_sglang_fields, SGLANG_FIELDS}, + worker::{Worker, WorkerError}, +}; + +#[derive(Debug)] +pub(crate) enum RequestBodyError { + Serialize(serde_json::Error), + Prepare(WorkerError), +} + +/// Serialize a typed request into the exact bytes the `Value`-mediated +/// pipeline (`to_value` → model rewrite → `prepare_request` → strip → +/// `to_vec`) produces. +pub(crate) fn serialize_request_body( + typed_req: &T, + canonical_model: Option<&str>, + worker: &dyn Worker, +) -> Result, RequestBodyError> { + if worker.mutates_request() { + return value_request_body(typed_req, canonical_model, worker); + } + + let bytes = to_vec_value_compatible(typed_req).map_err(RequestBodyError::Serialize)?; + let canonical_raw = canonical_model + .map(to_raw_value) + .transpose() + .map_err(RequestBodyError::Serialize)?; + + // A body the raw editor cannot parse (in practice: non-objects) takes + // the Value pipeline rather than skipping the hooks. + let mut body = match serde_json::from_slice::(&bytes) { + Ok(body) => body, + Err(_) => return value_request_body(typed_req, canonical_model, worker), + }; + if let Some(model) = canonical_raw.as_deref() { + body.set_model(model); + } + body.strip_default_sglang_fields(); + if body.mutated { + serde_json::to_vec(&body).map_err(RequestBodyError::Serialize) + } else { + Ok(bytes) + } +} + +/// The `Value` pipeline, kept for workers whose `prepare_request` edits the +/// body ([`Worker::mutates_request`]) and as the fallback for bodies +/// [`RawBody`] cannot parse. +fn value_request_body( + typed_req: &T, + canonical_model: Option<&str>, + worker: &dyn Worker, +) -> Result, RequestBodyError> { + let mut json_val = serde_json::to_value(typed_req).map_err(RequestBodyError::Serialize)?; + if let Some(canonical_model) = canonical_model { + super::set_request_model(&mut json_val, canonical_model); + } + let mut json_val = worker + .prepare_request(json_val) + .map_err(RequestBodyError::Prepare)?; + strip_default_sglang_fields(&mut json_val); + serde_json::to_vec(&json_val).map_err(RequestBodyError::Serialize) +} + +/// `serde_json::to_value` stores `f32` widened to `f64`, so the `Value` +/// pipeline has always emitted the widened decimal form. The plain writer +/// emits the shorter `f32` form instead; widen here to keep wire bytes +/// identical. +struct F32WideningFormatter; + +impl serde_json::ser::Formatter for F32WideningFormatter { + fn write_f32(&mut self, writer: &mut W, value: f32) -> std::io::Result<()> + where + W: ?Sized + std::io::Write, + { + self.write_f64(writer, f64::from(value)) + } +} + +fn to_vec_value_compatible(value: &T) -> Result, serde_json::Error> { + let mut buf = Vec::with_capacity(128); + let mut ser = serde_json::Serializer::with_formatter(&mut buf, F32WideningFormatter); + value.serialize(&mut ser)?; + Ok(buf) +} + +/// Top-level fields of a serialized request; values stay borrowed raw JSON. +struct RawBody<'a> { + fields: Vec<(String, &'a RawValue)>, + mutated: bool, +} + +impl<'a> RawBody<'a> { + /// Mirrors [`super::set_request_model`]: only replaces an existing field. + fn set_model(&mut self, canonical_model: &'a RawValue) { + if let Some((_, value)) = self.fields.iter_mut().find(|(name, _)| name == "model") { + *value = canonical_model; + self.mutated = true; + } + } + + /// Mirrors [`strip_default_sglang_fields`], including the swap-remove + /// ordering of `serde_json::Map::remove` under `preserve_order`. + fn strip_default_sglang_fields(&mut self) { + for field in SGLANG_FIELDS { + let index = self.fields.iter().position(|(name, value)| { + name == field && is_stripped_sglang_default(field, value.get()) + }); + if let Some(index) = index { + self.fields.swap_remove(index); + self.mutated = true; + } + } + } +} + +impl<'de> Deserialize<'de> for RawBody<'de> { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct RawBodyVisitor; + + impl<'de> Visitor<'de> for RawBodyVisitor { + type Value = RawBody<'de>; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str("a JSON object") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut fields = Vec::with_capacity(map.size_hint().unwrap_or(0)); + while let Some(entry) = map.next_entry::()? { + fields.push(entry); + } + Ok(RawBody { + fields, + mutated: false, + }) + } + } + + deserializer.deserialize_map(RawBodyVisitor) + } +} + +impl Serialize for RawBody<'_> { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let mut map = serializer.serialize_map(Some(self.fields.len()))?; + for (name, value) in &self.fields { + map.serialize_entry(name, value)?; + } + map.end() + } +} + +#[cfg(test)] +mod tests { + use openai_protocol::{chat::ChatCompletionRequest, generate::GenerateRequest}; + use serde_json::{json, Value}; + + use super::*; + use crate::{ + routers::http::set_request_model, + worker::{BasicWorker, BasicWorkerBuilder, WorkerType}, + }; + + fn worker() -> BasicWorker { + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .build() + } + + fn dp_worker() -> BasicWorker { + BasicWorkerBuilder::new("http://worker:8080") + .worker_type(WorkerType::Regular) + .dp_config(3, 8) + .build() + } + + /// The pre-existing pipeline, verbatim: the produced bytes are the wire + /// contract the fast path must reproduce. + fn value_path_bytes( + typed_req: &T, + canonical_model: Option<&str>, + worker: &dyn Worker, + ) -> Vec { + let mut json_val = serde_json::to_value(typed_req).unwrap(); + if let Some(canonical_model) = canonical_model { + set_request_model(&mut json_val, canonical_model); + } + let mut json_val = worker.prepare_request(json_val).unwrap(); + strip_default_sglang_fields(&mut json_val); + serde_json::to_vec(&json_val).unwrap() + } + + fn generate_request(mut extra: Value) -> GenerateRequest { + let mut body = json!({ + "model": "alias-model", + "input_ids": [101, 7592, 2088, 1010, 2129, 2024, 2017, 2651, 1029, + 102, 2003, 2023, 1037, 2200, 2146, 3793, 6251, 102], + "sampling_params": {"temperature": 0.7, "top_p": 0.9, "max_new_tokens": 32}, + "stream": false, + "rid": "req-1" + }); + body.as_object_mut() + .unwrap() + .append(extra.as_object_mut().unwrap()); + serde_json::from_value(body).unwrap() + } + + #[test] + fn plain_generate_body_is_byte_identical_with_value_path() { + let worker = worker(); + assert!(!worker.mutates_request()); + let req = generate_request(json!({})); + + let body = serialize_request_body(&req, None, &worker).unwrap(); + + assert_eq!(body, value_path_bytes(&req, None, &worker)); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert!(parsed.get("return_hidden_states").is_none()); + } + + #[test] + fn aliased_model_is_rewritten_to_canonical() { + let worker = worker(); + let req = generate_request(json!({})); + + let body = serialize_request_body(&req, Some("canonical-model"), &worker).unwrap(); + + assert_eq!( + body, + value_path_bytes(&req, Some("canonical-model"), &worker) + ); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(parsed["model"], "canonical-model"); + } + + #[test] + fn dp_aware_worker_still_gets_prepare_request() { + let worker = dp_worker(); + assert!(worker.mutates_request()); + let req = generate_request(json!({})); + + let body = serialize_request_body(&req, None, &worker).unwrap(); + + assert_eq!(body, value_path_bytes(&req, None, &worker)); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(parsed["data_parallel_rank"], 3); + } + + #[test] + fn explicit_sglang_defaults_are_stripped() { + let worker = worker(); + let req = generate_request(json!({ + "ignore_eos": false, + "top_k": null, + "separate_reasoning": true, + "no_stop_trim": true, + "priority": 5, + "min_p": 0.0 + })); + + let body = serialize_request_body(&req, None, &worker).unwrap(); + + assert_eq!(body, value_path_bytes(&req, None, &worker)); + let parsed: Value = serde_json::from_slice(&body).unwrap(); + assert!(parsed.get("ignore_eos").is_none()); + assert!(parsed.get("top_k").is_none()); + assert!(parsed.get("separate_reasoning").is_none()); + assert_eq!(parsed["no_stop_trim"], true); + assert_eq!(parsed["priority"], 5); + assert_eq!(parsed["min_p"], 0.0); + } + + #[test] + fn untouched_body_reuses_the_direct_serialization() { + let worker = worker(); + // `return_hidden_states: true` survives the strip, so nothing in this + // body needs editing. + let req = generate_request(json!({"return_hidden_states": true})); + + let body = serialize_request_body(&req, None, &worker).unwrap(); + + assert_eq!(body, value_path_bytes(&req, None, &worker)); + assert_eq!(body, to_vec_value_compatible(&req).unwrap()); + } + + #[test] + fn f32_fields_keep_the_value_path_widening() { + let worker = worker(); + let req = generate_request(json!({})); + + let body = String::from_utf8(serialize_request_body(&req, None, &worker).unwrap()).unwrap(); + + // `to_value` widens `f32` to `f64`; the plain writer would emit the + // shorter `0.7` and change wire bytes. + let widened = serde_json::to_string(&Value::from(0.7f32)).unwrap(); + assert_ne!(widened, "0.7"); + assert!(body.contains(&widened)); + } + + #[test] + fn chat_completion_body_is_byte_identical_with_value_path() { + let worker = worker(); + let req: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "alias-model", + "messages": [{"role": "user", "content": "hello"}], + "temperature": 0.7 + })) + .unwrap(); + + let plain = serialize_request_body(&req, None, &worker).unwrap(); + assert_eq!(plain, value_path_bytes(&req, None, &worker)); + let parsed: Value = serde_json::from_slice(&plain).unwrap(); + assert!(parsed.get("separate_reasoning").is_none()); + assert_eq!(parsed["skip_special_tokens"], true); + + let aliased = serialize_request_body(&req, Some("canonical-model"), &worker).unwrap(); + assert_eq!( + aliased, + value_path_bytes(&req, Some("canonical-model"), &worker) + ); + } + + #[test] + fn non_object_body_falls_back_to_the_value_pipeline() { + let worker = worker(); + let req = vec![1, 2, 3]; + + // The raw editor rejects the shape, so this exercises the fallback. + let direct = to_vec_value_compatible(&req).unwrap(); + assert!(serde_json::from_slice::(&direct).is_err()); + + let body = serialize_request_body(&req, Some("canonical-model"), &worker).unwrap(); + + assert_eq!( + body, + value_path_bytes(&req, Some("canonical-model"), &worker) + ); + assert_eq!(body, b"[1,2,3]"); + } +} diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 883f42348..704e2b30c 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -55,7 +55,7 @@ use crate::{ }, error::{self, extract_error_code_from_response}, grpc::utils::{error_type_from_status, route_to_endpoint}, - openai::strip_default_sglang_fields, + http::request_body::{serialize_request_body, RequestBodyError}, RouterTrait, }, worker::{AttachedBody, ConnectionMode, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType}, @@ -1005,32 +1005,27 @@ impl Router { let api_key = worker.api_key().cloned(); let endpoint_url = worker.endpoint_url(route); - let mut json_val = match serde_json::to_value(typed_req) { - Ok(j) => j, - Err(e) => { + let body = match serialize_request_body(typed_req, canonical_model, worker) { + Ok(body) => body, + Err(RequestBodyError::Serialize(e)) => { return error::bad_request( "serialization_failed", - format!("Convert into serde_json::Value failed: {e}"), + format!("Failed to serialize request body: {e}"), ); } - }; - - if let Some(canonical_model) = canonical_model { - super::set_request_model(&mut json_val, canonical_model); - } - - let mut json_val = match worker.prepare_request(json_val) { - Ok(prepared) => prepared, - Err(e) => { + Err(RequestBodyError::Prepare(e)) => { return error::bad_request( "request_preparation_failed", format!("Failed to prepare request: {e}"), ); } }; - strip_default_sglang_fields(&mut json_val); - let mut request_builder = self.client.post(&endpoint_url).json(&json_val); + let mut request_builder = self + .client + .post(&endpoint_url) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) + .body(body); request_builder = header_utils::apply_forwarded_request_headers( request_builder, diff --git a/model_gateway/src/routers/openai/mod.rs b/model_gateway/src/routers/openai/mod.rs index 0ac0a897a..d8819ecb9 100644 --- a/model_gateway/src/routers/openai/mod.rs +++ b/model_gateway/src/routers/openai/mod.rs @@ -15,5 +15,5 @@ mod provider; pub mod responses; mod router; -pub(crate) use provider::strip_default_sglang_fields; +pub(crate) use provider::{is_stripped_sglang_default, strip_default_sglang_fields, SGLANG_FIELDS}; pub use router::OpenAIRouter; diff --git a/model_gateway/src/routers/openai/provider/mod.rs b/model_gateway/src/routers/openai/provider/mod.rs index f4814e696..c9506d2c4 100644 --- a/model_gateway/src/routers/openai/provider/mod.rs +++ b/model_gateway/src/routers/openai/provider/mod.rs @@ -17,6 +17,6 @@ pub use openai::OpenAIProvider; pub use provider_trait::Provider; pub use registry::ProviderRegistry; pub use sglang::SGLangProvider; -pub(crate) use types::strip_default_sglang_fields; pub use types::ProviderError; +pub(crate) use types::{is_stripped_sglang_default, strip_default_sglang_fields, SGLANG_FIELDS}; pub use xai::XAIProvider; diff --git a/model_gateway/src/routers/openai/provider/tests.rs b/model_gateway/src/routers/openai/provider/tests.rs index abfa4db88..1a157f3ce 100644 --- a/model_gateway/src/routers/openai/provider/tests.rs +++ b/model_gateway/src/routers/openai/provider/tests.rs @@ -13,7 +13,10 @@ use openai_protocol::{ }; use serde_json::{json, to_value, Value}; -use super::{types::strip_default_sglang_fields, OpenAIProvider, Provider, XAIProvider}; +use super::{ + types::{is_stripped_sglang_default, strip_default_sglang_fields, SGLANG_FIELDS}, + OpenAIProvider, Provider, XAIProvider, +}; use crate::worker::Endpoint; /// Build a `ResponsesRequest` whose single input message carries every @@ -89,6 +92,26 @@ fn strip_default_sglang_fields_removes_false_and_null_values() { assert_eq!(payload.get("model"), Some(&json!("test-model"))); } +#[test] +fn raw_predicate_agrees_with_value_strip_for_every_field() { + for field in SGLANG_FIELDS { + for raw in ["null", "false", "true", "0", "1.5", "\"false\"", "[false]"] { + let value: Value = serde_json::from_str(raw).expect("literal parses"); + let mut fields = serde_json::Map::new(); + fields.insert((*field).to_string(), value); + let mut payload = Value::Object(fields); + strip_default_sglang_fields(&mut payload); + + let value_stripped = payload.get(*field).is_none(); + assert_eq!( + is_stripped_sglang_default(field, raw), + value_stripped, + "field={field} raw={raw}" + ); + } + } +} + #[test] fn openai_provider_passes_input_image_untouched() { let req = request_with_all_content_parts(); diff --git a/model_gateway/src/routers/openai/provider/types.rs b/model_gateway/src/routers/openai/provider/types.rs index c20287283..3375d1edd 100644 --- a/model_gateway/src/routers/openai/provider/types.rs +++ b/model_gateway/src/routers/openai/provider/types.rs @@ -52,6 +52,14 @@ pub(crate) fn strip_default_sglang_fields(payload: &mut Value) { } } +/// Raw-slice twin of [`strip_default_sglang_fields`]: decides whether a +/// [`SGLANG_FIELDS`] entry would be stripped, given the compact serde_json +/// rendering of its value. Must mirror the `Value` version above. +pub(crate) fn is_stripped_sglang_default(field: &str, raw_json: &str) -> bool { + matches!(raw_json, "null" | "false") + || (matches!(field, "separate_reasoning" | "stream_reasoning") && raw_json == "true") +} + #[derive(Error, Debug)] pub enum ProviderError { #[error("Unsupported endpoint: {0:?}")] diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index 30efd62a3..0c289bece 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -472,10 +472,19 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { /// /// When the worker has a `dp_rank`, injects `data_parallel_rank` /// into the request body. Otherwise returns the request unchanged. + /// + /// Any override that edits the body must also override + /// [`Worker::mutates_request`] to return `true`. fn prepare_request(&self, req: serde_json::Value) -> WorkerResult { self.metadata().prepare_request(req) } + /// Whether [`Worker::prepare_request`] rewrites the body. The HTTP proxy + /// path skips the `serde_json::Value` round-trip when this is `false`. + fn mutates_request(&self) -> bool { + self.metadata().mutates_request() + } + /// Get the model ID this worker serves. fn model_id(&self) -> &str { self.metadata().model_id() @@ -795,6 +804,11 @@ impl WorkerMetadata { } } + /// True when [`Self::prepare_request`] would modify the request. + pub fn mutates_request(&self) -> bool { + self.spec.dp_rank.is_some() + } + // ── Routing priorities / model lookup ─────────────────────────── /// Get the model ID this worker serves.