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/grpc/common/stages/worker_selection.rs
+14 -12
View File
@@ -61,15 +61,18 @@ impl PipelineStage for WorkerSelectionStage {
// For Harmony, use selection_text produced during Harmony encoding
// Otherwise, use original_text from regular preparation
let text = if prep.harmony_mode {
prep.selection_text.as_deref()
} else {
prep.original_text.as_deref()
let info = SelectWorkerInfo {
request_text: if prep.harmony_mode {
prep.selection_text.as_deref()
} else {
prep.original_text.as_deref()
},
routing_id: prep.routing_id.as_deref(),
};
let workers = match self.mode {
WorkerSelectionMode::Regular => {
match self.select_single_worker(ctx.input.model_id.as_deref(), text) {
match self.select_single_worker(ctx.input.model_id.as_deref(), &info) {
Some(w) => WorkerSelection::Single { worker: w },
None => {
error!(
@@ -86,7 +89,7 @@ impl PipelineStage for WorkerSelectionStage {
}
}
WorkerSelectionMode::PrefillDecode => {
match self.select_pd_pair(ctx.input.model_id.as_deref(), text) {
match self.select_pd_pair(ctx.input.model_id.as_deref(), &info) {
Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode },
None => {
error!(
@@ -120,7 +123,7 @@ impl WorkerSelectionStage {
fn select_single_worker(
&self,
model_id: Option<&str>,
text: Option<&str>,
info: &SelectWorkerInfo,
) -> Option<Arc<dyn Worker>> {
// Get workers for the specified model, filtered by connection mode
let workers = self.worker_registry.get_workers_filtered(
@@ -146,7 +149,7 @@ impl WorkerSelectionStage {
};
// Select worker using the policy
let idx = policy.select_worker(&available, &SelectWorkerInfo { request_text: text })?;
let idx = policy.select_worker(&available, info)?;
let selected = available[idx].clone();
// Record worker selection metric
@@ -163,7 +166,7 @@ impl WorkerSelectionStage {
fn select_pd_pair(
&self,
model_id: Option<&str>,
text: Option<&str>,
info: &SelectWorkerInfo,
) -> Option<(Arc<dyn Worker>, Arc<dyn Worker>)> {
let all_workers = self.worker_registry.get_workers_filtered(
model_id,
@@ -203,9 +206,8 @@ impl WorkerSelectionStage {
None => self.policy_registry.get_default_policy(),
};
let info = SelectWorkerInfo { request_text: text };
let prefill_idx = policy.select_worker(&available_prefill, &info)?;
let decode_idx = policy.select_worker(&available_decode, &info)?;
let prefill_idx = policy.select_worker(&available_prefill, info)?;
let decode_idx = policy.select_worker(&available_decode, info)?;
let model = model_id.unwrap_or("default");
let policy_name = policy.name();
sgl-model-gateway/src/routers/grpc/context.rs
+3
View File
@@ -94,6 +94,9 @@ pub struct PreparationOutput {
/// Original text (for chat) or resolved text (for generate)
pub original_text: Option<String>,
/// Routing ID for manual routing policy
pub routing_id: Option<String>,
/// Tokenized input
pub token_ids: Vec<u32>,
sgl-model-gateway/src/routers/grpc/harmony/stages/preparation.rs
+2
View File
@@ -123,6 +123,7 @@ impl HarmonyPreparationStage {
// Step 4: Store results
ctx.state.preparation = Some(PreparationOutput {
original_text: None,
routing_id: request.routing_id.clone(),
token_ids: build_output.input_ids,
processed_messages: None,
tool_constraints,
@@ -203,6 +204,7 @@ impl HarmonyPreparationStage {
// Step 4: Store results with constraint
ctx.state.preparation = Some(PreparationOutput {
original_text: None,
routing_id: request.routing_id.clone(),
token_ids: build_output.input_ids,
processed_messages: None,
tool_constraints: constraint,
sgl-model-gateway/src/routers/grpc/regular/responses/tool_loop.rs
+2
View File
@@ -492,6 +492,7 @@ pub(super) async fn execute_tool_loop(
top_k: current_request.top_k,
min_p: current_request.min_p,
repetition_penalty: current_request.repetition_penalty,
routing_id: current_request.routing_id.clone(),
};
// Continue to next iteration
@@ -1070,6 +1071,7 @@ async fn execute_tool_loop_streaming_internal(
top_k: current_request.top_k,
min_p: current_request.min_p,
repetition_penalty: current_request.repetition_penalty,
routing_id: current_request.routing_id.clone(),
};
continue;
sgl-model-gateway/src/routers/grpc/regular/stages/chat/preparation.rs
+1
View File
@@ -96,6 +96,7 @@ impl ChatPreparationStage {
// Store results in context
ctx.state.preparation = Some(PreparationOutput {
original_text: Some(processed_messages.text.clone()),
routing_id: request.routing_id.clone(),
token_ids,
processed_messages: Some(processed_messages),
tool_constraints: tool_call_constraint,
sgl-model-gateway/src/routers/grpc/regular/stages/embedding/preparation.rs
+3 -1
View File
@@ -47,8 +47,9 @@ impl PipelineStage for EmbeddingPreparationStage {
));
};
// Extract text from request
// Extract text and routing_id from request before borrowing ctx mutably
let text = request.extract_text_for_routing();
let routing_id = request.routing_id.clone();
if text.is_empty() {
return Err(error::bad_request(
"empty_input",
@@ -77,6 +78,7 @@ impl PipelineStage for EmbeddingPreparationStage {
// Store preparation output
ctx.state.preparation = Some(PreparationOutput {
original_text: Some(text),
routing_id,
token_ids,
processed_messages: None,
tool_constraints: None,
sgl-model-gateway/src/routers/grpc/regular/stages/generate/preparation.rs
+1
View File
@@ -68,6 +68,7 @@ impl GeneratePreparationStage {
ctx.state.preparation = Some(PreparationOutput {
original_text,
routing_id: request.routing_id.clone(),
token_ids,
processed_messages: None,
tool_constraints: None,
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(