feat(agentic router, 1/N): Add LoadBasedPolicy (#26480)

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() {

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`.

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<dyn 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 => {
             // 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}",
+        );
+    }
 }

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<Worker>]) -> Option<Arc<Worker>> {
+        workers
+            .iter()
+            .min_by_key(|w| w.active_load())
+            .map(Arc::clone)
+    }
+}
+
+impl Policy for LoadBasedPolicy {
+    fn select(&self, workers: &[Arc<Worker>], _ctx: &SelectionContext<'_>) -> Option<Arc<Worker>> {
+        Self::pick_min_load(workers)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
+
+    fn worker(id: &str) -> Arc<Worker> {
+        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
+        );
+    }
+}

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 {

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)