EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[router] Add Rerank Routing Logic in Regular Router (#10219)

Browse Source
This commit is contained in:
Frank Fang
2025-09-13 00:10:18 +08:00
committed by GitHub
parent efedbe6ca9
commit 4634fd5953
10 changed files with 545 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

89
sgl-router/src/protocols/spec.rs
View File

@@ -1891,7 +1891,7 @@ pub struct RerankResponse {
pub object: String,
/// Response ID
pub id: String,
pub id: Option<StringOrArray>,
/// Creation timestamp
pub created: i64,
@@ -1976,7 +1976,11 @@ impl RerankRequest {
}
impl RerankResponse {
pub fn new(results: Vec<RerankResult>, model: String, request_id: String) -> Self {
pub fn new(
results: Vec<RerankResult>,
model: String,
request_id: Option<StringOrArray>,
) -> Self {
RerankResponse {
results,
model,
@@ -2000,6 +2004,13 @@ impl RerankResponse {
pub fn apply_top_k(&mut self, k: usize) {
self.results.truncate(k);
}
/// Drop documents from results
pub fn drop_documents(&mut self) {
self.results.iter_mut().for_each(|result| {
result.document = None;
});
}
}
// ==================================================================
@@ -2268,12 +2279,15 @@ mod tests {
let response = RerankResponse::new(
results.clone(),
"test-model".to_string(),
"req-123".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
assert_eq!(response.results.len(), 2);
assert_eq!(response.model, "test-model");
assert_eq!(response.id, "req-123");
assert_eq!(
response.id,
Some(StringOrArray::String("req-123".to_string()))
);
assert_eq!(response.object, "rerank");
assert!(response.created > 0);
}
@@ -2287,8 +2301,11 @@ mod tests {
meta_info: None,
}];
let response =
RerankResponse::new(results, "test-model".to_string(), "req-123".to_string());
let response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
let serialized = serde_json::to_string(&response).unwrap();
let deserialized: RerankResponse = serde_json::from_str(&serialized).unwrap();
@@ -2322,8 +2339,11 @@ mod tests {
},
];
let mut response =
RerankResponse::new(results, "test-model".to_string(), "req-123".to_string());
let mut response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
response.sort_by_score();
@@ -2358,8 +2378,11 @@ mod tests {
},
];
let mut response =
RerankResponse::new(results, "test-model".to_string(), "req-123".to_string());
let mut response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
response.apply_top_k(2);
@@ -2377,14 +2400,36 @@ mod tests {
meta_info: None,
}];
let mut response =
RerankResponse::new(results, "test-model".to_string(), "req-123".to_string());
let mut response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
response.apply_top_k(5);
assert_eq!(response.results.len(), 1);
}
#[test]
fn test_rerank_response_drop_documents() {
let results = vec![RerankResult {
score: 0.8,
document: Some("doc1".to_string()),
index: 0,
meta_info: None,
}];
let mut response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
response.drop_documents();
assert_eq!(response.results[0].document, None);
}
// ==================================================================
// = RERANK RESULT TESTS =
// ==================================================================
@@ -2570,8 +2615,11 @@ mod tests {
meta_info: None,
}];
let mut response =
RerankResponse::new(results, "test-model".to_string(), "req-123".to_string());
let mut response = RerankResponse::new(
results,
"test-model".to_string(),
Some(StringOrArray::String("req-123".to_string())),
);
response.usage = Some(UsageInfo {
prompt_tokens: 100,
@@ -2645,18 +2693,7 @@ mod tests {
];
// Create response
let mut response = RerankResponse::new(
results,
request.model.clone(),
request
.rid
.as_ref()
.and_then(|r| match r {
StringOrArray::String(s) => Some(s.clone()),
StringOrArray::Array(arr) => arr.first().cloned(),
})
.unwrap_or_else(|| "unknown".to_string()),
);
let mut response = RerankResponse::new(results, request.model.clone(), request.rid.clone());
// Sort by score
response.sort_by_score();

6
sgl-router/src/routers/grpc/pd_router.rs
View File

@@ -301,7 +301,11 @@ impl RouterTrait for GrpcPDRouter {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
async fn route_rerank(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::RerankRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}

6
sgl-router/src/routers/grpc/router.rs
View File

@@ -234,7 +234,11 @@ impl RouterTrait for GrpcRouter {
(StatusCode::NOT_IMPLEMENTED).into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
async fn route_rerank(
&self,
_headers: Option<&HeaderMap>,
_body: &crate::protocols::spec::RerankRequest,
) -> Response {
(StatusCode::NOT_IMPLEMENTED).into_response()
}

6
sgl-router/src/routers/http/openai_router.rs
View File

@@ -2,7 +2,9 @@
use crate::config::CircuitBreakerConfig;
use crate::core::{CircuitBreaker, CircuitBreakerConfig as CoreCircuitBreakerConfig};
use crate::protocols::spec::{ChatCompletionRequest, CompletionRequest, GenerateRequest};
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, RerankRequest,
};
use async_trait::async_trait;
use axum::{
body::Body,
@@ -381,7 +383,7 @@ impl super::super::RouterTrait for OpenAIRouter {
.into_response()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: &RerankRequest) -> Response {
(
StatusCode::NOT_IMPLEMENTED,
"Rerank endpoint not implemented for OpenAI backend",

25
sgl-router/src/routers/http/pd_router.rs
View File

@@ -9,8 +9,8 @@ use crate::core::{
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::protocols::spec::{
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, ResponsesRequest,
StringOrArray, UserMessageContent,
ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, RerankRequest,
ResponsesRequest, StringOrArray, UserMessageContent,
};
use crate::routers::header_utils;
use crate::routers::{RouterTrait, WorkerManagement};
@@ -1946,8 +1946,25 @@ impl RouterTrait for PDRouter {
todo!()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
todo!()
async fn route_rerank(&self, headers: Option<&HeaderMap>, body: &RerankRequest) -> Response {
// Extract text for cache-aware routing
let req_text = if self.policies_need_request_text() {
Some(body.query.clone())
} else {
None
};
// Create context
let context = PDRequestContext {
route: "/v1/rerank",
batch_size: None,
is_stream: false,
return_logprob: false,
request_text: req_text,
};
// Execute with retry and bootstrap injection
self.execute_dual_dispatch(headers, body, context).await
}
async fn flush_cache(&self) -> Response {

45
sgl-router/src/routers/http/router.rs
View File

@@ -6,10 +6,12 @@ use crate::core::{
use crate::metrics::RouterMetrics;
use crate::policies::LoadBalancingPolicy;
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, GenerationRequest, ResponsesRequest,
ChatCompletionRequest, CompletionRequest, GenerateRequest, GenerationRequest, RerankRequest,
RerankResponse, RerankResult, ResponsesRequest,
};
use crate::routers::header_utils;
use crate::routers::{RouterTrait, WorkerManagement};
use axum::body::to_bytes;
use axum::{
body::Body,
extract::Request,
@@ -1124,6 +1126,25 @@ impl Router {
}
}
}
async fn build_rerank_response(
req: &RerankRequest,
response: Response,
) -> anyhow::Result<Response> {
let (_, response_body) = response.into_parts();
let body_bytes = to_bytes(response_body, usize::MAX).await?;
let rerank_results = serde_json::from_slice::<Vec<RerankResult>>(&body_bytes)?;
let mut rerank_response =
RerankResponse::new(rerank_results, req.model.clone(), req.rid.clone());
rerank_response.sort_by_score();
if let Some(top_k) = req.top_k {
rerank_response.apply_top_k(top_k);
}
if !req.return_documents {
rerank_response.drop_documents();
}
Ok(Json(rerank_response).into_response())
}
}
use async_trait::async_trait;
@@ -1223,8 +1244,26 @@ impl RouterTrait for Router {
todo!()
}
async fn route_rerank(&self, _headers: Option<&HeaderMap>, _body: Body) -> Response {
todo!()
async fn route_rerank(&self, headers: Option<&HeaderMap>, body: &RerankRequest) -> Response {
if let Err(e) = body.validate() {
return (StatusCode::BAD_REQUEST, e).into_response();
}
let response = self.route_typed_request(headers, body, "/v1/rerank").await;
if response.status().is_success() {
match Self::build_rerank_response(body, response).await {
Ok(rerank_response) => rerank_response,
Err(e) => {
error!("Failed to build rerank response: {}", e);
return (
StatusCode::INTERNAL_SERVER_ERROR,
"Failed to build rerank response".to_string(),
)
.into_response();
}
}
} else {
response
}
}
async fn flush_cache(&self) -> Response {

4
sgl-router/src/routers/mod.rs
View File

@@ -10,7 +10,7 @@ use axum::{
use std::fmt::Debug;
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, ResponsesRequest,
ChatCompletionRequest, CompletionRequest, GenerateRequest, RerankRequest, ResponsesRequest,
};
pub mod factory;
@@ -89,7 +89,7 @@ pub trait RouterTrait: Send + Sync + Debug + WorkerManagement {
async fn route_embeddings(&self, headers: Option<&HeaderMap>, body: Body) -> Response;
async fn route_rerank(&self, headers: Option<&HeaderMap>, body: Body) -> Response;
async fn route_rerank(&self, headers: Option<&HeaderMap>, body: &RerankRequest) -> Response;
/// Flush cache on all workers
async fn flush_cache(&self) -> Response;

24
sgl-router/src/server.rs
View File

@@ -3,7 +3,8 @@ use crate::logging::{self, LoggingConfig};
use crate::metrics::{self, PrometheusConfig};
use crate::middleware::TokenBucket;
use crate::protocols::spec::{
ChatCompletionRequest, CompletionRequest, GenerateRequest, ResponsesRequest,
ChatCompletionRequest, CompletionRequest, GenerateRequest, RerankRequest, ResponsesRequest,
V1RerankReqInput,
};
use crate::reasoning_parser::ParserFactory;
use crate::routers::{RouterFactory, RouterTrait};
@@ -152,6 +153,25 @@ async fn v1_completions(
state.router.route_completion(Some(&headers), &body).await
}
async fn rerank(
State(state): State<Arc<AppState>>,
headers: http::HeaderMap,
Json(body): Json<RerankRequest>,
) -> Response {
state.router.route_rerank(Some(&headers), &body).await
}
async fn v1_rerank(
State(state): State<Arc<AppState>>,
headers: http::HeaderMap,
Json(body): Json<V1RerankReqInput>,
) -> Response {
state
.router
.route_rerank(Some(&headers), &body.into())
.await
}
async fn v1_responses(
State(state): State<Arc<AppState>>,
headers: http::HeaderMap,
@@ -237,6 +257,8 @@ pub fn build_app(
.route("/generate", post(generate))
.route("/v1/chat/completions", post(v1_chat_completions))
.route("/v1/completions", post(v1_completions))
.route("/rerank", post(rerank))
.route("/v1/rerank", post(v1_rerank))
.route("/v1/responses", post(v1_responses))
.route_layer(axum::middleware::from_fn_with_state(
app_state.clone(),