From ad5bf6fddeac7c1b5871379ff828266d6f6f9a1b Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Thu, 20 Aug 2026 12:48:12 -0700 Subject: [PATCH] fix(grpc): release request payloads at dispatch via RequestLease Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../src/routers/common/request_lease.rs | 37 ++- .../grpc/common/stages/dispatch_metadata.rs | 3 + .../src/routers/grpc/common/stages/helpers.rs | 7 +- .../grpc/common/stages/request_execution.rs | 13 + model_gateway/src/routers/grpc/context.rs | 110 +++++++-- .../grpc/harmony/stages/request_building.rs | 3 +- .../harmony/stages/response_processing.rs | 3 +- model_gateway/src/routers/grpc/pipeline.rs | 225 +++++++++++++++++- .../src/routers/grpc/proto_wrapper.rs | 29 +++ .../grpc/regular/responses/streaming.rs | 4 + .../regular/stages/response_processing.rs | 17 +- model_gateway/src/routers/grpc/router.rs | 61 ++++- 12 files changed, 459 insertions(+), 53 deletions(-) diff --git a/model_gateway/src/routers/common/request_lease.rs b/model_gateway/src/routers/common/request_lease.rs index b06c702fae..6f7ba51250 100644 --- a/model_gateway/src/routers/common/request_lease.rs +++ b/model_gateway/src/routers/common/request_lease.rs @@ -69,6 +69,9 @@ enum SerializedBody { None, Single(Bytes), Legs(Bytes, Bytes), + /// Wire size only: the serialization lives outside the lease (typed + /// proto dispatch). + Sized(usize), } impl SerializedBody { @@ -79,6 +82,7 @@ impl SerializedBody { Self::None => 0, Self::Single(body) => body.len(), Self::Legs(prefill, decode) => prefill.len().max(decode.len()), + Self::Sized(len) => *len, } } } @@ -144,10 +148,16 @@ impl RequestLease { pub(crate) fn body(&self) -> Option { match &self.lock().body { SerializedBody::Single(body) => Some(body.clone()), - SerializedBody::None | SerializedBody::Legs(..) => None, + SerializedBody::None | SerializedBody::Legs(..) | SerializedBody::Sized(_) => None, } } + /// Record the upstream wire size for the release metric when the wire + /// serialization lives outside the lease (typed proto dispatch). + pub(crate) fn note_upstream_len(&self, len: usize) { + self.lock().body = SerializedBody::Sized(len); + } + /// 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. @@ -185,6 +195,31 @@ impl RequestLease { } } +/// Object-safe lease handle for pipelines that thread the release point +/// through a heterogeneous request context (the gRPC stage pipeline). +pub(crate) trait ErasedLease: Send + Sync { + /// True when the lease frees at dispatch (retries disabled). + fn releases_after_dispatch(&self) -> bool; + /// See [`RequestLease::note_upstream_len`]. + fn note_upstream_len(&self, len: usize); + /// See [`RequestLease::release_dispatch`]. + fn release_dispatch(&self); +} + +impl ErasedLease for RequestLease { + fn releases_after_dispatch(&self) -> bool { + self.release_point() == ReleasePoint::AfterDispatch + } + + fn note_upstream_len(&self, len: usize) { + RequestLease::note_upstream_len(self, len); + } + + fn release_dispatch(&self) { + RequestLease::release_dispatch(self); + } +} + /// Shared drop-probe idiom for release tests: a probed request type plus /// loopback stubs gated on the probe's weak count. #[cfg(test)] diff --git a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs index 73215c019d..b13877948c 100644 --- a/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs +++ b/model_gateway/src/routers/grpc/common/stages/dispatch_metadata.rs @@ -41,6 +41,9 @@ impl PipelineStage for DispatchMetadataStage { RequestType::Embedding(req) => req.model.clone(), RequestType::Classify(req) => req.model.clone(), RequestType::Messages(req) => req.model.clone(), + // Runs before request execution, so the payload is never released + // here; the canonical model ID is the same value regardless. + RequestType::Released(_) => ctx.input.model_id.clone(), }; let weight_version = ctx diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index 7de73ffe91..2ea3dded96 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -71,9 +71,10 @@ impl SamplingDefaultsMask { min_p: true, repetition_penalty: true, }), - RequestType::Responses(_) | RequestType::Embedding(_) | RequestType::Classify(_) => { - None - } + RequestType::Responses(_) + | RequestType::Embedding(_) + | RequestType::Classify(_) + | RequestType::Released(_) => None, } } 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 8129b5f827..d10fda852e 100644 --- a/model_gateway/src/routers/grpc/common/stages/request_execution.rs +++ b/model_gateway/src/routers/grpc/common/stages/request_execution.rs @@ -195,6 +195,19 @@ impl PipelineStage for RequestExecutionStage { sub_requests, )); + // Dispatch-phase release: response processing reads only the view + // captured at request building, so the context sheds its payload + // handle before the send. The lease then frees the parsed request + // now (retries disabled) or keeps it for replay until the retry + // window closes. + if let Some(lease) = ctx.input.request_lease.take() { + ctx.input.request_type.release_payload(); + if lease.releases_after_dispatch() { + lease.note_upstream_len(execution_plan.wire_len()); + } + lease.release_dispatch(); + } + // Extract dispatch metadata for tracing span let dispatch = ctx.state.dispatch.as_ref(); let request_id = dispatch.map(|d| d.request_id.as_str()).unwrap_or("unknown"); diff --git a/model_gateway/src/routers/grpc/context.rs b/model_gateway/src/routers/grpc/context.rs index 36a816fa88..dc65ad9ce5 100644 --- a/model_gateway/src/routers/grpc/context.rs +++ b/model_gateway/src/routers/grpc/context.rs @@ -33,6 +33,7 @@ use super::{ }; use crate::{ middleware::TenantRequestMeta, + routers::common::request_lease::ErasedLease, worker::{RuntimeType, Worker, WorkerLoadGuard, WorkerRegistry}, }; @@ -53,12 +54,19 @@ pub(crate) struct RequestInput { pub headers: Option, /// Canonical model ID used after aliases are resolved at request entry. pub model_id: String, + /// Captured at construction so it survives payload release. + pub streaming: bool, pub tenant_request_meta: Option, /// Shared across every retry attempt of one logical request so /// `RateLimitReserveStage` reserves at most once. `None` for endpoints /// that haven't opted into tenant rate limiting yet (Responses, /// embeddings, classify). pub rate_limit_cell: Option>, + /// Dispatch-phase owner of the parsed request (the router retry loop's + /// handle). When set, `RequestExecutionStage` drops the context's own + /// payload handle at dispatch and releases the lease; response stages + /// then read only the pre-extracted request view. + pub request_lease: Option>, } /// Request type variants @@ -71,9 +79,58 @@ pub(crate) enum RequestType { Embedding(Arc), Classify(Arc), Messages(Arc), + /// Payload dropped at dispatch; only the original kind survives for + /// post-dispatch stage dispatching. + Released(RequestKind), +} + +/// Payload-free discriminant of [`RequestType`]. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum RequestKind { + Chat, + Generate, + Completion, + Responses, + Embedding, + Classify, + Messages, +} + +impl std::fmt::Display for RequestKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let name = match self { + Self::Chat => "Chat", + Self::Generate => "Generate", + Self::Completion => "Completion", + Self::Responses => "Responses", + Self::Embedding => "Embedding", + Self::Classify => "Classify", + Self::Messages => "Messages", + }; + write!(f, "{name}") + } } impl RequestType { + pub fn kind(&self) -> RequestKind { + match self { + Self::Chat(_) => RequestKind::Chat, + Self::Generate(_) => RequestKind::Generate, + Self::Completion(_) => RequestKind::Completion, + Self::Responses(_) => RequestKind::Responses, + Self::Embedding(_) => RequestKind::Embedding, + Self::Classify(_) => RequestKind::Classify, + Self::Messages(_) => RequestKind::Messages, + Self::Released(kind) => *kind, + } + } + + /// Drop the context's payload handle at dispatch. The lease (or the + /// router loop) is the remaining owner; the typed accessors panic past + /// this point by construction. + pub fn release_payload(&mut self) { + *self = Self::Released(self.kind()); + } /// Overwrite the request's own `model` field. /// /// Callers hold the request behind an `Arc` that the retry loop also @@ -94,6 +151,7 @@ impl RequestType { Self::Embedding(request) => replace(&mut Arc::make_mut(request).model, model_id), Self::Classify(request) => replace(&mut Arc::make_mut(request).model, model_id), Self::Messages(request) => replace(&mut Arc::make_mut(request).model, model_id), + Self::Released(_) => {} } } @@ -107,22 +165,14 @@ impl RequestType { Self::Embedding(r) => r.rid.as_deref(), Self::Classify(r) => r.rid.as_deref(), Self::Messages(r) => r.rid.as_deref(), - Self::Responses(_) => None, + Self::Responses(_) | Self::Released(_) => None, } } } impl std::fmt::Display for RequestType { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Chat(_) => write!(f, "Chat"), - Self::Generate(_) => write!(f, "Generate"), - Self::Completion(_) => write!(f, "Completion"), - Self::Responses(_) => write!(f, "Responses"), - Self::Embedding(_) => write!(f, "Embedding"), - Self::Classify(_) => write!(f, "Classify"), - Self::Messages(_) => write!(f, "Messages"), - } + self.kind().fmt(f) } } @@ -281,6 +331,19 @@ impl ExecutionPlan { }, } } + + /// Serialized wire size of the built request(s), for the release metric. + pub(crate) fn wire_len(&self) -> usize { + match self { + Self::Single(request) => request.wire_len(), + Self::PrefillDecode(request) | Self::EncodePrefillDecode { request } => { + request.wire_len() + } + Self::Batch { requests, .. } => { + requests.iter().map(ProtoGenerateRequest::wire_len).sum() + } + } + } } /// Output from preparation stage (Step 1) @@ -491,7 +554,7 @@ pub(crate) struct ResponseState { pub skip_special_tokens: Option, /// Response-phase view of the request, set by request building so - /// response processing never reads the payload after dispatch. + /// response processing never reads the (possibly released) payload. pub request_view: Option, /// Execution result (streams from workers) @@ -528,13 +591,25 @@ impl RequestContext { model_id.push_str(&canonical_model_id); request_type.set_model(&model_id); } + let streaming = match &request_type { + RequestType::Chat(req) => req.stream, + RequestType::Generate(req) => req.stream, + RequestType::Completion(req) => req.stream, + RequestType::Responses(req) => req.stream.unwrap_or(false), + RequestType::Messages(req) => req.stream.unwrap_or(false), + // Embeddings and classification never stream. + RequestType::Embedding(_) | RequestType::Classify(_) => false, + RequestType::Released(_) => false, + }; Self { input: RequestInput { request_type, headers, model_id, + streaming, tenant_request_meta: None, rate_limit_cell: None, + request_lease: None, }, components, state: ProcessingState::default(), @@ -701,17 +776,10 @@ impl RequestContext { } } - /// Check if request is streaming + /// Check if request is streaming (captured at construction, so valid + /// after payload release). pub fn is_streaming(&self) -> bool { - match &self.input.request_type { - RequestType::Chat(req) => req.stream, - RequestType::Generate(req) => req.stream, - RequestType::Completion(req) => req.stream, - RequestType::Responses(req) => req.stream.unwrap_or(false), - RequestType::Messages(req) => req.stream.unwrap_or(false), - RequestType::Embedding(_) => false, // Embeddings are never streaming - RequestType::Classify(_) => false, // Classification is never streaming - } + self.input.streaming } /// Get the cached tokenizer, cloning the Arc (cheap 8-byte clone) diff --git a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs index 536e30b489..842dcd904b 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs @@ -100,7 +100,8 @@ impl PipelineStage for HarmonyRequestBuildingStage { | RequestType::Completion(_) | RequestType::Embedding(_) | RequestType::Classify(_) - | RequestType::Messages(_)) => { + | RequestType::Messages(_) + | RequestType::Released(_)) => { error!( function = "HarmonyRequestBuildingStage::execute", request_type = %request_type, diff --git a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs index 810cb5b41b..9a3fe73009 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs @@ -170,7 +170,8 @@ impl PipelineStage for HarmonyResponseProcessingStage { | RequestType::Completion(_) | RequestType::Embedding(_) | RequestType::Classify(_) - | RequestType::Messages(_)) => { + | RequestType::Messages(_) + | RequestType::Released(_)) => { error!( function = "HarmonyResponseProcessingStage::execute", request_type = %request_type, diff --git a/model_gateway/src/routers/grpc/pipeline.rs b/model_gateway/src/routers/grpc/pipeline.rs index 842b3fc5ff..37facec2ec 100644 --- a/model_gateway/src/routers/grpc/pipeline.rs +++ b/model_gateway/src/routers/grpc/pipeline.rs @@ -54,7 +54,7 @@ use crate::{ observability::metrics::{bool_to_static_str, metrics_labels, Metrics}, policies::PolicyRegistry, rate_limit::{RateLimitManager, UsageSettlement}, - routers::error, + routers::{common::request_lease::ErasedLease, error}, worker::WorkerRegistry, }; @@ -453,12 +453,14 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, rate_limit_cell: Option>, + request_lease: Option>, ) -> Response { let start = Instant::now(); let streaming = request.stream; let mut ctx = RequestContext::for_chat(request, headers, model_id, components); ctx.input.tenant_request_meta = tenant_request_meta; ctx.input.rate_limit_cell = rate_limit_cell; + ctx.input.request_lease = request_lease; let model = ctx.input.model_id.clone(); // Record request start @@ -552,12 +554,14 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, rate_limit_cell: Option>, + request_lease: Option>, ) -> Response { let start = Instant::now(); let streaming = request.stream; let mut ctx = RequestContext::for_generate(request, headers, model_id, components); ctx.input.tenant_request_meta = tenant_request_meta; ctx.input.rate_limit_cell = rate_limit_cell; + ctx.input.request_lease = request_lease; let model_id = ctx.input.model_id.clone(); // Record request start @@ -654,12 +658,14 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, rate_limit_cell: Option>, + request_lease: Option>, ) -> Response { let start = Instant::now(); let streaming = request.stream; let mut ctx = RequestContext::for_completion(request, headers, model_id, components); ctx.input.tenant_request_meta = tenant_request_meta; ctx.input.rate_limit_cell = rate_limit_cell; + ctx.input.request_lease = request_lease; let model = ctx.input.model_id.clone(); Metrics::record_router_request( @@ -957,12 +963,14 @@ impl RequestPipeline { components: Arc, tenant_request_meta: Option, rate_limit_cell: Option>, + request_lease: Option>, ) -> Response { let start = Instant::now(); let streaming = request.stream.unwrap_or(false); let mut ctx = RequestContext::for_messages(request, headers, model_id, components); ctx.input.tenant_request_meta = tenant_request_meta; ctx.input.rate_limit_cell = rate_limit_cell; + ctx.input.request_lease = request_lease; let model = ctx.input.model_id.clone(); // Record request start @@ -1637,6 +1645,7 @@ mod request_release_tests { use super::*; use crate::{ config::types::PolicyConfig, + routers::common::request_lease::{ReleasePoint, RequestLease, RoutingDerivatives}, worker::{BasicWorkerBuilder, ConnectionMode, RuntimeType, WorkerType}, }; @@ -1648,13 +1657,20 @@ mod request_release_tests { Pin> + Send>>; /// TokenSpeed stub gated on the parsed request's drop probe: it withholds - /// its tokens until the probe reaches zero strong references or a - /// deadline passes, recording the outcome in `released`. An ungated stub - /// (no probe) answers immediately -- used for the PD prefill leg. + /// its tokens (or, with `gate_rpc`, the generate RPC itself) until the + /// probe reaches zero strong references or a deadline passes, recording + /// the outcome in `released`. An ungated stub (no probe) answers + /// immediately -- used for the PD prefill leg. `fail_first` makes the + /// first generate call return UNAVAILABLE, for retry-replay tests; every + /// call's input token ids are recorded in `seen_input_ids`. #[derive(Clone, Default)] struct GatedScheduler { probe: Option>, + gate_rpc: bool, released: Arc, + fail_first: bool, + calls: Arc, + seen_input_ids: Arc>>>, } impl GatedScheduler { @@ -1711,9 +1727,22 @@ mod request_release_tests { &self, request: TonicRequest, ) -> Result, Status> { - let request_id = request.into_inner().request_id; + let request = request.into_inner(); + self.seen_input_ids + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(request.tokenized.map(|t| t.input_ids).unwrap_or_default()); + if self.fail_first && self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + return Err(Status::unavailable("release-test induced failure")); + } + if self.gate_rpc { + if let Some(probe) = &self.probe { + Self::await_probe(probe, &self.released).await; + } + } + let request_id = request.request_id; let (tx, rx) = tokio::sync::mpsc::channel(8); - let probe = self.probe.clone(); + let probe = (!self.gate_rpc).then(|| self.probe.clone()).flatten(); let released = Arc::clone(&self.released); tokio::spawn(async move { if let Some(probe) = probe { @@ -1913,9 +1942,18 @@ mod request_release_tests { pipeline: RequestPipeline, components: Arc, request: Arc, + lease: Option>, ) -> bytes::Bytes { let response = pipeline - .execute_completion(request, None, MODEL.to_string(), components, None, None) + .execute_completion( + request, + None, + MODEL.to_string(), + components, + None, + None, + lease, + ) .await; assert_eq!(response.status(), http::StatusCode::OK); axum::body::to_bytes(response.into_body(), usize::MAX) @@ -1942,7 +1980,7 @@ mod request_release_tests { let pipeline = completion_pipeline(&worker_registry, Mode::Regular); let components = components(worker_registry).await; - let body = run_and_drain(pipeline, components, request).await; + let body = run_and_drain(pipeline, components, request, None).await; assert!( released.load(Ordering::SeqCst), @@ -1972,7 +2010,7 @@ mod request_release_tests { let pipeline = completion_pipeline(&worker_registry, Mode::PrefillDecode); let components = components(worker_registry).await; - let body = run_and_drain(pipeline, components, request).await; + let body = run_and_drain(pipeline, components, request, None).await; assert!( released.load(Ordering::SeqCst), @@ -1981,6 +2019,175 @@ mod request_release_tests { let body = String::from_utf8_lossy(&body); assert!(body.contains("data: [DONE]"), "stream must finish: {body}"); } + + /// Retries disabled: the parsed request must be freed at dispatch. The + /// stub refuses to answer the generate RPC until the probe frees. + #[tokio::test] + async fn disabled_retries_release_parsed_request_before_upstream_responds() { + let request = completion_request(false); + let released = Arc::new(AtomicBool::new(false)); + let port = spawn_stub(GatedScheduler { + probe: Some(Arc::downgrade(&request)), + gate_rpc: true, + released: Arc::clone(&released), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, port, WorkerType::Regular); + let pipeline = completion_pipeline(&worker_registry, Mode::Regular); + let components = components(worker_registry).await; + + let lease = Arc::new(RequestLease::new( + request, + RoutingDerivatives::default(), + ReleasePoint::AfterDispatch, + )); + let attempt = lease.with_view(|view| Arc::clone(view.request)); + let response = pipeline + .execute_completion( + attempt, + None, + MODEL.to_string(), + components, + None, + None, + Some(lease as Arc), + ) + .await; + + assert_eq!(response.status(), http::StatusCode::OK); + assert!( + released.load(Ordering::SeqCst), + "the parsed request must be freed before the upstream answers" + ); + } + + /// grpc_pd twin of the dispatch-release probe: both legs' RPCs answer + /// only after the parsed request is freed. + #[tokio::test] + async fn pd_disabled_retries_release_parsed_request_before_upstream_responds() { + let request = completion_request(false); + let released = Arc::new(AtomicBool::new(false)); + let prefill_port = spawn_stub(GatedScheduler::default()).await; + let decode_port = spawn_stub(GatedScheduler { + probe: Some(Arc::downgrade(&request)), + gate_rpc: true, + released: Arc::clone(&released), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, prefill_port, WorkerType::Prefill); + register_worker(&worker_registry, decode_port, WorkerType::Decode); + let pipeline = completion_pipeline(&worker_registry, Mode::PrefillDecode); + let components = components(worker_registry).await; + + let lease = Arc::new(RequestLease::new( + request, + RoutingDerivatives::default(), + ReleasePoint::AfterDispatch, + )); + let attempt = lease.with_view(|view| Arc::clone(view.request)); + let response = pipeline + .execute_completion( + attempt, + None, + MODEL.to_string(), + components, + None, + None, + Some(lease as Arc), + ) + .await; + + assert_eq!(response.status(), http::StatusCode::OK); + assert!( + released.load(Ordering::SeqCst), + "the parsed request must be freed before the decode leg answers" + ); + } + + /// Retries enabled (AtRetryClose): a failed first dispatch must leave the + /// request intact, the second attempt must send identical token ids, and + /// the retry-window close (lease drop) frees it. + #[tokio::test] + async fn enabled_retries_replay_identical_token_ids() { + let request = completion_request(false); + let probe = Arc::downgrade(&request); + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let port = spawn_stub(GatedScheduler { + fail_first: true, + seen_input_ids: Arc::clone(&seen), + ..Default::default() + }) + .await; + + let worker_registry = Arc::new(WorkerRegistry::new()); + register_worker(&worker_registry, port, WorkerType::Regular); + let pipeline = completion_pipeline(&worker_registry, Mode::Regular); + let components = components(worker_registry).await; + + let lease = Arc::new(RequestLease::new( + request, + RoutingDerivatives::default(), + ReleasePoint::AtRetryClose, + )); + + let attempt = lease.with_view(|view| Arc::clone(view.request)); + let response = pipeline + .execute_completion( + attempt, + None, + MODEL.to_string(), + components.clone(), + None, + None, + Some(lease.clone() as Arc), + ) + .await; + assert!( + !response.status().is_success(), + "first dispatch is induced to fail" + ); + assert!(probe.upgrade().is_some(), "request must survive for replay"); + + let attempt = lease.with_view(|view| Arc::clone(view.request)); + let response = pipeline + .execute_completion( + attempt, + None, + MODEL.to_string(), + components, + None, + None, + Some(lease.clone() as Arc), + ) + .await; + assert_eq!(response.status(), http::StatusCode::OK); + + { + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 2, "503 then 200 must mean two attempts"); + assert_eq!( + seen[0], seen[1], + "the retry must replay identical input ids" + ); + assert!( + !seen[0].is_empty(), + "attempts must carry the tokenized prompt" + ); + } + + drop(lease); + assert_eq!( + probe.strong_count(), + 0, + "retry-window close must free the request" + ); + } } #[cfg(test)] diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index 374cb224c5..bf3bfc422c 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -1096,6 +1096,14 @@ pub enum ProtoRequest { } impl ProtoRequest { + /// Serialized wire size, for the release metric. + pub fn wire_len(&self) -> usize { + match self { + Self::Generate(request) => request.wire_len(), + Self::Embed(request) => request.wire_len(), + } + } + /// Get request ID from either variant pub fn request_id(&self) -> &str { match self { @@ -1301,6 +1309,18 @@ impl ProtoGenerateRequest { self.clone() } + /// Serialized wire size, for the release metric. + pub fn wire_len(&self) -> usize { + use prost::Message; + match self { + Self::Sglang(req) => req.encoded_len(), + Self::Vllm(req) => req.encoded_len(), + Self::Trtllm(req) => req.encoded_len(), + Self::Mlx(req) => req.encoded_len(), + Self::TokenSpeed(req) => req.encoded_len(), + } + } + /// Drop raw multimodal encoder tensors while keeping item metadata. /// /// Used by the EPD prefill leg: multimodal embeddings arrive from encode workers, @@ -2094,6 +2114,15 @@ pub enum ProtoEmbedRequest { } impl ProtoEmbedRequest { + /// Serialized wire size, for the release metric. + pub fn wire_len(&self) -> usize { + use prost::Message; + match self { + Self::Sglang(req) => req.encoded_len(), + Self::Vllm(req) => req.encoded_len(), + } + } + /// Get SGLang variant #[expect( clippy::panic, diff --git a/model_gateway/src/routers/grpc/regular/responses/streaming.rs b/model_gateway/src/routers/grpc/regular/responses/streaming.rs index fa81d57006..81451847c8 100644 --- a/model_gateway/src/routers/grpc/regular/responses/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/responses/streaming.rs @@ -98,6 +98,8 @@ pub(super) async fn convert_chat_stream_to_responses_stream( // (out of scope for this phase; see RequestPipeline::build's // stage-insertion comment). None, + // No lease: the tool loop re-reads the request across iterations. + None, ) .await; @@ -596,6 +598,8 @@ async fn execute_tool_loop_streaming_internal( Some(params.tenant_request_meta.clone()), // Responses endpoint hasn't opted into tenant rate limiting yet. None, + // No lease: the tool loop re-reads the request across iterations. + None, ) .await; diff --git a/model_gateway/src/routers/grpc/regular/stages/response_processing.rs b/model_gateway/src/routers/grpc/regular/stages/response_processing.rs index 95d21a710f..fef1297075 100644 --- a/model_gateway/src/routers/grpc/regular/stages/response_processing.rs +++ b/model_gateway/src/routers/grpc/regular/stages/response_processing.rs @@ -14,7 +14,7 @@ use crate::routers::{ error, grpc::{ common::stages::PipelineStage, - context::{RequestContext, RequestType}, + context::{RequestContext, RequestKind}, regular::{processor, streaming}, }, }; @@ -43,18 +43,19 @@ impl ChatGenerateResponseProcessingStage { #[async_trait] impl PipelineStage for ChatGenerateResponseProcessingStage { async fn execute(&self, ctx: &mut RequestContext) -> Result, Response> { - match &ctx.input.request_type { - RequestType::Chat(_) => self.chat_stage.execute(ctx).await, - RequestType::Generate(_) => self.generate_stage.execute(ctx).await, - request_type => { + // Dispatch on the kind, which survives payload release. + match ctx.input.request_type.kind() { + RequestKind::Chat => self.chat_stage.execute(ctx).await, + RequestKind::Generate => self.generate_stage.execute(ctx).await, + request_kind => { error!( function = "ChatGenerateResponseProcessingStage::execute", - request_type = %request_type, - "{request_type} should not reach this stage" + request_type = %request_kind, + "{request_kind} should not reach this stage" ); Err(error::internal_error( "wrong_pipeline", - format!("{request_type} should use its dedicated pipeline"), + format!("{request_kind} should use its dedicated pipeline"), )) } } diff --git a/model_gateway/src/routers/grpc/router.rs b/model_gateway/src/routers/grpc/router.rs index 3e3a27ff5b..a7a7ec2187 100644 --- a/model_gateway/src/routers/grpc/router.rs +++ b/model_gateway/src/routers/grpc/router.rs @@ -41,7 +41,10 @@ use crate::{ middleware::TenantRequestMeta, observability::metrics::{metrics_labels, Metrics}, routers::{ - common::retry::{is_retryable_response, RetryExecutor}, + common::{ + request_lease::{ErasedLease, ReleasePoint, RequestLease, RoutingDerivatives}, + retry::{is_retryable_response, RetryExecutor}, + }, error, RouterTrait, }, worker::{WorkerRegistry, WorkerType}, @@ -569,7 +572,6 @@ impl GrpcRouter { let model_id_cloned = self.resolve_canonical_model_id(model_id); let mut canonical_body = body; canonical_body.model = model_id_cloned.clone(); - let request = Arc::new(canonical_body); let headers_cloned = headers.cloned(); let components = self.shared_components.clone(); let tenant_meta_cloned = tenant_meta.clone(); @@ -577,16 +579,32 @@ impl GrpcRouter { let retry_config = self.resolve_retry_config_for_canonical(&model_id_cloned); + // The lease owns the parsed request for the dispatch phase; its + // release point encodes the retry policy. The pipeline releases it at + // dispatch — except on the Harmony pipeline, whose response + // processing still reads the request after dispatch. + let lease = Arc::new(RequestLease::new( + Arc::new(canonical_body), + RoutingDerivatives::default(), + ReleasePoint::from_retry_config(&retry_config), + )); + let pipeline_lease: Option> = if is_harmony { + None + } else { + Some(lease.clone()) + }; + let response = RetryExecutor::execute_response_with_retry( &retry_config, // Operation: execute pipeline (creates fresh context each attempt) |_attempt| { - let request = Arc::clone(&request); + let request = lease.with_view(|view| Arc::clone(view.request)); let headers = headers_cloned.clone(); let model_id = model_id_cloned.clone(); let components = Arc::clone(&components); let tenant_meta = tenant_meta_cloned.clone(); let rate_limit_cell = Arc::clone(&rate_limit_cell); + let request_lease = pipeline_lease.clone(); async move { pipeline .execute_chat( @@ -596,6 +614,7 @@ impl GrpcRouter { components, Some(tenant_meta), Some(rate_limit_cell), + request_lease, ) .await } @@ -641,7 +660,6 @@ impl GrpcRouter { let model_id_cloned = self.resolve_canonical_model_id(model_id); let mut canonical_body = body; canonical_body.model = model_id_cloned.clone(); - let request = Arc::new(canonical_body); let headers_cloned = headers.cloned(); let components = self.shared_components.clone(); let tenant_meta_cloned = tenant_meta.clone(); @@ -650,16 +668,24 @@ impl GrpcRouter { let retry_config = self.resolve_retry_config_for_canonical(&model_id_cloned); + // Dispatch-phase owner of the parsed request; see route_chat_impl. + let lease = Arc::new(RequestLease::new( + Arc::new(canonical_body), + RoutingDerivatives::default(), + ReleasePoint::from_retry_config(&retry_config), + )); + let response = RetryExecutor::execute_response_with_retry( &retry_config, // Operation: execute pipeline (creates fresh context each attempt) |_attempt| { - let request = Arc::clone(&request); + let request = lease.with_view(|view| Arc::clone(view.request)); let headers = headers_cloned.clone(); let model_id = model_id_cloned.clone(); let components = Arc::clone(&components); let tenant_meta = tenant_meta_cloned.clone(); let rate_limit_cell = Arc::clone(&rate_limit_cell); + let request_lease: Arc = lease.clone(); async move { pipeline .execute_generate( @@ -669,6 +695,7 @@ impl GrpcRouter { components, Some(tenant_meta), Some(rate_limit_cell), + Some(request_lease), ) .await } @@ -809,7 +836,6 @@ impl GrpcRouter { let model_id_cloned = self.resolve_canonical_model_id(model_id); let mut canonical_body = body; canonical_body.model = model_id_cloned.clone(); - let request = Arc::new(canonical_body); let headers_cloned = headers.cloned(); let components = self.shared_components.clone(); let tenant_meta_cloned = tenant_meta.clone(); @@ -818,15 +844,23 @@ impl GrpcRouter { let retry_config = self.resolve_retry_config_for_canonical(&model_id_cloned); + // Dispatch-phase owner of the parsed request; see route_chat_impl. + let lease = Arc::new(RequestLease::new( + Arc::new(canonical_body), + RoutingDerivatives::default(), + ReleasePoint::from_retry_config(&retry_config), + )); + let response = RetryExecutor::execute_response_with_retry( &retry_config, |_attempt| { - let request = Arc::clone(&request); + let request = lease.with_view(|view| Arc::clone(view.request)); let headers = headers_cloned.clone(); let model_id = model_id_cloned.clone(); let components = Arc::clone(&components); let tenant_meta = tenant_meta_cloned.clone(); let rate_limit_cell = Arc::clone(&rate_limit_cell); + let request_lease: Arc = lease.clone(); async move { pipeline .execute_messages( @@ -836,6 +870,7 @@ impl GrpcRouter { components, Some(tenant_meta), Some(rate_limit_cell), + Some(request_lease), ) .await } @@ -877,7 +912,6 @@ impl GrpcRouter { let model_id_cloned = self.resolve_canonical_model_id(model_id); let mut canonical_body = body; canonical_body.model = model_id_cloned.clone(); - let request = Arc::new(canonical_body); let headers_cloned = headers.cloned(); let components = self.shared_components.clone(); let tenant_meta_cloned = tenant_meta.clone(); @@ -886,15 +920,23 @@ impl GrpcRouter { let retry_config = self.resolve_retry_config_for_canonical(&model_id_cloned); + // Dispatch-phase owner of the parsed request; see route_chat_impl. + let lease = Arc::new(RequestLease::new( + Arc::new(canonical_body), + RoutingDerivatives::default(), + ReleasePoint::from_retry_config(&retry_config), + )); + let response = RetryExecutor::execute_response_with_retry( &retry_config, |_attempt| { - let request = Arc::clone(&request); + let request = lease.with_view(|view| Arc::clone(view.request)); let headers = headers_cloned.clone(); let model_id = model_id_cloned.clone(); let components = Arc::clone(&components); let tenant_meta = tenant_meta_cloned.clone(); let rate_limit_cell = Arc::clone(&rate_limit_cell); + let request_lease: Arc = lease.clone(); async move { pipeline .execute_completion( @@ -904,6 +946,7 @@ impl GrpcRouter { components, Some(tenant_meta), Some(rate_limit_cell), + Some(request_lease), ) .await }