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

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;