chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[smg] cleanup router RAII guards (#16560)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-09 08:39:53 +08:00
committed by GitHub
parent 1bc7aa5801
commit 9e3a032ad6
9 changed files with 79 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
sgl-model-gateway/src/core/mod.rs
View File

@@ -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};

94
sgl-model-gateway/src/core/worker.rs
View File

@@ -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<WorkerLoadGuard>,
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<T> {
inner: Body,
_guard: WorkerLoadGuard,
_attached: T,
}
/// Body wrapper that holds multiple WorkerLoadGuards (for dual prefill/decode)
struct MultiGuardedBody {
inner: Body,
_guards: Vec<WorkerLoadGuard>,
impl<T> AttachedBody<T> {
pub fn new(inner: Body, attached: T) -> Self {
Self {
inner,
_attached: attached,
}
}
}
impl http_body::Body for GuardedBody {
impl<T: Send + Unpin + 'static> AttachedBody<T> {
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<T: Send + Unpin + 'static> http_body::Body for AttachedBody<T> {
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<Option<Result<http_body::Frame<Self::Data>, 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<Option<Result<http_body::Frame<Self::Data>, 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 {

40
sgl-model-gateway/src/routers/grpc/context.rs
View File

@@ -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 {

15
sgl-model-gateway/src/routers/grpc/harmony/stages/response_processing.rs
View File

@@ -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,
};

17
sgl-model-gateway/src/routers/grpc/regular/stages/chat/response_processing.rs
View File

@@ -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,
};

17
sgl-model-gateway/src/routers/grpc/regular/stages/generate/response_processing.rs
View File

@@ -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,
};

4
sgl-model-gateway/src/routers/http/pd_router.rs
View File

@@ -835,7 +835,7 @@ impl PDRouter {
prefill: Arc<dyn Worker>,
decode: Arc<dyn Worker>,
) -> 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

4
sgl-model-gateway/src/routers/http/router.rs
View File

@@ -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
}

15
sgl-model-gateway/tests/load_guard_raii_test.rs
View File

@@ -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();