diff --git a/experimental/sgl-router/src/config/cli.rs b/experimental/sgl-router/src/config/cli.rs index 6957e383b..41361eee8 100644 --- a/experimental/sgl-router/src/config/cli.rs +++ b/experimental/sgl-router/src/config/cli.rs @@ -573,6 +573,19 @@ mod tests { ); } + /// `--policy load_based` parses to the load-based selector. + #[test] + fn parses_load_based_policy() { + let c = into_config_owned(with_model(&[ + "--worker-urls", + "http://10.0.0.1:30000", + "--policy", + "load_based", + ])) + .unwrap(); + assert_eq!(c.model.policy, PolicyKind::LoadBased); + } + /// clap rejects `--cb-threshold 0` because the field is `NonZeroU32`. #[test] fn rejects_zero_cb_threshold() { diff --git a/experimental/sgl-router/src/config/types.rs b/experimental/sgl-router/src/config/types.rs index 5cff09325..0af45461f 100644 --- a/experimental/sgl-router/src/config/types.rs +++ b/experimental/sgl-router/src/config/types.rs @@ -70,7 +70,7 @@ impl Default for ActiveLoadConfig { /// policy factory. /// /// Accepted on the CLI (`--policy`) as `round_robin` / `random` / -/// `power_of_two` / `cache_aware_zmq`. +/// `power_of_two` / `load_based` / `cache_aware_zmq`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)] pub enum PolicyKind { #[default] @@ -80,6 +80,9 @@ pub enum PolicyKind { Random, #[value(name = "power_of_two")] PowerOfTwo, + /// Selects the currently least-loaded worker. + #[value(name = "load_based")] + LoadBased, /// Cache-aware routing fed by SGLang's ZMQ KV-cache event publisher. /// Requires the model to have a tokenizer loaded; cache_aware tuning /// lives on `ModelConfig::cache_aware`. diff --git a/experimental/sgl-router/src/policies/factory.rs b/experimental/sgl-router/src/policies/factory.rs index 6bd2bcb00..271614313 100644 --- a/experimental/sgl-router/src/policies/factory.rs +++ b/experimental/sgl-router/src/policies/factory.rs @@ -6,6 +6,7 @@ use crate::discovery::ModelId; use crate::policies::{ cache_aware_zmq::CacheAwareZmqPolicy, kv_events::{BlockSizeOracle, HashTree}, + load_based::LoadBasedPolicy, power_of_two::PowerOfTwoChoicesPolicy, random::RandomPolicy, round_robin::RoundRobinPolicy, @@ -32,6 +33,7 @@ pub fn build_policy( PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), PolicyKind::Random => Arc::new(RandomPolicy::new()), PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), + PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), PolicyKind::CacheAwareZmq => { let cache_cfg = model.cache_aware.unwrap_or_default(); Arc::new(CacheAwareZmqPolicy::new( @@ -54,6 +56,7 @@ pub fn build_policy_kind_only(kind: PolicyKind) -> Arc { PolicyKind::RoundRobin => Arc::new(RoundRobinPolicy::new()), PolicyKind::Random => Arc::new(RandomPolicy::new()), PolicyKind::PowerOfTwo => Arc::new(PowerOfTwoChoicesPolicy::new()), + PolicyKind::LoadBased => Arc::new(LoadBasedPolicy::new()), PolicyKind::CacheAwareZmq => { // Provide an empty tree + empty tokenizer registry + fresh // oracle so the test policy is constructible. Production @@ -144,6 +147,7 @@ mod tests { let _ = build_policy_kind_only(PolicyKind::RoundRobin); let _ = build_policy_kind_only(PolicyKind::Random); let _ = build_policy_kind_only(PolicyKind::PowerOfTwo); + let _ = build_policy_kind_only(PolicyKind::LoadBased); let _ = build_policy_kind_only(PolicyKind::CacheAwareZmq); } @@ -173,4 +177,18 @@ mod tests { "expected CacheAwareZmqPolicy debug repr, got: {dbg}", ); } + + #[test] + fn load_based_builds_via_factory() { + let cfg = cfg_with_model("modelA", PolicyKind::LoadBased); + let tree = Arc::new(HashTree::new()); + let tokenizers = Arc::new(TokenizerRegistry::default()); + let reg = build_registry(&cfg, tree, tokenizers, BlockSizeOracle::new()).unwrap(); + let p = reg.get(&ModelId("modelA".into())).unwrap(); + let dbg = format!("{p:?}"); + assert!( + dbg.contains("LoadBasedPolicy"), + "expected LoadBasedPolicy debug repr, got: {dbg}", + ); + } } diff --git a/experimental/sgl-router/src/policies/load_based.rs b/experimental/sgl-router/src/policies/load_based.rs new file mode 100644 index 000000000..6767ad12f --- /dev/null +++ b/experimental/sgl-router/src/policies/load_based.rs @@ -0,0 +1,70 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use crate::policies::{Policy, SelectionContext}; +use crate::workers::Worker; +use std::sync::Arc; + +/// Deterministic load-based policy. +/// +/// Chooses the candidate with the lowest current `Worker::active_load`. +/// Ties follow the candidate slice order, which is registry-dependent. +#[derive(Debug, Default)] +pub struct LoadBasedPolicy; + +impl LoadBasedPolicy { + pub fn new() -> Self { + Self + } + + pub fn pick_min_load(workers: &[Arc]) -> Option> { + workers + .iter() + .min_by_key(|w| w.active_load()) + .map(Arc::clone) + } +} + +impl Policy for LoadBasedPolicy { + fn select(&self, workers: &[Arc], _ctx: &SelectionContext<'_>) -> Option> { + Self::pick_min_load(workers) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec}; + + fn worker(id: &str) -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}:30000"), + mode: WorkerMode::Plain, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: None, + })) + } + + #[test] + fn empty_returns_none() { + let policy = LoadBasedPolicy::new(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + assert!(policy.select(&[], &ctx).is_none()); + } + + #[test] + fn picks_lowest_active_load() { + let policy = LoadBasedPolicy::new(); + let model = ModelId("tiny".into()); + let ctx = SelectionContext::new(&model, None); + let w0 = worker("w0"); + let w1 = worker("w1"); + let _g0 = w0.load_guard(); + assert_eq!( + policy.select(&[w0, Arc::clone(&w1)], &ctx).unwrap().id, + w1.id + ); + } +} diff --git a/experimental/sgl-router/src/policies/mod.rs b/experimental/sgl-router/src/policies/mod.rs index f30652c81..c5fa0243a 100644 --- a/experimental/sgl-router/src/policies/mod.rs +++ b/experimental/sgl-router/src/policies/mod.rs @@ -5,6 +5,7 @@ pub mod active_load; pub mod cache_aware_zmq; pub mod factory; pub mod kv_events; +pub mod load_based; pub mod power_of_two; pub mod random; pub mod registry; @@ -26,6 +27,7 @@ use std::sync::Arc; pub struct SelectionContext<'a> { model: &'a ModelId, request_body: Option<&'a [u8]>, + routing_key: Option<&'a str>, } impl<'a> SelectionContext<'a> { @@ -33,6 +35,19 @@ impl<'a> SelectionContext<'a> { Self { model, request_body, + routing_key: None, + } + } + + pub fn with_routing_key( + model: &'a ModelId, + request_body: Option<&'a [u8]>, + routing_key: Option<&'a str>, + ) -> Self { + Self { + model, + request_body, + routing_key, } } @@ -43,6 +58,10 @@ impl<'a> SelectionContext<'a> { pub fn request_body(&self) -> Option<&[u8]> { self.request_body } + + pub fn routing_key(&self) -> Option<&str> { + self.routing_key + } } pub trait Policy: Send + Sync + std::fmt::Debug { diff --git a/experimental/sgl-router/tests/e2e/chat_completions/test_load_based_policy.py b/experimental/sgl-router/tests/e2e/chat_completions/test_load_based_policy.py new file mode 100644 index 000000000..4683e572d --- /dev/null +++ b/experimental/sgl-router/tests/e2e/chat_completions/test_load_based_policy.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Real-GPU acceptance coverage for ``policy = "load_based"``.""" + +from __future__ import annotations + +import re +import threading +import time + +import httpx +import pytest +from infra.gateway import Gateway +from infra.model_pool import spawn_worker +from infra.model_specs import get_model_spec + +_ACTIVE_RE = re.compile( + r'^sgl_router_active_load\{worker_url="([^"]+)",kind="prefill_tokens"\}\s+(-?\d+)' +) +_REQ_TOTAL_RE = re.compile( + r"^sgl_router_requests_total\{([^}]*)\}\s+(\d+(?:\.\d+)?)\s*$" +) +_LABEL_RE = re.compile(r'(\w+)="([^"]*)"') + + +def _chat( + router_url: str, model_id: str, prompt: str, max_tokens: int +) -> httpx.Response: + return httpx.post( + f"{router_url}/v1/chat/completions", + json={ + "model": model_id, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "ignore_eos": True, + "stream": False, + }, + timeout=240.0, + ) + + +def _active_prefill_loads(router_url: str) -> dict[str, int]: + resp = httpx.get(f"{router_url}/metrics", timeout=5.0) + resp.raise_for_status() + loads: dict[str, int] = {} + for line in resp.text.splitlines(): + match = _ACTIVE_RE.match(line) + if match: + loads[match.group(1)] = int(match.group(2)) + return loads + + +def _success_counts(router_url: str) -> dict[str, int]: + resp = httpx.get(f"{router_url}/metrics", timeout=5.0) + resp.raise_for_status() + counts: dict[str, int] = {} + for line in resp.text.splitlines(): + match = _REQ_TOTAL_RE.match(line) + if not match: + continue + labels = dict(_LABEL_RE.findall(match.group(1))) + if labels.get("outcome") != "success": + continue + worker_url = labels.get("worker_url") + if worker_url: + counts[worker_url] = counts.get(worker_url, 0) + int(float(match.group(2))) + return counts + + +def _wait_for_busy_worker(router_url: str, worker_urls: list[str]) -> str: + deadline = time.time() + 60.0 + while time.time() < deadline: + loads = _active_prefill_loads(router_url) + busy = [url for url in worker_urls if loads.get(url, 0) > 0] + if busy: + return busy[0] + time.sleep(0.2) + raise AssertionError( + f"no worker became busy; last loads={_active_prefill_loads(router_url)}" + ) + + +def _single_success_delta(before: dict[str, int], after: dict[str, int]) -> str: + deltas = { + url: after.get(url, 0) - before.get(url, 0) for url in set(before) | set(after) + } + winners = [url for url, delta in deltas.items() if delta == 1] + assert len(winners) == 1, f"expected exactly one success delta, got {deltas}" + return winners[0] + + +@pytest.mark.real_gpu +@pytest.mark.slow +def test_load_based_routes_to_the_cooler_worker( + router_binary, # noqa: ARG001 - fixture forces release-binary presence + gpu_allocator, +) -> None: + spec = get_model_spec("qwen3-0.6b") + gpus = gpu_allocator.acquire(2) + try: + with ( + spawn_worker("qwen3-0.6b", gpu_ids=[gpus[0]]) as worker_a, + spawn_worker("qwen3-0.6b", gpu_ids=[gpus[1]]) as worker_b, + Gateway() as router, + ): + worker_urls = [worker_a.url, worker_b.url] + router.start_regular( + model_id=spec["model"], + tokenizer_path=spec["model"], + worker_urls=worker_urls, + policy="load_based", + timeout=120.0, + ) + + errors: list[BaseException] = [] + + def long_request() -> None: + try: + resp = _chat( + router.base_url, + spec["model"], + "Write a long numbered list of routing test facts.", + 1024, + ) + assert resp.status_code == 200, resp.text + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + thread = threading.Thread(target=long_request, daemon=True) + thread.start() + busy_worker = _wait_for_busy_worker(router.base_url, worker_urls) + + before = _success_counts(router.base_url) + resp = _chat( + router.base_url, + spec["model"], + "Answer with one short sentence.", + 8, + ) + assert resp.status_code == 200, resp.text + routed_worker = _single_success_delta( + before, _success_counts(router.base_url) + ) + assert routed_worker != busy_worker + + thread.join(timeout=240.0) + assert not thread.is_alive(), "long request did not finish" + assert not errors, errors + finally: + gpu_allocator.release(gpus)