diff --git a/sgl-model-gateway/src/core/mod.rs b/sgl-model-gateway/src/core/mod.rs index 587ec44cd..055812c80 100644 --- a/sgl-model-gateway/src/core/mod.rs +++ b/sgl-model-gateway/src/core/mod.rs @@ -34,8 +34,8 @@ pub use job_queue::{Job, JobQueue, JobQueueConfig}; pub use model_card::{ModelCard, ProviderType}; pub use retry::{is_retryable_status, RetryExecutor}; pub use worker::{ - attach_guards_to_response, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker, - WorkerLoadGuard, WorkerType, + AttachedBody, BasicWorker, ConnectionMode, HealthConfig, RuntimeType, Worker, WorkerLoadGuard, + WorkerType, }; pub use worker_builder::{BasicWorkerBuilder, DPAwareWorkerBuilder}; pub use worker_manager::{LoadMonitor, WorkerManager}; diff --git a/sgl-model-gateway/src/core/worker.rs b/sgl-model-gateway/src/core/worker.rs index 036e27597..0cb986f4a 100644 --- a/sgl-model-gateway/src/core/worker.rs +++ b/sgl-model-gateway/src/core/worker.rs @@ -1015,28 +1015,6 @@ impl WorkerLoadGuard { worker.increment_load(); Self { worker } } - - /// Attach this guard to a Response, tying the guard's lifetime to the response body. - /// - /// When the response body is fully consumed or dropped (e.g., client disconnects), - /// the guard is dropped and worker load is decremented automatically. - /// - /// This is the proper RAII pattern for SSE/streaming responses where the handler - /// returns immediately but the stream continues in a background task. - pub fn attach_to_response( - self, - response: axum::response::Response, - ) -> axum::response::Response { - let (parts, body) = response.into_parts(); - - // Wrap body with guard - guard drops when body drops - let guarded_body = GuardedBody { - inner: body, - _guard: self, - }; - - axum::response::Response::from_parts(parts, Body::new(guarded_body)) - } } impl Drop for WorkerLoadGuard { @@ -1045,65 +1023,45 @@ impl Drop for WorkerLoadGuard { } } -/// Attach multiple guards to a Response (for dual prefill/decode workers) -pub fn attach_guards_to_response( - guards: Vec, - response: axum::response::Response, -) -> axum::response::Response { - let (parts, body) = response.into_parts(); - - let guarded_body = MultiGuardedBody { - inner: body, - _guards: guards, - }; - - axum::response::Response::from_parts(parts, Body::new(guarded_body)) -} - -/// Body wrapper that holds a WorkerLoadGuard +/// Body wrapper that holds an attached value. /// /// When this body is dropped (stream ends or client disconnects), -/// the guard is dropped, decrementing worker load. -struct GuardedBody { +/// the attached value is dropped automatically. This is useful for RAII guards +/// like WorkerLoadGuard that need to be tied to a response body's lifetime. +pub struct AttachedBody { inner: Body, - _guard: WorkerLoadGuard, + _attached: T, } -/// Body wrapper that holds multiple WorkerLoadGuards (for dual prefill/decode) -struct MultiGuardedBody { - inner: Body, - _guards: Vec, +impl AttachedBody { + pub fn new(inner: Body, attached: T) -> Self { + Self { + inner, + _attached: attached, + } + } } -impl http_body::Body for GuardedBody { +impl AttachedBody { + pub fn wrap_response( + response: axum::response::Response, + attached: T, + ) -> axum::response::Response { + let (parts, body) = response.into_parts(); + axum::response::Response::from_parts(parts, Body::new(Self::new(body, attached))) + } +} + +impl http_body::Body for AttachedBody { type Data = bytes::Bytes; type Error = axum::Error; fn poll_frame( - mut self: std::pin::Pin<&mut Self>, + self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>, ) -> std::task::Poll, Self::Error>>> { - std::pin::Pin::new(&mut self.inner).poll_frame(cx) - } - - fn is_end_stream(&self) -> bool { - self.inner.is_end_stream() - } - - fn size_hint(&self) -> http_body::SizeHint { - self.inner.size_hint() - } -} - -impl http_body::Body for MultiGuardedBody { - type Data = bytes::Bytes; - type Error = axum::Error; - - fn poll_frame( - mut self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll, Self::Error>>> { - std::pin::Pin::new(&mut self.inner).poll_frame(cx) + let this = self.get_mut(); + std::pin::Pin::new(&mut this.inner).poll_frame(cx) } fn is_end_stream(&self) -> bool { diff --git a/sgl-model-gateway/src/routers/grpc/context.rs b/sgl-model-gateway/src/routers/grpc/context.rs index 49a2404fd..7a6a8eab8 100644 --- a/sgl-model-gateway/src/routers/grpc/context.rs +++ b/sgl-model-gateway/src/routers/grpc/context.rs @@ -13,7 +13,7 @@ use super::{ proto_wrapper::{ProtoEmbedComplete, ProtoRequest, ProtoStream}, }; use crate::{ - core::{attach_guards_to_response, Worker, WorkerLoadGuard}, + core::{Worker, WorkerLoadGuard}, protocols::{ chat::{ChatCompletionRequest, ChatCompletionResponse}, classify::{ClassifyRequest, ClassifyResponse}, @@ -158,47 +158,29 @@ pub(crate) struct DispatchMetadata { /// Load guards for worker load tracking /// Automatically decrements load when dropped pub(crate) enum LoadGuards { - Single(WorkerLoadGuard), + Single { + _guard: WorkerLoadGuard, + }, Dual { - prefill: WorkerLoadGuard, - decode: WorkerLoadGuard, + _prefill: WorkerLoadGuard, + _decode: WorkerLoadGuard, }, } impl From<&WorkerSelection> for LoadGuards { fn from(selection: &WorkerSelection) -> Self { match selection { - WorkerSelection::Single { worker } => { - LoadGuards::Single(WorkerLoadGuard::new(worker.clone())) - } + WorkerSelection::Single { worker } => LoadGuards::Single { + _guard: WorkerLoadGuard::new(worker.clone()), + }, WorkerSelection::Dual { prefill, decode } => LoadGuards::Dual { - prefill: WorkerLoadGuard::new(prefill.clone()), - decode: WorkerLoadGuard::new(decode.clone()), + _prefill: WorkerLoadGuard::new(prefill.clone()), + _decode: WorkerLoadGuard::new(decode.clone()), }, } } } -impl LoadGuards { - /// Attach these load guards to a Response, tying their lifetime to the response body. - /// - /// When the response body is fully consumed or dropped (e.g., client disconnects), - /// the guards are dropped and worker load is decremented automatically. - /// - /// This is the proper RAII pattern for SSE/streaming responses. - pub fn attach_to_response( - self, - response: axum::response::Response, - ) -> axum::response::Response { - let guards = match self { - LoadGuards::Single(guard) => vec![guard], - LoadGuards::Dual { prefill, decode } => vec![prefill, decode], - }; - - attach_guards_to_response(guards, response) - } -} - /// Response processing state (Step 6) #[derive(Default)] pub(crate) struct ResponseState { diff --git a/sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs b/sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs index 2a237ab1a..c0346c9ad 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs @@ -7,11 +7,14 @@ use axum::response::Response; use tracing::error; use super::super::{HarmonyResponseProcessor, HarmonyStreamingProcessor}; -use crate::routers::{ - error, - grpc::{ - common::stages::PipelineStage, - context::{FinalResponse, RequestContext, RequestType}, +use crate::{ + core::AttachedBody, + routers::{ + error, + grpc::{ + common::stages::PipelineStage, + context::{FinalResponse, RequestContext, RequestType}, + }, }, }; @@ -81,7 +84,7 @@ impl PipelineStage for HarmonyResponseProcessingStage { // Attach load guards to response body for proper RAII lifecycle let response = match ctx.state.load_guards.take() { - Some(guards) => guards.attach_to_response(response), + Some(guards) => AttachedBody::wrap_response(response, guards), None => response, }; diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs index 044c2063d..8d2a5f933 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs @@ -9,12 +9,15 @@ use async_trait::async_trait; use axum::response::Response; use tracing::error; -use crate::routers::{ - error, - grpc::{ - common::stages::PipelineStage, - context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, +use crate::{ + core::AttachedBody, + routers::{ + error, + grpc::{ + common::stages::PipelineStage, + context::{FinalResponse, RequestContext}, + regular::{processor, streaming}, + }, }, }; @@ -100,7 +103,7 @@ impl ChatResponseProcessingStage { // Attach load guards to response body for proper RAII lifecycle let response = match ctx.state.load_guards.take() { - Some(guards) => guards.attach_to_response(response), + Some(guards) => AttachedBody::wrap_response(response, guards), None => response, }; diff --git a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs index 62093268d..d0a55d4ea 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs @@ -6,12 +6,15 @@ use async_trait::async_trait; use axum::response::Response; use tracing::error; -use crate::routers::{ - error, - grpc::{ - common::stages::PipelineStage, - context::{FinalResponse, RequestContext}, - regular::{processor, streaming}, +use crate::{ + core::AttachedBody, + routers::{ + error, + grpc::{ + common::stages::PipelineStage, + context::{FinalResponse, RequestContext}, + regular::{processor, streaming}, + }, }, }; @@ -100,7 +103,7 @@ impl GenerateResponseProcessingStage { // Attach load guards to response body for proper RAII lifecycle let response = match ctx.state.load_guards.take() { - Some(guards) => guards.attach_to_response(response), + Some(guards) => AttachedBody::wrap_response(response, guards), None => response, }; diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 8fb73b05c..78a40a30d 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -835,7 +835,7 @@ impl PDRouter { prefill: Arc, decode: Arc, ) -> Response { - use crate::core::attach_guards_to_response; + use crate::core::AttachedBody; let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); @@ -885,7 +885,7 @@ impl PDRouter { // Attach load guards to response body for proper RAII lifecycle // Guards are dropped when response body is consumed or client disconnects let guards = vec![WorkerLoadGuard::new(prefill), WorkerLoadGuard::new(decode)]; - attach_guards_to_response(guards, response) + AttachedBody::wrap_response(response, guards) } // Helper to process non-streaming decode response with logprob merging diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 347ef095e..f76b0372e 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -16,7 +16,7 @@ use crate::{ app_context::AppContext, config::types::RetryConfig, core::{ - is_retryable_status, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard, + is_retryable_status, AttachedBody, ConnectionMode, RetryExecutor, Worker, WorkerLoadGuard, WorkerRegistry, WorkerType, UNKNOWN_MODEL_ID, }, observability::{ @@ -629,7 +629,7 @@ impl Router { // Attach load guard to response body for proper RAII lifecycle // Guard is dropped when response body is consumed or client disconnects if let Some(guard) = load_guard { - response = guard.attach_to_response(response); + response = AttachedBody::wrap_response(response, guard); } response } diff --git a/sgl-model-gateway/tests/load_guard_raii_test.rs b/sgl-model-gateway/tests/load_guard_raii_test.rs index 7be1bd03f..3f7634307 100644 --- a/sgl-model-gateway/tests/load_guard_raii_test.rs +++ b/sgl-model-gateway/tests/load_guard_raii_test.rs @@ -11,7 +11,7 @@ use axum::{body::Body, response::Response}; use bytes::Bytes; use futures_util::StreamExt; use http_body_util::BodyExt; -use smg::core::{attach_guards_to_response, BasicWorkerBuilder, Worker, WorkerLoadGuard}; +use smg::core::{AttachedBody, BasicWorkerBuilder, Worker, WorkerLoadGuard}; use tokio::sync::mpsc; use tokio_stream::wrappers::UnboundedReceiverStream; @@ -40,7 +40,7 @@ async fn test_guard_dropped_when_response_body_consumed() { let guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); - let guarded_response = guard.attach_to_response(response); + let guarded_response = AttachedBody::wrap_response(response, guard); // Load should still be 1 (guard is in the body) assert_eq!(worker.load(), 1); @@ -65,7 +65,7 @@ async fn test_guard_dropped_when_response_dropped_without_consumption() { let guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); - let _guarded_response = guard.attach_to_response(response); + let _guarded_response = AttachedBody::wrap_response(response, guard); // Load is still 1 assert_eq!(worker.load(), 1); @@ -89,7 +89,7 @@ async fn test_streaming_guard_dropped_when_stream_ends() { let guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); - let guarded_response = guard.attach_to_response(response); + let guarded_response = AttachedBody::wrap_response(response, guard); // Spawn a task to consume the response let worker_clone = worker.clone(); @@ -136,7 +136,7 @@ async fn test_streaming_guard_dropped_on_client_disconnect() { let guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); - let guarded_response = guard.attach_to_response(response); + let guarded_response = AttachedBody::wrap_response(response, guard); // Start consuming but drop early (simulate client disconnect) { @@ -176,8 +176,7 @@ async fn test_multiple_guards_all_dropped() { assert_eq!(worker1.load(), 1); assert_eq!(worker2.load(), 1); - // Attach both guards using attach_guards_to_response - let _response = attach_guards_to_response(vec![guard1, guard2], response); + let _response = AttachedBody::wrap_response(response, vec![guard1, guard2]); // Both loads are 1 assert_eq!(worker1.load(), 1); @@ -201,7 +200,7 @@ async fn test_guard_with_empty_body() { let guard = WorkerLoadGuard::new(worker.clone()); assert_eq!(worker.load(), 1); - let guarded_response = guard.attach_to_response(response); + let guarded_response = AttachedBody::wrap_response(response, guard); // Consume empty body let body = guarded_response.into_body();