minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] Restore response streaming by optimizing WASM middleware buffering (#16804)

Browse Source
This commit is contained in:
Praneth Paruchuri
2026-01-09 09:29:22 -08:00
committed by GitHub
parent 8ef5b90528
commit bd1afeb568
3 changed files with 193 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/Cargo.toml
+4
View File
@@ -149,6 +149,10 @@ tonic-v12 = { version = "0.12.3", package = "tonic" }
serial_test = "3.0" serial_test = "3.0"
rsa = { version = "0.9", features = ["sha2"] } rsa = { version = "0.9", features = ["sha2"] }
[[bench]]
name = "wasm_middleware_latency"
harness = false
path = "benches/wasm_middleware_latency.rs"
[[bench]] [[bench]]
name = "request_processing" name = "request_processing"
harness = false harness = false
sgl-model-gateway/benches/wasm_middleware_latency.rs
+108
View File
@@ -0,0 +1,108 @@
use std::sync::Arc;
use axum::{
body::Body,
http::{HeaderMap, Request, Response, StatusCode},
middleware,
response::IntoResponse,
};
use criterion::{criterion_group, criterion_main, Criterion};
use http_body_util::BodyExt;
use smg::{
app_context::AppContext, config::RouterConfig, middleware::wasm_middleware,
protocols::chat::ChatCompletionRequest, routers::RouterTrait, server::AppState,
};
use tokio::runtime::Runtime;
use tower::{Layer, Service};
#[derive(Debug)]
struct MockRouter;
#[async_trait::async_trait]
impl RouterTrait for MockRouter {
fn as_any(&self) -> &dyn std::any::Any {
self
}
async fn route_chat(
&self,
_headers: Option<&HeaderMap>,
_body: &ChatCompletionRequest,
_model_id: Option<&str>,
) -> Response<Body> {
StatusCode::OK.into_response()
}
fn router_type(&self) -> &'static str {
"mock"
}
}
/// Mock service that simulates a streaming response with a 500ms delay.
async fn mock_next_streaming(_req: Request<Body>) -> Response<Body> {
let (tx, rx) = tokio::sync::mpsc::channel(16);
tokio::spawn(async move {
// Send first chunk immediately
let _ = tx
.send(Ok::<_, std::io::Error>(bytes::Bytes::from("chunk 1 ")))
.await;
// Simulate generation delay
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
// Send final chunk
let _ = tx
.send(Ok::<_, std::io::Error>(bytes::Bytes::from("chunk 2")))
.await;
});
Response::new(Body::from_stream(
tokio_stream::wrappers::ReceiverStream::new(rx),
))
}
fn bench_wasm_middleware_buffering(c: &mut Criterion) {
let rt = Runtime::new().unwrap();
// Setup AppContext with WASM enabled
let config = RouterConfig::builder().enable_wasm(true).build_unchecked();
let context = rt.block_on(AppContext::from_config(config, 30)).unwrap();
let app_state = Arc::new(AppState {
router: Arc::new(MockRouter),
context: Arc::new(context),
concurrency_queue_tx: None,
router_manager: None,
});
c.bench_function("wasm_middleware_pre_fix_latency", |b| {
b.iter(|| {
rt.block_on(async {
let req = Request::builder()
.uri("/v1/chat/completions")
.body(Body::empty())
.unwrap();
// Create the service by applying the middleware layer to the mock streamer
let mut service =
middleware::from_fn_with_state(app_state.clone(), wasm_middleware).layer(
tower::service_fn(|req: Request<Body>| async move {
Ok::<_, std::convert::Infallible>(mock_next_streaming(req).await)
}),
);
// Explicitly poll the service
let response: Response<Body> =
service.call(req).await.expect("Middleware service failed");
// Measure how long it takes to receive the FIRST frame
let mut body = response.into_body();
let _first_frame = body.frame().await;
});
});
});
}
criterion_group! {
name = benches;
config = Criterion::default().sample_size(10);
targets = bench_wasm_middleware_buffering
}
criterion_main!(benches);
sgl-model-gateway/src/middleware.rs
+81 -82
View File
@@ -773,95 +773,91 @@ pub async fn wasm_middleware(
} }
}; };
// Extract request body once before processing modules let response = if modules_on_request.is_empty() {
let method = request.method().clone(); next.run(request).await
let uri = request.uri().clone(); } else {
let mut headers = request.headers().clone(); // Extract request body once before processing modules
let max_body_size = wasm_manager.get_max_body_size(); let method = request.method().clone();
let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await { let uri = request.uri().clone();
Ok(bytes) => bytes.to_vec(), let mut headers = request.headers().clone();
Err(e) => { let max_body_size = wasm_manager.get_max_body_size();
error!("Failed to read request body: {}", e); let body_bytes = match axum::body::to_bytes(request.into_body(), max_body_size).await {
// Create a minimal request with empty body for error recovery Ok(bytes) => bytes.to_vec(),
let error_request = Request::builder() Err(e) => {
.uri(uri) error!("Failed to read request body: {}", e);
.body(Body::empty()) // Create a minimal request with empty body for error recovery
.unwrap_or_else(|_| Request::new(Body::empty())); let error_request = Request::builder()
return Ok(next.run(error_request).await); .uri(uri)
} .body(Body::empty())
}; .unwrap_or_else(|_| Request::new(Body::empty()));
return Ok(next.run(error_request).await);
// Process each OnRequest module }
let mut modified_body = body_bytes;
// Pre-compute strings once before the loop to avoid repeated allocations
let method_str = method.to_string();
let path_str = uri.path().to_string();
let query_str = uri.query().unwrap_or("").to_string();
for module in modules_on_request {
// Build WebAssembly request from collected data
let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
let wasm_request = WasmRequest {
method: method_str.clone(),
path: path_str.clone(),
query: query_str.clone(),
headers: wasm_headers,
body: modified_body.clone(),
request_id: request_id.clone(),
now_epoch_ms: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_else(|_| {
// Fallback to 0 if system time is before UNIX_EPOCH
// This should never happen in practice, but provides a safe fallback
Duration::from_millis(0)
})
.as_millis() as u64,
}; };
// Execute WASM component // Process each OnRequest module
let action = match wasm_manager let mut modified_body = body_bytes;
.execute_module_for_attach_point(
&module,
on_request_attach_point.clone(),
WasmComponentInput::MiddlewareRequest(wasm_request),
)
.await
{
Some(action) => action,
None => continue, // Continue to next module on error
};
// Process action // Pre-compute strings once before the loop to avoid repeated allocations
match action { let method_str = method.to_string();
Action::Continue => { let path_str = uri.path().to_string();
// Continue to next module or request processing let query_str = uri.query().unwrap_or("").to_string();
}
Action::Reject(status) => { for module in modules_on_request {
// Immediately reject the request // Build WebAssembly request from collected data
return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST)); let wasm_headers = build_wasm_headers_from_axum_headers(&headers);
} let wasm_request = WasmRequest {
Action::Modify(modify) => { method: method_str.clone(),
// Apply modifications to headers and body path: path_str.clone(),
apply_modify_action_to_headers(&mut headers, &modify); query: query_str.clone(),
// Apply body_replace headers: wasm_headers,
if let Some(body_bytes) = modify.body_replace { body: modified_body.clone(),
modified_body = body_bytes; request_id: request_id.clone(),
now_epoch_ms: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_else(|_| Duration::from_millis(0))
.as_millis() as u64,
};
// Execute WASM component
let action = match wasm_manager
.execute_module_for_attach_point(
&module,
on_request_attach_point.clone(),
WasmComponentInput::MiddlewareRequest(wasm_request),
)
.await
{
Some(action) => action,
None => continue, // Continue to next module on error
};
// Process action
match action {
Action::Continue => {}
Action::Reject(status) => {
return Err(StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST));
}
Action::Modify(modify) => {
// Apply modifications to headers and body
apply_modify_action_to_headers(&mut headers, &modify);
if let Some(body_bytes) = modify.body_replace {
modified_body = body_bytes;
}
} }
} }
} }
}
// Reconstruct request with modifications // Reconstruct request with modifications
let mut final_request = Request::builder() let mut final_request = Request::builder()
.method(method) .method(method)
.uri(uri) .uri(uri)
.body(Body::from(modified_body)) .body(Body::from(modified_body))
.unwrap_or_else(|_| Request::new(Body::empty())); .unwrap_or_else(|_| Request::new(Body::empty()));
*final_request.headers_mut() = headers; *final_request.headers_mut() = headers;
// Continue with request processing // Continue with request processing
let response = next.run(final_request).await; next.run(final_request).await
};
// ===== OnResponse Phase ===== // ===== OnResponse Phase =====
let on_response_attach_point = let on_response_attach_point =
@@ -875,10 +871,13 @@ pub async fn wasm_middleware(
return Ok(response); return Ok(response);
} }
}; };
if modules_on_response.is_empty() {
return Ok(response);
}
// Extract response data once before processing modules // Extract response data once before processing modules
let mut status = response.status(); let mut status = response.status();
let mut headers = response.headers().clone(); let mut headers = response.headers().clone();
let max_body_size = wasm_manager.get_max_body_size();
let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await { let mut body_bytes = match axum::body::to_bytes(response.into_body(), max_body_size).await {
Ok(bytes) => bytes.to_vec(), Ok(bytes) => bytes.to_vec(),
Err(e) => { Err(e) => {