diff --git a/experimental/sgl-router/POLICY_DESIGN.md b/experimental/sgl-router/POLICY_DESIGN.md new file mode 100644 index 000000000..53cec7bc5 --- /dev/null +++ b/experimental/sgl-router/POLICY_DESIGN.md @@ -0,0 +1,575 @@ +# Engine selection: target design + +This document describes the target architecture for engine selection. It defines +responsibilities and behavior; the interface and configuration examples are +sketches, not a specification of the current API or CLI. Implementation status is +listed at the end. + +## Principles + +1. **Order compatible buckets by token length first.** `BucketResolver` returns + all matching buckets, smallest capacity first, without inspecting workers or policies. +2. **The bucket owns plain versus PD engine selection.** `Bucket::pick_engines` + calls its one plain group or both prefill and decode groups, returning a + complete selection. Both PD engines come from that same bucket. +3. **An engine group owns membership and its policy.** `EngineGroup::pick` + filters live workers by model, health, stage, and membership, invokes the + policy, and validates that its result belongs to the exact candidate set. +4. **A policy owns selection, fallback, admission, and its state dependencies.** + Construction injects the load, KV, or affinity handles it needs. Request-time + arguments contain facts and candidates. A policy cannot choose another bucket + or cross a PD role boundary. +5. **The handler owns bucket fallback and dispatch.** It calls `pick_engines` + on each bucket and dispatches only after a complete selection succeeds. + Missing candidates or admission rejection advance to the next bucket, where + all required engines are selected again. Invalid signals or policy results stop routing. + +An engine is represented by `Worker`. Registry role labels remain authoritative +when filtering candidates. Groups reference worker IDs rather than owning live +workers. Group policy instances are reused across requests. + +### Data ownership + +| Type | Owns | +| --- | --- | +| `BucketResolver` | A model's bucket collection, length filtering, and ordering | +| `Bucket` | ID, length constraints, rank, groups, and complete plain/PD engine selection | +| `EngineGroup` | Engine membership, attached policy, and engine selection | +| `WorkerRegistry` | Live workers, model membership, health, and role | +| Request handler | Request preparation, ordered bucket attempts, HTTP errors, and dispatch | + +`worker_ids: None` means every healthy engine serving the requested model and +role. An explicit empty set means no engines. `EngineGroup::new(policy)` creates +a catch-all membership group. `BucketGroups::Pd` requires both groups, making +partial or mixed plain/PD bucket configurations unrepresentable. + +## 1. Code organization + +```text +src/ + buckets.rs Bucket resolution and engine-group selection + policies/ + mod.rs Policy contract and construction + admission.rs Per-engine acceptance checks + cache_aware.rs Cache selection and local/remote adapter + session_aware.rs Session selection and fallback + sticky.rs Routing-key selection and fallback + least_load.rs Least-load selection + power_of_two.rs Pair sampling and stage-aware comparison + random.rs Random selection + round_robin.rs Rotation + state/ + mod.rs Shared exports + kv_events/ Local index, subscriptions, hashing, wire format + load_monitor/ Engine reports and router-local request accounting + affinity_store.rs Assignments, expiry, and atomic updates + server/ + app_context.rs Shared service lifecycle and policy wiring + routes/chat.rs Request preparation, PD coordination, dispatch +``` + +During migration, `buckets_reorg.rs` and `policies_reorg/` implement this design +beside the live `policies/` path. The layout above is the target after switchover. +Shared state already lives directly under `src/state/`. + +Dependencies flow from engine groups to policies, and from policies and admission +to shared state. State does not depend on bucket ordering or concrete policy +strategies. Small shared helpers are sufficient; no generic score-composition, +tier executor, or separate selection framework is required. + +## 2. Responsibilities and request flow + +`BucketResolver::resolve(input_tokens, expected_peak_tokens)` returns an ordered +list of compatible bucket references (possibly empty), or an invalid-signal error. +It does not receive a stage or a load view, resolve live engines, or invoke policies. +The handler iterates this list until a bucket supplies the complete engine selection. + +`Bucket::pick_engines(workers, request)` accepts a `BucketRequest` +of prepared routing facts, invokes the required groups, and returns `BucketPick` +(one plain pick or a complete P/D pair). Failures retain their stage. This API +has no HTTP headers, `AppContext`, or forwarding dependency. + +`EngineGroup::pick(workers, request)` resolves healthy workers for the request's +model and stage, intersects them with its membership, sorts them by stable ID, +and invokes its attached policy. It rejects foreign results, including a newly +allocated worker with the same ID as a candidate. + +```text +chat_completions (reorg configured): prepare tokens and expected peak + | + v +BucketResolver::resolve: ordered length-compatible buckets + | + v +For each bucket: bucket.pick_engines(...) + | + +-- BucketGroups::Plain + | plain.pick() -> one plain engine + | + +-- BucketGroups::Pd + prefill.pick() -> P engine + decode.pick() -> D engine from the same bucket + | + +-- empty group / admission rejection -> try next bucket (repeat all picks) + +-- invalid signal / configuration / foreign pick -> return error + | + v +Complete selection -> forward_chat_request: acquire guards, attach PD bootstrap, forward response +``` + +The handler extracts token facts and header keys once into `BucketRequest`. +The bucket creates a stage-specific `PickRequest` for each group call, supplying +its own ID and the role associated with that group. Input length, expected peak, +token IDs, and session/routing keys pass through. Policies obtain observations +from their own shared-state handles; buckets and handlers do not provide load, +KV, or affinity services on each call. PD uses separate group policies, but never +independently resolves a decode bucket. Decode selection +failure discards that tentative prefill choice and advances to the next bucket +on missing candidates or admission rejection. No forwarding guards are acquired +and no prefill request is sent until both picks in one bucket succeed. +PD compatibility constraints beyond model and role remain follow-up work. + +There is one endpoint: `POST /v1/chat/completions`. `AppContext::chat_routing` +chooses its implementation: + +- `ChatRouting::Legacy` (default) uses the existing policies and bucket selector. +- `ChatRouting::Reorg(HashMap)` uses the new bucket and + policy interfaces, with explicit model-specific resolvers. + +Callers set this field before building the router. A missing model in the reorg +map returns 404, without falling back to legacy routing. This PR adds the +programmatic configuration switch; CLI/configuration factory construction and +the remaining production policies remain follow-ups. Power-of-two is implemented +for explicit attachments; the default serving path remains legacy. + +Both implementations reuse request preparation (including sampling validation +and tokenization), forwarding, streaming, middleware, and the 32 MiB body limit. +The reorg implementation requests tokenization for length matching, retaining +the existing body-size estimate when tokenization is unavailable. + +## 3. Bucket resolution + +1. Validate that a known expected peak is at least the input length. +2. Keep buckets whose inclusive input-token range contains the input length. +3. Check the bucket context capacity against input plus requested output when + known, or against input length when the output budget is unknown. +4. Sort by ascending input capacity (the lesser of the input upper bound and + context capacity). Unbounded capacities sort last. Break ties by ascending + bucket rank, then ID, and return the entire ordered list. +5. The handler calls each bucket's `pick_engines` until one supplies its complete + selection. A failed PD attempt never contributes an engine to a later pair. + +The handler checks addition overflow when computing the expected peak. +An empty bucket list becomes a 400 `NoMatchingBucket` response. After exhausting +the list, accumulated admission rejection details produce a selection failure +(503); if there were no admission rejections, the last unavailable stage produces +a stage-specific 503. Invalid policy signals/configuration or out-of-candidate +picks stop the pass immediately with an internal error. Successful engine +selection ends the pass; forwarding errors do not restart bucket iteration. + +Token ranges and rank belong to the bucket, not its engine groups. For a PD +bucket, both groups share this one request-length decision. Policy fallback on +a cache/affinity miss stays within that group's candidates. There is no second +pass with relaxed admission and no post-policy substitution. + +SLO ordering, cache lookup, and session/sticky policies are follow-ups. Their +integration must preserve bucket-first selection and the same-bucket PD rule. +Cross-bucket affinity probing is not part of this interface. Session/routing +keys still pass through `PickRequest` for policies operating inside the selected +group; unsupported legacy modes need explicit migration decisions before the +standard serving path switches. + +## 4. Policy and admission contracts + +### Policy + +All policies implement one asynchronous, object-safe `Policy::pick` interface. +A boxed future supports the remote prefix indexer; local policies can return an +immediately ready result. + +```rust +pub trait Policy: Send + Sync + std::fmt::Debug { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result>; +} +``` + +Policies own their required state handles and read observations locally. Nested +fallback uses `pick_fallback(engines, request)`, which calls the fallback's `pick`; +the fallback reads its own state. There is no shared observation context or cache. +Buckets and HTTP handlers supply only candidates and request facts. + +`Pick` identifies one engine and a selection reason for metrics and tracing. +`PickRequest` carries model, stage, selected bucket ID, input and optional +expected peak counts, optional token IDs, and session/routing keys. It contains +no HTTP body, bucket resolver, state handles, snapshots, or backend configuration. + +| Outcome | Meaning | +| --- | --- | +| `Pick` | The chosen engine belongs to the supplied set and passed admission | +| `NoMatchingBucket` | No bucket supports the requested length | +| `NoCandidates` | No eligible member or policy selection miss | +| `NoAdmissibleEngine` | A policy exhausted its candidates; includes rejection reasons | +| `AdmissionRejected` | The chosen engine failed after-selection admission | +| `InvalidSignal` / `InvalidConfiguration` | Invalid policy input or configuration | +| `OutsideCandidates` | Policy returned an engine outside its exact candidate set | + +### Admission + +Admission evaluates acceptance. It does not rank engines, choose replacements, +change buckets, or mutate affinity. + +| Check | Acceptance rule | +| --- | --- | +| `AllowAll` | Add no acceptance constraint | +| `CapacityAdmission` | Projected running requests and KV tokens fit reported capacity | +| `PendingPrefillAdmission` | Waiting uncached tokens plus incoming uncached work fit the budget | +| `InFlightLimitAdmission` | Router-local in-flight requests are below the limit | +| `QueueLimitAdmission` | Engine-reported waiting requests are below the limit | +| `AllOfAdmission` | Every attached check allows the request | + +`EngineAdmission::check(engine, request, load)` checks one engine and returns +`Allow`, `Reject(reason)`, or an error for invalid inputs. Policies attach the +checker directly as `Arc`. There is no placement setting, +filtering wrapper, or before/after API; each policy decides where checking +belongs in its selection algorithm. The `load` argument is an +`Option<&EngineReportedWorkerLoad>` retained by the policy for this engine, including +request counts, token usage, capacity, and the report timestamp. `None` means +no usable observation, never zero load; each check defines its missing-data +behavior. Other required state handles belong to the checker. + +Power-of-two first selects an engine, then calls admission exactly once on that +engine. A rejection returns `AdmissionRejected` to the bucket loop; it does not +resample, choose the other sampled engine, or run a policy fallback. No candidates +returns `NoCandidates` without invoking admission. A single candidate is selected +directly; otherwise two distinct candidates are sampled uniformly, and the one +with lower stage pressure wins. A complete tie keeps the first sampled engine. + +Power-of-two reuses the existing pure pressure-comparison functions. Plain and +prefill stages compare estimated prefill queue time when both reports provide it, +then waiting uncached tokens, waiting requests, and running requests. Decode +compares waiting requests, running requests, KV usage fraction, then used KV tokens. +Reported-pressure ties use router-local active requests. If either sampled engine +lacks a fresh, complete native report with valid capacity, both are compared by +router-local active requests instead. Basic reports from older publishers are +still passed to admission when fresh, but do not supply native pressure metrics. + +Prepare the signals needed by admission before checking. A pending-prefill check +uses per-engine uncached work when a prefix is known, and full input otherwise. +Decode capacity uses the expected peak sequence length when available, including +on a cache hit. Power-of-two retains the selected engine's load record from +selection and passes it to admission without another snapshot. A single candidate +still has its load read for admission, even though selection needs no comparison. +Neither the bucket nor HTTP handler supplies observations. Concrete load-aware +acceptance rules and additional cache-specific admission signals remain follow-up +work. Synchronous checks do not fetch telemetry over the network themselves. + +`AllowAll` is the default for new explicit policy attachments. It leaves health, +role, membership, and policy preferences in force. Migrated configurations must +retain their existing capacity and configured budget checks; see compatibility +below. Each check defines its missing-data behavior. Unknown load is not zero; +the existing capacity and pending-prefill checks allow requests without a fresh, +complete native report. + +The cache policy's `worker_queue_limit` is a **soft preference**, not +`QueueLimitAdmission`. Saturation handling can reconsider a queued engine, but +cannot bypass attached hard admission. + +Admission checks observe capacity; they do not reserve it. Concurrent requests +may pass against the same observation. Strict reservations would require a +separate mechanism. + +## 5. Concrete policies + +| Policy | Selection behavior | +| --- | --- | +| `RoundRobinPolicy` | Rotate over candidates using a cursor owned by this policy instance | +| `RandomPolicy` | Choose uniformly from candidates | +| `PowerOfTwoPolicy` | Sample two distinct candidates when possible and choose the lower-pressure engine using the stage's load comparison | +| `LeastLoadPolicy` (`load_based`) | Choose the least loaded engine; preserve tie-breaking, telemetry fallback, and recent-dispatch correction | +| `SessionAwarePolicy` | Reuse an admitted session binding; use power-of-two for new or keyless sessions | +| `StickyPolicy` | Reuse an admitted routing-key binding; use the configured fallback for new or missing keys | +| `CacheAwarePolicy` | Prefer a usable prefix under cache and pressure rules; use a load-based fallback on a miss | + +Session and sticky policies do not create assignments for missing keys. A +binding outside the candidates cannot win. A missing binding may invoke policy +fallback within the group; hard admission rejection remains an error. + +Sticky fallback supports `round_robin`, `random`, `power_of_two`, and `load_based`, +with round-robin as the default. Nested fallbacks use `AllowAll`; the owning +policy explicitly checks the engine returned by its fallback. + +### Cache-aware behavior + +The architecture preserves the cache algorithm within the supplied candidate +set. Its responsibilities are: + +1. Look up prefix ownership through the local index or remote indexer. +2. Apply minimum matched-token and optional ratio thresholds. +3. Bound candidates using prefix/pressure ordering and the configured minimum, + ratio, and maximum worker counts. +4. Apply the soft queue gate and saturation rules, and call admission explicitly + as required by the cache policy's candidate-selection algorithm. +5. Choose among usable prefix holders using uncached work, the switch margin, + and the pressure guard. +6. On a miss, run the load fallback, preferring engines admitted by the soft + queue gate when available. + +Candidate limits and saturation observations use only the selected bucket's +role-group candidates. +Saturation pinning must still pass hard admission. + +The target load fallback supports power-of-k sampling through +`--min-load-choices`, default 2. When k covers the group, choose the exact minimum. +Preserve queue-tier preference and avoid sorting with a pairwise pressure +comparator that does not define a total ordering. + +Memoize the prefix lookup once per request, including remote I/O. Each policy +restricts those matches to its own candidates. A memoized lookup does not imply +that another bucket or stage's cache selection and admission have already run. +Any optimization that skips those steps must establish that the previous result +applies to the current group. + +## 6. Shared state and construction + +Application wiring starts shared services once. Policy construction validates +configuration and passes the required handles to each policy and admission +implementation. For example, `PowerOfTwoPolicy::new(Arc)` +retains the application's shared load table. KV-aware and affinity-aware policies +receive their corresponding shared handles when implemented. Policies with no +state dependency require none. Policy instances do not create duplicate +subscriptions, polling loops, indexes, or remote-client concurrency limits. + +Requirements come from all configured policies, their admission checks, and +nested fallbacks. This includes tokenization, affinity-header extraction, load +observations, and dispatch timestamps. A role-group override must receive its +required settings even when the model's default uses a different policy. + +### Load state + +`state/load_monitor/` owns engine reports and existing router-local request +accounting. Power-of-two owns an `Arc` and captures a snapshot +locally for each nonempty selection attempt. Its `pick` method selects the engine, +then passes that engine's borrowed load record directly to admission. +No snapshot or observation is added to `Pick`, `PickRequest`, +or the bucket interface, and no shared observation context is threaded through +policies or fallbacks. + +The existing snapshot reader preserves rank aggregation, freshness, and capacity +fields. Missing, stale, or rank-incomplete reports yield `None`, not zero load. +A new pick reads current state. Admission reuses the selected observation even if +reports change after selection; it neither recaptures nor reserves capacity. +Fallback policies read their own state and do not share snapshots with callers. + +Snapshot capture still scans the full table; an engine-scoped reader can be added +if profiling justifies it. Power-of-two reuses the legacy prefill/decode pressure +comparisons, including router-local fallback. Concrete load-aware admission remains +in #40271. Further shared load interpretation and correction for dispatches since +the report remain follow-ups; these must preserve source, freshness, and available +measurements without adding another +independent in-flight counter. Load and cache observations are not an atomic global +snapshot. Preserve request-guard cleanup. + +### Cache state + +`state/kv_events/` owns local subscriptions, event application, hashing, and the +radix-tree index. The tree stores prefix ownership and storage tiers, not KV +tensors. Eviction, invalidation, and worker removal update this shared state. + +A small `CacheSource` adapter beside the cache policy normalizes local and remote +lookups into common prefix results. Preserve model/cache namespaces, required +rank information, and hash/block-size handling. Convert block counts to token +counts only with a valid conversion. + +A confirmed miss and an unavailable backend both allow cache fallback, but +remain distinguishable in diagnostics. Invalid queries and configuration errors +remain explicit errors. Remote indexing and load-only deployments do not need +a duplicate local cache tree. + +### Affinity state + +`AffinityStore` owns scoped assignments, idle expiry, worker invalidation, and +atomic updates. Policies decide when to reuse or replace an assignment. + +Concurrent first assignments must converge on an effective binding that is +still a candidate and passes admission. If a concurrent binding returns a +different engine, the policy must check that engine before returning it. +Reconcile conflicts with bounded retry. + +Create or replace a binding only after admission succeeds. A binding records +preferred placement, not successful execution, so it may remain if later PD +selection or dispatch fails. It must not increment dispatch accounting. + +## 7. Configuration and compatibility + +This conceptual example shows the target attachment model. It is not copyable +current CLI/JSON syntax; existing bucket field names need not change. + +```yaml +buckets: + - id: short-context + rank: 10 + limits: {min: 0, max: 4096} + max_context_tokens: 8192 + groups: + pd: + prefill: + worker_ids: [P1, P2] + policy: + type: cache_aware + admission: {type: capacity} + decode: + worker_ids: [D1, D2] + policy: + type: power_of_two + admission: {type: capacity} + + - id: long-context + rank: 20 + limits: {min: 0, max: 131072} + max_context_tokens: 131072 + groups: + pd: + prefill: + worker_ids: [P3, P4] + policy: + type: cache_aware + admission: {type: capacity} + decode: + worker_ids: [D3, D4] + policy: + type: power_of_two + admission: {type: capacity} +``` + +A request with 4k input tokens and a 16k expected peak cannot fit the short +bucket's context capacity. It selects the long bucket and both of its P/D groups. +With a known peak of 8k or less, the same input selects both groups of the short bucket. + +The planned factory validates unique nonempty bucket IDs, token ranges, and +role-compatible membership. Each bucket is either plain or PD. The selected +bucket invokes its required groups through `pick_engines`. The existing worker registry +still rejects mixed plain and PD engines within one model; this PR preserves +that constraint. The engine group's model and stage filters apply on every pick. + +Legacy `BucketSpec` represents a single role-specific membership set. Migration +must explicitly associate prefill and decode specs into complete PD buckets; +never infer those associations from matching rank or similar names. Translation +of role-specific ranges/ranks into bucket-level constraints needs explicit +validation and is deferred with the configuration factory. + +Retained settings keep their meanings, defaults, units, and validation unless a +change is listed below. Policy-specific tuning applies to role groups using that +policy. Reject unsupported settings and incompatible combinations at startup; +do not accept and ignore them. + +### Retained behavior + +- Keep `load_based` as the CLI name for `LeastLoadPolicy`. +- Preserve explicit engine membership and context constraints. Bucket-level + ranges and ordering replace independent per-stage selection. + Restore existing SLO behavior in the separate SLO PR before serving switchover; + legacy routing continues to support SLOs during this skeleton-only phase. +- Preserve cache-provider selection, endpoint validation, query timeout and + concurrency limits, and unavailable-backend fallback. +- Preserve cache thresholds and tuning: the 1,024-token default minimum hit, + optional ratio gate, candidate bounds, switch margin, pressure guard, soft + queue limit, and saturation floor. +- Preserve session and sticky headers, idle timeouts, eviction cadence, and the + four sticky fallback choices. Global modes need a bucket-first migration design. +- Translate `--filter overloaded` and `--max-in-flight` into + `InFlightLimitAdmission`, composed with other checks through `AllOfAdmission`. +- Preserve configured capacity, pending-prefill, and in-flight checks, including + their missing-report behavior. Power-of-two applies admission to its selected + engine; other policies explicitly place checks in their selection logic. + Other paths use `AllowAll` unless a check is configured. Never silently discard + a configured budget. + +Listener and shutdown configuration, discovery, worker health and circuit +breakers, tokenizer loading, request timeouts, sampling overrides, and logging +remain outside the policy redesign. Preserve their existing behavior and +validation, including dispatch-time breaker probes and request cancellation. + +### Deliberate changes + +| Behavior | Target | +| --- | --- | +| Power-of-two admission | Check only the chosen engine; rejection advances to the next bucket | +| Round-robin cursor | One cursor per role-group policy instance | +| Capacity exhaustion | Try the next compatible bucket; return accumulated rejection details if all fail | +| Primary/backup proposals and post-policy substitution | Removed; each policy returns one engine | +| Session affinity | Reuse admitted bindings; remove primary/backup pressure escape | +| Omitted `--affinity-mode` | Admitted-binding reuse replaces the former soft-mode default | +| Pressure-guard tuning | Applies to cache-aware selection; reject session-only use | +| Policy attachment | Explicit role-group policy overrides the applicable model/stage default | + +Reject these dropped options explicitly: + +- `--policy fused_score` and `--policy score_policy`, including `--fuse` terms + and weights. +- `--decode-policy legacy_host_affinity`. +- `--stable-pair`. +- `--affinity-mode soft`; only strict admitted-binding reuse remains. +- `--filter prefix_cache` and `--prefix-cache-min-share`. The removed prefix-share + filter is not equivalent to the cache-aware minimum-hit gate. + +## 8. Dispatch and failure handling + +A successful pick means admission passed against the observed state. Health and +capacity can change before dispatch. The handler owns network operations, retry +rules, and accounting for the engines actually dispatched to. + +Selection fallback advances through the ordered compatible buckets before any +network dispatch. Every PD attempt picks both engines from that bucket. Transport +retry integration remains separate work; a forwarding failure does not resume +the bucket loop or silently replace a successful policy pick. + +For PD, acquire and release accounting for the actual stages and clean up +partial setup on failure. Policy selection does not own the PD request lifetime. +Selection metrics must not imply that dispatch or execution succeeded. + +## Implementation status + +This PR adds the side-by-side interfaces in `src/buckets_reorg.rs` and +`src/policies_reorg/`, plus a configurable bucket-first implementation behind +`chat_completions`. The live `src/policies/` path remains the default. + +Implemented here: + +- `BucketResolver::resolve` returns all length-compatible buckets in capacity/rank/ID order. +- `Bucket::pick_engines` owns plain/PD orchestration and stage-specific policy + requests; `BucketRequest` carries prepared facts and `BucketPick` retains picks. +- `Bucket` owns input limits, context capacity, rank, and plain-or-PD groups. +- `EngineGroup::pick` owns live candidate filtering, policy invocation, and + exact candidate validation, without cross-bucket fallback. +- `Policy::pick`, within-group fallback interface, per-engine `EngineAdmission::check`, + and `AllowAll`. Power-of-two samples two distinct engines, compares stage pressure, + and checks its selected engine with no replacement on rejection. +- Policy-owned load dependency and local observations. Power-of-two passes the + selected engine's load record directly to admission, without another snapshot. + `PickRequest`, `Pick`, and bucket APIs carry no load observations. +- The reorg chat implementation iterates resolved buckets, calls `pick_engines`, + advances on empty candidates/admission rejection, and + dispatches only after one complete selection. Exhaustion retains admission reasons. +- `AppContext::chat_routing` configures legacy versus reorg routing on the same + endpoint and carries the reorg model-resolver map. + +Follow-up order: concrete admission (#40271), then bucket SLO ordering +in a separate PR, followed by remaining policies and production configuration. + +Not yet implemented in the reorg path: + +- Other concrete policies and capacity/in-flight admission checks. +- SLO estimates, targets, and bucket preference ordering. +- CLI/configuration parsing, validation, and model-specific construction. + The YAML above is illustrative; reorg resolvers are installed in code. +- Session modes, prefix memoization, and cache-aware selection. +- Shared load interpretation, dispatch correction, and policy-specific + dispatch-timestamp requirements. +- PD compatibility filtering, retry integration, and legacy-route switchover. + +The preceding policy sections describe target behavior for those follow-ups; +they do not claim those capabilities are present in this PR. diff --git a/experimental/sgl-router/src/buckets_reorg.rs b/experimental/sgl-router/src/buckets_reorg.rs new file mode 100644 index 000000000..2bf39509c --- /dev/null +++ b/experimental/sgl-router/src/buckets_reorg.rs @@ -0,0 +1,238 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Order buckets by request length; each bucket owns the groups used to pick engines. +//! +//! ```text +//! BucketResolver (one model's buckets) +//! -> Bucket (token limits, context capacity, rank) +//! -> Plain: one EngineGroup +//! -> PD: prefill + decode EngineGroups +//! -> each EngineGroup: worker membership + its own Policy +//! ``` +//! +//! - [`BucketResolver::resolve`] returns length-compatible buckets in preference order. +//! - [`EngineGroup::pick`] filters live workers by model, health, stage, and membership, +//! then calls [`Policy::pick`] and validates the returned engine. +//! +//! The handler tries buckets in order, advancing on missing candidates or admission +//! rejection. Both P/D picks must succeed in the same bucket before dispatch. +//! [`WorkerRegistry`] owns live workers; groups reference their IDs. Policies own +//! their load/KV/affinity dependencies and share observations within each attempt. + +use std::collections::HashSet; +use std::sync::Arc; + +use crate::discovery::{ModelId, WorkerId}; +use crate::policies_reorg::{Pick, PickError, PickRequest, Policy, Stage}; +use crate::workers::WorkerRegistry; + +#[derive(Debug, Clone, Copy, Default)] +pub struct TokenLimits { + pub min: Option, + pub max: Option, +} + +impl TokenLimits { + fn fits(&self, tokens: u64) -> bool { + self.min.is_none_or(|min| tokens >= min) && self.max.is_none_or(|max| tokens <= max) + } +} + +#[derive(Debug)] +pub struct EngineGroup { + /// `None` includes all registered engines matching the request's model and role. + pub worker_ids: Option>, + pub policy: Arc, +} + +impl EngineGroup { + pub fn new(policy: Arc) -> Self { + Self { + worker_ids: None, + policy, + } + } + + /// Resolve live members and invoke this group's policy. Never changes buckets. + pub async fn pick( + &self, + workers: &WorkerRegistry, + request: &PickRequest<'_>, + ) -> Result { + let mut engines: Vec<_> = workers + .healthy_workers_for(request.model) + .into_iter() + .filter(|engine| engine.mode() == request.stage) + .filter(|engine| { + self.worker_ids + .as_ref() + .is_none_or(|ids| ids.contains(&engine.id)) + }) + .collect(); + // Stable order so cursor-based policies see a consistent candidate list. + engines.sort_by(|left, right| left.id.0.cmp(&right.id.0)); + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + let pick = self.policy.pick(&engines, request).await?; + if !engines + .iter() + .any(|engine| Arc::ptr_eq(engine, &pick.engine)) + { + return Err(PickError::OutsideCandidates(pick.engine.id.clone())); + } + Ok(pick) + } +} + +/// A bucket serves a request on one plain engine or on its own P/D groups. +#[derive(Debug)] +pub enum BucketGroups { + Plain(EngineGroup), + Pd { + prefill: EngineGroup, + decode: EngineGroup, + }, +} + +/// Prepared request facts shared by all bucket attempts. The bucket supplies +/// its ID and each group's stage when calling policies. +#[derive(Debug)] +pub struct BucketRequest<'a> { + pub model: &'a ModelId, + pub input_tokens: u64, + pub expected_peak_tokens: Option, + pub token_ids: Option<&'a [u32]>, + pub session_key: Option<&'a str>, + pub routing_key: Option<&'a str>, +} + +/// A complete selection from one bucket. For plain serving, `prefill` is the +/// plain engine and `decode` is absent; PD supplies both picks. +#[derive(Debug)] +pub struct BucketPick { + pub prefill: Pick, + pub decode: Option, +} + +#[derive(Debug)] +pub struct Bucket { + pub id: String, + /// Break ties between equally sized buckets; lower ranks win. + pub rank: u32, + /// Inclusive input-token range used to choose the bucket. + pub limits: TokenLimits, + /// Full sequence capacity, checked against the expected peak when known. + pub max_context_tokens: Option, + pub groups: BucketGroups, +} + +impl Bucket { + pub fn new(id: impl Into, groups: BucketGroups) -> Self { + Self { + id: id.into(), + rank: 0, + limits: TokenLimits::default(), + max_context_tokens: None, + groups, + } + } + + /// Select this bucket's plain engine or complete P/D pair, without dispatching. + /// A failed group reports its stage; the caller may then try another bucket. + pub async fn pick_engines( + &self, + workers: &WorkerRegistry, + request: &BucketRequest<'_>, + ) -> Result { + let (prefill, decode) = match &self.groups { + BucketGroups::Plain(group) => ( + self.pick_from_group(group, Stage::Plain, workers, request) + .await?, + None, + ), + BucketGroups::Pd { prefill, decode } => { + let prefill = self + .pick_from_group(prefill, Stage::Prefill, workers, request) + .await?; + let decode = self + .pick_from_group(decode, Stage::Decode, workers, request) + .await?; + (prefill, Some(decode)) + } + }; + Ok(BucketPick { prefill, decode }) + } + + /// Scope the request to this bucket and role, then ask the group for one engine. + async fn pick_from_group( + &self, + group: &EngineGroup, + stage: Stage, + workers: &WorkerRegistry, + request: &BucketRequest<'_>, + ) -> Result { + let request = PickRequest { + model: request.model, + stage, + bucket: &self.id, + input_tokens: request.input_tokens, + expected_peak_tokens: request.expected_peak_tokens, + token_ids: request.token_ids, + session_key: request.session_key, + routing_key: request.routing_key, + }; + group + .pick(workers, &request) + .await + .map_err(|error| (stage, error)) + } + + fn fits(&self, input_tokens: u64, expected_peak_tokens: Option) -> bool { + self.limits.fits(input_tokens) + && self + .max_context_tokens + .is_none_or(|max| expected_peak_tokens.unwrap_or(input_tokens) <= max) + } + + fn input_capacity(&self) -> u64 { + self.limits + .max + .unwrap_or(u64::MAX) + .min(self.max_context_tokens.unwrap_or(u64::MAX)) + } +} + +/// Model-specific bucket configuration. Selection does not inspect engine state. +#[derive(Debug, Default)] +pub struct BucketResolver { + pub buckets: Vec, +} + +impl BucketResolver { + pub fn new(buckets: Vec) -> Self { + Self { buckets } + } + + /// Return all length-compatible buckets, ordered by input capacity, rank, and ID. + /// The caller tries their groups in order until a complete engine selection succeeds. + pub fn resolve( + &self, + input_tokens: u64, + expected_peak_tokens: Option, + ) -> Result, PickError> { + if expected_peak_tokens.is_some_and(|tokens| tokens < input_tokens) { + return Err(PickError::InvalidSignal( + "expected peak tokens are below input length".into(), + )); + } + let mut buckets: Vec<_> = self + .buckets + .iter() + .filter(|bucket| bucket.fits(input_tokens, expected_peak_tokens)) + .collect(); + buckets.sort_by_key(|bucket| (bucket.input_capacity(), bucket.rank, &bucket.id)); + Ok(buckets) + } +} diff --git a/experimental/sgl-router/src/lib.rs b/experimental/sgl-router/src/lib.rs index bb7a1252e..530eaa1bb 100644 --- a/experimental/sgl-router/src/lib.rs +++ b/experimental/sgl-router/src/lib.rs @@ -8,10 +8,12 @@ pub const VERSION: &str = env!("CARGO_PKG_VERSION"); +pub mod buckets_reorg; pub mod config; pub mod discovery; pub mod health; pub mod policies; +pub mod policies_reorg; pub mod proxy; pub mod server; pub mod state; diff --git a/experimental/sgl-router/src/policies_reorg/admission.rs b/experimental/sgl-router/src/policies_reorg/admission.rs new file mode 100644 index 000000000..f63896b38 --- /dev/null +++ b/experimental/sgl-router/src/policies_reorg/admission.rs @@ -0,0 +1,42 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::fmt::Debug; + +use crate::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; +use crate::workers::Worker; + +use super::{PickError, PickRequest}; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum Decision { + Allow, + Reject(String), +} + +/// Checks one engine using the load observation retained by selection. +/// `None` means no usable load observation, never zero load. Each check defines +/// its missing-data behavior and owns any other state handles it needs. +/// Each policy decides when to check an engine and how to handle rejection. +pub trait EngineAdmission: Send + Sync + Debug { + fn check( + &self, + engine: &Worker, + request: &PickRequest<'_>, + load: Option<&EngineReportedWorkerLoad>, + ) -> Result; +} + +#[derive(Debug)] +pub struct AllowAll; + +impl EngineAdmission for AllowAll { + fn check( + &self, + _: &Worker, + _: &PickRequest<'_>, + _: Option<&EngineReportedWorkerLoad>, + ) -> Result { + Ok(Decision::Allow) + } +} diff --git a/experimental/sgl-router/src/policies_reorg/mod.rs b/experimental/sgl-router/src/policies_reorg/mod.rs new file mode 100644 index 000000000..34e5ca265 --- /dev/null +++ b/experimental/sgl-router/src/policies_reorg/mod.rs @@ -0,0 +1,105 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +//! Side-by-side implementation of POLICY_DESIGN.md. Chat routing can opt into +//! this interface through AppContext; `policies` remains the default. + +pub mod admission; +pub mod power_of_two; + +use std::fmt::Debug; +use std::sync::Arc; + +use futures::future::BoxFuture; + +use crate::discovery::{ModelId, WorkerId}; +use crate::workers::Worker; + +pub use crate::discovery::WorkerMode as Stage; + +/// Request facts for engine selection. Policies own their shared-state handles. +#[derive(Debug, Clone, Copy)] +pub struct PickRequest<'a> { + pub model: &'a ModelId, + pub stage: Stage, + pub bucket: &'a str, + pub input_tokens: u64, + pub expected_peak_tokens: Option, + pub token_ids: Option<&'a [u32]>, + pub session_key: Option<&'a str>, + pub routing_key: Option<&'a str>, +} + +impl<'a> PickRequest<'a> { + pub fn new(model: &'a ModelId, stage: Stage, input_tokens: u64) -> Self { + Self { + model, + stage, + bucket: "", + input_tokens, + expected_peak_tokens: None, + token_ids: None, + session_key: None, + routing_key: None, + } + } +} + +#[derive(Debug, Clone)] +pub struct Pick { + pub engine: Arc, + pub reason: &'static str, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Rejection { + pub engine: WorkerId, + pub reason: String, +} + +#[derive(Debug, thiserror::Error)] +pub enum PickError { + #[error("no bucket matches the request length")] + NoMatchingBucket, + #[error("no candidates")] + NoCandidates, + #[error("no admissible engine: {0:?}")] + NoAdmissibleEngine(Vec), + #[error("selected engine rejected: {0:?}")] + AdmissionRejected(Rejection), + #[error("invalid signal: {0}")] + InvalidSignal(String), + #[error("invalid configuration: {0}")] + InvalidConfiguration(String), + #[error("policy selected an engine outside its candidates: {0:?}")] + OutsideCandidates(WorkerId), +} + +/// Returns one admitted engine from exactly the supplied candidates. +/// Implementations receive shared load, KV, and affinity handles at construction; +/// they obtain their own observations rather than asking callers to supply them. +pub trait Policy: Send + Sync + Debug { + /// Read required state locally and pass the selected engine's observations to admission. + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result>; + + /// Runs on a miss within the same candidates; never on an admission rejection. + fn fallback(&self) -> Option<&dyn Policy> { + None + } + + /// Delegate within the same candidates; the fallback reads its own state. + fn pick_fallback<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + match self.fallback() { + Some(fallback) => fallback.pick(engines, request), + None => Box::pin(async { Err(PickError::NoCandidates) }), + } + } +} diff --git a/experimental/sgl-router/src/policies_reorg/power_of_two.rs b/experimental/sgl-router/src/policies_reorg/power_of_two.rs new file mode 100644 index 000000000..4d6e7b06a --- /dev/null +++ b/experimental/sgl-router/src/policies_reorg/power_of_two.rs @@ -0,0 +1,79 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Instant; + +use futures::future::BoxFuture; +use rand::Rng; + +use crate::policies::admission::{compare_decode_pressure, compare_prefill_pressure}; +use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable; +use crate::workers::Worker; + +use super::admission::{AllowAll, Decision, EngineAdmission}; +use super::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; + +/// Samples two distinct engines and selects the one with lower stage pressure. +/// Checks admission only on the selected engine; rejection never resamples. +#[derive(Debug)] +pub struct PowerOfTwoPolicy { + /// Shared application state; snapshots are local to each pick. + engine_load: Arc, + pub admission: Arc, +} + +impl PowerOfTwoPolicy { + pub fn new(engine_load: Arc) -> Self { + Self { + engine_load, + admission: Arc::new(AllowAll), + } + } +} + +impl Policy for PowerOfTwoPolicy { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + // Selection and admission use the same load observation. + let load = self.engine_load.capture_snapshot(Instant::now()); + let engine = match engines { + [engine] => Arc::clone(engine), + _ => { + let mut rng = rand::thread_rng(); + let i = rng.gen_range(0..engines.len()); + let mut j = rng.gen_range(0..engines.len() - 1); + if j >= i { + j += 1; + } + let (left, right) = (&engines[i], &engines[j]); + let pressure = match request.stage { + Stage::Plain | Stage::Prefill => { + compare_prefill_pressure(left, right, Some(&load)) + } + Stage::Decode => compare_decode_pressure(left, right, Some(&load)), + }; + Arc::clone(if pressure.is_gt() { right } else { left }) + } + }; + let engine_load = load.fresh_load_for_url(&engine.url); + if let Decision::Reject(reason) = self.admission.check(&engine, request, engine_load)? { + return Err(PickError::AdmissionRejected(Rejection { + engine: engine.id.clone(), + reason, + })); + } + Ok(Pick { + engine, + reason: "power_of_two", + }) + }) + } +} diff --git a/experimental/sgl-router/src/server/app_context.rs b/experimental/sgl-router/src/server/app_context.rs index eaeb3da1a..4502bef32 100644 --- a/experimental/sgl-router/src/server/app_context.rs +++ b/experimental/sgl-router/src/server/app_context.rs @@ -1,7 +1,9 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors // SPDX-License-Identifier: Apache-2.0 +use crate::buckets_reorg::BucketResolver; use crate::config::Config; +use crate::discovery::ModelId; use crate::policies::buckets::BucketSelector; use crate::policies::prefix_provider::RadixTreePrefixProvider; @@ -14,6 +16,7 @@ use crate::state::load_monitor::engine_reported_load::EngineReportedLoadTable; use crate::state::load_monitor::router_inflight_load::RouterInflightLoadRegistry; use crate::tokenizer::TokenizerRegistry; use crate::workers::WorkerRegistry; +use std::collections::HashMap; use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::Arc; @@ -26,6 +29,15 @@ const READINESS_NOT_READY: u8 = 0; const READINESS_READY: u8 = 1; const READINESS_DRAINING: u8 = 2; +/// Routing implementation used by the standard chat-completions endpoint. +/// Reorg configuration is installed explicitly until its CLI factory is available. +#[derive(Debug, Default)] +pub enum ChatRouting { + #[default] + Legacy, + Reorg(HashMap), +} + pub struct AppContext { pub config: Config, pub tokenizers: Arc, @@ -34,6 +46,8 @@ pub struct AppContext { pub policies: Arc, /// Converts static Bucket configuration into request candidate domains. pub bucket_selector: Arc, + /// Select legacy policies or model-specific bucket-first routing. + pub chat_routing: ChatRouting, /// Per-worker active-load bookkeeping shared by the proxy, policies, /// timeout janitor, and metrics. pub router_inflight_load: Arc, @@ -106,6 +120,7 @@ impl AppContext { registry, policies, bucket_selector, + chat_routing: ChatRouting::Legacy, router_inflight_load, metrics, prefix_index: None, @@ -195,6 +210,7 @@ impl AppContext { registry: Arc::new(WorkerRegistry::default()), policies: Arc::new(PolicyRegistry::default()), bucket_selector: Arc::new(BucketSelector::new(None)), + chat_routing: ChatRouting::Legacy, router_inflight_load: RouterInflightLoadRegistry::with_defaults(), metrics: MetricsRegistry::new(), prefix_index: None, diff --git a/experimental/sgl-router/src/server/routes/chat.rs b/experimental/sgl-router/src/server/routes/chat.rs index 7d4d72673..008b73522 100644 --- a/experimental/sgl-router/src/server/routes/chat.rs +++ b/experimental/sgl-router/src/server/routes/chat.rs @@ -3,6 +3,7 @@ mod forward; mod preparation; +mod reorg; use crate::config::{SessionAffinityMode, DEFAULT_MIN_LOAD_CHOICES}; use crate::discovery::{ModelId, WorkerMode}; @@ -11,7 +12,7 @@ use crate::policies::selection::{ select_decode_peer, select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs, }; use crate::policies::{ExternalPrefixSignal, Policy}; -use crate::server::app_context::AppContext; +use crate::server::app_context::{AppContext, ChatRouting}; use crate::server::error::ApiError; use crate::server::metrics::PolicySelectionFailureReason; use crate::state::kv_events::{compute_block_hashes, compute_block_hashes_bigram}; @@ -38,6 +39,19 @@ pub async fn chat_completions( State(ctx): State>, headers: HeaderMap, body: Bytes, +) -> Result, ApiError> { + match &ctx.chat_routing { + ChatRouting::Legacy => chat_completions_legacy(&ctx, headers, body).await, + ChatRouting::Reorg(resolvers) => { + reorg::chat_completions(&ctx, resolvers, headers, body).await + } + } +} + +async fn chat_completions_legacy( + ctx: &AppContext, + headers: HeaderMap, + body: Bytes, ) -> Result, ApiError> { let start = Instant::now(); let mut fields = parse_routing_fields(&body)?; @@ -59,11 +73,11 @@ pub async fn chat_completions( .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?; let request = - PreparedChatRequest::prepare(&ctx, model, fields, body, policy.needs_request_tokens())?; + PreparedChatRequest::prepare(ctx, model, fields, body, policy.needs_request_tokens())?; // Pick a plain worker, or a prefill worker followed by a decode peer in PD mode. let workers = select_workers( - &ctx, + ctx, &request, &headers, policy.as_ref(), @@ -73,7 +87,7 @@ pub async fn chat_completions( .await?; // PD sends to both workers and returns the decode response. - forward_chat_request(&ctx, request, workers, headers, start).await + forward_chat_request(ctx, request, workers, headers, start).await } fn pool_error(error: PdResolveError, model: &ModelId) -> ApiError { diff --git a/experimental/sgl-router/src/server/routes/chat/reorg.rs b/experimental/sgl-router/src/server/routes/chat/reorg.rs new file mode 100644 index 000000000..eb872d7d9 --- /dev/null +++ b/experimental/sgl-router/src/server/routes/chat/reorg.rs @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use super::forward::{forward_chat_request, SelectedWorkers}; +use super::nonempty_header; +use super::preparation::{parse_routing_fields, PreparedChatRequest}; +use crate::buckets_reorg::{BucketRequest, BucketResolver}; +use crate::discovery::ModelId; +use crate::policies_reorg::{PickError, Stage}; +use crate::server::app_context::AppContext; +use crate::server::error::ApiError; +use axum::body::Body; +use axum::http::{HeaderMap, Response}; +use bytes::Bytes; +use std::collections::HashMap; +use std::time::Instant; + +/// Bucket-first implementation selected by `AppContext::chat_routing`. +pub(super) async fn chat_completions( + ctx: &AppContext, + resolvers: &HashMap, + headers: HeaderMap, + body: Bytes, +) -> Result, ApiError> { + let start = Instant::now(); + let mut fields = parse_routing_fields(&body)?; + let model = ModelId( + fields + .model + .take() + .ok_or_else(|| ApiError::BadRequest("missing `model` field".into()))?, + ); + let resolver = resolvers + .get(&model) + .ok_or_else(|| ApiError::ModelNotFound(model.0.clone()))?; + // Length-based routing needs tokenization even for load-only group policies. + let request = PreparedChatRequest::prepare(ctx, model, fields, body, true)?; + let input_tokens = request.input_token_count as u64; + let expected_peak_tokens = request + .max_output_tokens + .map(|output| { + input_tokens.checked_add(output).ok_or_else(|| { + ApiError::BadRequest("input and output token counts overflow".into()) + }) + }) + .transpose()?; + + let buckets = resolver + .resolve(input_tokens, expected_peak_tokens) + .map_err(|error| selection_error(error, &request.model, None))?; + if buckets.is_empty() { + return Err(selection_error( + PickError::NoMatchingBucket, + &request.model, + None, + )); + } + let bucket_request = BucketRequest { + model: &request.model, + input_tokens, + expected_peak_tokens, + token_ids: request.tokens.as_ref().map(|tokens| tokens.ids.as_slice()), + session_key: ctx + .config + .model + .affinity + .as_ref() + .and_then(|config| nonempty_header(&headers, &config.session_id_header)), + routing_key: ctx + .config + .model + .sticky + .as_ref() + .and_then(|config| nonempty_header(&headers, &config.header_name)), + }; + let mut rejections: Option> = None; + let mut missing_stage = None; + for bucket in buckets { + match bucket.pick_engines(&ctx.registry, &bucket_request).await { + Ok(picks) => { + // Dispatch only after this bucket supplies the entire plain or PD selection. + let workers = SelectedWorkers { + prefill: picks.prefill.engine, + decode: picks.decode.map(|pick| pick.engine), + track_dispatch_timestamps: false, + }; + return forward_chat_request(ctx, request, workers, headers, start).await; + } + Err((stage, error)) => { + tracing::debug!(bucket = %bucket.id, ?stage, %error, "bucket selection failed"); + match error { + PickError::NoCandidates => missing_stage = Some(stage), + PickError::NoAdmissibleEngine(reasons) => { + rejections.get_or_insert_with(Vec::new).extend(reasons); + } + PickError::AdmissionRejected(reason) => { + rejections.get_or_insert_with(Vec::new).push(reason); + } + error => return Err(selection_error(error, &request.model, Some(stage))), + } + } + } + } + // Preserve admission exhaustion even if a later bucket has no candidates. + let error = match rejections { + Some(reasons) => PickError::NoAdmissibleEngine(reasons), + None => PickError::NoCandidates, + }; + Err(selection_error(error, &request.model, missing_stage)) +} + +fn selection_error(error: PickError, model: &ModelId, stage: Option) -> ApiError { + tracing::warn!(%model, ?stage, %error, "reorg selection failed"); + match error { + PickError::NoMatchingBucket => { + ApiError::BadRequest("no bucket supports the requested token length".into()) + } + PickError::NoCandidates => match stage { + Some(Stage::Prefill) => ApiError::NoPrefillWorkersAvailable { + model: model.0.clone(), + }, + Some(Stage::Decode) => ApiError::NoDecodeWorkersAvailable { + model: model.0.clone(), + }, + _ => ApiError::NoHealthyWorkers { + model: model.0.clone(), + }, + }, + PickError::NoAdmissibleEngine(_) | PickError::AdmissionRejected(_) => { + ApiError::PolicySelectionFailed { + model: model.0.clone(), + } + } + error => ApiError::Internal(error.into()), + } +} diff --git a/experimental/sgl-router/tests/component/main.rs b/experimental/sgl-router/tests/component/main.rs index d2a8461f9..1a5accf8b 100644 --- a/experimental/sgl-router/tests/component/main.rs +++ b/experimental/sgl-router/tests/component/main.rs @@ -10,5 +10,8 @@ mod discovery; mod health; mod policies; +mod policies_reorg; +mod policies_reorg_load; +mod policies_reorg_power_of_two; mod tokenizer; mod workers; diff --git a/experimental/sgl-router/tests/component/policies_reorg.rs b/experimental/sgl-router/tests/component/policies_reorg.rs new file mode 100644 index 000000000..65ac82cf9 --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg.rs @@ -0,0 +1,486 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; + +use futures::future::BoxFuture; +use sgl_router::buckets_reorg::{ + Bucket, BucketGroups, BucketRequest, BucketResolver, EngineGroup, TokenLimits, +}; +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission}; +use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; +use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; +use sgl_router::workers::{Worker, WorkerRegistry}; + +#[derive(Debug)] +struct TestPolicy { + admission: Arc, + result: Option>, + miss: bool, + invalid: bool, + calls: Mutex>, +} + +impl Default for TestPolicy { + fn default() -> Self { + Self { + admission: Arc::new(AllowAll), + result: None, + miss: false, + invalid: false, + calls: Mutex::new(Vec::new()), + } + } +} + +impl Policy for TestPolicy { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.calls.lock().unwrap().push(request.bucket.to_owned()); + if self.invalid { + return Err(PickError::InvalidSignal("test signal".into())); + } + if self.miss { + return Err(PickError::NoCandidates); + } + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + let engine = self.result.clone().unwrap_or_else(|| engines[0].clone()); + if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? { + return Err(PickError::AdmissionRejected(Rejection { + engine: engine.id.clone(), + reason, + })); + } + Ok(Pick { + engine, + reason: "test", + }) + }) + } +} + +#[derive(Debug)] +struct Reject(&'static str); + +impl EngineAdmission for Reject { + fn check( + &self, + engine: &Worker, + _: &PickRequest<'_>, + _: Option<&EngineReportedWorkerLoad>, + ) -> Result { + Ok(if engine.id.0 == self.0 { + Decision::Reject("full".into()) + } else { + Decision::Allow + }) + } +} + +fn spec(id: &str, mode: Stage, model: &str) -> WorkerSpec { + WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}"), + mode, + model_ids: vec![ModelId(model.into())], + bootstrap_port: None, + } +} + +fn registry() -> Arc { + let workers = Arc::new(WorkerRegistry::default()); + for (id, mode, model) in [ + ("a", Stage::Plain, "m"), + ("b", Stage::Plain, "m"), + ("unhealthy", Stage::Plain, "m"), + ("other", Stage::Plain, "other"), + ("p", Stage::Prefill, "pd"), + ("d", Stage::Decode, "pd"), + ] { + workers.add(spec(id, mode, model)).unwrap(); + } + let unhealthy = workers.get(&WorkerId("unhealthy".into())).unwrap(); + for _ in 0..3 { + unhealthy.breaker.record_failure(); + } + workers +} + +fn group(members: &[&str], policy: Arc) -> EngineGroup { + EngineGroup { + worker_ids: Some(members.iter().map(|id| WorkerId((*id).into())).collect()), + policy, + } +} + +fn bucket(id: &str, max: Option, policy: Arc) -> Bucket { + let mut bucket = Bucket::new(id, BucketGroups::Plain(EngineGroup::new(policy))); + bucket.limits.max = max; + bucket +} + +#[tokio::test] +async fn groups_isolate_model_health_stage_and_membership() { + let workers = registry(); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let group = EngineGroup::new(Arc::new(TestPolicy::default())); + assert_eq!( + group.pick(&workers, &request).await.unwrap().engine.id.0, + "a" + ); + workers.remove(&WorkerId("a".into())); + assert_eq!( + group.pick(&workers, &request).await.unwrap().engine.id.0, + "b" + ); + workers.remove(&WorkerId("b".into())); + assert!(matches!( + group.pick(&workers, &request).await, + Err(PickError::NoCandidates) + )); + + let pd = ModelId("pd".into()); + for (stage, expected) in [(Stage::Prefill, "p"), (Stage::Decode, "d")] { + let request = PickRequest::new(&pd, stage, 10); + let group = self::group( + &["p", "d", "other", "unhealthy"], + Arc::new(TestPolicy::default()), + ); + assert_eq!( + group.pick(&workers, &request).await.unwrap().engine.id.0, + expected + ); + } + let empty = self::group(&[], Arc::new(TestPolicy::default())); + assert!(matches!( + empty.pick(&workers, &request).await, + Err(PickError::NoCandidates) + )); +} + +#[test] +fn resolve_orders_all_length_fits_by_capacity_rank_and_id() { + let policy = Arc::new(TestPolicy::default()); + let mut z = bucket("z", Some(20), policy.clone()); + z.rank = 1; + let mut a = bucket("a", Some(20), policy.clone()); + a.rank = 1; + let mut later = bucket("later", Some(20), policy.clone()); + later.rank = 2; + let mut min = bucket("min", Some(15), policy.clone()); + min.limits.min = Some(11); + let resolver = BucketResolver::new(vec![ + bucket("catch-all", None, policy.clone()), + bucket("too-small", Some(9), policy), + z, + later, + a, + min, + ]); + assert_eq!( + resolver + .resolve(10, None) + .unwrap() + .iter() + .map(|bucket| bucket.id.as_str()) + .collect::>(), + ["a", "z", "later", "catch-all"] + ); + assert_eq!(resolver.resolve(11, None).unwrap()[0].id, "min"); + assert_eq!(resolver.resolve(15, None).unwrap()[0].id, "min"); + assert_eq!(resolver.resolve(20, None).unwrap()[0].id, "a"); + assert_eq!(resolver.resolve(21, None).unwrap()[0].id, "catch-all"); +} + +#[test] +fn context_capacity_checks_peak_when_known_and_input_otherwise() { + let policy = Arc::new(TestPolicy::default()); + let mut short = bucket("short", None, policy.clone()); + short.max_context_tokens = Some(20); + let mut long = bucket("long", None, policy); + long.max_context_tokens = Some(30); + let resolver = BucketResolver::new(vec![long, short]); + assert_eq!(resolver.resolve(10, None).unwrap()[0].id, "short"); + assert_eq!(resolver.resolve(10, Some(20)).unwrap()[0].id, "short"); + assert_eq!(resolver.resolve(10, Some(21)).unwrap()[0].id, "long"); + assert!(resolver.resolve(10, Some(31)).unwrap().is_empty()); + assert!(resolver.resolve(31, None).unwrap().is_empty()); + assert!(matches!( + resolver.resolve(10, Some(9)), + Err(PickError::InvalidSignal(_)) + )); + assert!(BucketResolver::default() + .resolve(1, None) + .unwrap() + .is_empty()); +} + +#[tokio::test] +async fn selected_pd_bucket_owns_both_memberships_and_policies() { + let workers = registry(); + workers.add(spec("p2", Stage::Prefill, "pd")).unwrap(); + workers.add(spec("d2", Stage::Decode, "pd")).unwrap(); + let model = ModelId("pd".into()); + let prefill_policy = Arc::new(TestPolicy::default()); + let decode_policy = Arc::new(TestPolicy::default()); + let resolver = BucketResolver::new(vec![Bucket::new( + "shared", + BucketGroups::Pd { + prefill: group(&["p2", "d", "a"], prefill_policy.clone()), + decode: group(&["d2", "p", "other"], decode_policy.clone()), + }, + )]); + let bucket = resolver.resolve(10, Some(20)).unwrap()[0]; + let request = BucketRequest { + model: &model, + input_tokens: 10, + expected_peak_tokens: Some(20), + token_ids: None, + session_key: None, + routing_key: None, + }; + let picks = bucket.pick_engines(&workers, &request).await.unwrap(); + assert_eq!(picks.prefill.engine.id.0, "p2"); + assert_eq!(picks.decode.unwrap().engine.id.0, "d2"); + assert_eq!(*prefill_policy.calls.lock().unwrap(), ["shared"]); + assert_eq!(*decode_policy.calls.lock().unwrap(), ["shared"]); +} + +#[tokio::test] +async fn resolver_includes_empty_groups_without_invoking_policies() { + let workers = registry(); + let model = ModelId("m".into()); + let policy = Arc::new(TestPolicy::default()); + let mut empty = Bucket::new( + "empty", + BucketGroups::Plain(group(&["missing"], policy.clone())), + ); + empty.limits = TokenLimits { + min: None, + max: Some(10), + }; + let resolver = BucketResolver::new(vec![empty, bucket("available", Some(20), policy.clone())]); + let buckets = resolver.resolve(10, None).unwrap(); + assert_eq!( + buckets + .iter() + .map(|bucket| bucket.id.as_str()) + .collect::>(), + ["empty", "available"] + ); + let bucket = buckets[0]; + let BucketGroups::Plain(group) = &bucket.groups else { + panic!("expected plain") + }; + let request = PickRequest { + bucket: &bucket.id, + ..PickRequest::new(&model, Stage::Plain, 10) + }; + assert!(matches!( + group.pick(&workers, &request).await, + Err(PickError::NoCandidates) + )); + assert!(policy.calls.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn group_propagates_rejections_misses_and_invalid_signals() { + let workers = registry(); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let rejected = group( + &["a"], + Arc::new(TestPolicy { + admission: Arc::new(Reject("a")), + ..Default::default() + }), + ); + assert!(matches!( + rejected.pick(&workers, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + let miss = group( + &["a"], + Arc::new(TestPolicy { + miss: true, + ..Default::default() + }), + ); + assert!(matches!( + miss.pick(&workers, &request).await, + Err(PickError::NoCandidates) + )); + let invalid = group( + &["a"], + Arc::new(TestPolicy { + invalid: true, + ..Default::default() + }), + ); + assert!(matches!( + invalid.pick(&workers, &request).await, + Err(PickError::InvalidSignal(_)) + )); +} + +#[tokio::test] +async fn group_rejects_foreign_pick_even_with_same_worker_id() { + let workers = registry(); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let group = group( + &["a"], + Arc::new(TestPolicy { + result: Some(Arc::new(Worker::new(spec("a", Stage::Plain, "m")))), + ..Default::default() + }), + ); + assert!(matches!( + group.pick(&workers, &request).await, + Err(PickError::OutsideCandidates(_)) + )); +} + +#[tokio::test] +async fn selected_engine_rejection_does_not_try_an_alternative() { + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let engines = [ + Arc::new(Worker::new(spec("a", Stage::Plain, "m"))), + Arc::new(Worker::new(spec("b", Stage::Plain, "m"))), + ]; + let policy = TestPolicy { + admission: Arc::new(Reject("a")), + ..Default::default() + }; + assert!(matches!( + policy.pick(&engines, &request).await, + Err(PickError::AdmissionRejected(_)) + )); + assert!(matches!( + policy.pick(&[], &request).await, + Err(PickError::NoCandidates) + )); +} + +#[tokio::test] +async fn bucket_scopes_plain_pick_and_preserves_request_facts() { + #[derive(Debug)] + struct InspectRequest; + + impl Policy for InspectRequest { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + assert_eq!(request.model.0, "m"); + assert_eq!(request.bucket, "plain-bucket"); + assert_eq!(request.stage, Stage::Plain); + assert_eq!(request.input_tokens, 2); + assert_eq!(request.expected_peak_tokens, Some(12)); + assert_eq!(request.token_ids, Some([7, 9].as_slice())); + assert_eq!(request.session_key, Some("session")); + assert_eq!(request.routing_key, Some("routing")); + Ok(Pick { + engine: engines[0].clone(), + reason: "inspected", + }) + }) + } + } + + let workers = registry(); + let model = ModelId("m".into()); + let bucket = Bucket::new( + "plain-bucket", + BucketGroups::Plain(group(&["b"], Arc::new(InspectRequest))), + ); + let request = BucketRequest { + model: &model, + input_tokens: 2, + expected_peak_tokens: Some(12), + token_ids: Some(&[7, 9]), + session_key: Some("session"), + routing_key: Some("routing"), + }; + let picks = bucket.pick_engines(&workers, &request).await.unwrap(); + assert_eq!(picks.prefill.engine.id.0, "b"); + assert_eq!(picks.prefill.reason, "inspected"); + assert!(picks.decode.is_none()); +} + +#[tokio::test] +async fn power_of_two_checks_selected_engine_and_propagates_rejection_without_fallback() { + use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; + use sgl_router::state::load_monitor::engine_reported_load::EngineReportedLoadTable; + + #[derive(Debug)] + struct Check { + calls: Mutex>, + reject: bool, + invalid: bool, + } + + impl EngineAdmission for Check { + fn check( + &self, + engine: &Worker, + _: &PickRequest<'_>, + _: Option<&EngineReportedWorkerLoad>, + ) -> Result { + self.calls.lock().unwrap().push(engine.id.clone()); + if self.invalid { + Err(PickError::InvalidSignal("admission input".into())) + } else if self.reject { + Ok(Decision::Reject("full".into())) + } else { + Ok(Decision::Allow) + } + } + } + + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let engine = Arc::new(Worker::new(spec("a", Stage::Plain, "m"))); + let other = Arc::new(Worker::new(spec("b", Stage::Plain, "m"))); + let _busy = other.load_guard(); + for engines in [vec![engine.clone()], vec![other.clone(), engine.clone()]] { + for (reject, invalid) in [(false, false), (true, false), (false, true)] { + let check = Arc::new(Check { + calls: Mutex::new(Vec::new()), + reject, + invalid, + }); + let mut policy = PowerOfTwoPolicy::new(EngineReportedLoadTable::new()); + policy.admission = check.clone(); + assert!(matches!( + policy.pick(&[], &request).await, + Err(PickError::NoCandidates) + )); + assert!(check.calls.lock().unwrap().is_empty()); + let result = policy.pick(&engines, &request).await; + if invalid { + assert!(matches!(result, Err(PickError::InvalidSignal(_)))); + } else if reject { + assert!(matches!(result, Err(PickError::AdmissionRejected(reason)) + if reason.engine == engine.id && reason.reason == "full")); + } else { + assert!(Arc::ptr_eq(&result.unwrap().engine, &engine)); + } + assert_eq!( + check.calls.lock().unwrap().as_slice(), + std::slice::from_ref(&engine.id) + ); + } + } +} diff --git a/experimental/sgl-router/tests/component/policies_reorg_load.rs b/experimental/sgl-router/tests/component/policies_reorg_load.rs new file mode 100644 index 000000000..e73cbb313 --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_load.rs @@ -0,0 +1,138 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies_reorg::admission::{Decision, EngineAdmission}; +use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; +use sgl_router::policies_reorg::{PickError, PickRequest, Policy, Stage}; +use sgl_router::state::load_monitor::engine_reported_load::{ + EngineReportedLoadTable, EngineReportedWorkerLoad, LoadStat, +}; +use sgl_router::workers::Worker; + +const URL: &str = "http://engine"; + +fn report(table: &EngineReportedLoadTable, rank: u32, running: u64, at: Instant) { + table.set( + URL, + rank, + LoadStat { + num_running_reqs: running, + num_waiting_reqs: 2, + num_tokens: 30, + max_total_num_tokens: 100, + native_cache: None, + }, + at, + ); +} + +fn engine() -> Arc { + Arc::new(Worker::new(WorkerSpec { + id: WorkerId("a".into()), + url: URL.into(), + mode: Stage::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })) +} + +#[derive(Debug)] +struct ObserveAdmission { + table: Arc, + observations: Mutex>>, +} + +impl EngineAdmission for ObserveAdmission { + fn check( + &self, + engine: &Worker, + _: &PickRequest<'_>, + load: Option<&EngineReportedWorkerLoad>, + ) -> Result { + assert_eq!(engine.url, URL); + // A new report arriving after selection must not change the observation + // supplied to admission. The next pick should read the new report. + report(&self.table, 0, 99, Instant::now()); + self.observations.lock().unwrap().push(load.cloned()); + Ok(Decision::Allow) + } +} + +#[tokio::test] +async fn selected_load_reaches_admission_and_next_pick_reads_fresh_state() { + let table = EngineReportedLoadTable::new(); + let first_at = Instant::now(); + report(&table, 0, 1, first_at); + report(&table, 1, 3, first_at); + table.mark_expected_rank(URL, 0); + table.mark_expected_rank(URL, 1); + let admission = Arc::new(ObserveAdmission { + table: table.clone(), + observations: Mutex::default(), + }); + let mut policy = PowerOfTwoPolicy::new(table); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + let alternative = Arc::new(Worker::new(WorkerSpec { + id: WorkerId("b".into()), + url: "http://other".into(), + mode: Stage::Plain, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })); + // These old-format reports lack native pressure metrics, so selection uses + // local active counts for both candidates and chooses the second engine. + let _busy = alternative.load_guard(); + let engines = [alternative, engine()]; + for _ in 0..2 { + let pick = policy.pick(&engines, &request).await.unwrap(); + assert!(Arc::ptr_eq(&pick.engine, &engines[1])); + } + let observations = admission.observations.lock().unwrap(); + assert_eq!(observations.len(), 2); + assert_eq!( + observations[0], + Some(EngineReportedWorkerLoad { + num_running_reqs: 4, + num_waiting_reqs: 4, + num_tokens: 60, + max_total_num_tokens: 200, + captured_at: first_at, + }) + ); + assert_eq!(observations[1].as_ref().unwrap().num_running_reqs, 102); +} + +#[tokio::test] +async fn missing_stale_and_incomplete_reports_reach_admission_as_unknown() { + for case in ["missing", "stale", "incomplete"] { + let table = EngineReportedLoadTable::new(); + match case { + "missing" => {} + "stale" => report(&table, 0, 1, Instant::now() - Duration::from_secs(3600)), + "incomplete" => { + report(&table, 0, 1, Instant::now()); + table.mark_expected_rank(URL, 0); + table.mark_expected_rank(URL, 1); + } + _ => unreachable!(), + } + let admission = Arc::new(ObserveAdmission { + table: table.clone(), + observations: Mutex::default(), + }); + let mut policy = PowerOfTwoPolicy::new(table); + policy.admission = admission.clone(); + let model = ModelId("m".into()); + policy + .pick(&[engine()], &PickRequest::new(&model, Stage::Plain, 10)) + .await + .unwrap(); + assert_eq!(*admission.observations.lock().unwrap(), [None], "{case}"); + } +} diff --git a/experimental/sgl-router/tests/component/policies_reorg_power_of_two.rs b/experimental/sgl-router/tests/component/policies_reorg_power_of_two.rs new file mode 100644 index 000000000..6b0a04db3 --- /dev/null +++ b/experimental/sgl-router/tests/component/policies_reorg_power_of_two.rs @@ -0,0 +1,203 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use sgl_router::discovery::{ModelId, WorkerId, WorkerSpec}; +use sgl_router::policies_reorg::power_of_two::PowerOfTwoPolicy; +use sgl_router::policies_reorg::{PickRequest, Policy, Stage}; +use sgl_router::state::load_monitor::engine_reported_load::{ + EngineReportedLoadTable, LoadStat, NativeCacheRankLoad, +}; +use sgl_router::workers::Worker; + +fn engine(id: &str, stage: Stage, active: usize) -> Arc { + let worker = Arc::new(Worker::new(WorkerSpec { + id: WorkerId(id.into()), + url: format!("http://{id}"), + mode: stage, + model_ids: vec![ModelId("m".into())], + bootstrap_port: None, + })); + worker.active_requests.store(active, Ordering::Relaxed); + worker +} + +fn load(running: u64, waiting: u64, tokens: u64, capacity: u64, pending: u64) -> LoadStat { + LoadStat { + num_running_reqs: running, + num_waiting_reqs: waiting, + num_tokens: tokens, + max_total_num_tokens: capacity, + native_cache: Some(NativeCacheRankLoad { + num_waiting_uncached_tokens: pending, + num_total_tokens: tokens, + max_running_requests: 100, + total_prefill_uncached_tokens: 0, + total_prefill_busy_us: 0, + }), + } +} + +async fn assert_winner( + policy: &PowerOfTwoPolicy, + engines: &[Arc], + stage: Stage, + expected: &Arc, +) { + let model = ModelId("m".into()); + let request = PickRequest::new(&model, stage, 10); + // With exactly two candidates the winner is independent of sample order. + for _ in 0..16 { + let pick = policy.pick(engines, &request).await.unwrap(); + assert!(Arc::ptr_eq(&pick.engine, expected), "stage: {stage:?}"); + assert_eq!(pick.reason, "power_of_two"); + } +} + +#[tokio::test] +async fn plain_and_prefill_use_pending_work_while_decode_uses_request_pressure() { + for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] { + let engines = [engine("a", stage, 100), engine("b", stage, 0)]; + let table = EngineReportedLoadTable::new(); + table.set(&engines[0].url, 0, load(5, 8, 10, 100, 1), Instant::now()); + table.set(&engines[1].url, 0, load(1, 1, 10, 100, 100), Instant::now()); + let expected = if stage == Stage::Decode { 1 } else { 0 }; + assert_winner( + &PowerOfTwoPolicy::new(table), + &engines, + stage, + &engines[expected], + ) + .await; + } +} + +#[tokio::test] +async fn prefill_uses_estimated_queue_time_only_when_both_engines_have_rates() { + for both_have_rates in [true, false] { + let engines = [ + engine("a", Stage::Prefill, 0), + engine("b", Stage::Prefill, 0), + ]; + let table = EngineReportedLoadTable::new(); + for (i, worker) in engines.iter().enumerate() { + let mut report = load(1, 1, 10, 100, if i == 0 { 10 } else { 20 }); + if i == 0 || both_have_rates { + table.set(&worker.url, 0, report.clone(), Instant::now()); + } + let native = report.native_cache.as_mut().unwrap(); + native.total_prefill_uncached_tokens = if i == 0 { 100 } else { 1000 }; + native.total_prefill_busy_us = 1_000_000; + table.set(&worker.url, 0, report, Instant::now()); + } + // B has more queued tokens, but its higher throughput gives a shorter + // estimated queue. Without B's rate, compare queued tokens for both. + let expected = usize::from(both_have_rates); + assert_winner( + &PowerOfTwoPolicy::new(table), + &engines, + Stage::Prefill, + &engines[expected], + ) + .await; + } +} + +#[tokio::test] +async fn decode_orders_by_waiting_running_kv_fraction_then_tokens() { + let cases = [ + (load(50, 1, 90, 100, 0), load(1, 2, 1, 100, 0)), + (load(1, 1, 90, 100, 0), load(2, 1, 1, 100, 0)), + (load(1, 1, 100, 1000, 0), load(1, 1, 20, 100, 0)), + (load(1, 1, 10, 100, 0), load(1, 1, 100, 1000, 0)), + ]; + for (left, right) in cases { + let engines = [ + engine("a", Stage::Decode, 100), + engine("b", Stage::Decode, 0), + ]; + let table = EngineReportedLoadTable::new(); + table.set(&engines[0].url, 0, left, Instant::now()); + table.set(&engines[1].url, 0, right, Instant::now()); + assert_winner( + &PowerOfTwoPolicy::new(table), + &engines, + Stage::Decode, + &engines[0], + ) + .await; + } +} + +#[tokio::test] +async fn unusable_telemetry_falls_back_to_local_load_for_both_candidates() { + for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] { + for case in [ + "missing", + "stale", + "incomplete", + "old_publisher", + "unknown_capacity", + ] { + let engines = [engine("a", stage, 1), engine("b", stage, 5)]; + let table = EngineReportedLoadTable::new(); + // A's high reported pressure must not be compared to B's local + // count or to a fabricated zero for its unavailable telemetry. + table.set( + &engines[0].url, + 0, + load(90, 90, 90, 100, 900), + Instant::now(), + ); + let mut right = load(0, 0, 0, 100, 0); + let mut at = Instant::now(); + match case { + "missing" => {} + "stale" => at -= Duration::from_secs(3600), + "incomplete" => { + table.mark_expected_rank(&engines[1].url, 0); + table.mark_expected_rank(&engines[1].url, 1); + } + "old_publisher" => right.native_cache = None, + "unknown_capacity" => right.max_total_num_tokens = 0, + _ => unreachable!(), + } + if case != "missing" { + table.set(&engines[1].url, 0, right, at); + } + assert_winner(&PowerOfTwoPolicy::new(table), &engines, stage, &engines[0]).await; + } + } +} + +#[tokio::test] +async fn equal_reported_pressure_uses_local_active_load_as_tiebreaker() { + for stage in [Stage::Plain, Stage::Prefill, Stage::Decode] { + let engines = [engine("a", stage, 5), engine("b", stage, 1)]; + let table = EngineReportedLoadTable::new(); + for worker in &engines { + table.set(&worker.url, 0, load(1, 1, 10, 100, 10), Instant::now()); + } + assert_winner(&PowerOfTwoPolicy::new(table), &engines, stage, &engines[1]).await; + } +} + +#[tokio::test] +async fn multiple_candidates_never_select_the_unique_busiest_engine() { + let engines: Vec<_> = (0..8) + .map(|i| engine(&i.to_string(), Stage::Plain, i)) + .collect(); + let policy = PowerOfTwoPolicy::new(EngineReportedLoadTable::new()); + let model = ModelId("m".into()); + let request = PickRequest::new(&model, Stage::Plain, 10); + for _ in 0..64 { + let pick = policy.pick(&engines, &request).await.unwrap(); + // Every distinct pair has an engine less busy than the last candidate. + assert!(engines[..7] + .iter() + .any(|engine| Arc::ptr_eq(engine, &pick.engine))); + } +} diff --git a/experimental/sgl-router/tests/proxy/chat_routing.rs b/experimental/sgl-router/tests/proxy/chat_routing.rs index 32f11a73a..28bf7c718 100644 --- a/experimental/sgl-router/tests/proxy/chat_routing.rs +++ b/experimental/sgl-router/tests/proxy/chat_routing.rs @@ -21,6 +21,8 @@ use std::sync::Arc; use std::time::Duration; use tower::ServiceExt; +mod reorg; + const TEST_TIMEOUT: Duration = Duration::from_secs(5); fn config_for(_worker_url: &str) -> Config { diff --git a/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs new file mode 100644 index 000000000..013ea6d6d --- /dev/null +++ b/experimental/sgl-router/tests/proxy/chat_routing/reorg.rs @@ -0,0 +1,523 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 The SGLang Authors +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::common::mock_worker::MockWorker; +use futures::future::BoxFuture; +use sgl_router::buckets_reorg::{Bucket, BucketGroups, BucketResolver, EngineGroup}; +use sgl_router::policies::PolicyRegistry; +use sgl_router::policies_reorg::admission::{AllowAll, Decision, EngineAdmission}; +use sgl_router::policies_reorg::{Pick, PickError, PickRequest, Policy, Rejection, Stage}; +use sgl_router::server::app_context::ChatRouting; +use sgl_router::state::load_monitor::engine_reported_load::EngineReportedWorkerLoad; +use std::sync::Mutex; + +type PickCall = (String, Stage, u64, Option); + +#[derive(Debug)] +struct FirstPolicy { + calls: Mutex>, + admission: Arc, + miss: bool, + invalid: bool, +} + +impl Default for FirstPolicy { + fn default() -> Self { + Self { + admission: Arc::new(AllowAll), + miss: false, + invalid: false, + calls: Mutex::new(Vec::new()), + } + } +} + +impl Policy for FirstPolicy { + fn pick<'a>( + &'a self, + engines: &'a [Arc], + request: &'a PickRequest<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + self.calls.lock().unwrap().push(( + request.bucket.to_owned(), + request.stage, + request.input_tokens, + request.expected_peak_tokens, + )); + if self.invalid { + return Err(PickError::InvalidSignal("invalid policy input".into())); + } + if self.miss { + return Err(PickError::NoCandidates); + } + if engines.is_empty() { + return Err(PickError::NoCandidates); + } + let engine = engines[0].clone(); + if let Decision::Reject(reason) = self.admission.check(&engine, request, None)? { + return Err(PickError::AdmissionRejected(Rejection { + engine: engine.id.clone(), + reason, + })); + } + Ok(Pick { + engine, + reason: "test", + }) + }) + } +} + +#[derive(Debug)] +struct RejectAll; + +impl EngineAdmission for RejectAll { + fn check( + &self, + _: &Worker, + _: &PickRequest<'_>, + _: Option<&EngineReportedWorkerLoad>, + ) -> Result { + Ok(Decision::Reject("full".into())) + } +} + +fn rejecting_policy() -> Arc { + Arc::new(FirstPolicy { + admission: Arc::new(RejectAll), + ..Default::default() + }) +} + +fn group(id: &str, policy: Arc) -> EngineGroup { + EngineGroup { + worker_ids: Some([WorkerId(id.into())].into_iter().collect()), + policy, + } +} + +fn context(workers: &[(&str, Stage, &MockWorker)], buckets: Vec) -> Arc { + let config = config_for(""); + let registry = Arc::new(WorkerRegistry::default()); + for &(id, mode, worker) in workers { + registry + .add(WorkerSpec { + id: WorkerId(id.into()), + url: worker.url.clone(), + mode, + model_ids: vec![ModelId("tiny".into())], + bootstrap_port: Some(8998), + }) + .unwrap(); + } + let tokenizers = Arc::new(TokenizerRegistry::load_from_config(&config).unwrap()); + let mut ctx = AppContext::new( + config, + tokenizers, + Arc::new(Proxy::new(TEST_TIMEOUT).unwrap()), + registry, + // The reorg route must not require the legacy policy registry. + Arc::new(PolicyRegistry::default()), + ); + ctx.chat_routing = ChatRouting::Reorg( + [(ModelId("tiny".into()), BucketResolver::new(buckets))] + .into_iter() + .collect(), + ); + Arc::new(ctx) +} + +fn request(value: serde_json::Value) -> Request { + Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(&value).unwrap())) + .unwrap() +} + +fn body(content: &str) -> serde_json::Value { + serde_json::json!({"model": "tiny", "messages": [{"role": "user", "content": content}]}) +} + +#[tokio::test] +async fn length_selects_plain_bucket_before_engine_selection() { + let short_worker = MockWorker::start(vec![]).await; + let long_worker = MockWorker::start(vec![]).await; + let policy = Arc::new(FirstPolicy::default()); + let mut short = Bucket::new("short", BucketGroups::Plain(group("short", policy.clone()))); + short.limits.max = Some(4); + let long = Bucket::new("long", BucketGroups::Plain(group("long", policy.clone()))); + let ctx = context( + &[ + ("short", Stage::Plain, &short_worker), + ("long", Stage::Plain, &long_worker), + ], + vec![long, short], + ); + let app = build_router(ctx); + let response = app.clone().oneshot(request(body("hi"))).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(short_worker.captured.lock().unwrap().last_body.is_some()); + assert!(long_worker.captured.lock().unwrap().last_body.is_none()); + + let response = app + .oneshot(request(body(&"hello ".repeat(30)))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(long_worker.captured.lock().unwrap().last_body.is_some()); + let calls = policy.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!((&*calls[0].0, calls[0].1), ("short", Stage::Plain)); + assert_eq!((&*calls[1].0, calls[1].1), ("long", Stage::Plain)); +} + +#[tokio::test] +async fn pd_picks_both_groups_from_selected_bucket_and_shares_bootstrap() { + let prefill = MockWorker::start(vec![]).await; + let decode = MockWorker::start(vec![]).await; + let policy = Arc::new(FirstPolicy::default()); + let mut selected = Bucket::new( + "selected", + BucketGroups::Pd { + prefill: group("p", policy.clone()), + decode: group("d", policy.clone()), + }, + ); + selected.limits.max = Some(100); + let other_policy = Arc::new(FirstPolicy::default()); + let other = Bucket::new( + "other", + BucketGroups::Pd { + prefill: group("p", other_policy.clone()), + decode: group("d", other_policy.clone()), + }, + ); + let ctx = context( + &[ + ("p", Stage::Prefill, &prefill), + ("d", Stage::Decode, &decode), + ], + vec![other, selected], + ); + let mut body = body("hello"); + body["max_completion_tokens"] = 10.into(); + let response = build_router(ctx).oneshot(request(body)).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(response.headers()["x-sgl-decode-url"], decode.url); + let _ = response.into_body().collect().await.unwrap(); + tokio::time::timeout(TEST_TIMEOUT, async { + while prefill.captured.lock().unwrap().last_body.is_none() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let p: serde_json::Value = + serde_json::from_slice(prefill.captured.lock().unwrap().last_body.as_ref().unwrap()) + .unwrap(); + let d: serde_json::Value = + serde_json::from_slice(decode.captured.lock().unwrap().last_body.as_ref().unwrap()) + .unwrap(); + assert!(p["bootstrap_room"].is_number()); + assert_eq!(p["bootstrap_room"], d["bootstrap_room"]); + let calls = policy.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!((&*calls[0].0, calls[0].1), ("selected", Stage::Prefill)); + assert_eq!((&*calls[1].0, calls[1].1), ("selected", Stage::Decode)); + assert_eq!(calls[0].3, Some(calls[0].2 + 10)); + assert_eq!(calls[1].3, calls[0].3); + assert!(other_policy.calls.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn missing_decode_in_all_buckets_does_not_dispatch_prefill() { + let prefill = MockWorker::start(vec![]).await; + let decode = MockWorker::start(vec![]).await; + let policy = Arc::new(FirstPolicy::default()); + let mut selected = Bucket::new( + "selected", + BucketGroups::Pd { + prefill: group("p", policy.clone()), + decode: group("missing", policy.clone()), + }, + ); + selected.limits.max = Some(100); + let other = Bucket::new( + "other", + BucketGroups::Pd { + prefill: group("p", policy.clone()), + decode: group("also-missing", policy.clone()), + }, + ); + let ctx = context( + &[ + ("p", Stage::Prefill, &prefill), + ("d", Stage::Decode, &decode), + ], + vec![other, selected], + ); + let response = build_router(ctx.clone()) + .oneshot(request(body("hi"))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response.headers()["x-router-error-code"], + "no_decode_workers_available" + ); + assert!(prefill.captured.lock().unwrap().last_body.is_none()); + assert!(decode.captured.lock().unwrap().last_body.is_none()); + assert_eq!(policy.calls.lock().unwrap().len(), 2); + assert_eq!(ctx.router_inflight_load.inflight_count(), 0); + assert_eq!( + ctx.registry + .get(&WorkerId("p".into())) + .unwrap() + .router_inflight_load(), + 0 + ); +} + +#[tokio::test] +async fn rejects_unsupported_length_unknown_model_and_overflow_before_policy() { + let worker = MockWorker::start(vec![]).await; + let policy = Arc::new(FirstPolicy::default()); + let mut bucket = Bucket::new("short", BucketGroups::Plain(group("w", policy.clone()))); + bucket.max_context_tokens = Some(4); + let app = build_router(context(&[("w", Stage::Plain, &worker)], vec![bucket])); + let mut long = body("hi"); + long["max_tokens"] = 100.into(); + let mut overflow = body("hi"); + overflow["max_tokens"] = u64::MAX.into(); + let mut unknown = body("hi"); + unknown["model"] = "unknown".into(); + for (body, status) in [ + (long, StatusCode::BAD_REQUEST), + (overflow, StatusCode::BAD_REQUEST), + (unknown, StatusCode::NOT_FOUND), + (serde_json::json!({}), StatusCode::BAD_REQUEST), + ] { + let response = app.clone().oneshot(request(body)).await.unwrap(); + assert_eq!(response.status(), status); + } + assert!(policy.calls.lock().unwrap().is_empty()); + assert!(worker.captured.lock().unwrap().last_body.is_none()); +} + +#[tokio::test] +async fn streaming_uses_existing_forwarder() { + let worker = MockWorker::start(vec!["data: {\"choices\":[]}\n\n", "data: [DONE]\n\n"]).await; + let policy = Arc::new(FirstPolicy::default()); + let bucket = Bucket::new("plain", BucketGroups::Plain(group("w", policy))); + let app = build_router(context(&[("w", Stage::Plain, &worker)], vec![bucket])); + let mut body = body("hi"); + body["stream"] = true.into(); + let response = app.oneshot(request(body)).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers()["content-type"] + .to_str() + .unwrap() + .starts_with("text/event-stream")); + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + assert!(String::from_utf8_lossy(&bytes).contains("data: [DONE]")); +} + +#[tokio::test] +async fn reorg_route_keeps_chat_body_limit() { + let app = build_router(context(&[], vec![])); + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from(vec![b' '; MAX_CHAT_BODY_BYTES + 1])) + .unwrap(); + assert_eq!( + app.oneshot(request).await.unwrap().status(), + StatusCode::PAYLOAD_TOO_LARGE + ); +} + +#[tokio::test] +async fn plain_fallback_skips_empty_missed_and_rejected_buckets_then_stops_on_success() { + let rejected_worker = MockWorker::start(vec![]).await; + let winner = MockWorker::start(vec![]).await; + let skipped = Arc::new(FirstPolicy::default()); + let rejected = rejecting_policy(); + let missed = Arc::new(FirstPolicy { + miss: true, + ..Default::default() + }); + let accepted = Arc::new(FirstPolicy::default()); + let buckets = vec![ + Bucket::new( + "a-empty", + BucketGroups::Plain(group("missing", skipped.clone())), + ), + Bucket::new( + "b-miss", + BucketGroups::Plain(group("rejected", missed.clone())), + ), + Bucket::new( + "c-rejected", + BucketGroups::Plain(group("rejected", rejected.clone())), + ), + Bucket::new( + "d-winner", + BucketGroups::Plain(group("winner", accepted.clone())), + ), + Bucket::new( + "e-unused", + BucketGroups::Plain(group("winner", skipped.clone())), + ), + ]; + let ctx = context( + &[ + ("rejected", Stage::Plain, &rejected_worker), + ("winner", Stage::Plain, &winner), + ], + buckets, + ); + let response = build_router(ctx) + .oneshot(request(body("hi"))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + assert!(rejected_worker.captured.lock().unwrap().last_body.is_none()); + assert!(winner.captured.lock().unwrap().last_body.is_some()); + assert!(skipped.calls.lock().unwrap().is_empty()); + assert_eq!(missed.calls.lock().unwrap().len(), 1); + assert_eq!(rejected.calls.lock().unwrap().len(), 1); + assert_eq!(accepted.calls.lock().unwrap()[0].0, "d-winner"); +} + +#[tokio::test] +async fn decode_failure_retries_both_groups_in_next_bucket_without_dispatching_first_prefill() { + // Cover an empty decode group and rejection of a selected decode engine. + for reject_decode in [false, true] { + let first_prefill = MockWorker::start(vec![]).await; + let second_prefill = MockWorker::start(vec![]).await; + let decode = MockWorker::start(vec![]).await; + let first = Arc::new(FirstPolicy::default()); + let rejected = rejecting_policy(); + let accepted = Arc::new(FirstPolicy::default()); + let buckets = vec![ + Bucket::new( + "a-first", + BucketGroups::Pd { + prefill: group("p1", first.clone()), + decode: group(if reject_decode { "d" } else { "missing" }, rejected), + }, + ), + Bucket::new( + "b-second", + BucketGroups::Pd { + prefill: group("p2", accepted.clone()), + decode: group("d", accepted.clone()), + }, + ), + ]; + let ctx = context( + &[ + ("p1", Stage::Prefill, &first_prefill), + ("p2", Stage::Prefill, &second_prefill), + ("d", Stage::Decode, &decode), + ], + buckets, + ); + let response = build_router(ctx.clone()) + .oneshot(request(body("hi"))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = response.into_body().collect().await.unwrap(); + tokio::time::timeout(TEST_TIMEOUT, async { + while second_prefill.captured.lock().unwrap().last_body.is_none() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert!(first_prefill.captured.lock().unwrap().last_body.is_none()); + assert!(decode.captured.lock().unwrap().last_body.is_some()); + assert_eq!( + ctx.registry + .get(&WorkerId("p1".into())) + .unwrap() + .router_inflight_load(), + 0 + ); + assert_eq!(first.calls.lock().unwrap().len(), 1); + let calls = accepted.calls.lock().unwrap(); + assert_eq!(calls.len(), 2); + assert_eq!((&*calls[0].0, calls[0].1), ("b-second", Stage::Prefill)); + assert_eq!((&*calls[1].0, calls[1].1), ("b-second", Stage::Decode)); + } +} + +#[tokio::test] +async fn admission_exhaustion_is_preserved_when_later_buckets_are_empty() { + let worker = MockWorker::start(vec![]).await; + let first = rejecting_policy(); + let second = rejecting_policy(); + let empty = Arc::new(FirstPolicy::default()); + let ctx = context( + &[("w", Stage::Plain, &worker)], + vec![ + Bucket::new("a-first", BucketGroups::Plain(group("w", first.clone()))), + Bucket::new("b-second", BucketGroups::Plain(group("w", second.clone()))), + Bucket::new( + "c-empty", + BucketGroups::Plain(group("missing", empty.clone())), + ), + ], + ); + let response = build_router(ctx.clone()) + .oneshot(request(body("hi"))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!( + response.headers()["x-router-error-code"], + "policy_selection_failed" + ); + assert_eq!(first.calls.lock().unwrap().len(), 1); + assert_eq!(second.calls.lock().unwrap().len(), 1); + assert!(empty.calls.lock().unwrap().is_empty()); + assert!(worker.captured.lock().unwrap().last_body.is_none()); + assert_eq!(ctx.router_inflight_load.inflight_count(), 0); +} + +#[tokio::test] +async fn invalid_policy_signal_stops_bucket_iteration() { + let worker = MockWorker::start(vec![]).await; + let invalid = Arc::new(FirstPolicy { + invalid: true, + ..Default::default() + }); + let later = Arc::new(FirstPolicy::default()); + let ctx = context( + &[("w", Stage::Plain, &worker)], + vec![ + Bucket::new( + "a-invalid", + BucketGroups::Plain(group("w", invalid.clone())), + ), + Bucket::new("b-later", BucketGroups::Plain(group("w", later.clone()))), + ], + ); + let response = build_router(ctx) + .oneshot(request(body("hi"))) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR); + assert_eq!(invalid.calls.lock().unwrap().len(), 1); + assert!(later.calls.lock().unwrap().is_empty()); + assert!(worker.captured.lock().unwrap().last_body.is_none()); +}