[sgl-router] refactor - layout BucketResolver, Bucket, EngineGroup and implement PowerOfTwo (#40241)
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
co-authored by
Claude Fable 5.1
Claude Fable 5.1
parent
4a9dc5c4af
commit
aedda8377e
@@ -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<ModelId, BucketResolver>)` 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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>>;
|
||||
}
|
||||
```
|
||||
|
||||
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<dyn EngineAdmission>`. 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<EngineReportedLoadTable>)`
|
||||
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<EngineReportedLoadTable>` 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.
|
||||
@@ -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<u64>,
|
||||
pub max: Option<u64>,
|
||||
}
|
||||
|
||||
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<HashSet<WorkerId>>,
|
||||
pub policy: Arc<dyn Policy>,
|
||||
}
|
||||
|
||||
impl EngineGroup {
|
||||
pub fn new(policy: Arc<dyn Policy>) -> 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<Pick, PickError> {
|
||||
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<u64>,
|
||||
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<Pick>,
|
||||
}
|
||||
|
||||
#[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<u64>,
|
||||
pub groups: BucketGroups,
|
||||
}
|
||||
|
||||
impl Bucket {
|
||||
pub fn new(id: impl Into<String>, 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<BucketPick, (Stage, PickError)> {
|
||||
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<Pick, (Stage, PickError)> {
|
||||
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<u64>) -> 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<Bucket>,
|
||||
}
|
||||
|
||||
impl BucketResolver {
|
||||
pub fn new(buckets: Vec<Bucket>) -> 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<u64>,
|
||||
) -> Result<Vec<&Bucket>, 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)
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<Decision, PickError>;
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AllowAll;
|
||||
|
||||
impl EngineAdmission for AllowAll {
|
||||
fn check(
|
||||
&self,
|
||||
_: &Worker,
|
||||
_: &PickRequest<'_>,
|
||||
_: Option<&EngineReportedWorkerLoad>,
|
||||
) -> Result<Decision, PickError> {
|
||||
Ok(Decision::Allow)
|
||||
}
|
||||
}
|
||||
@@ -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<u64>,
|
||||
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<Worker>,
|
||||
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<Rejection>),
|
||||
#[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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>>;
|
||||
|
||||
/// 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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>> {
|
||||
match self.fallback() {
|
||||
Some(fallback) => fallback.pick(engines, request),
|
||||
None => Box::pin(async { Err(PickError::NoCandidates) }),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<EngineReportedLoadTable>,
|
||||
pub admission: Arc<dyn EngineAdmission>,
|
||||
}
|
||||
|
||||
impl PowerOfTwoPolicy {
|
||||
pub fn new(engine_load: Arc<EngineReportedLoadTable>) -> Self {
|
||||
Self {
|
||||
engine_load,
|
||||
admission: Arc::new(AllowAll),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Policy for PowerOfTwoPolicy {
|
||||
fn pick<'a>(
|
||||
&'a self,
|
||||
engines: &'a [Arc<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>> {
|
||||
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",
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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<ModelId, BucketResolver>),
|
||||
}
|
||||
|
||||
pub struct AppContext {
|
||||
pub config: Config,
|
||||
pub tokenizers: Arc<TokenizerRegistry>,
|
||||
@@ -34,6 +46,8 @@ pub struct AppContext {
|
||||
pub policies: Arc<PolicyRegistry>,
|
||||
/// Converts static Bucket configuration into request candidate domains.
|
||||
pub bucket_selector: Arc<BucketSelector>,
|
||||
/// 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<RouterInflightLoadRegistry>,
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Arc<AppContext>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Result<Response<Body>, 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<Response<Body>, 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 {
|
||||
|
||||
@@ -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<ModelId, BucketResolver>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Result<Response<Body>, 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<Vec<_>> = 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<Stage>) -> 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()),
|
||||
}
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<dyn EngineAdmission>,
|
||||
result: Option<Arc<Worker>>,
|
||||
miss: bool,
|
||||
invalid: bool,
|
||||
calls: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>> {
|
||||
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<Decision, PickError> {
|
||||
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<WorkerRegistry> {
|
||||
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<dyn Policy>) -> EngineGroup {
|
||||
EngineGroup {
|
||||
worker_ids: Some(members.iter().map(|id| WorkerId((*id).into())).collect()),
|
||||
policy,
|
||||
}
|
||||
}
|
||||
|
||||
fn bucket(id: &str, max: Option<u64>, policy: Arc<dyn Policy>) -> 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::<Vec<_>>(),
|
||||
["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::<Vec<_>>(),
|
||||
["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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>> {
|
||||
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<Vec<WorkerId>>,
|
||||
reject: bool,
|
||||
invalid: bool,
|
||||
}
|
||||
|
||||
impl EngineAdmission for Check {
|
||||
fn check(
|
||||
&self,
|
||||
engine: &Worker,
|
||||
_: &PickRequest<'_>,
|
||||
_: Option<&EngineReportedWorkerLoad>,
|
||||
) -> Result<Decision, PickError> {
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<Worker> {
|
||||
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<EngineReportedLoadTable>,
|
||||
observations: Mutex<Vec<Option<EngineReportedWorkerLoad>>>,
|
||||
}
|
||||
|
||||
impl EngineAdmission for ObserveAdmission {
|
||||
fn check(
|
||||
&self,
|
||||
engine: &Worker,
|
||||
_: &PickRequest<'_>,
|
||||
load: Option<&EngineReportedWorkerLoad>,
|
||||
) -> Result<Decision, PickError> {
|
||||
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}");
|
||||
}
|
||||
}
|
||||
@@ -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<Worker> {
|
||||
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<Worker>],
|
||||
stage: Stage,
|
||||
expected: &Arc<Worker>,
|
||||
) {
|
||||
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)));
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<u64>);
|
||||
|
||||
#[derive(Debug)]
|
||||
struct FirstPolicy {
|
||||
calls: Mutex<Vec<PickCall>>,
|
||||
admission: Arc<dyn EngineAdmission>,
|
||||
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<Worker>],
|
||||
request: &'a PickRequest<'a>,
|
||||
) -> BoxFuture<'a, Result<Pick, PickError>> {
|
||||
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<Decision, PickError> {
|
||||
Ok(Decision::Reject("full".into()))
|
||||
}
|
||||
}
|
||||
|
||||
fn rejecting_policy() -> Arc<FirstPolicy> {
|
||||
Arc::new(FirstPolicy {
|
||||
admission: Arc::new(RejectAll),
|
||||
..Default::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn group(id: &str, policy: Arc<FirstPolicy>) -> EngineGroup {
|
||||
EngineGroup {
|
||||
worker_ids: Some([WorkerId(id.into())].into_iter().collect()),
|
||||
policy,
|
||||
}
|
||||
}
|
||||
|
||||
fn context(workers: &[(&str, Stage, &MockWorker)], buckets: Vec<Bucket>) -> Arc<AppContext> {
|
||||
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<Body> {
|
||||
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());
|
||||
}
|
||||
Reference in New Issue
Block a user