From 32c6d8e1e1574db95890b3fdf12a87988504d89a Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Thu, 20 Aug 2026 09:21:28 -0700 Subject: [PATCH] refactor(router): unify request-buffer lifetime under RequestLease Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/routers/common/mod.rs | 28 + .../src/routers/common/request_lease.rs | 351 +++++++++++++ model_gateway/src/routers/http/pd_router.rs | 480 ++++++++---------- model_gateway/src/routers/http/router.rs | 221 ++------ 4 files changed, 648 insertions(+), 432 deletions(-) create mode 100644 model_gateway/src/routers/common/request_lease.rs diff --git a/model_gateway/src/routers/common/mod.rs b/model_gateway/src/routers/common/mod.rs index 78f867c4f..d72c86fc1 100644 --- a/model_gateway/src/routers/common/mod.rs +++ b/model_gateway/src/routers/common/mod.rs @@ -17,6 +17,9 @@ //! guard, shared by the HTTP and gRPC selection paths //! - [`worker_selection`] — per-request worker-selection helpers used //! by every routing path (regular, PD, fallback, external provider) +//! - [`request_lease`] — dispatch-phase owner of a request's parsed +//! body, routing derivatives and serialized upstream bytes, with a +//! retry-aware release point //! - [`retry`] — generic async retry executor + backoff calculator, //! used by every router for transport-level retries. Has zero //! coupling to the `Worker` trait — it lived in `worker/` for @@ -30,6 +33,31 @@ pub mod openai_bridge; pub mod overload; pub mod persistence_utils; pub mod realtime; +pub mod request_lease; pub mod retry; pub mod sse; pub mod worker_selection; + +/// Threshold above which upstream request bodies are sent as one-shot +/// streams. A streamed body's `try_clone()` is `None` in every layer that +/// would otherwise hold a refcount until response headers (the stale-conn +/// resend guard, reqwest's internal retry and redirect clones), so the +/// allocation frees at upload completion instead of living for the whole +/// generation on buffered upstreams. The explicit Content-Length keeps h1 +/// framing non-chunked; h2 is unaffected. +pub(crate) const STREAM_UPSTREAM_BODY_OVER: usize = 1 << 20; + +pub(crate) fn attach_sized_body( + builder: reqwest::RequestBuilder, + body: bytes::Bytes, +) -> reqwest::RequestBuilder { + let len = body.len(); + let builder = builder.header(reqwest::header::CONTENT_LENGTH, len); + if len >= STREAM_UPSTREAM_BODY_OVER { + builder.body(reqwest::Body::wrap_stream(futures::stream::once( + async move { Ok::<_, std::convert::Infallible>(body) }, + ))) + } else { + builder.body(body) + } +} diff --git a/model_gateway/src/routers/common/request_lease.rs b/model_gateway/src/routers/common/request_lease.rs new file mode 100644 index 000000000..b06c702fa --- /dev/null +++ b/model_gateway/src/routers/common/request_lease.rs @@ -0,0 +1,351 @@ +//! Dispatch-phase ownership of request memory, shared by every router. + +use std::sync::{Mutex, MutexGuard, PoisonError}; + +use bytes::Bytes; + +use crate::{config::types::RetryConfig, observability::metrics::Metrics}; + +/// When a [`RequestLease`] lets go of the parsed request. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum ReleasePoint { + /// Retries disabled: release at dispatch, once the upstream bytes exist. + AfterDispatch, + /// Retries enabled: keep the request for replay; released when the lease + /// drops, at retry-window close (first non-retryable response). + AtRetryClose, +} + +impl ReleasePoint { + pub(crate) fn from_retry_config(config: &RetryConfig) -> Self { + if config.max_retries.max(1) <= 1 { + Self::AfterDispatch + } else { + Self::AtRetryClose + } + } +} + +/// Routing inputs derived from the request body before dispatch. +#[derive(Default)] +pub(crate) struct RoutingDerivatives { + pub tokens: Option>, + pub text: Option, + pub rid_key: Option, +} + +/// Borrowed view of the leased request and its routing derivatives. +#[derive(Clone, Copy)] +pub(crate) struct LeaseView<'a, T> { + pub request: &'a T, + pub tokens: Option<&'a [u32]>, + pub text: Option<&'a str>, + pub rid_key: Option<&'a str>, +} + +/// Single owner of a request's dispatch-phase memory: the parsed request, +/// its routing derivatives, and the memoized serialized upstream body. +/// +/// Invariant: the parsed request and its derivatives live exactly until the +/// lease's release point — [`ReleasePoint::AfterDispatch`] frees them the +/// moment the upstream bytes exist, [`ReleasePoint::AtRetryClose`] keeps them +/// for retry replay until the lease drops. +pub(crate) struct RequestLease { + inner: Mutex>, + release: ReleasePoint, +} + +struct Inner { + held: Option>, + body: SerializedBody, +} + +struct Held { + request: T, + routing: RoutingDerivatives, +} + +enum SerializedBody { + None, + Single(Bytes), + Legs(Bytes, Bytes), +} + +impl SerializedBody { + /// One upstream rendering of the request; PD legs differ only in + /// injected fields, so the larger leg stands for both. + fn released_len(&self) -> usize { + match self { + Self::None => 0, + Self::Single(body) => body.len(), + Self::Legs(prefill, decode) => prefill.len().max(decode.len()), + } + } +} + +impl RequestLease { + pub(crate) fn new(request: T, routing: RoutingDerivatives, release: ReleasePoint) -> Self { + Self { + inner: Mutex::new(Inner { + held: Some(Held { request, routing }), + body: SerializedBody::None, + }), + release, + } + } + + pub(crate) fn release_point(&self) -> ReleasePoint { + self.release + } + + /// Run `f` over the leased request and derivatives. Valid only before + /// release; the closure keeps the borrow synchronous by construction. + pub(crate) fn with_view(&self, f: impl FnOnce(LeaseView<'_, T>) -> R) -> R { + let inner = self.lock(); + f(Self::view(&inner)) + } + + /// Serialize the upstream body from the leased request, memoizing the + /// bytes for [`Self::body`]. Calling again (a later retry attempt whose + /// worker may shape the body differently) refreshes the memo. + pub(crate) fn serialize_with( + &self, + f: impl FnOnce(LeaseView<'_, T>) -> Result, E>, + ) -> Result { + let mut inner = self.lock(); + let body = Bytes::from(f(Self::view(&inner))?); + inner.body = SerializedBody::Single(body.clone()); + Ok(body) + } + + /// Two-leg (PD) variant of [`Self::serialize_with`]: both leg bodies come + /// from one pass over the leased request, and the closure's intermediate + /// trees die with it. + pub(crate) fn serialize_legs_with( + &self, + f: impl FnOnce(LeaseView<'_, T>) -> Result<(Vec, Vec), E>, + ) -> Result<(Bytes, Bytes), E> { + let mut inner = self.lock(); + let (prefill, decode) = f(Self::view(&inner))?; + let legs = (Bytes::from(prefill), Bytes::from(decode)); + inner.body = SerializedBody::Legs(legs.0.clone(), legs.1.clone()); + Ok(legs) + } + + /// Memoized upstream body from the last single-leg serialization; the + /// handle is cheap and safe to hand to (re)dispatch attempts. + #[cfg_attr( + not(test), + expect( + dead_code, + reason = "accessor for dispatch paths that resend an unchanged memoized body" + ) + )] + pub(crate) fn body(&self) -> Option { + match &self.lock().body { + SerializedBody::Single(body) => Some(body.clone()), + SerializedBody::None | SerializedBody::Legs(..) => None, + } + } + + /// Under `AfterDispatch`, free the parsed request and derivatives and + /// count the serialized size as released early; under `AtRetryClose` a + /// no-op — release happens when the lease drops. + pub(crate) fn release_dispatch(&self) { + if self.release != ReleasePoint::AfterDispatch { + return; + } + let mut inner = self.lock(); + if inner.held.take().is_some() { + Metrics::record_request_buffers_released_early(inner.body.released_len()); + } + // In-flight sends hold their own handles to the serialized bytes. + inner.body = SerializedBody::None; + } + + fn lock(&self) -> MutexGuard<'_, Inner> { + self.inner.lock().unwrap_or_else(PoisonError::into_inner) + } + + fn view<'a>(inner: &'a Inner) -> LeaseView<'a, T> { + #[expect( + clippy::expect_used, + reason = "using a lease after release_dispatch is a dispatch-order bug, not a runtime condition" + )] + let held = inner + .held + .as_ref() + .expect("request lease used after release"); + LeaseView { + request: &held.request, + tokens: held.routing.tokens.as_deref(), + text: held.routing.text.as_deref(), + rid_key: held.routing.rid_key.as_deref(), + } + } +} + +/// Shared drop-probe idiom for release tests: a probed request type plus +/// loopback stubs gated on the probe's weak count. +#[cfg(test)] +pub(crate) mod test_probe { + use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Weak, + }, + time::Duration, + }; + + use axum::http::header::CONTENT_TYPE; + use openai_protocol::common::GenerationRequest; + + /// Typed request carrying a drop probe; tests watch the `Arc` count to + /// observe exactly when a router frees the parsed body. + #[derive(serde::Serialize)] + pub(crate) struct DropProbeRequest { + pub text: String, + #[serde(skip)] + pub _probe: Arc<()>, + } + + impl GenerationRequest for DropProbeRequest { + fn is_stream(&self) -> bool { + false + } + + fn get_model(&self) -> Option<&str> { + None + } + + fn extract_text_for_routing(&self) -> String { + self.text.clone() + } + } + + /// Loopback POST /generate stub that answers `{}` only after every probe + /// clone outside the test is gone (or after a deadline, leaving + /// `released` false). + #[expect( + clippy::disallowed_methods, + reason = "test stub server lives for the duration of the test process" + )] + pub(crate) async fn spawn_release_gated_stub(probe: Weak<()>) -> (String, Arc) { + let released = Arc::new(AtomicBool::new(false)); + let flag = Arc::clone(&released); + let app = axum::Router::new().route( + "/generate", + axum::routing::post(move || { + let probe = probe.clone(); + let flag = Arc::clone(&flag); + async move { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while probe.strong_count() > 1 && tokio::time::Instant::now() < deadline { + tokio::time::sleep(Duration::from_millis(2)).await; + } + flag.store(probe.strong_count() <= 1, Ordering::SeqCst); + ([(CONTENT_TYPE, "application/json")], "{}") + } + }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + (format!("http://{addr}"), released) + } + + /// Loopback POST /generate stub answering `{}` immediately. + #[expect( + clippy::disallowed_methods, + reason = "test stub server lives for the duration of the test process" + )] + pub(crate) async fn spawn_immediate_stub() -> String { + let app = axum::Router::new().route( + "/generate", + axum::routing::post(|| async { ([(CONTENT_TYPE, "application/json")], "{}") }), + ); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + format!("http://{addr}") + } +} + +#[cfg(test)] +mod tests { + use std::{convert::Infallible, sync::Arc}; + + use super::*; + + fn lease(release: ReleasePoint) -> (RequestLease>, std::sync::Weak<()>) { + let probe = Arc::new(()); + let weak = Arc::downgrade(&probe); + ( + RequestLease::new(probe, RoutingDerivatives::default(), release), + weak, + ) + } + + #[test] + fn after_dispatch_release_frees_request_and_memo() { + let (lease, probe) = lease(ReleasePoint::AfterDispatch); + let body = lease + .serialize_with(|_| Ok::<_, Infallible>(b"abcd".to_vec())) + .unwrap(); + assert_eq!(probe.strong_count(), 1); + assert_eq!(lease.body(), Some(body.clone())); + + lease.release_dispatch(); + + assert_eq!(probe.strong_count(), 0, "release must free the request"); + assert_eq!(lease.body(), None, "release must drop the memo handle"); + assert_eq!(body.as_ref(), b"abcd", "caller handles stay valid"); + } + + #[test] + fn at_retry_close_keeps_request_until_drop() { + let (lease, probe) = lease(ReleasePoint::AtRetryClose); + lease + .serialize_with(|_| Ok::<_, Infallible>(b"abcd".to_vec())) + .unwrap(); + + lease.release_dispatch(); + + assert_eq!(probe.strong_count(), 1, "request must survive for replay"); + assert!(lease.body().is_some()); + + drop(lease); + assert_eq!(probe.strong_count(), 0, "drop closes the lease"); + } + + #[test] + fn serialize_refreshes_the_memo_per_attempt() { + let (lease, _probe) = lease(ReleasePoint::AtRetryClose); + lease + .serialize_with(|_| Ok::<_, Infallible>(b"first".to_vec())) + .unwrap(); + lease + .serialize_with(|_| Ok::<_, Infallible>(b"second".to_vec())) + .unwrap(); + assert_eq!(lease.body().as_deref(), Some(b"second".as_slice())); + } + + #[test] + fn legs_serialize_from_one_view_and_release_together() { + let (lease, probe) = lease(ReleasePoint::AfterDispatch); + let (prefill, decode) = lease + .serialize_legs_with(|_| Ok::<_, Infallible>((b"pp".to_vec(), b"dddd".to_vec()))) + .unwrap(); + assert_eq!(lease.body(), None, "legs are not a single-body memo"); + + lease.release_dispatch(); + + assert_eq!(probe.strong_count(), 0); + assert_eq!(prefill.as_ref(), b"pp"); + assert_eq!(decode.as_ref(), b"dddd"); + } +} diff --git a/model_gateway/src/routers/http/pd_router.rs b/model_gateway/src/routers/http/pd_router.rs index b20ed882b..22bce3081 100644 --- a/model_gateway/src/routers/http/pd_router.rs +++ b/model_gateway/src/routers/http/pd_router.rs @@ -35,7 +35,8 @@ use crate::{ policies::{LoadBalancingPolicy, PolicyRegistry, SelectWorkerInfo}, routers::{ common::{ - header_utils, overload, + attach_sized_body, header_utils, overload, + request_lease::{ReleasePoint, RequestLease, RoutingDerivatives}, retry::{is_retryable_response, RetryExecutor}, sse::{SseEncoder, SSE_CHANNEL_BUFFER}, }, @@ -52,6 +53,8 @@ use crate::{ /// Split so the overload shed keeps its own error code, message and counter /// instead of being reworded as a circuit-breaker failure by /// [`PDRouter::handle_server_selection_error`]. +type PdPair = (Arc, Arc); + #[derive(Debug)] enum PdSelectionFailure { /// Every worker on one leg is vetoed: a ready-made, already-counted 503. @@ -76,30 +79,10 @@ struct PDRequestContext<'a> { batch_size: Option, is_stream: bool, return_logprob: bool, - request_text: Option, - routing_tokens: Option>, - rid_key: Option, model_id: &'a str, headers: Option, } -/// Per-attempt view of the parsed request. `Shared` keeps it alive across -/// retry attempts for replay; `Owned` (retries disabled) is dropped by the -/// attempt as soon as its JSON tree exists. -enum AttemptRequest<'a, T> { - Shared(&'a T), - Owned(T), -} - -impl AttemptRequest<'_, T> { - fn req(&self) -> &T { - match self { - Self::Shared(req) => req, - Self::Owned(req) => req, - } - } -} - impl PDRouter { async fn proxy_to_first_prefill_worker( &self, @@ -343,6 +326,7 @@ impl PDRouter { &self, headers: Option<&HeaderMap>, original_request: T, + routing: RoutingDerivatives, mut context: PDRequestContext<'_>, ) -> Response { let start_time = Instant::now(); @@ -372,43 +356,39 @@ impl PDRouter { .as_ref() .unwrap_or(&self.retry_config); - let response = if retry_config.max_retries.max(1) <= 1 { - // Retries disabled: one dispatch that owns the parsed request and - // releases it as soon as its JSON tree exists. + // The lease owns the parsed request and its routing derivatives for + // the dispatch phase; its release point encodes the retry policy. + let lease = RequestLease::new( + original_request, + routing, + ReleasePoint::from_retry_config(retry_config), + ); + + let response = if lease.release_point() == ReleasePoint::AfterDispatch { + // Retries disabled: one dispatch; the lease frees the parsed + // request as soon as its serialized legs exist. let res = self - .execute_dual_dispatch_attempt( - 0, - headers, - AttemptRequest::Owned(original_request), - context, - ) + .execute_dual_dispatch_attempt(0, headers, &lease, context) .await; // Mirror the retry executor's exhaustion accounting for a - // retryable status that gets no retry. + // retryable response that gets no retry. if is_retryable_response(&res) { Metrics::record_worker_retries_exhausted(metrics_labels::WORKER_PREFILL, endpoint); Metrics::record_worker_retries_exhausted(metrics_labels::WORKER_DECODE, endpoint); } res } else { - // Arc-share the request across attempts; it stays alive for replay - // until the retry window closes (first non-retryable response). - let shared_request = Arc::new(original_request); + // The lease keeps the request alive for replay until the retry + // window closes (first non-retryable response). + let lease = &lease; RetryExecutor::execute_response_with_retry( retry_config, { move |attempt: u32| { - // Clone Arc (cheap reference count increment) instead of cloning the entire request - let shared_request = Arc::clone(&shared_request); let context = context.clone(); async move { - self.execute_dual_dispatch_attempt( - attempt, - headers, - AttemptRequest::Shared(&shared_request), - context, - ) - .await + self.execute_dual_dispatch_attempt(attempt, headers, lease, context) + .await } } }, @@ -458,28 +438,28 @@ impl PDRouter { response } - /// One PD dispatch attempt: select the pair, materialize the per-leg - /// JSON bodies, dispatch, and record per-attempt worker outcomes. + /// One PD dispatch attempt: select the pair, lease-serialize the + /// per-leg bodies, dispatch, and record per-attempt worker outcomes. async fn execute_dual_dispatch_attempt( &self, attempt: u32, headers: Option<&HeaderMap>, - request: AttemptRequest<'_, T>, - mut context: PDRequestContext<'_>, + lease: &RequestLease, + context: PDRequestContext<'_>, ) -> Response { - let (prefill, decode) = match self - .select_pd_pair( - context.request_text.as_deref(), - context.routing_tokens.as_deref(), - context.rid_key.as_deref(), + let selected = lease.with_view(|view| { + self.select_pd_pair( + view.text, + view.tokens, + view.rid_key, context.model_id, context.headers.as_ref(), ) - .await - { + }); + let (prefill, decode) = match selected { Ok(pair) => pair, Err(e) => { - return Self::handle_server_selection_error(e); + return Self::handle_server_selection_error(*e); } }; @@ -490,96 +470,118 @@ impl PDRouter { decode.url() ); - let mut json_request = match serde_json::to_value(request.req()) { - Ok(v) => v, - Err(e) => return Self::handle_serialization_error(e), - }; - // The JSON tree is all dispatch needs from here on: an owned request - // (retries disabled) is freed now, before the upstream send. - drop(request); - // The prefill and decode workers only know the - // canonical name, so forward that, not the alias the - // client sent. - super::set_request_model(&mut json_request, context.model_id); - - json_request = match Self::inject_bootstrap_into_value( - json_request, - prefill.as_ref(), - context.batch_size, - ) { - Ok(v) => v, - Err(e) => { - Metrics::record_pd_bootstrap_failure(); - return Self::handle_serialization_error(e); - } - }; + // Dispatch-time re-check of both legs, the same one the regular HTTP + // and gRPC paths take just before their load guards. + if let Some(shed) = overload::shed_if_worker_overloaded(prefill.as_ref(), context.model_id) + .or_else(|| overload::shed_if_worker_overloaded(decode.as_ref(), context.model_id)) + { + return shed; + } - let mut prefill_json_request = json_request.clone(); - let mut decode_json_request = json_request; + // Keyed-load accounting uses the same effective key as selection: + // rid-derived first, header fallback. Built before the lease releases. + let load_guards = lease.with_view(|view| { + let key = view + .rid_key + .or_else(|| self.policy_registry.sticky_header_key(headers)); + vec![ + WorkerLoadGuard::with_key(prefill.clone(), key), + WorkerLoadGuard::with_key(decode.clone(), key), + ] + }); - let mut prefill_rank = prefill.dp_rank().map(|rank| rank as isize); - let mut decode_rank = decode.dp_rank().map(|rank| rank as isize); + let legs = lease.serialize_legs_with(|view| -> Result<(Vec, Vec), Box> { + let mut json_request = serde_json::to_value(view.request) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?; + // The prefill and decode workers only know the canonical name, so + // forward that, not the alias the client sent. + super::set_request_model(&mut json_request, context.model_id); + + json_request = Self::inject_bootstrap_into_value( + json_request, + prefill.as_ref(), + context.batch_size, + ) + .map_err(|e| { + Metrics::record_pd_bootstrap_failure(); + Self::handle_serialization_error(e) + })?; - let dp_rank_policy_opt = self.policy_registry.get_dp_rank_policy(); - if let Some(dp_rank_policy) = dp_rank_policy_opt.as_ref() { - let estimated_cost: isize = match ( - context.routing_tokens.as_deref(), - context.request_text.as_ref(), - ) { - (Some(tokens), _) => (tokens.len() as isize).max(1), - (None, Some(text)) => { - // Calculate token count using a simple heuristic - // In a real implementation, we would use the tokenizer - // For now, use a simple words-to-tokens ratio - let word_count = text.split_whitespace().count(); - // Assume average 1.3 tokens per word - let token_count = (word_count as f64 * 1.3).ceil() as isize; - token_count.max(1) + let mut prefill_json_request = json_request.clone(); + let mut decode_json_request = json_request; + + let mut prefill_rank = prefill.dp_rank().map(|rank| rank as isize); + let mut decode_rank = decode.dp_rank().map(|rank| rank as isize); + + let dp_rank_policy_opt = self.policy_registry.get_dp_rank_policy(); + if let Some(dp_rank_policy) = dp_rank_policy_opt.as_ref() { + let estimated_cost: isize = match (view.tokens, view.text) { + (Some(tokens), _) => (tokens.len() as isize).max(1), + (None, Some(text)) => { + // Calculate token count using a simple heuristic + // In a real implementation, we would use the tokenizer + // For now, use a simple words-to-tokens ratio + let word_count = text.split_whitespace().count(); + // Assume average 1.3 tokens per word + let token_count = (word_count as f64 * 1.3).ceil() as isize; + token_count.max(1) + } + (None, None) => 1, // Use at least 1 to avoid no-op + }; + let policy_prefill_rank = + dp_rank_policy.select_dp_rank(prefill.as_ref(), estimated_cost); + let policy_decode_rank = + dp_rank_policy.select_dp_rank(decode.as_ref(), estimated_cost); + if let Some(rank) = policy_prefill_rank { + prefill_rank = Some(rank); + } + if let Some(rank) = policy_decode_rank { + decode_rank = Some(rank); } - (None, None) => 1, // Use at least 1 to avoid no-op - }; - let policy_prefill_rank = - dp_rank_policy.select_dp_rank(prefill.as_ref(), estimated_cost); - let policy_decode_rank = dp_rank_policy.select_dp_rank(decode.as_ref(), estimated_cost); - if let Some(rank) = policy_prefill_rank { - prefill_rank = Some(rank); - } - if let Some(rank) = policy_decode_rank { - decode_rank = Some(rank); } - } - if let Some(p_rank) = prefill_rank { - Self::inject_dp_rank_to_json(&mut prefill_json_request, p_rank, "routed_dp_rank"); - Self::inject_dp_rank_to_json( - &mut decode_json_request, - p_rank, - "disagg_prefill_dp_rank", - ); - } - if let Some(d_rank) = decode_rank { - Self::inject_dp_rank_to_json(&mut decode_json_request, d_rank, "routed_dp_rank"); - } - if prefill_rank.is_some() || decode_rank.is_some() { - debug!( - "PD selected DP ranks prefill={:?} decode={:?}", - prefill_rank, decode_rank - ); - } + if let Some(p_rank) = prefill_rank { + Self::inject_dp_rank_to_json(&mut prefill_json_request, p_rank, "routed_dp_rank"); + Self::inject_dp_rank_to_json( + &mut decode_json_request, + p_rank, + "disagg_prefill_dp_rank", + ); + } + if let Some(d_rank) = decode_rank { + Self::inject_dp_rank_to_json(&mut decode_json_request, d_rank, "routed_dp_rank"); + } + if prefill_rank.is_some() || decode_rank.is_some() { + debug!( + "PD selected DP ranks prefill={:?} decode={:?}", + prefill_rank, decode_rank + ); + } - // Selection and cost estimation are done with these; the dispatch and - // response relay must not pin the routing text or token vector. - context.request_text = None; - context.routing_tokens = None; + Ok(( + serde_json::to_vec(&prefill_json_request) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?, + serde_json::to_vec(&decode_json_request) + .map_err(|e| Box::new(Self::handle_serialization_error(e)))?, + )) + }); + let (prefill_body, decode_body) = match legs { + Ok(pair) => pair, + Err(response) => return *response, + }; + // The serialized legs are all dispatch needs; the lease frees the + // parsed request and its routing derivatives now when retries are + // disabled. + lease.release_dispatch(); let response = self .execute_dual_dispatch_internal( headers, - prefill_json_request, - decode_json_request, + (prefill_body, decode_body), context, Arc::clone(&prefill), Arc::clone(&decode), + load_guards, ) .await; @@ -722,28 +724,13 @@ impl PDRouter { async fn execute_dual_dispatch_internal( &self, headers: Option<&HeaderMap>, - prefill_json_request: Value, - decode_json_request: Value, + leg_bodies: (Bytes, Bytes), context: PDRequestContext<'_>, prefill: Arc, decode: Arc, + load_guards: Vec, ) -> Response { - let effective_key = context - .rid_key - .as_deref() - .or_else(|| self.policy_registry.sticky_header_key(headers)); - // Dispatch-time re-check of both legs, the same one the regular HTTP - // and gRPC paths take just before their load guards. - if let Some(shed) = overload::shed_if_worker_overloaded(prefill.as_ref(), context.model_id) - .or_else(|| overload::shed_if_worker_overloaded(decode.as_ref(), context.model_id)) - { - return shed; - } - - let load_guards = vec![ - WorkerLoadGuard::with_key(prefill.clone(), effective_key), - WorkerLoadGuard::with_key(decode.clone(), effective_key), - ]; + let (prefill_body, decode_body) = leg_bodies; let mut headers_with_trace = headers.cloned().unwrap_or_default(); inject_trace_context_http(&mut headers_with_trace); @@ -754,7 +741,7 @@ impl PDRouter { &self.client, prefill.as_ref(), context.route, - &prefill_json_request, + prefill_body, headers, false, ); @@ -762,14 +749,10 @@ impl PDRouter { &self.client, decode.as_ref(), context.route, - &decode_json_request, + decode_body, headers, false, ); - // `.json()` already serialized both trees into the builders; free - // them before the sends instead of at scope end after the response. - drop(prefill_json_request); - drop(decode_json_request); // Send both requests concurrently and wait for both // Note: Using borrowed references avoids heap allocation @@ -929,18 +912,14 @@ impl PDRouter { prefill_policy.needs_request_text() || decode_policy.needs_request_text() } - #[expect( - clippy::unused_async, - reason = "async for API consistency; callers await uniformly" - )] - async fn select_pd_pair( + fn select_pd_pair( &self, request_text: Option<&str>, tokens: Option<&[u32]>, rid_key: Option<&str>, model_id: &str, headers: Option<&HeaderMap>, - ) -> Result<(Arc, Arc), PdSelectionFailure> { + ) -> Result> { debug!("Selecting PD pair: model_id={:?}", model_id); let is_unknown_model = model_id == UNKNOWN_MODEL_ID; @@ -995,7 +974,7 @@ impl PDRouter { "prefill", crate::policies::WorkerLeg::Prefill, ) - .map_err(|e| Self::leg_failure(&prefill_workers, model_id, e))?; + .map_err(|e| Box::new(Self::leg_failure(&prefill_workers, model_id, e)))?; let decode = self .pick_worker_by_policy_arc( @@ -1009,7 +988,7 @@ impl PDRouter { "decode", crate::policies::WorkerLeg::Decode, ) - .map_err(|e| Self::leg_failure(&decode_workers, model_id, e))?; + .map_err(|e| Box::new(Self::leg_failure(&decode_workers, model_id, e)))?; // Record worker selection metrics (Layer 3) let model = model_id; @@ -1314,12 +1293,17 @@ impl PDRouter { client: &Client, worker: &dyn Worker, route: &'static str, - json_request: &Value, + body: Bytes, headers: Option<&HeaderMap>, connection_close: bool, ) -> reqwest::RequestBuilder { let endpoint_url = worker.endpoint_url(route); - let mut request = client.post(endpoint_url).json(json_request); + let mut request = attach_sized_body( + client + .post(endpoint_url) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")), + bytes::Bytes::from(body), + ); if connection_close { request = request.header("Connection", "close"); } @@ -1463,21 +1447,21 @@ impl RouterTrait for PDRouter { // Note: This endpoint actually causes the model to generate tokens, so we only test one pair // Select a random worker pair using the policy - let (prefill, decode) = match self - .select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None) - .await + let (prefill, decode) = match self.select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None) { Ok(pair) => pair, // A deep probe that generates gets the same answer routing does: // an all-vetoed fleet fails the probe, exactly as an all-circuit- // broken one already did. - Err(PdSelectionFailure::Shed(shed)) => return shed, - Err(PdSelectionFailure::Unavailable(e)) => { - return error::service_unavailable( - "no_healthy_worker_pair", - format!("No healthy worker pair available: {e}"), - ); - } + Err(failure) => match *failure { + PdSelectionFailure::Shed(shed) => return shed, + PdSelectionFailure::Unavailable(e) => { + return error::service_unavailable( + "no_healthy_worker_pair", + format!("No healthy worker pair available: {e}"), + ); + } + }, }; let prefill_url = format!("{}/health_generate", prefill.url()); @@ -1579,22 +1563,25 @@ impl RouterTrait for PDRouter { let batch_size = Self::get_generate_batch_size(&body); + let routing = RoutingDerivatives { + tokens: routing_tokens, + text: request_text, + rid_key: self + .policy_registry + .derive_rid_key(body.rid()) + .map(str::to_string), + }; let context = PDRequestContext { route: "/generate", batch_size, is_stream, return_logprob, - request_text, - routing_tokens, - rid_key: self - .policy_registry - .derive_rid_key(body.rid()) - .map(str::to_string), model_id, headers: headers.cloned(), }; - self.execute_dual_dispatch(headers, body, context).await + self.execute_dual_dispatch(headers, body, routing, context) + .await } async fn route_chat( @@ -1616,22 +1603,25 @@ impl RouterTrait for PDRouter { // Calculate batch size let batch_size = Self::get_chat_batch_size(&body); + let routing = RoutingDerivatives { + tokens: None, + text: request_text, + rid_key: self + .policy_registry + .derive_rid_key(body.rid()) + .map(str::to_string), + }; let context = PDRequestContext { route: "/v1/chat/completions", batch_size, is_stream, return_logprob, - request_text, - routing_tokens: None, - rid_key: self - .policy_registry - .derive_rid_key(body.rid()) - .map(str::to_string), model_id, headers: headers.cloned(), }; - self.execute_dual_dispatch(headers, body, context).await + self.execute_dual_dispatch(headers, body, routing, context) + .await } async fn route_completion( @@ -1656,22 +1646,25 @@ impl RouterTrait for PDRouter { // Calculate batch size let batch_size = Self::get_completion_batch_size(&body); + let routing = RoutingDerivatives { + tokens: None, + text: request_text, + rid_key: self + .policy_registry + .derive_rid_key(body.rid()) + .map(str::to_string), + }; let context = PDRequestContext { route: "/v1/completions", batch_size, is_stream, return_logprob, - request_text, - routing_tokens: None, - rid_key: self - .policy_registry - .derive_rid_key(body.rid()) - .map(str::to_string), model_id, headers: headers.cloned(), }; - self.execute_dual_dispatch(headers, body, context).await + self.execute_dual_dispatch(headers, body, routing, context) + .await } async fn route_rerank( @@ -1688,22 +1681,25 @@ impl RouterTrait for PDRouter { None }; + let routing = RoutingDerivatives { + tokens: None, + text: req_text, + rid_key: self + .policy_registry + .derive_rid_key(body.rid()) + .map(str::to_string), + }; let context = PDRequestContext { route: "/v1/rerank", batch_size: None, is_stream: false, return_logprob: false, - request_text: req_text, - routing_tokens: None, - rid_key: self - .policy_registry - .derive_rid_key(body.rid()) - .map(str::to_string), model_id, headers: headers.cloned(), }; - self.execute_dual_dispatch(headers, body, context).await + self.execute_dual_dispatch(headers, body, routing, context) + .await } fn router_type(&self) -> &'static str { @@ -1834,7 +1830,7 @@ mod tests { &router.client, &worker, "/generate", - &json!({"text": "hello"}), + Bytes::from(r#"{"text":"hello"}"#), None, false, ) @@ -1870,9 +1866,7 @@ mod tests { .worker_registry .register_or_replace(Arc::from(decode_worker)); - let result = router - .select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None) - .await; + let result = router.select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None); assert!(result.is_ok()); let (prefill, _decode) = result.unwrap(); @@ -1898,14 +1892,12 @@ mod tests { let (prefill, decode) = router .select_pd_pair(None, None, None, "GLM-5.2-Coding", None) - .await .expect("alias should select a PD pair"); assert_eq!(prefill.url(), "http://prefill"); assert_eq!(decode.url(), "http://decode"); assert!(router .select_pd_pair(None, None, None, "GLM-5.2-Unknown", None) - .await .is_err()); } @@ -1913,9 +1905,7 @@ mod tests { async fn test_empty_worker_lists() { let router = create_test_pd_router(); - let result = router - .select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None) - .await; + let result = router.select_pd_pair(None, None, None, UNKNOWN_MODEL_ID, None); assert!(result.is_err()); // No workers at all is the pre-existing unavailable string, not a shed: @@ -1980,9 +1970,6 @@ mod tests { batch_size: None, is_stream: true, return_logprob: false, - request_text: None, - routing_tokens: None, - rid_key: None, model_id: UNKNOWN_MODEL_ID, headers: None, }; @@ -2013,57 +2000,21 @@ mod tests { ); } - /// PD request with a drop probe (see the regular router's twin test): - /// with retries disabled the parsed request must be freed once its JSON - /// tree exists, before either upstream leg answers. The decode stub - /// refuses to respond until the probe's only remaining holder is the - /// test itself. + /// PD twin of the regular router's release test: with retries disabled + /// the lease must free the parsed request once the serialized legs + /// exist, before either upstream leg answers. The decode stub refuses to + /// respond until the probe's only remaining holder is the test itself. #[tokio::test] async fn pd_disabled_retries_release_parsed_request_before_upstream_responds() { - use std::sync::atomic::{AtomicBool, Ordering}; - - #[derive(Serialize)] - struct DropProbeRequest { - text: String, - #[serde(skip)] - _probe: Arc<()>, - } + use std::sync::atomic::Ordering; - #[expect( - clippy::disallowed_methods, - reason = "test stub servers live for the duration of the test process" - )] - async fn spawn_stub(gate: Option<(std::sync::Weak<()>, Arc)>) -> String { - let app = axum::Router::new().route( - "/generate", - axum::routing::post(move || { - let gate = gate.clone(); - async move { - if let Some((probe, flag)) = gate { - let deadline = - tokio::time::Instant::now() + std::time::Duration::from_secs(5); - while probe.strong_count() > 1 && tokio::time::Instant::now() < deadline - { - tokio::time::sleep(std::time::Duration::from_millis(2)).await; - } - flag.store(probe.strong_count() <= 1, Ordering::SeqCst); - } - "{}" - } - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - format!("http://{addr}") - } + use crate::routers::common::request_lease::test_probe::{ + spawn_immediate_stub, spawn_release_gated_stub, DropProbeRequest, + }; let probe = Arc::new(()); - let released = Arc::new(AtomicBool::new(false)); - let prefill_url = spawn_stub(None).await; - let decode_url = spawn_stub(Some((Arc::downgrade(&probe), Arc::clone(&released)))).await; + let prefill_url = spawn_immediate_stub().await; + let (decode_url, released) = spawn_release_gated_stub(Arc::downgrade(&probe)).await; let router = create_test_pd_router(); let router = PDRouter { @@ -2093,9 +2044,6 @@ mod tests { batch_size: None, is_stream: false, return_logprob: false, - request_text: None, - routing_tokens: None, - rid_key: None, model_id: UNKNOWN_MODEL_ID, headers: None, }; @@ -2104,7 +2052,9 @@ mod tests { _probe: Arc::clone(&probe), }; - let response = router.execute_dual_dispatch(None, request, context).await; + let response = router + .execute_dual_dispatch(None, request, RoutingDerivatives::default(), context) + .await; assert_eq!(response.status(), StatusCode::OK); assert!( diff --git a/model_gateway/src/routers/http/router.rs b/model_gateway/src/routers/http/router.rs index 522b8d105..2de124058 100644 --- a/model_gateway/src/routers/http/router.rs +++ b/model_gateway/src/routers/http/router.rs @@ -53,11 +53,12 @@ use crate::{ policies::{PolicyRegistry, SelectWorkerInfo}, routers::{ common::{ - header_utils, overload, + attach_sized_body, header_utils, overload, realtime::{ rest::forward_realtime_rest, webrtc, webrtc::handle_realtime_webrtc, ws::handle_realtime_ws, RealtimeLabels, RealtimeRegistry, }, + request_lease::{ReleasePoint, RequestLease, RoutingDerivatives}, retry::{is_retryable_response, is_retryable_status, RetryExecutor}, sse::SSE_CHANNEL_BUFFER, worker_selection::{SelectWorkerRequest, WorkerSelector}, @@ -84,65 +85,6 @@ const STREAMED_BODY_STALLED: &str = "request_body_stalled"; const STREAMED_BODY_TOO_LARGE: &str = "request_body_too_large"; const STREAMED_BODY_ABORTED: &str = "request_body_aborted"; -/// Per-attempt view of the parsed request and its routing derivatives. -/// -/// `Borrowed` keeps them alive across retry attempts for replay; `Owned` -/// (retries disabled) is released by the dispatch path as soon as the -/// serialized upstream body exists, so neither the send wait nor the response -/// relay pins the parsed request. -enum AttemptPayload<'a, T> { - Borrowed { - req: &'a T, - text: Option<&'a str>, - tokens: Option<&'a [u32]>, - rid_key: Option<&'a str>, - }, - Owned { - req: T, - text: Option, - tokens: Option>, - rid_key: Option, - }, -} - -impl AttemptPayload<'_, T> { - fn req(&self) -> &T { - match self { - Self::Borrowed { req, .. } => req, - Self::Owned { req, .. } => req, - } - } - - fn text(&self) -> Option<&str> { - match self { - Self::Borrowed { text, .. } => *text, - Self::Owned { text, .. } => text.as_deref(), - } - } - - fn tokens(&self) -> Option<&[u32]> { - match self { - Self::Borrowed { tokens, .. } => *tokens, - Self::Owned { tokens, .. } => tokens.as_deref(), - } - } - - fn rid_key(&self) -> Option<&str> { - match self { - Self::Borrowed { rid_key, .. } => *rid_key, - Self::Owned { rid_key, .. } => rid_key.as_deref(), - } - } - - /// Consume the payload once the upstream body is serialized. - /// `upstream_bytes` sizes the early-release metric for `Owned`. - fn release(self, upstream_bytes: usize) { - if matches!(self, Self::Owned { .. }) { - Metrics::record_request_buffers_released_early(upstream_bytes); - } - } -} - /// Regular router that uses injected load balancing policies pub struct Router { worker_registry: Arc, @@ -428,6 +370,10 @@ impl Router { let text = routing_tokens .is_none() .then(|| typed_req.extract_text_for_routing()); + let rid_key = self + .policy_registry + .derive_rid_key(typed_req.rid()) + .map(str::to_string); // Resolve once, here, so every registry, policy and metrics lookup // below is keyed by the canonical model ID. Only `get_by_model` // understands aliases; retry configs, hash rings and policies do not, @@ -453,23 +399,25 @@ impl Router { .as_ref() .unwrap_or(&self.retry_config); - let response = if retry_config.max_retries.max(1) <= 1 { - // Retries disabled: one dispatch that owns the parsed request and - // its routing derivatives, releasing them the moment the upstream - // bytes are serialized instead of holding them for the response. - let rid_key = self - .policy_registry - .derive_rid_key(typed_req.rid()) - .map(str::to_string); + // The lease owns the parsed request and its routing derivatives for + // the dispatch phase; its release point encodes the retry policy. + let lease = RequestLease::new( + typed_req, + RoutingDerivatives { + tokens: routing_tokens, + text, + rid_key, + }, + ReleasePoint::from_retry_config(retry_config), + ); + + let response = if lease.release_point() == ReleasePoint::AfterDispatch { + // Retries disabled: one dispatch; the lease frees the parsed + // request the moment the upstream bytes are serialized. let res = self .route_typed_request_once( headers, - AttemptPayload::Owned { - req: typed_req, - text, - tokens: routing_tokens, - rid_key, - }, + &lease, route, model_id, canonical_model.as_deref(), @@ -482,28 +430,22 @@ impl Router { extract_error_code_from_response(&res), ); // Mirror the retry executor's exhaustion accounting for a - // retryable status that gets no retry. + // retryable response that gets no retry. if is_retryable_response(&res) { Metrics::record_worker_retries_exhausted(metrics_labels::WORKER_REGULAR, endpoint); } res } else { - let rid_key = self.policy_registry.derive_rid_key(typed_req.rid()); RetryExecutor::execute_response_with_retry( retry_config, - // operation per attempt; the parsed request stays alive for - // replay until the retry window closes (first non-retryable - // response). + // operation per attempt; the lease keeps the request alive + // for replay until the retry window closes (first + // non-retryable response). |_: u32| async { let res = self .route_typed_request_once( headers, - AttemptPayload::Borrowed { - req: &typed_req, - text: text.as_deref(), - tokens: routing_tokens.as_deref(), - rid_key, - }, + &lease, route, model_id, canonical_model.as_deref(), @@ -566,19 +508,15 @@ impl Router { async fn route_typed_request_once( &self, headers: Option<&HeaderMap>, - payload: AttemptPayload<'_, T>, + lease: &RequestLease, route: &'static str, model_id: &str, canonical_model: Option<&str>, is_stream: bool, ) -> Response { - let worker = match self.select_worker_for_model( - model_id, - payload.text(), - payload.tokens(), - headers, - payload.rid_key(), - ) { + let worker = match lease.with_view(|view| { + self.select_worker_for_model(model_id, view.text, view.tokens, headers, view.rid_key) + }) { Some(w) => w, None => { // Distinguish "no workers for this model" from "workers exist but unavailable" @@ -620,12 +558,13 @@ impl Router { // Keyed-load accounting uses the same effective key as selection: // rid-derived first, header fallback. - let load_guard = WorkerLoadGuard::with_key( - worker.clone(), - payload - .rid_key() - .or_else(|| self.policy_registry.sticky_header_key(headers)), - ); + let load_guard = lease.with_view(|view| { + WorkerLoadGuard::with_key( + worker.clone(), + view.rid_key + .or_else(|| self.policy_registry.sticky_header_key(headers)), + ) + }); // Note: Using borrowed reference avoids heap allocation events::RequestSentEvent { url: worker.url() }.emit(); @@ -633,13 +572,14 @@ impl Router { inject_trace_context_http(&mut headers_with_trace); let headers = Some(&headers_with_trace); - let response = match serialize_request_body(payload.req(), canonical_model, worker.as_ref()) - { + let response = match lease.serialize_with(|view| { + serialize_request_body(view.request, canonical_model, worker.as_ref()) + }) { Ok(body) => { - // Past this point dispatch needs only the serialized bytes: an - // owned payload (retries disabled) frees the parsed request and - // its routing token/text derivatives before the upstream send. - payload.release(body.len()); + // Past this point dispatch needs only the serialized bytes; + // the lease frees the parsed request and its routing + // derivatives now when retries are disabled. + lease.release_dispatch(); self.send_serialized_request( headers, body, @@ -1235,7 +1175,7 @@ impl Router { async fn send_serialized_request( &self, headers: Option<&HeaderMap>, - body: Vec, + body: Bytes, route: &'static str, worker: &dyn Worker, is_stream: bool, @@ -1244,11 +1184,12 @@ impl Router { let api_key = worker.api_key().cloned(); let endpoint_url = worker.endpoint_url(route); - let mut request_builder = self - .client - .post(&endpoint_url) - .header(CONTENT_TYPE, HeaderValue::from_static("application/json")) - .body(body); + let mut request_builder = attach_sized_body( + self.client + .post(&endpoint_url) + .header(CONTENT_TYPE, HeaderValue::from_static("application/json")), + bytes::Bytes::from(body), + ); request_builder = header_utils::apply_forwarded_request_headers( request_builder, @@ -2173,7 +2114,7 @@ impl RouterTrait for Router { mod tests { use std::{ net::SocketAddr, - sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}, + sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}, }; use openai_protocol::worker::HealthCheckConfig; @@ -2182,6 +2123,7 @@ mod tests { use crate::{ config::types::{PolicyConfig, RoutingKeyOverrideConfig}, policies::CacheAwarePolicy, + routers::common::request_lease::test_probe::{spawn_release_gated_stub, DropProbeRequest}, worker::BasicWorkerBuilder, }; @@ -2939,61 +2881,6 @@ mod tests { .is_err()); } - /// Typed request with a drop probe: the test watches the `Arc` count to - /// observe exactly when the router frees the parsed body. - #[derive(serde::Serialize)] - struct DropProbeRequest { - text: String, - #[serde(skip)] - _probe: Arc<()>, - } - - impl GenerationRequest for DropProbeRequest { - fn is_stream(&self) -> bool { - false - } - - fn get_model(&self) -> Option<&str> { - None - } - - fn extract_text_for_routing(&self) -> String { - self.text.clone() - } - } - - /// Upstream stub that answers only after every probe clone outside the - /// test is gone (or after a deadline, leaving `released` false). - #[expect( - clippy::disallowed_methods, - reason = "test stub server lives for the duration of the test process" - )] - async fn spawn_release_gated_stub(probe: std::sync::Weak<()>) -> (String, Arc) { - let released = Arc::new(AtomicBool::new(false)); - let flag = Arc::clone(&released); - let app = axum::Router::new().route( - "/generate", - axum::routing::post(move || { - let probe = probe.clone(); - let flag = Arc::clone(&flag); - async move { - let deadline = tokio::time::Instant::now() + Duration::from_secs(5); - while probe.strong_count() > 1 && tokio::time::Instant::now() < deadline { - tokio::time::sleep(Duration::from_millis(2)).await; - } - flag.store(probe.strong_count() <= 1, AtomicOrdering::SeqCst); - ([(CONTENT_TYPE, "application/json")], "{}") - } - }), - ); - let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); - let addr = listener.local_addr().unwrap(); - tokio::spawn(async move { - axum::serve(listener, app).await.unwrap(); - }); - (format!("http://{addr}"), released) - } - /// With retries disabled the parsed request must be freed at dispatch: /// the upstream stub refuses to answer until the probe's only remaining /// holder is the test itself.