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

Add manual routing policy for router (#15586)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-25 11:27:36 +08:00
committed by GitHub
parent 1ba897f330
commit 45adad37d0
35 changed files with 1128 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/routers/http/pd_router.rs
+25 -19
View File
@@ -26,7 +26,7 @@ use crate::{
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::{LoadBalancingPolicy, PolicyRegistry, SelectWorkerInfo},
policies::{LoadBalancingPolicy, PolicyRegistry},
protocols::{
chat::{ChatCompletionRequest, ChatMessage, MessageContent},
common::{InputIds, StringOrArray},
@@ -58,6 +58,7 @@ struct PDRequestContext<'a> {
is_stream: bool,
return_logprob: bool,
request_text: Option<String>,
routing_id: Option<String>,
model_id: Option<&'a str>,
}
@@ -303,7 +304,11 @@ impl PDRouter {
let context = context.clone();
async move {
let (prefill, decode) = match self
.select_pd_pair(context.request_text.as_deref(), context.model_id)
.select_pd_pair(
context.request_text.as_deref(),
context.routing_id.as_deref(),
context.model_id,
)
.await
{
Ok(pair) => pair,
@@ -691,6 +696,7 @@ impl PDRouter {
async fn select_pd_pair(
&self,
request_text: Option<&str>,
routing_id: Option<&str>,
model_id: Option<&str>,
) -> Result<(Arc<dyn Worker>, Arc<dyn Worker>), String> {
let effective_model_id = if !self.enable_igw { None } else { model_id };
@@ -725,19 +731,16 @@ impl PDRouter {
let prefill_policy = self.policy_registry.get_prefill_policy();
let decode_policy = self.policy_registry.get_decode_policy();
let prefill = Self::pick_worker_by_policy_arc(
&prefill_workers,
&*prefill_policy,
let info = crate::policies::SelectWorkerInfo {
request_text,
"prefill",
)?;
routing_id,
};
let decode = Self::pick_worker_by_policy_arc(
&decode_workers,
&*decode_policy,
request_text,
"decode",
)?;
let prefill =
Self::pick_worker_by_policy_arc(&prefill_workers, &*prefill_policy, &info, "prefill")?;
let decode =
Self::pick_worker_by_policy_arc(&decode_workers, &*decode_policy, &info, "decode")?;
// Record worker selection metrics (Layer 3)
let model = model_id.unwrap_or("default");
@@ -760,7 +763,7 @@ impl PDRouter {
fn pick_worker_by_policy_arc(
workers: &[Arc<dyn Worker>],
policy: &dyn LoadBalancingPolicy,
request_text: Option<&str>,
info: &crate::policies::SelectWorkerInfo,
worker_type: &str,
) -> Result<Arc<dyn Worker>, String> {
if workers.is_empty() {
@@ -784,7 +787,7 @@ impl PDRouter {
}
let selected_idx = policy
.select_worker(&available_workers, &SelectWorkerInfo { request_text })
.select_worker(&available_workers, info)
.ok_or_else(|| {
format!(
"Policy {} failed to select a {} worker",
@@ -1120,7 +1123,7 @@ impl RouterTrait for PDRouter {
// Note: This endpoint actually causes the model to generate tokens, so we only test one pair
// Select a random worker pair using the policy
let (prefill, decode) = match self.select_pd_pair(None, None).await {
let (prefill, decode) = match self.select_pd_pair(None, None, None).await {
Ok(pair) => pair,
Err(e) => {
return error::service_unavailable(
@@ -1242,6 +1245,7 @@ impl RouterTrait for PDRouter {
is_stream,
return_logprob,
request_text,
routing_id: body.routing_id.clone(),
model_id,
};
@@ -1283,6 +1287,7 @@ impl RouterTrait for PDRouter {
is_stream,
return_logprob,
request_text,
routing_id: body.routing_id.clone(),
model_id,
};
@@ -1316,6 +1321,7 @@ impl RouterTrait for PDRouter {
is_stream,
return_logprob,
request_text,
routing_id: body.routing_id.clone(),
model_id,
};
@@ -1328,7 +1334,6 @@ impl RouterTrait for PDRouter {
body: &RerankRequest,
model_id: Option<&str>,
) -> Response {
// Extract text for cache-aware routing
let req_text = if self.policies_need_request_text() {
Some(body.query.clone())
} else {
@@ -1341,6 +1346,7 @@ impl RouterTrait for PDRouter {
is_stream: false,
return_logprob: false,
request_text: req_text,
routing_id: body.routing_id.clone(),
model_id,
};
@@ -1405,7 +1411,7 @@ mod tests {
router.worker_registry.register(Arc::from(healthy_worker));
router.worker_registry.register(Arc::from(decode_worker));
let result = router.select_pd_pair(None, None).await;
let result = router.select_pd_pair(None, None, None).await;
assert!(result.is_ok());
let (prefill, _decode) = result.unwrap();
@@ -1418,7 +1424,7 @@ mod tests {
async fn test_empty_worker_lists() {
let router = create_test_pd_router();
let result = router.select_pd_pair(None, None).await;
let result = router.select_pd_pair(None, None, None).await;
assert!(result.is_err());
assert!(result.unwrap_err().contains("No prefill workers available"));
sgl-model-gateway/src/routers/http/router.rs
+13 -18
View File
@@ -27,7 +27,7 @@ use crate::{
metrics::{bool_to_static_str, metrics_labels, Metrics},
otel_trace::inject_trace_context_http,
},
policies::{PolicyRegistry, SelectWorkerInfo},
policies::PolicyRegistry,
protocols::{
chat::ChatCompletionRequest,
classify::ClassifyRequest,
@@ -39,7 +39,7 @@ use crate::{
responses::{ResponsesGetParams, ResponsesRequest},
},
routers::{
error::{self, extract_error_code_from_response},
error,
grpc::utils::{error_type_from_status, route_to_endpoint},
header_utils, RouterTrait,
},
@@ -140,7 +140,7 @@ impl Router {
fn select_worker_for_model(
&self,
model_id: Option<&str>,
text: Option<&str>,
info: &crate::policies::SelectWorkerInfo,
) -> Option<Arc<dyn Worker>> {
let effective_model_id = if !self.enable_igw { None } else { model_id };
@@ -168,7 +168,7 @@ impl Router {
None => self.policy_registry.get_default_policy(),
};
let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?;
let idx = policy.select_worker(&available, info)?;
// Record worker selection metric (Layer 3)
Metrics::record_worker_selection(
@@ -191,6 +191,11 @@ impl Router {
let start = Instant::now();
let is_stream = typed_req.is_stream();
let text = typed_req.extract_text_for_routing();
let routing_id = typed_req.get_routing_id().map(|s| s.to_string());
let info = crate::policies::SelectWorkerInfo {
request_text: Some(&text),
routing_id: routing_id.as_deref(),
};
let model = model_id.unwrap_or("default");
let endpoint = route_to_endpoint(route);
@@ -208,18 +213,8 @@ impl Router {
&self.retry_config,
// operation per attempt
|_: u32| async {
let res = self
.route_typed_request_once(headers, typed_req, route, model_id, is_stream, &text)
.await;
// Need to be outside `route_typed_request_once` because that function has multiple return paths
Metrics::record_router_upstream_response(
metrics_labels::ROUTER_HTTP,
res.status().as_u16(),
extract_error_code_from_response(&res),
);
res
self.route_typed_request_once(headers, typed_req, route, model_id, is_stream, &info)
.await
},
// should_retry predicate
|res, _attempt| is_retryable_status(res.status()),
@@ -267,9 +262,9 @@ impl Router {
route: &'static str,
model_id: Option<&str>,
is_stream: bool,
text: &str,
info: &crate::policies::SelectWorkerInfo<'_>,
) -> Response {
let worker = match self.select_worker_for_model(model_id, Some(text)) {
let worker = match self.select_worker_for_model(model_id, info) {
Some(w) => w,
None => {
return error::service_unavailable(