[Router] Add --worker-queue-limit: stop sending cache-affinity traffic to a queueing worker (#39168)
Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Kangyan Zhou
Claude Opus 5
parent
3f8eb35ead
commit
c1f5b4736a
@@ -163,6 +163,16 @@ pub struct Cli {
|
||||
/// Maximum uncached-work difference that pressure may override.
|
||||
#[arg(long)]
|
||||
pub cache_switch_margin_tokens: Option<u64>,
|
||||
/// Queue gate for cache affinity: a worker whose engine reports at least
|
||||
/// this many waiting (queued) requests cannot win a selection on cache
|
||||
/// affinity. The request goes to another worker holding the same prefix,
|
||||
/// or failing that to the least-loaded worker that is not queueing; when
|
||||
/// every worker is queueing the least-loaded worker overall keeps the
|
||||
/// fleet routable. Unset disables the gate. Requires
|
||||
/// `--policy cache_aware`; scale with the engine's `--dp-size` because
|
||||
/// the published queue sums across a worker's DP ranks.
|
||||
#[arg(long)]
|
||||
pub worker_queue_limit: Option<u64>,
|
||||
|
||||
// ---- score composition ----
|
||||
/// Policies to sum, spelled exactly as `--policy` spells them and each
|
||||
@@ -349,7 +359,13 @@ impl Cli {
|
||||
|| self.cache_candidate_min_workers.is_some()
|
||||
|| self.cache_candidate_ratio.is_some()
|
||||
|| self.cache_candidate_max_workers.is_some()
|
||||
|| self.cache_switch_margin_tokens.is_some();
|
||||
|| self.cache_switch_margin_tokens.is_some()
|
||||
|| self.worker_queue_limit.is_some();
|
||||
// Value checks before the policy check: a value that is wrong under
|
||||
// every policy should say so, rather than pointing at --policy.
|
||||
if self.worker_queue_limit == Some(0) {
|
||||
return Err(anyhow!("--worker-queue-limit must be at least 1"));
|
||||
}
|
||||
if tuned_cache_candidates && self.policy != PolicyKind::CacheAware {
|
||||
return Err(anyhow!(
|
||||
"cache candidate tuning flags require --policy cache_aware"
|
||||
@@ -573,6 +589,7 @@ impl Cli {
|
||||
cache_switch_margin_tokens: self
|
||||
.cache_switch_margin_tokens
|
||||
.unwrap_or(d.cache_switch_margin_tokens),
|
||||
worker_queue_limit: self.worker_queue_limit.or(d.worker_queue_limit),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
@@ -1817,6 +1834,58 @@ mod tests {
|
||||
assert!(err.contains("--stable-pair"), "got: {err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn worker_queue_limit_requires_cache_aware_and_a_positive_value() {
|
||||
let config = cfg_of(
|
||||
"--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \
|
||||
--worker-queue-limit 4",
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
config
|
||||
.model
|
||||
.affinity
|
||||
.expect("cache-aware needs affinity config")
|
||||
.worker_queue_limit,
|
||||
Some(4)
|
||||
);
|
||||
|
||||
// Unset, the gate is disabled.
|
||||
let defaults =
|
||||
cfg_of("--policy cache_aware --kv-indexer-endpoint http://indexer:50051").unwrap();
|
||||
assert_eq!(
|
||||
defaults
|
||||
.model
|
||||
.affinity
|
||||
.expect("default affinity config")
|
||||
.worker_queue_limit,
|
||||
None
|
||||
);
|
||||
|
||||
let err = cfg_of("--policy power_of_two --worker-queue-limit 4")
|
||||
.expect_err("the gate only governs cache-affinity selection")
|
||||
.to_string();
|
||||
assert!(
|
||||
err.contains("cache candidate tuning flags require --policy cache_aware"),
|
||||
"got: {err}"
|
||||
);
|
||||
|
||||
let err = cfg_of(
|
||||
"--policy cache_aware --kv-indexer-endpoint http://indexer:50051 \
|
||||
--worker-queue-limit 0",
|
||||
)
|
||||
.expect_err("a zero limit would reject every queue reading")
|
||||
.to_string();
|
||||
assert!(err.contains("--worker-queue-limit"), "got: {err}");
|
||||
|
||||
// A zero limit is wrong under every policy, so the value error must
|
||||
// win over the policy error rather than being masked by it.
|
||||
let err = cfg_of("--policy power_of_two --worker-queue-limit 0")
|
||||
.expect_err("a zero limit is rejected regardless of policy")
|
||||
.to_string();
|
||||
assert!(err.contains("--worker-queue-limit"), "got: {err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_candidate_cli_rejects_invalid_bounds() {
|
||||
for (args, expected) in [
|
||||
|
||||
@@ -453,6 +453,29 @@ pub struct AffinityConfig {
|
||||
pub cache_candidate_ratio: f64,
|
||||
pub cache_candidate_max_workers: usize,
|
||||
pub cache_switch_margin_tokens: u64,
|
||||
/// Queue gate (`--worker-queue-limit`): a worker whose engine reports at
|
||||
/// least this many *waiting* requests cannot win a selection on cache
|
||||
/// affinity — the request goes to another worker holding the same
|
||||
/// prefix, or failing that to the least-loaded worker that is not
|
||||
/// queueing. `None` disables the gate.
|
||||
///
|
||||
/// Gating on the queue rather than on total depth is what makes this
|
||||
/// targeted: `num_waiting_reqs` IS the question the request cares about
|
||||
/// — will I sit behind other work before my prefill starts — whereas
|
||||
/// depth only proxies it, and proxies it badly (an engine can queue at
|
||||
/// 7-8 running on long-prompt traffic, far below its running cap, so a
|
||||
/// depth threshold either fires on healthy busy workers or misses the
|
||||
/// workers actually making requests wait).
|
||||
///
|
||||
/// The gate reads the engine-published load sample and fails OPEN on a
|
||||
/// worker with no fresh sample: the router-side in-flight counter cannot
|
||||
/// separate a running request from a waiting one, so there is no honest
|
||||
/// substitute to compare the limit against.
|
||||
///
|
||||
/// Note the firing point scales with `dp_size`: the sample sums `waiting`
|
||||
/// across a worker's DP ranks while a request lands on one of them, so
|
||||
/// scale the limit with `--dp-size` on DP-attention deployments.
|
||||
pub worker_queue_limit: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for AffinityConfig {
|
||||
@@ -475,6 +498,7 @@ impl Default for AffinityConfig {
|
||||
cache_candidate_ratio: 0.05,
|
||||
cache_candidate_max_workers: 32,
|
||||
cache_switch_margin_tokens: 1_024,
|
||||
worker_queue_limit: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,8 +6,13 @@
|
||||
//! Native Cache-Aware uses monitor-backed admission only when every expected
|
||||
//! DP rank has a fresh, complete #34608 ZMQ sample. Otherwise it falls back to
|
||||
//! Router-local load.
|
||||
//!
|
||||
//! The optional queue gate (`--worker-queue-limit`) is the one criterion here
|
||||
//! that reads [`EngineWorkerLoad::num_waiting_reqs`] rather than the native
|
||||
//! monitor fields: a worker already making requests wait cannot win on cache
|
||||
//! affinity. It fails open on a missing sample — see [`queue_gate_admits`].
|
||||
|
||||
use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
|
||||
use crate::policies::engine_load::{EngineLoadSnapshot, EngineWorkerLoad, NativeCacheWorkerLoad};
|
||||
use crate::policies::power_of_two::select_with_snapshot;
|
||||
use crate::policies::{CacheCandidate, CacheCandidateProposal, GuardHints, SelectionProposal};
|
||||
use crate::workers::Worker;
|
||||
@@ -124,32 +129,125 @@ pub struct FinalDecision {
|
||||
pub struct CacheCandidateResolution {
|
||||
pub decision: Option<FinalDecision>,
|
||||
pub prefill_pressure_source: &'static str,
|
||||
/// Candidates actually put through KV-capacity / pending-prefill
|
||||
/// admission. Queue-gate rejections are excluded: the gate runs first and
|
||||
/// they are never evaluated, so counting them here would silently deflate
|
||||
/// the rejected/evaluated ratio whenever the gate is armed.
|
||||
pub admission_evaluated_candidates: u64,
|
||||
/// Candidates rejected by KV-capacity / pending-prefill admission. Does
|
||||
/// not include queue-gate rejections — those are counted separately so a
|
||||
/// busy fleet never reads as a capacity problem.
|
||||
pub admission_rejected_candidates: u64,
|
||||
/// Candidates rejected by the queue gate: their engine queue is at or
|
||||
/// over `--worker-queue-limit`. The gate runs BEFORE hard admission, so
|
||||
/// these candidates were never evaluated against capacity and are
|
||||
/// disjoint from `admission_rejected_candidates` by construction.
|
||||
pub queue_gate_rejected_candidates: u64,
|
||||
/// Deepest matched prefix, in blocks, among the queue-gate-rejected
|
||||
/// candidates — the locality a diversion gave up. 0 when the gate
|
||||
/// rejected nothing.
|
||||
pub queue_gate_best_rejected_blocks: u32,
|
||||
/// True when the gate removed EVERY candidate and nowhere in the fleet is
|
||||
/// unqueued, so the winner came from the ungated candidate set instead.
|
||||
/// The mirror of `range_fallback`'s second tier: diversion buys nothing
|
||||
/// here, so discarding the prefix would be pure loss.
|
||||
pub queue_gate_fell_back: bool,
|
||||
/// True when no worker in the fleet has room before the gate. Computed
|
||||
/// once here so the route handler does not re-derive the gate over the
|
||||
/// fleet; always false when the gate is disabled.
|
||||
pub fleet_all_queued: bool,
|
||||
pub pressure_guard_compared_pairs: u64,
|
||||
pub pressure_guard_overrides: u64,
|
||||
}
|
||||
|
||||
/// The queue gate, in one place. Both subtle decisions live here: the
|
||||
/// boundary is `<` (a worker AT the limit is already making this request
|
||||
/// wait), and an unknown queue ADMITS. The gate reads
|
||||
/// [`EngineWorkerLoad::num_waiting_reqs`] because that is what the request
|
||||
/// cares about, and the router-side in-flight counter cannot separate a
|
||||
/// running request from a waiting one — so there is no honest substitute,
|
||||
/// and the gate fails open rather than comparing the limit against a
|
||||
/// different quantity.
|
||||
pub(crate) fn queue_gate_admits(
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
worker: &Worker,
|
||||
limit: Option<u64>,
|
||||
) -> bool {
|
||||
let Some(limit) = limit else {
|
||||
return true;
|
||||
};
|
||||
snapshot
|
||||
.fresh_load_for_url(&worker.url)
|
||||
.is_none_or(|load| load.num_waiting_reqs < limit)
|
||||
}
|
||||
|
||||
/// True when the gate has provably nowhere unqueued to divert to: every
|
||||
/// worker in `fleet` has a fresh sample at or over the limit. An unset limit,
|
||||
/// an empty fleet, or a single worker with no fresh sample all make this
|
||||
/// false — an unknown queue is not a proven full one.
|
||||
pub(crate) fn fleet_is_all_queued(
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
fleet: &[Arc<Worker>],
|
||||
limit: Option<u64>,
|
||||
) -> bool {
|
||||
limit.is_some()
|
||||
&& !fleet.is_empty()
|
||||
&& fleet
|
||||
.iter()
|
||||
.all(|worker| !queue_gate_admits(snapshot, worker, limit))
|
||||
}
|
||||
|
||||
/// Selects a worker from bounded cache candidates and records guard coverage.
|
||||
pub fn resolve_cache_candidates(
|
||||
proposal: &CacheCandidateProposal,
|
||||
request_input_tokens: u64,
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
fleet: &[Arc<Worker>],
|
||||
) -> CacheCandidateResolution {
|
||||
let queue_limit = proposal.worker_queue_limit;
|
||||
let fleet_all_queued = fleet_is_all_queued(snapshot, fleet, queue_limit);
|
||||
// The queue gate runs before hard admission so a busy worker never
|
||||
// pollutes the capacity-rejection counters, and so the diverted-overlap
|
||||
// audit below sees exactly the candidates the gate removed. One pass:
|
||||
// the gate answer per candidate is what splits the set, so asking twice
|
||||
// would re-read the snapshot for every candidate on every request.
|
||||
let mut evaluated: Vec<&CacheCandidate> = Vec::with_capacity(proposal.candidates.len());
|
||||
let mut queue_gate_rejected_candidates = 0u64;
|
||||
let mut queue_gate_best_rejected_blocks = 0u32;
|
||||
for candidate in &proposal.candidates {
|
||||
if queue_gate_admits(snapshot, &candidate.worker, queue_limit) {
|
||||
evaluated.push(candidate);
|
||||
} else {
|
||||
queue_gate_rejected_candidates += 1;
|
||||
queue_gate_best_rejected_blocks =
|
||||
queue_gate_best_rejected_blocks.max(candidate.matched_prefix_blocks);
|
||||
}
|
||||
}
|
||||
// Second tier, mirroring `range_fallback`: when the gate removed every
|
||||
// candidate AND nowhere in the fleet is unqueued, diversion cannot dodge
|
||||
// a wait, so returning no decision would trade the whole prefix for
|
||||
// nothing. Re-admit the ungated set. While an unqueued worker still
|
||||
// exists the gate keeps its teeth and the request leaves the prefix.
|
||||
let queue_gate_fell_back =
|
||||
evaluated.is_empty() && queue_gate_rejected_candidates > 0 && fleet_all_queued;
|
||||
if queue_gate_fell_back {
|
||||
evaluated.extend(proposal.candidates.iter());
|
||||
}
|
||||
// Built over the candidates that actually reach admission: a gated-out
|
||||
// candidate missing native monitor data would otherwise break the
|
||||
// lookup's full-coverage check and silently downgrade every pressure
|
||||
// comparison (and `prefill_pressure_source`) to router-local.
|
||||
let loads = FreshLoadLookup::new(
|
||||
Some(snapshot),
|
||||
proposal
|
||||
.candidates
|
||||
.iter()
|
||||
.map(|candidate| &candidate.worker),
|
||||
evaluated.iter().copied().map(|candidate| &candidate.worker),
|
||||
);
|
||||
let admitted: Vec<&CacheCandidate> = proposal
|
||||
.candidates
|
||||
.iter()
|
||||
let admission_evaluated_candidates = evaluated.len() as u64;
|
||||
let admitted: Vec<&CacheCandidate> = evaluated
|
||||
.into_iter()
|
||||
.filter(|candidate| is_cache_candidate_admitted(candidate, request_input_tokens, &loads))
|
||||
.collect();
|
||||
let admission_rejected_candidates =
|
||||
proposal.candidates.len().saturating_sub(admitted.len()) as u64;
|
||||
admission_evaluated_candidates.saturating_sub(admitted.len() as u64);
|
||||
let Some(work_floor) = admitted
|
||||
.iter()
|
||||
.copied()
|
||||
@@ -158,8 +256,12 @@ pub fn resolve_cache_candidates(
|
||||
return CacheCandidateResolution {
|
||||
decision: None,
|
||||
prefill_pressure_source: loads.prefill_pressure_source(),
|
||||
admission_evaluated_candidates: proposal.candidates.len() as u64,
|
||||
admission_evaluated_candidates,
|
||||
admission_rejected_candidates,
|
||||
queue_gate_rejected_candidates,
|
||||
queue_gate_best_rejected_blocks,
|
||||
queue_gate_fell_back,
|
||||
fleet_all_queued,
|
||||
pressure_guard_compared_pairs: 0,
|
||||
pressure_guard_overrides: 0,
|
||||
};
|
||||
@@ -201,8 +303,12 @@ pub fn resolve_cache_candidates(
|
||||
load_snapshot_version: snapshot.version,
|
||||
}),
|
||||
prefill_pressure_source: loads.prefill_pressure_source(),
|
||||
admission_evaluated_candidates: proposal.candidates.len() as u64,
|
||||
admission_evaluated_candidates,
|
||||
admission_rejected_candidates,
|
||||
queue_gate_rejected_candidates,
|
||||
queue_gate_best_rejected_blocks,
|
||||
queue_gate_fell_back,
|
||||
fleet_all_queued,
|
||||
pressure_guard_compared_pairs,
|
||||
pressure_guard_overrides,
|
||||
}
|
||||
@@ -213,27 +319,30 @@ pub fn resolve_prefill(
|
||||
proposal: &SelectionProposal,
|
||||
request_input_tokens: u64,
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
queue_limit: Option<u64>,
|
||||
) -> Option<FinalDecision> {
|
||||
resolve_prefill_admitted(range, proposal, request_input_tokens, snapshot).or_else(|| {
|
||||
if !contains_worker(range, &proposal.primary) {
|
||||
return None;
|
||||
}
|
||||
let backup = proposal
|
||||
.backup
|
||||
.as_ref()
|
||||
.filter(|worker| contains_worker(range, worker))
|
||||
.cloned();
|
||||
let legal = legal_prefill_candidates(range, proposal);
|
||||
let selected = select_with_snapshot(&legal, Some(snapshot))?;
|
||||
Some(FinalDecision {
|
||||
selected,
|
||||
primary: Arc::clone(&proposal.primary),
|
||||
backup,
|
||||
reason: DecisionReason::CapacityFallbackPowerOfTwo,
|
||||
candidate_range_id: range.id.to_string(),
|
||||
load_snapshot_version: snapshot.version,
|
||||
})
|
||||
})
|
||||
resolve_prefill_admitted(range, proposal, request_input_tokens, snapshot, queue_limit).or_else(
|
||||
|| {
|
||||
if !contains_worker(range, &proposal.primary) {
|
||||
return None;
|
||||
}
|
||||
let backup = proposal
|
||||
.backup
|
||||
.as_ref()
|
||||
.filter(|worker| contains_worker(range, worker))
|
||||
.cloned();
|
||||
let legal = legal_prefill_candidates(range, proposal);
|
||||
let selected = select_with_snapshot(&legal, Some(snapshot))?;
|
||||
Some(FinalDecision {
|
||||
selected,
|
||||
primary: Arc::clone(&proposal.primary),
|
||||
backup,
|
||||
reason: DecisionReason::CapacityFallbackPowerOfTwo,
|
||||
candidate_range_id: range.id.to_string(),
|
||||
load_snapshot_version: snapshot.version,
|
||||
})
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
/// Resolves prefill admission without overcommitting a full candidate range.
|
||||
@@ -242,6 +351,7 @@ pub fn resolve_prefill_admitted(
|
||||
proposal: &SelectionProposal,
|
||||
request_input_tokens: u64,
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
queue_limit: Option<u64>,
|
||||
) -> Option<FinalDecision> {
|
||||
if !contains_worker(range, &proposal.primary) {
|
||||
return None;
|
||||
@@ -251,11 +361,20 @@ pub fn resolve_prefill_admitted(
|
||||
.as_ref()
|
||||
.filter(|worker| contains_worker(range, worker))
|
||||
.cloned();
|
||||
// The queue gate demotes an admitted primary or backup exactly as it
|
||||
// demotes a cache candidate: a worker already making requests wait must
|
||||
// not win on proposal position alone, or the fallback the gate diverted
|
||||
// to hands the request straight back to it. Demotion is not rejection —
|
||||
// `range_fallback` below is two-tier, and its second tier returns the
|
||||
// least-pressured admitted worker (the demoted one included) when every
|
||||
// admitted worker is queueing, so an all-queueing fleet still routes.
|
||||
let primary_admitted = is_proposal_worker_eligible(proposal, &proposal.primary)
|
||||
&& is_prefill_admitted(range, &proposal.primary, request_input_tokens, snapshot);
|
||||
&& is_prefill_admitted(range, &proposal.primary, request_input_tokens, snapshot)
|
||||
&& queue_gate_admits(snapshot, &proposal.primary, queue_limit);
|
||||
let backup_admitted = backup.as_ref().is_some_and(|worker| {
|
||||
is_proposal_worker_eligible(proposal, worker)
|
||||
&& is_prefill_admitted(range, worker, request_input_tokens, snapshot)
|
||||
&& queue_gate_admits(snapshot, worker, queue_limit)
|
||||
});
|
||||
|
||||
let (selected, reason) = match (primary_admitted, backup.as_ref(), backup_admitted) {
|
||||
@@ -275,7 +394,7 @@ pub fn resolve_prefill_admitted(
|
||||
(false, Some(backup), true) => (Arc::clone(backup), DecisionReason::BackupPrimaryAdmission),
|
||||
_ => {
|
||||
let legal = legal_prefill_candidates(range, proposal);
|
||||
range_fallback(range, &legal, request_input_tokens, snapshot)?
|
||||
range_fallback(range, &legal, request_input_tokens, snapshot, queue_limit)?
|
||||
}
|
||||
};
|
||||
Some(FinalDecision {
|
||||
@@ -493,7 +612,7 @@ fn materially_more_pressured(
|
||||
/// candidate sets use Router-local active load to preserve ordering.
|
||||
pub(crate) struct FreshLoadLookup<'a> {
|
||||
by_worker_id: HashMap<String, &'a NativeCacheWorkerLoad>,
|
||||
basic_by_worker_id: HashMap<String, &'a crate::policies::engine_load::EngineWorkerLoad>,
|
||||
basic_by_worker_id: HashMap<String, &'a EngineWorkerLoad>,
|
||||
local_active_by_worker_id: HashMap<String, usize>,
|
||||
compare_engine: bool,
|
||||
compare_basic_engine: bool,
|
||||
@@ -662,6 +781,7 @@ fn range_fallback(
|
||||
legal: &[Arc<Worker>],
|
||||
request_input_tokens: u64,
|
||||
snapshot: &EngineLoadSnapshot,
|
||||
queue_limit: Option<u64>,
|
||||
) -> Option<(Arc<Worker>, DecisionReason)> {
|
||||
let admitted = legal
|
||||
.iter()
|
||||
@@ -669,9 +789,38 @@ fn range_fallback(
|
||||
.filter(|worker| is_prefill_admitted(range, worker, request_input_tokens, snapshot))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
let loads = FreshLoadLookup::new(Some(snapshot), admitted.iter());
|
||||
// Two-tier under the queue gate: least-pressured worker that is not
|
||||
// queueing, and only when every admitted worker is queueing,
|
||||
// least-pressured overall. Both tiers are load-bearing. Pressure ranks by
|
||||
// queue tokens / depth while the gate reads queue length, and those
|
||||
// disagree exactly where the gate earns its keep — a shallow worker with
|
||||
// a backlog is the fleet minimum BY PRESSURE, so a single-tier fallback
|
||||
// hands the request straight back to the cache home the gate just
|
||||
// rejected. The second tier is what keeps an all-queueing fleet routable
|
||||
// instead of failing every request.
|
||||
let pool = match queue_limit {
|
||||
// Gate disabled: the tiers coincide, so filtering would only clone
|
||||
// the vector.
|
||||
None => admitted,
|
||||
Some(_) => {
|
||||
let unqueued = admitted
|
||||
.iter()
|
||||
.filter(|worker| queue_gate_admits(snapshot, worker, queue_limit))
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
if unqueued.is_empty() {
|
||||
admitted
|
||||
} else {
|
||||
unqueued
|
||||
}
|
||||
}
|
||||
};
|
||||
// Scoped to the pool actually ranked: an admitted-but-gated worker
|
||||
// missing native monitor data would otherwise downgrade the comparison
|
||||
// for the whole unqueued tier to router-local.
|
||||
let loads = FreshLoadLookup::new(Some(snapshot), pool.iter());
|
||||
loads
|
||||
.min_by_pressure_key(admitted, FreshLoadLookup::compare_prefill_keys)
|
||||
.min_by_pressure_key(pool, FreshLoadLookup::compare_prefill_keys)
|
||||
.map(|worker| (worker, DecisionReason::RangeFallback))
|
||||
}
|
||||
|
||||
@@ -866,11 +1015,12 @@ mod tests {
|
||||
&range,
|
||||
&SelectionProposal::primary(Arc::clone(&full)),
|
||||
20,
|
||||
&loads
|
||||
&loads,
|
||||
None,
|
||||
)
|
||||
.is_some());
|
||||
assert_eq!(
|
||||
resolve_prefill(&range, &SelectionProposal::primary(full), 20, &loads)
|
||||
resolve_prefill(&range, &SelectionProposal::primary(full), 20, &loads, None)
|
||||
.expect("fallback selects the admitted worker")
|
||||
.selected
|
||||
.id,
|
||||
@@ -896,8 +1046,14 @@ mod tests {
|
||||
(&filtered, 0, 0, 0, 100),
|
||||
]);
|
||||
|
||||
let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &loads)
|
||||
.expect("capacity exhaustion must degrade within the legal domain");
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
None,
|
||||
)
|
||||
.expect("capacity exhaustion must degrade within the legal domain");
|
||||
|
||||
assert!(matches!(
|
||||
decision.selected.id.0.as_str(),
|
||||
@@ -918,8 +1074,14 @@ mod tests {
|
||||
.expect("the opposite snapshot has the same legal workers");
|
||||
assert_eq!(opposite_decision.id, backup.id);
|
||||
|
||||
let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &explicit)
|
||||
.expect("capacity exhaustion must degrade to Power-of-Two");
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&explicit,
|
||||
None,
|
||||
)
|
||||
.expect("capacity exhaustion must degrade to Power-of-Two");
|
||||
|
||||
assert_eq!(decision.selected.id, primary.id);
|
||||
assert_eq!(decision.load_snapshot_version, explicit.version);
|
||||
@@ -951,6 +1113,7 @@ mod tests {
|
||||
worker: Arc::clone(&congested),
|
||||
matched_prefix_tokens: 90,
|
||||
uncached_tokens: 10,
|
||||
matched_prefix_blocks: 9,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
},
|
||||
@@ -958,6 +1121,7 @@ mod tests {
|
||||
worker: Arc::clone(&idle),
|
||||
matched_prefix_tokens: 80,
|
||||
uncached_tokens: 20,
|
||||
matched_prefix_blocks: 8,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
},
|
||||
@@ -967,13 +1131,14 @@ mod tests {
|
||||
pressure_abs_threshold_tokens: 100,
|
||||
pressure_abs_threshold_ms: None,
|
||||
pressure_rel_threshold: 1.5,
|
||||
worker_queue_limit: None,
|
||||
};
|
||||
let loads = snapshot(&[
|
||||
(&congested, 1, 1_000, 10, 10_000),
|
||||
(&idle, 1, 10, 10, 10_000),
|
||||
]);
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads);
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &[]);
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
@@ -987,4 +1152,401 @@ mod tests {
|
||||
assert_eq!(resolution.pressure_guard_compared_pairs, 1);
|
||||
assert_eq!(resolution.pressure_guard_overrides, 1);
|
||||
}
|
||||
|
||||
fn candidate(
|
||||
worker: &Arc<Worker>,
|
||||
uncached_tokens: u64,
|
||||
matched_blocks: u32,
|
||||
) -> CacheCandidate {
|
||||
CacheCandidate {
|
||||
worker: Arc::clone(worker),
|
||||
matched_prefix_tokens: 100 - uncached_tokens,
|
||||
uncached_tokens,
|
||||
matched_prefix_blocks: matched_blocks,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn queue_gate_proposal(
|
||||
candidates: Vec<CacheCandidate>,
|
||||
limit: Option<u64>,
|
||||
) -> CacheCandidateProposal {
|
||||
CacheCandidateProposal {
|
||||
candidates,
|
||||
worker_queue_limit: limit,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_diverts_off_an_owner_over_the_limit() {
|
||||
let owner = worker("owner");
|
||||
let other = worker("other");
|
||||
let proposal = queue_gate_proposal(
|
||||
vec![
|
||||
// The owner holds the deeper prefix and would always win
|
||||
// without the gate.
|
||||
candidate(&owner, 10, 9),
|
||||
candidate(&other, 60, 4),
|
||||
],
|
||||
Some(4),
|
||||
);
|
||||
// The owner is at the limit (4 waiting >= 4); the other is idle.
|
||||
let loads = snapshot(&[(&owner, 1, 4, 10, 10_000), (&other, 1, 0, 10, 10_000)]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&other)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
.expect("the unqueued candidate remains admitted")
|
||||
.selected
|
||||
.id,
|
||||
other.id
|
||||
);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 1);
|
||||
assert_eq!(resolution.queue_gate_best_rejected_blocks, 9);
|
||||
// A gate rejection must not read as a capacity rejection.
|
||||
assert_eq!(resolution.admission_rejected_candidates, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_disabled_keeps_the_deepest_owner() {
|
||||
let owner = worker("owner");
|
||||
let other = worker("other");
|
||||
let proposal = queue_gate_proposal(
|
||||
vec![candidate(&owner, 10, 9), candidate(&other, 60, 4)],
|
||||
None,
|
||||
);
|
||||
let loads = snapshot(&[(&owner, 1, 40, 10, 10_000), (&other, 1, 0, 10, 10_000)]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&other)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
.expect("no gate, owner wins")
|
||||
.selected
|
||||
.id,
|
||||
owner.id
|
||||
);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_missing_snapshot_admits() {
|
||||
let unknown = worker("unknown");
|
||||
let proposal = queue_gate_proposal(vec![candidate(&unknown, 10, 9)], Some(4));
|
||||
// The worker never published a load sample: the gate fails open.
|
||||
let loads = snapshot(&[]);
|
||||
let fleet = vec![Arc::clone(&unknown)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
.expect("an unknown queue admits")
|
||||
.selected
|
||||
.id,
|
||||
unknown.id
|
||||
);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_exhausted_candidates_returns_no_decision_with_audit() {
|
||||
// The only owner is queueing but a non-owner is idle, so diversion
|
||||
// still buys something: the request must leave the prefix.
|
||||
let owner = worker("owner");
|
||||
let idle_stranger = worker("idle_stranger");
|
||||
let proposal = queue_gate_proposal(vec![candidate(&owner, 10, 9)], Some(4));
|
||||
let loads = snapshot(&[
|
||||
(&owner, 1, 9, 10, 10_000),
|
||||
(&idle_stranger, 0, 0, 10, 10_000),
|
||||
]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&idle_stranger)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert!(resolution.decision.is_none());
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 1);
|
||||
assert_eq!(resolution.queue_gate_best_rejected_blocks, 9);
|
||||
assert_eq!(resolution.admission_rejected_candidates, 0);
|
||||
assert_eq!(resolution.admission_evaluated_candidates, 0);
|
||||
assert!(!resolution.fleet_all_queued);
|
||||
assert!(!resolution.queue_gate_fell_back);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_admits_one_below_the_limit() {
|
||||
// Pins the admit side of the `<` boundary: `limit - 1` waiting must
|
||||
// still win on affinity, or the gate fires a request early.
|
||||
let owner = worker("owner");
|
||||
let other = worker("other");
|
||||
let proposal = queue_gate_proposal(
|
||||
vec![candidate(&owner, 10, 9), candidate(&other, 60, 4)],
|
||||
Some(4),
|
||||
);
|
||||
let loads = snapshot(&[(&owner, 1, 3, 10, 10_000), (&other, 1, 0, 10, 10_000)]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&other)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
.expect("one below the limit is not queueing")
|
||||
.selected
|
||||
.id,
|
||||
owner.id
|
||||
);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 0);
|
||||
assert_eq!(resolution.admission_evaluated_candidates, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_keeps_the_prefix_when_the_whole_fleet_is_queueing() {
|
||||
// Every owner is over the limit AND so is every other worker, so a
|
||||
// diversion could not dodge a wait. Discarding the prefix would be
|
||||
// pure loss: the ungated tier re-admits the owners.
|
||||
let owner = worker("owner");
|
||||
let shallow_owner = worker("shallow_owner");
|
||||
let proposal = queue_gate_proposal(
|
||||
vec![candidate(&owner, 10, 9), candidate(&shallow_owner, 60, 4)],
|
||||
Some(4),
|
||||
);
|
||||
let loads = snapshot(&[
|
||||
(&owner, 1, 9, 10, 10_000),
|
||||
(&shallow_owner, 1, 5, 10, 10_000),
|
||||
]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&shallow_owner)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert_eq!(
|
||||
resolution
|
||||
.decision
|
||||
.expect("an all-queueing fleet must keep its prefix")
|
||||
.selected
|
||||
.id,
|
||||
owner.id,
|
||||
"the deepest prefix must win once diversion buys nothing"
|
||||
);
|
||||
assert!(resolution.queue_gate_fell_back);
|
||||
assert!(resolution.fleet_all_queued);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 2);
|
||||
assert_eq!(
|
||||
resolution.admission_evaluated_candidates, 2,
|
||||
"the ungated tier puts every candidate through capacity admission"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_saturation_survives_a_capacity_exhausted_re_admission() {
|
||||
// The audit tuple the decision label is read from, in the one case
|
||||
// that used to lose the saturation signal: the fleet is queueing, the
|
||||
// ungated tier re-admits the owners, and they then fail hard
|
||||
// admission, so there is no winner. `fleet_all_queued` must still be
|
||||
// set on the way out, because the label is keyed on the fleet being
|
||||
// saturated and not on where the request finally landed.
|
||||
let owner = worker("owner");
|
||||
let shallow_owner = worker("shallow_owner");
|
||||
let proposal = queue_gate_proposal(
|
||||
vec![candidate(&owner, 10, 9), candidate(&shallow_owner, 60, 4)],
|
||||
Some(4),
|
||||
);
|
||||
// Queueing AND out of KV: used == capacity on both.
|
||||
let loads = snapshot(&[
|
||||
(&owner, 1, 9, 10_000, 10_000),
|
||||
(&shallow_owner, 1, 5, 10_000, 10_000),
|
||||
]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&shallow_owner)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100_000, &loads, &fleet);
|
||||
|
||||
assert!(
|
||||
resolution.decision.is_none(),
|
||||
"a capacity-exhausted fleet cannot produce a winner"
|
||||
);
|
||||
assert!(resolution.fleet_all_queued);
|
||||
assert!(resolution.queue_gate_fell_back);
|
||||
assert_eq!(resolution.queue_gate_rejected_candidates, 2);
|
||||
assert_eq!(
|
||||
resolution.admission_evaluated_candidates, 2,
|
||||
"re-admission is what makes a zero-evaluated saturated audit unreachable"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_fleet_saturation_needs_a_fresh_sample_everywhere() {
|
||||
// A worker with no fresh sample has an unknown queue, not a proven
|
||||
// full one, so the fleet is not saturated and the gate keeps diverting.
|
||||
let owner = worker("owner");
|
||||
let silent = worker("silent");
|
||||
let proposal = queue_gate_proposal(vec![candidate(&owner, 10, 9)], Some(4));
|
||||
let loads = snapshot(&[(&owner, 1, 9, 10, 10_000)]);
|
||||
let fleet = vec![Arc::clone(&owner), Arc::clone(&silent)];
|
||||
|
||||
let resolution = resolve_cache_candidates(&proposal, 100, &loads, &fleet);
|
||||
|
||||
assert!(resolution.decision.is_none());
|
||||
assert!(!resolution.fleet_all_queued);
|
||||
assert!(!resolution.queue_gate_fell_back);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_fallback_prefers_an_unqueued_worker_over_a_shallower_queueing_one() {
|
||||
// The primary fails KV-capacity admission so selection reaches the
|
||||
// range fallback. The queueing worker is the fallback minimum BY
|
||||
// PRESSURE (1 waiting uncached token), so a single-tier fallback
|
||||
// would pick it — handing the request straight back to a worker the
|
||||
// gate just rejected.
|
||||
let primary = worker("primary");
|
||||
let shallow_queued = worker("shallow_queued");
|
||||
let busy_unqueued = worker("busy_unqueued");
|
||||
let workers = vec![
|
||||
Arc::clone(&primary),
|
||||
Arc::clone(&shallow_queued),
|
||||
Arc::clone(&busy_unqueued),
|
||||
];
|
||||
let proposal = SelectionProposal::primary(Arc::clone(&primary));
|
||||
let load = |waiting: u64, waiting_uncached: u64, total: u64, max_total: u64| {
|
||||
NativeCacheWorkerLoad {
|
||||
num_running_reqs: 0,
|
||||
num_waiting_reqs: waiting,
|
||||
num_waiting_uncached_tokens: waiting_uncached,
|
||||
num_used_tokens: total,
|
||||
num_total_tokens: total,
|
||||
max_total_num_tokens: max_total,
|
||||
max_running_requests: 64,
|
||||
prefill_throughput_tokens_per_s: None,
|
||||
estimated_prefill_queue_ms: None,
|
||||
captured_at: Instant::now(),
|
||||
}
|
||||
};
|
||||
// Queue depth and pressure disagree by construction: shallow_queued
|
||||
// waits 5 (over the limit) behind 1 uncached token, busy_unqueued
|
||||
// waits 3 (under the limit) behind 1000 uncached tokens. The primary
|
||||
// is KV-full, so it is not admitted at all.
|
||||
let loads = EngineLoadSnapshot::from_native_cache_workers(
|
||||
7,
|
||||
[
|
||||
(primary.url.clone(), load(0, 0, 10_000, 10_000)),
|
||||
(shallow_queued.url.clone(), load(5, 1, 10, 10_000)),
|
||||
(busy_unqueued.url.clone(), load(3, 1_000, 10, 10_000)),
|
||||
]
|
||||
.into_iter()
|
||||
.collect(),
|
||||
);
|
||||
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
Some(4),
|
||||
)
|
||||
.expect("an admitted worker exists");
|
||||
|
||||
assert_eq!(
|
||||
decision.selected.id, busy_unqueued.id,
|
||||
"the fallback must prefer the unqueued worker even though the \
|
||||
queueing one is the pressure minimum"
|
||||
);
|
||||
assert_eq!(decision.reason, DecisionReason::RangeFallback);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_demotes_a_queueing_primary_that_capacity_would_admit() {
|
||||
// The whole point of the gate is that the cache-affinity fallback
|
||||
// must not hand the request back to a queueing worker. The primary
|
||||
// here has plenty of KV capacity, so without the gate the
|
||||
// `(true, _, _)` arm returns it unconditionally and the gate is
|
||||
// bypassed on the single most common path.
|
||||
let queued_primary = worker("queued_primary");
|
||||
let unqueued = worker("unqueued");
|
||||
let workers = vec![Arc::clone(&queued_primary), Arc::clone(&unqueued)];
|
||||
let proposal = SelectionProposal::primary(Arc::clone(&queued_primary));
|
||||
let loads = snapshot(&[
|
||||
(&queued_primary, 0, 6, 10, 10_000),
|
||||
(&unqueued, 0, 0, 10, 10_000),
|
||||
]);
|
||||
|
||||
let ungated = resolve_prefill_admitted(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
None,
|
||||
)
|
||||
.expect("without a limit the primary is admitted");
|
||||
assert_eq!(ungated.selected.id, queued_primary.id);
|
||||
assert_eq!(ungated.reason, DecisionReason::Primary);
|
||||
|
||||
let gated = resolve_prefill_admitted(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
Some(4),
|
||||
)
|
||||
.expect("the unqueued worker takes over");
|
||||
assert_eq!(gated.selected.id, unqueued.id);
|
||||
assert_eq!(gated.reason, DecisionReason::RangeFallback);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn queue_gate_still_routes_when_the_only_admitted_worker_is_queueing() {
|
||||
// Demotion must never become rejection: with every admitted worker
|
||||
// over the limit, the second fallback tier keeps the request routable.
|
||||
let only = worker("only");
|
||||
let workers = vec![Arc::clone(&only)];
|
||||
let proposal = SelectionProposal::primary(Arc::clone(&only));
|
||||
let loads = snapshot(&[(&only, 0, 9, 10, 10_000)]);
|
||||
|
||||
let decision = resolve_prefill_admitted(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
Some(4),
|
||||
)
|
||||
.expect("an all-queueing fleet must still route");
|
||||
|
||||
assert_eq!(decision.selected.id, only.id);
|
||||
assert_eq!(decision.reason, DecisionReason::RangeFallback);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn range_fallback_keeps_an_all_queueing_fleet_routable() {
|
||||
let primary = worker("primary");
|
||||
let left = worker("left");
|
||||
let right = worker("right");
|
||||
let workers = vec![Arc::clone(&primary), Arc::clone(&left), Arc::clone(&right)];
|
||||
let proposal = SelectionProposal::primary(Arc::clone(&primary));
|
||||
// The primary is KV-full; both fallback workers are over the limit.
|
||||
// The second tier takes the least-pressured one instead of failing
|
||||
// the request.
|
||||
let loads = snapshot(&[
|
||||
(&primary, 0, 0, 10_000, 10_000),
|
||||
(&left, 0, 9, 10, 10_000),
|
||||
(&right, 0, 5, 10, 10_000),
|
||||
]);
|
||||
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
Some(4),
|
||||
)
|
||||
.expect("an all-queueing fleet must still route");
|
||||
|
||||
assert_eq!(decision.selected.id, right.id);
|
||||
assert_eq!(decision.reason, DecisionReason::RangeFallback);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -55,10 +55,12 @@ impl CacheAwarePolicy {
|
||||
if entry.matched_prefix_blocks == 0 || !seen.insert(worker.id.clone()) {
|
||||
continue;
|
||||
}
|
||||
let matched_prefix_blocks =
|
||||
cap_matched_prefix_blocks(signal.query_blocks, entry.matched_prefix_blocks);
|
||||
let matched_prefix_tokens = estimate_matched_prefix_tokens(
|
||||
input_tokens,
|
||||
signal.query_blocks,
|
||||
entry.matched_prefix_blocks,
|
||||
matched_prefix_blocks,
|
||||
);
|
||||
if !self.passes_cache_gate(input_tokens, matched_prefix_tokens) {
|
||||
continue;
|
||||
@@ -67,6 +69,7 @@ impl CacheAwarePolicy {
|
||||
worker: Arc::clone(worker),
|
||||
matched_prefix_tokens,
|
||||
uncached_tokens: input_tokens.saturating_sub(matched_prefix_tokens),
|
||||
matched_prefix_blocks,
|
||||
candidate_range_id: ctx.candidate_range_id().to_string(),
|
||||
max_pending_prefill_tokens: None,
|
||||
});
|
||||
@@ -106,6 +109,7 @@ impl CacheAwarePolicy {
|
||||
pressure_abs_threshold_tokens: self.config.pressure_abs_threshold_tokens,
|
||||
pressure_abs_threshold_ms: self.config.pressure_abs_threshold_ms,
|
||||
pressure_rel_threshold: self.config.pressure_rel_threshold,
|
||||
worker_queue_limit: self.config.worker_queue_limit,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -189,13 +193,24 @@ impl Policy for CacheAwarePolicy {
|
||||
}
|
||||
}
|
||||
|
||||
/// Caps an indexer-supplied matched-block count at the blocks the query
|
||||
/// actually asked about: a query cannot match more blocks than it contains.
|
||||
/// Both the token estimate and the diverted-overlap histogram read the capped
|
||||
/// value, so the clamp lives here rather than at each use.
|
||||
fn cap_matched_prefix_blocks(query_blocks: usize, matched_prefix_blocks: u32) -> u32 {
|
||||
matched_prefix_blocks.min(u32::try_from(query_blocks).unwrap_or(u32::MAX))
|
||||
}
|
||||
|
||||
fn estimate_matched_prefix_tokens(
|
||||
input_tokens: u64,
|
||||
query_blocks: usize,
|
||||
matched_prefix_blocks: u32,
|
||||
) -> u64 {
|
||||
let matched_prefix_blocks = u64::from(cap_matched_prefix_blocks(
|
||||
query_blocks,
|
||||
matched_prefix_blocks,
|
||||
));
|
||||
let query_blocks = u64::try_from(query_blocks).unwrap_or(u64::MAX).max(1);
|
||||
let matched_prefix_blocks = u64::from(matched_prefix_blocks).min(query_blocks);
|
||||
input_tokens.saturating_mul(matched_prefix_blocks) / query_blocks
|
||||
}
|
||||
|
||||
@@ -207,4 +222,13 @@ mod tests {
|
||||
fn matched_token_estimate_caps_untrusted_block_count() {
|
||||
assert_eq!(estimate_matched_prefix_tokens(80, 8, 99), 80);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn matched_block_cap_is_shared_by_the_estimate_and_the_candidate() {
|
||||
// One clamp, two readers: the histogram must never see a block count
|
||||
// the token estimate would have thrown away.
|
||||
assert_eq!(cap_matched_prefix_blocks(8, 99), 8);
|
||||
assert_eq!(cap_matched_prefix_blocks(8, 3), 3);
|
||||
assert_eq!(cap_matched_prefix_blocks(0, 3), 0);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -332,6 +332,11 @@ pub struct CacheCandidate {
|
||||
pub worker: Arc<Worker>,
|
||||
pub matched_prefix_tokens: u64,
|
||||
pub uncached_tokens: u64,
|
||||
/// Matched prefix length in blocks, as reported by the prefix signal.
|
||||
/// Selection reads `matched_prefix_tokens`; the block count exists for
|
||||
/// observability (the diverted-overlap histogram reads against the
|
||||
/// tree/indexer block domain).
|
||||
pub matched_prefix_blocks: u32,
|
||||
/// Domain containing this candidate.
|
||||
pub candidate_range_id: String,
|
||||
/// Optional pending prefill limit checked against `E`.
|
||||
@@ -347,6 +352,10 @@ pub struct CacheCandidateProposal {
|
||||
pub pressure_abs_threshold_tokens: u64,
|
||||
pub pressure_abs_threshold_ms: Option<f64>,
|
||||
pub pressure_rel_threshold: f64,
|
||||
/// Queue gate: a candidate whose engine reports at least this many
|
||||
/// waiting requests cannot win on cache affinity. `None` disables the
|
||||
/// gate. See [`crate::config::AffinityConfig::worker_queue_limit`].
|
||||
pub worker_queue_limit: Option<u64>,
|
||||
}
|
||||
|
||||
/// Prefill proposal returned as either a pair or a Cache-Aware candidate set.
|
||||
@@ -661,6 +670,7 @@ mod tests {
|
||||
worker: Arc::clone(&hot),
|
||||
matched_prefix_tokens: 75,
|
||||
uncached_tokens: 25,
|
||||
matched_prefix_blocks: 3,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
}],
|
||||
@@ -771,8 +781,14 @@ mod tests {
|
||||
},
|
||||
),
|
||||
]);
|
||||
let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &loads)
|
||||
.expect("the admitted backup must become Final P");
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&loads,
|
||||
None,
|
||||
)
|
||||
.expect("the admitted backup must become Final P");
|
||||
assert_eq!(decision.selected.id, backup.id);
|
||||
policy.commit_prefill_selection(&ctx, proposal.kind, &decision.selected);
|
||||
|
||||
@@ -1261,6 +1277,7 @@ mod tests {
|
||||
worker: Arc::clone(worker),
|
||||
matched_prefix_tokens,
|
||||
uncached_tokens,
|
||||
matched_prefix_blocks: 0,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens,
|
||||
}
|
||||
@@ -1297,7 +1314,7 @@ mod tests {
|
||||
),
|
||||
]);
|
||||
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads)
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads, &[])
|
||||
.decision
|
||||
.expect("a later admitted cache match must survive");
|
||||
|
||||
@@ -1345,7 +1362,7 @@ mod tests {
|
||||
),
|
||||
]);
|
||||
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads)
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads, &[])
|
||||
.decision
|
||||
.expect("all admitted candidates must participate in the tournament");
|
||||
|
||||
@@ -1371,7 +1388,7 @@ mod tests {
|
||||
},
|
||||
)]);
|
||||
assert!(
|
||||
resolve_cache_candidates(&proposal, 100, &pending_allows)
|
||||
resolve_cache_candidates(&proposal, 100, &pending_allows, &[])
|
||||
.decision
|
||||
.is_some(),
|
||||
"pending admission must project E=20, not L=100"
|
||||
@@ -1387,7 +1404,7 @@ mod tests {
|
||||
},
|
||||
)]);
|
||||
assert!(
|
||||
resolve_cache_candidates(&proposal, 100, &kv_rejects)
|
||||
resolve_cache_candidates(&proposal, 100, &kv_rejects, &[])
|
||||
.decision
|
||||
.is_none(),
|
||||
"KV safety must conservatively project the complete input L=100"
|
||||
@@ -1425,7 +1442,7 @@ mod tests {
|
||||
),
|
||||
]);
|
||||
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads)
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads, &[])
|
||||
.decision
|
||||
.unwrap();
|
||||
assert_eq!(decision.selected.id, congested.id);
|
||||
@@ -1462,7 +1479,7 @@ mod tests {
|
||||
),
|
||||
]);
|
||||
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads)
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads, &[])
|
||||
.decision
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
@@ -1515,7 +1532,7 @@ mod tests {
|
||||
),
|
||||
]);
|
||||
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads)
|
||||
let decision = resolve_cache_candidates(&proposal, 100, &loads, &[])
|
||||
.decision
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
@@ -1564,7 +1581,7 @@ mod tests {
|
||||
let range = CandidateRange::global(&workers);
|
||||
let proposal = SelectionProposal::with_backup(Arc::clone(&primary), Arc::clone(&backup));
|
||||
|
||||
let decision = resolve_prefill(&range, &proposal, 32, &snapshot)
|
||||
let decision = resolve_prefill(&range, &proposal, 32, &snapshot, None)
|
||||
.expect("an admitted backup must be selected");
|
||||
|
||||
assert_eq!(decision.selected.id, backup.id);
|
||||
@@ -1582,6 +1599,7 @@ mod tests {
|
||||
&SelectionProposal::primary(Arc::clone(&primary)),
|
||||
1_000_000,
|
||||
&snapshot,
|
||||
None,
|
||||
)
|
||||
.expect("disabled reporting must preserve the healthy registry candidate");
|
||||
|
||||
@@ -1614,8 +1632,14 @@ mod tests {
|
||||
]);
|
||||
let proposal = SelectionProposal::with_backup(primary, backup);
|
||||
|
||||
let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 80, &snapshot)
|
||||
.expect("both candidates fit capacity");
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
80,
|
||||
&snapshot,
|
||||
None,
|
||||
)
|
||||
.expect("both candidates fit capacity");
|
||||
|
||||
assert_eq!(decision.reason, DecisionReason::Primary);
|
||||
}
|
||||
@@ -1658,8 +1682,14 @@ mod tests {
|
||||
]);
|
||||
let proposal = SelectionProposal::with_backup(primary, backup);
|
||||
|
||||
let decision = resolve_prefill(&CandidateRange::global(&workers), &proposal, 32, &snapshot)
|
||||
.expect("an admitted range fallback must be selected");
|
||||
let decision = resolve_prefill(
|
||||
&CandidateRange::global(&workers),
|
||||
&proposal,
|
||||
32,
|
||||
&snapshot,
|
||||
None,
|
||||
)
|
||||
.expect("an admitted range fallback must be selected");
|
||||
|
||||
assert_eq!(decision.selected.id, fallback.id);
|
||||
assert_eq!(decision.reason, DecisionReason::RangeFallback);
|
||||
|
||||
@@ -747,8 +747,9 @@ mod tests {
|
||||
(&ws[2], 0, 0, 0, 4_096),
|
||||
]);
|
||||
|
||||
let decision = resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot)
|
||||
.expect("capacity exhaustion must degrade inside the filtered domain");
|
||||
let decision =
|
||||
resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot, None)
|
||||
.expect("capacity exhaustion must degrade inside the filtered domain");
|
||||
assert!(matches!(decision.selected.id.0.as_str(), "a" | "b"));
|
||||
}
|
||||
|
||||
@@ -783,8 +784,9 @@ mod tests {
|
||||
assert_eq!(proposal.primary.id, ws[2].id);
|
||||
|
||||
let snapshot = EngineLoadSnapshot::default();
|
||||
let decision = resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot)
|
||||
.expect("an eligible escape worker exists");
|
||||
let decision =
|
||||
resolve_prefill(&CandidateRange::global(&ws), &proposal, 32, &snapshot, None)
|
||||
.expect("an eligible escape worker exists");
|
||||
assert_ne!(decision.selected.id, ws[2].id);
|
||||
assert!(matches!(decision.selected.id.0.as_str(), "a" | "b"));
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@
|
||||
//! The module owns the decision and reports why; it does not own the HTTP
|
||||
//! response. Mapping a failed selection onto a status code stays in the route.
|
||||
|
||||
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::config::{DecodePolicyKind, PolicyKind, SessionAffinityMode};
|
||||
@@ -38,7 +39,7 @@ use crate::policies::engine_load::EngineLoadSnapshot;
|
||||
use crate::policies::{
|
||||
ExternalPrefixSignal, Policy, PrefillProposal, ProposalKind, SelectionContext,
|
||||
};
|
||||
use crate::server::metrics::{MetricsRegistry, PolicySelectionFailureReason};
|
||||
use crate::server::metrics::{CacheAwareDecision, MetricsRegistry, PolicySelectionFailureReason};
|
||||
use crate::workers::Worker;
|
||||
|
||||
/// Everything one prefill selection reads. Collaborators first, then the
|
||||
@@ -67,6 +68,49 @@ pub(crate) struct PrefillSelectionInputs<'a> {
|
||||
/// The configured mode. Without Bucket partitioning all modes reduce to
|
||||
/// the single global domain, and the ladder applies that reduction itself.
|
||||
pub session_affinity_mode: SessionAffinityMode,
|
||||
/// `--worker-queue-limit`. `None` disables the queue gate entirely.
|
||||
pub worker_queue_limit: Option<u64>,
|
||||
}
|
||||
|
||||
/// The queue-gate blind warn is sampled: it fires on a per-request path, and
|
||||
/// the condition (gate configured, zero fresh engine load samples fleet-wide)
|
||||
/// is steady-state, so 1-in-64 is plenty to surface it without log flooding.
|
||||
const QUEUE_GATE_BLIND_LOG_SAMPLE: u64 = 64;
|
||||
static QUEUE_GATE_BLIND_LOG_COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
|
||||
/// Maps the queue-gate audit of a Cache-Aware selection that produced no
|
||||
/// winner onto its decision label. Pure so every boundary is pinned by unit
|
||||
/// tests rather than inferred from the ladder that calls it:
|
||||
///
|
||||
/// - `queue_gate_rejected == 0` means the gate took nothing out, so whatever
|
||||
/// emptied the candidate set was not the gate.
|
||||
/// - `fleet_all_queued` is asked BEFORE the capacity question, because
|
||||
/// `all_queued` is keyed on the fleet being saturated and not on where the
|
||||
/// request landed. Asking capacity first drops the saturation signal in the
|
||||
/// worst case there is: a queueing fleet whose owners are also out of KV
|
||||
/// books a plain `cache_miss`, and a fully saturated fleet reads as a
|
||||
/// healthy one — the exact blind spot the label exists to remove.
|
||||
/// - `admission_evaluated == 0` is then what makes the GATE, rather than KV
|
||||
/// capacity, the reason an unsaturated fleet's candidate set came back
|
||||
/// empty. Without it a capacity exhaustion books as a gate diversion
|
||||
/// whenever one owner happens to be queueing.
|
||||
fn cache_aware_fallback_decision(
|
||||
queue_gate_rejected: u64,
|
||||
admission_evaluated: u64,
|
||||
fleet_all_queued: bool,
|
||||
) -> CacheAwareDecision {
|
||||
if queue_gate_rejected == 0 {
|
||||
return CacheAwareDecision::CacheMiss;
|
||||
}
|
||||
if fleet_all_queued {
|
||||
return CacheAwareDecision::AllQueued;
|
||||
}
|
||||
if admission_evaluated > 0 {
|
||||
// Capacity, not the gate: owners survived the gate and then failed
|
||||
// capacity admission.
|
||||
return CacheAwareDecision::CacheMiss;
|
||||
}
|
||||
CacheAwareDecision::CacheWorkerQueued
|
||||
}
|
||||
|
||||
/// Runs the prefill selection ladder. `Err` carries the reason the last rung
|
||||
@@ -85,6 +129,7 @@ pub(crate) fn select_prefill_worker(
|
||||
tps_slo: inputs.tps_slo,
|
||||
},
|
||||
failure_reason: PolicySelectionFailureReason::ProposalEmpty,
|
||||
cache_gate_audit: None,
|
||||
};
|
||||
let selected = selector.run();
|
||||
selected.ok_or(selector.failure_reason)
|
||||
@@ -98,6 +143,12 @@ struct Selector<'a> {
|
||||
inputs: &'a PrefillSelectionInputs<'a>,
|
||||
bucket_request: BucketRequest,
|
||||
failure_reason: PolicySelectionFailureReason,
|
||||
/// Queue-gate audit of a Cache-Aware resolution that produced no winner:
|
||||
/// (gate-rejected candidates, candidates that reached capacity admission,
|
||||
/// fleet saturation, deepest rejected prefix). `None` when there was no
|
||||
/// resolution at all — no load snapshot, or no candidate proposal — which
|
||||
/// reads as a plain miss because nothing was gated out.
|
||||
cache_gate_audit: Option<(u64, u64, bool, u32)>,
|
||||
}
|
||||
|
||||
impl<'a> Selector<'a> {
|
||||
@@ -117,6 +168,7 @@ impl<'a> Selector<'a> {
|
||||
|
||||
// Cache-Aware resolves one bounded global candidate set and returns a final winner.
|
||||
let cache_winner = self.cache_winner();
|
||||
let cache_winner_hit = cache_winner.is_some();
|
||||
|
||||
let global_affinity_probe = use_global_affinity_probe
|
||||
.then(|| {
|
||||
@@ -147,7 +199,7 @@ impl<'a> Selector<'a> {
|
||||
// Rebuild the backup inside the primary's own Bucket.
|
||||
.and_then(|domain| self.select_in_domain(&domain, true, false, false));
|
||||
|
||||
cache_winner.or_else(|| {
|
||||
let selected = cache_winner.or_else(|| {
|
||||
// Materializing the normal domains clones the member list of every
|
||||
// Bucket, so build them only on the rung that actually reads them.
|
||||
let prefill_domains = || {
|
||||
@@ -173,7 +225,32 @@ impl<'a> Selector<'a> {
|
||||
self.select_domains(&prefill_domains(), true, true)
|
||||
}
|
||||
}
|
||||
})
|
||||
});
|
||||
// Only a selection that resolved a worker books a decision. A ladder
|
||||
// that ran out of rungs is a 503, already counted by
|
||||
// `sgl_router_policy_selection_failures_total`; booking it here too
|
||||
// would break the documented sum and, worse, let a request that
|
||||
// reached nothing book `cache_worker_queued` and contribute to
|
||||
// `sgl_router_diverted_overlap_blocks` — a diversion that never
|
||||
// arrived is not evidence about what the gate traded away.
|
||||
if inputs.policy_kind == PolicyKind::CacheAware && !cache_winner_hit && selected.is_some() {
|
||||
let (rejected, evaluated, fleet_all_queued, blocks) =
|
||||
self.cache_gate_audit.unwrap_or((0, 0, false, 0));
|
||||
let decision = cache_aware_fallback_decision(rejected, evaluated, fleet_all_queued);
|
||||
if matches!(decision, CacheAwareDecision::CacheWorkerQueued) {
|
||||
// A real diversion: an unqueued destination existed and the
|
||||
// gate gave up `blocks` of matched prefix to reach it. The
|
||||
// histogram is the evidence for whether the gate is trading
|
||||
// large cached prefixes for short waits.
|
||||
inputs
|
||||
.metrics
|
||||
.observe_diverted_overlap_blocks(&inputs.model_id.0, u64::from(blocks));
|
||||
}
|
||||
inputs
|
||||
.metrics
|
||||
.record_cache_aware_decision(&inputs.model_id.0, decision);
|
||||
}
|
||||
selected
|
||||
}
|
||||
|
||||
/// The per-request `SelectionContext` every rung starts from.
|
||||
@@ -206,8 +283,42 @@ impl<'a> Selector<'a> {
|
||||
return None;
|
||||
};
|
||||
let bounded_candidate_count = proposal.candidates.len();
|
||||
let cache_decision =
|
||||
resolve_cache_candidates(&proposal, inputs.request_input_tokens, snapshot);
|
||||
// The queue gate reads the engine-published load sample and fails open
|
||||
// per worker. When NO worker has a fresh sample the gate is inert
|
||||
// fleet-wide and nothing would say so: `cache_worker_queued` sitting at
|
||||
// 0 is indistinguishable from a healthy fleet. Warn (sampled) — a fleet
|
||||
// that never advertised a load port must not silently disable the gate.
|
||||
if inputs.worker_queue_limit.is_some()
|
||||
&& !inputs.workers.is_empty()
|
||||
&& inputs
|
||||
.workers
|
||||
.iter()
|
||||
.all(|worker| snapshot.fresh_load_for_url(&worker.url).is_none())
|
||||
&& QUEUE_GATE_BLIND_LOG_COUNTER
|
||||
.fetch_add(1, AtomicOrdering::Relaxed)
|
||||
.is_multiple_of(QUEUE_GATE_BLIND_LOG_SAMPLE)
|
||||
{
|
||||
tracing::warn!(
|
||||
model = %inputs.model_id,
|
||||
worker_queue_limit = inputs.worker_queue_limit,
|
||||
workers = inputs.workers.len(),
|
||||
"--worker-queue-limit is set but no worker has a fresh engine load \
|
||||
sample, so the queue gate is inert. Check that engines advertise a \
|
||||
load port and publish LoadStat",
|
||||
);
|
||||
}
|
||||
let cache_decision = resolve_cache_candidates(
|
||||
&proposal,
|
||||
inputs.request_input_tokens,
|
||||
snapshot,
|
||||
inputs.workers,
|
||||
);
|
||||
self.cache_gate_audit = Some((
|
||||
cache_decision.queue_gate_rejected_candidates,
|
||||
cache_decision.admission_evaluated_candidates,
|
||||
cache_decision.fleet_all_queued,
|
||||
cache_decision.queue_gate_best_rejected_blocks,
|
||||
));
|
||||
inputs
|
||||
.metrics
|
||||
.record_cache_admission_evaluations(cache_decision.admission_evaluated_candidates);
|
||||
@@ -246,6 +357,18 @@ impl<'a> Selector<'a> {
|
||||
inputs
|
||||
.metrics
|
||||
.record_policy_decision("cache_aware", "cache_candidate");
|
||||
inputs.metrics.record_cache_aware_decision(
|
||||
&inputs.model_id.0,
|
||||
if cache_decision.queue_gate_fell_back {
|
||||
// The gate removed every owner and nowhere in the fleet is
|
||||
// unqueued, so the prefix was kept rather than traded for a
|
||||
// wait that cannot be dodged. Booked as saturation, never as
|
||||
// a plain hit.
|
||||
CacheAwareDecision::AllQueued
|
||||
} else {
|
||||
CacheAwareDecision::CacheHit
|
||||
},
|
||||
);
|
||||
Some(decision.selected)
|
||||
}
|
||||
|
||||
@@ -317,6 +440,7 @@ impl<'a> Selector<'a> {
|
||||
&proposal,
|
||||
inputs.request_input_tokens,
|
||||
snapshot,
|
||||
inputs.worker_queue_limit,
|
||||
)
|
||||
} else {
|
||||
resolve_prefill_admitted(
|
||||
@@ -324,6 +448,7 @@ impl<'a> Selector<'a> {
|
||||
&proposal,
|
||||
inputs.request_input_tokens,
|
||||
snapshot,
|
||||
inputs.worker_queue_limit,
|
||||
)
|
||||
};
|
||||
let Some(decision) = decision else {
|
||||
@@ -521,17 +646,20 @@ fn projected_decode_kv_tokens(input_tokens: u64, max_output_tokens: Option<u64>)
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
prefill_policy_reason, projected_decode_kv_tokens, select_decode_peer,
|
||||
select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
|
||||
cache_aware_fallback_decision, prefill_policy_reason, projected_decode_kv_tokens,
|
||||
select_decode_peer, select_prefill_worker, DecodeSelectionInputs, PrefillSelectionInputs,
|
||||
};
|
||||
use crate::config::{DecodePolicyKind, PolicyKind, SessionAffinityMode};
|
||||
use crate::config::{AffinityConfig, DecodePolicyKind, PolicyKind, SessionAffinityMode};
|
||||
use crate::discovery::{ModelId, WorkerId, WorkerMode, WorkerSpec};
|
||||
use crate::policies::admission::{resolve_prefill_admitted, CandidateRange, DecisionReason};
|
||||
use crate::policies::buckets::BucketSelector;
|
||||
use crate::policies::cache_aware::CacheAwarePolicy;
|
||||
use crate::policies::engine_load::{EngineLoadSnapshot, NativeCacheWorkerLoad};
|
||||
use crate::policies::power_of_two::PowerOfTwoChoicesPolicy;
|
||||
use crate::policies::{Policy, ProposalKind, SelectionProposal};
|
||||
use crate::server::metrics::{MetricsRegistry, PolicySelectionFailureReason};
|
||||
use crate::policies::{ExternalPrefixSignal, Policy, ProposalKind, SelectionProposal};
|
||||
use crate::server::metrics::{
|
||||
CacheAwareDecision, MetricsRegistry, PolicySelectionFailureReason,
|
||||
};
|
||||
use crate::workers::Worker;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
@@ -573,6 +701,62 @@ mod tests {
|
||||
)
|
||||
}
|
||||
|
||||
/// `(worker, waiting requests, tokens already held, published KV
|
||||
/// capacity)`. The queue gate reads `num_waiting_reqs`, which the plain
|
||||
/// [`snapshot`] fixture pins at zero.
|
||||
fn queued_snapshot(entries: &[(&Arc<Worker>, u64, u64, u64)]) -> EngineLoadSnapshot {
|
||||
EngineLoadSnapshot::from_native_cache_workers(
|
||||
7,
|
||||
entries
|
||||
.iter()
|
||||
.map(|(worker, waiting, used, capacity)| {
|
||||
(
|
||||
worker.url.clone(),
|
||||
NativeCacheWorkerLoad {
|
||||
num_running_reqs: 1,
|
||||
num_waiting_reqs: *waiting,
|
||||
num_waiting_uncached_tokens: *waiting,
|
||||
num_used_tokens: *used,
|
||||
num_total_tokens: *used,
|
||||
max_total_num_tokens: *capacity,
|
||||
max_running_requests: 64,
|
||||
prefill_throughput_tokens_per_s: None,
|
||||
estimated_prefill_queue_ms: None,
|
||||
captured_at: Instant::now(),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
/// An indexer hit placing `matched_prefix_blocks` on each named worker.
|
||||
fn prefix_signal(matches: &[(&Arc<Worker>, u32)], query_blocks: usize) -> ExternalPrefixSignal {
|
||||
ExternalPrefixSignal {
|
||||
outcome: sgl_kv_indexer::PrefixOutcome::Matched {
|
||||
matches: matches
|
||||
.iter()
|
||||
.map(|(worker, blocks)| sgl_kv_indexer::PrefixMatch {
|
||||
matched_prefix_blocks: *blocks,
|
||||
worker_id: worker.id.0.clone(),
|
||||
address: worker.url.clone(),
|
||||
})
|
||||
.collect(),
|
||||
best_prefix_blocks: matches.iter().map(|(_, blocks)| *blocks).max().unwrap_or(0),
|
||||
},
|
||||
query_blocks,
|
||||
}
|
||||
}
|
||||
|
||||
fn booked_decisions(metrics: &MetricsRegistry) -> Vec<String> {
|
||||
metrics
|
||||
.render()
|
||||
.lines()
|
||||
.filter(|line| line.starts_with("sgl_router_cache_aware_decisions_total{"))
|
||||
.map(str::to_owned)
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Power-of-two prefill, Bucket partitioning off, no session affinity —
|
||||
/// the single global domain, which is what isolates the ladder's rungs.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
@@ -599,6 +783,7 @@ mod tests {
|
||||
external_prefix: None,
|
||||
load_snapshot,
|
||||
workers,
|
||||
worker_queue_limit: None,
|
||||
ttft_slo_ms: None,
|
||||
tps_slo: None,
|
||||
session_affinity_mode: SessionAffinityMode::Bucket,
|
||||
@@ -651,6 +836,150 @@ mod tests {
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_failed_cache_aware_selection_books_no_decision() {
|
||||
// The documented invariant on `sgl_router_cache_aware_decisions_total`
|
||||
// is one decision per selection that RESOLVES a worker, so the labels
|
||||
// sum to the cache-aware rate less the 503s. A ladder that ran out of
|
||||
// rungs is already counted by the failure counter; booking it here too
|
||||
// would break that sum and let a request that reached nothing feed
|
||||
// `sgl_router_diverted_overlap_blocks`.
|
||||
let policy = PowerOfTwoChoicesPolicy::new();
|
||||
let buckets = BucketSelector::new(None);
|
||||
let metrics = MetricsRegistry::new();
|
||||
let model = ModelId("model".into());
|
||||
let workers: Vec<Arc<Worker>> = Vec::new();
|
||||
let loads = snapshot(&[]);
|
||||
let mut inputs = prefill_inputs(
|
||||
&policy,
|
||||
&buckets,
|
||||
&metrics,
|
||||
&model,
|
||||
&workers,
|
||||
Some(&loads),
|
||||
64,
|
||||
);
|
||||
inputs.policy_kind = PolicyKind::CacheAware;
|
||||
|
||||
assert!(select_prefill_worker(&inputs).is_err());
|
||||
|
||||
let rendered = metrics.render();
|
||||
let booked: Vec<&str> = rendered
|
||||
.lines()
|
||||
.filter(|line| line.starts_with("sgl_router_cache_aware_decisions_total{"))
|
||||
.collect();
|
||||
assert!(
|
||||
booked.is_empty(),
|
||||
"a 503 selection must book no cache-aware decision, got {booked:?}"
|
||||
);
|
||||
}
|
||||
|
||||
/// The wiring between `resolve_cache_candidates`' audit and the decision
|
||||
/// label, which the resolver test and the pure-mapper test each cover only
|
||||
/// one side of. This is the end the `selected.is_some()` guard could
|
||||
/// silently kill.
|
||||
#[test]
|
||||
fn a_saturated_capacity_exhausted_fleet_books_all_queued_through_the_ladder() {
|
||||
let owner = worker("owner");
|
||||
let shallow_owner = worker("shallow-owner");
|
||||
let workers = vec![Arc::clone(&owner), Arc::clone(&shallow_owner)];
|
||||
// Both owners are over the limit (saturation) AND out of KV, so the
|
||||
// re-admitted set yields no winner and the ladder falls through to the
|
||||
// capacity fallback. Before the ordering fix this booked `cache_miss`.
|
||||
let loads = queued_snapshot(&[
|
||||
(&owner, 9, 10_000, 10_000),
|
||||
(&shallow_owner, 5, 10_000, 10_000),
|
||||
]);
|
||||
let signal = prefix_signal(&[(&owner, 9), (&shallow_owner, 4)], 10);
|
||||
let config = AffinityConfig {
|
||||
worker_queue_limit: Some(4),
|
||||
..Default::default()
|
||||
};
|
||||
let policy = CacheAwarePolicy::new(config);
|
||||
let buckets = BucketSelector::new(None);
|
||||
let metrics = MetricsRegistry::new();
|
||||
let model = ModelId("model".into());
|
||||
let mut inputs = prefill_inputs(
|
||||
&policy,
|
||||
&buckets,
|
||||
&metrics,
|
||||
&model,
|
||||
&workers,
|
||||
Some(&loads),
|
||||
100_000,
|
||||
);
|
||||
inputs.policy_kind = PolicyKind::CacheAware;
|
||||
inputs.worker_queue_limit = Some(4);
|
||||
inputs.external_prefix = Some(&signal);
|
||||
|
||||
assert!(
|
||||
select_prefill_worker(&inputs).is_ok(),
|
||||
"a saturated fleet must still route"
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
booked_decisions(&metrics),
|
||||
vec![
|
||||
r#"sgl_router_cache_aware_decisions_total{model_id="model",decision="all_queued"} 1"#
|
||||
],
|
||||
"saturation must survive a capacity-exhausted re-admission"
|
||||
);
|
||||
}
|
||||
|
||||
/// The other side of the same wiring: an unsaturated fleet where the gate
|
||||
/// really did divert, which must book the diversion AND the prefix depth
|
||||
/// it gave up.
|
||||
#[test]
|
||||
fn a_real_diversion_books_cache_worker_queued_and_its_overlap_depth() {
|
||||
let owner = worker("owner");
|
||||
let idle = worker("idle");
|
||||
let workers = vec![Arc::clone(&owner), Arc::clone(&idle)];
|
||||
// The only prefix owner is queueing; a non-owner is idle, so a
|
||||
// diversion can dodge the wait and the fleet is NOT saturated.
|
||||
// `cache_affinity_min_matched_tokens` defaults to 1024, so the request
|
||||
// has to be large enough for a 7/10-block match to clear it, or the
|
||||
// candidate never reaches the gate at all and this books a plain miss.
|
||||
let loads = queued_snapshot(&[(&owner, 9, 10, 10_000_000), (&idle, 0, 10, 10_000_000)]);
|
||||
let signal = prefix_signal(&[(&owner, 7)], 10);
|
||||
let config = AffinityConfig {
|
||||
worker_queue_limit: Some(4),
|
||||
..Default::default()
|
||||
};
|
||||
let policy = CacheAwarePolicy::new(config);
|
||||
let buckets = BucketSelector::new(None);
|
||||
let metrics = MetricsRegistry::new();
|
||||
let model = ModelId("model".into());
|
||||
let mut inputs = prefill_inputs(
|
||||
&policy,
|
||||
&buckets,
|
||||
&metrics,
|
||||
&model,
|
||||
&workers,
|
||||
Some(&loads),
|
||||
100_000,
|
||||
);
|
||||
inputs.policy_kind = PolicyKind::CacheAware;
|
||||
inputs.worker_queue_limit = Some(4);
|
||||
inputs.external_prefix = Some(&signal);
|
||||
|
||||
let selected = select_prefill_worker(&inputs).expect("an idle worker exists");
|
||||
|
||||
assert_eq!(selected.id, idle.id, "the gate must divert off the prefix");
|
||||
assert_eq!(
|
||||
booked_decisions(&metrics),
|
||||
vec![
|
||||
r#"sgl_router_cache_aware_decisions_total{model_id="model",decision="cache_worker_queued"} 1"#
|
||||
]
|
||||
);
|
||||
// The depth given up is the evidence the histogram exists for, and a
|
||||
// diversion that never arrived must never reach it.
|
||||
let rendered = metrics.render();
|
||||
assert!(
|
||||
rendered.contains(r#"sgl_router_diverted_overlap_blocks_count{model_id="model"} 1"#),
|
||||
"a real diversion must observe its overlap depth, got:\n{rendered}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_saturated_fleet_still_routes_through_the_capacity_fallback() {
|
||||
let full = worker("full");
|
||||
@@ -667,6 +996,7 @@ mod tests {
|
||||
&SelectionProposal::with_backup(Arc::clone(&full), Arc::clone(&also_full)),
|
||||
64,
|
||||
&loads,
|
||||
None,
|
||||
)
|
||||
.is_none(),
|
||||
"fixture must saturate every worker so the strict rung admits none",
|
||||
@@ -842,4 +1172,49 @@ mod tests {
|
||||
"cache_candidate"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_aware_fallback_decision_needs_the_gate_to_have_emptied_the_set() {
|
||||
// The trap this pins: on an UNSATURATED fleet, one owner queueing
|
||||
// while the others exhaust KV capacity is a CAPACITY problem, not a
|
||||
// gate diversion. Only a gate that removed every owner leaves zero
|
||||
// candidates evaluated.
|
||||
assert!(matches!(
|
||||
cache_aware_fallback_decision(1, 3, false),
|
||||
CacheAwareDecision::CacheMiss
|
||||
));
|
||||
// Nothing gated out at all: a plain miss, saturated or not.
|
||||
assert!(matches!(
|
||||
cache_aware_fallback_decision(0, 0, false),
|
||||
CacheAwareDecision::CacheMiss
|
||||
));
|
||||
assert!(matches!(
|
||||
cache_aware_fallback_decision(0, 4, true),
|
||||
CacheAwareDecision::CacheMiss
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_aware_fallback_decision_separates_diversion_from_saturation() {
|
||||
// Gate removed every owner and somewhere unqueued exists: a real
|
||||
// diversion off the prefix. `resolve_cache_candidates` leaves
|
||||
// `evaluated` at zero here because its second tier does not fire on
|
||||
// an unsaturated fleet.
|
||||
assert!(matches!(
|
||||
cache_aware_fallback_decision(2, 0, false),
|
||||
CacheAwareDecision::CacheWorkerQueued
|
||||
));
|
||||
// Saturation, in the shape the resolver actually produces: the
|
||||
// second tier re-admitted the gated-out owners, so `evaluated` is
|
||||
// NON-zero, and they then failed hard admission. Booking the
|
||||
// capacity outcome here would drop the saturation signal exactly
|
||||
// where it matters — hence saturation is asked first. Pinning
|
||||
// `(2, 0, true)` instead would assert a state the resolver cannot
|
||||
// reach: re-admission and the `AllQueued` precondition are the same
|
||||
// condition, so an empty `evaluated` never survives it.
|
||||
assert!(matches!(
|
||||
cache_aware_fallback_decision(2, 2, true),
|
||||
CacheAwareDecision::AllQueued
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -39,8 +39,42 @@
|
||||
//! | `sgl_router_cache_pressure_guard_compared_total` | Counter | — |
|
||||
//! | `sgl_router_cache_pressure_guard_override_total` | Counter | — |
|
||||
//! | `sgl_router_cache_monitor_decisions_total` | Counter | `source` |
|
||||
//! | `sgl_router_cache_aware_decisions_total` | Counter | `model_id`, `decision` |
|
||||
//! | `sgl_router_diverted_overlap_blocks` | Histogram | `model_id` |
|
||||
//! | `sgl_router_ingress_tokenize_errors_total` | Counter | `model_id` |
|
||||
//!
|
||||
//! `sgl_router_cache_aware_decisions_total` records exactly one decision per
|
||||
//! cache-aware prefill selection that resolves a worker, so the labels sum to
|
||||
//! the cache-aware request rate less the selections that ended in a 503 (see
|
||||
//! `sgl_router_policy_selection_failures_total` for those) and ratios between
|
||||
//! them are meaningful:
|
||||
//!
|
||||
//! - `cache_hit` — a prefix owner won the selection. Note this includes a
|
||||
//! PARTIAL gate diversion: when the gate removed the deepest owner but a
|
||||
//! shallower one survived, an owner still won, so the request books here
|
||||
//! and contributes nothing to `sgl_router_diverted_overlap_blocks`.
|
||||
//! - `cache_miss` — no usable prefix owner (tree miss, or every owner
|
||||
//! rejected by hard capacity admission). A tree miss books here even on a
|
||||
//! saturated fleet: with no prefix owner the gate never fired, so there was
|
||||
//! no affinity to keep or trade. `all_queued` is the saturation signal for
|
||||
//! traffic the gate ACTED on, not a fleet-wide saturation gauge — read
|
||||
//! engine queue depth for that.
|
||||
//! - `cache_worker_queued` — the queue gate (`--worker-queue-limit`) removed
|
||||
//! every owner while an unqueued destination still existed, so the request
|
||||
//! was diverted off its prefix. The matched-prefix depth it gave up is in
|
||||
//! `sgl_router_diverted_overlap_blocks` — read it against the overlap of
|
||||
//! all selections: a diverted curve skewing high means the gate is trading
|
||||
//! large cached prefixes for short waits.
|
||||
//! - `all_queued` — the queue gate removed every owner AND every worker in
|
||||
//! the prefill fleet is queueing, so no diversion could dodge a wait. This
|
||||
//! is the saturation signal for traffic the gate acted on, keyed on
|
||||
//! saturation rather than on where the request landed: usually the request
|
||||
//! kept its prefix, but when the re-admitted owners are also out of KV
|
||||
//! capacity it lands off-owner and still books here. Reporting that case as `cache_miss` would hide the
|
||||
//! saturation in the one state where it matters most. It deliberately
|
||||
//! does NOT spell `cache_hit*`: a `decision=~"cache_hit.*"` hit-rate query
|
||||
//! must not absorb it, or a fully saturated fleet reads as a healthy one.
|
||||
//!
|
||||
//! The four `sgl_router_worker*` gauges and `sgl_router_workers` are sampled
|
||||
//! at scrape time from the live [`crate::workers::WorkerRegistry`] (passed to
|
||||
//! [`MetricsRegistry::render_with_workers`]) rather than pushed — there is no
|
||||
@@ -87,6 +121,14 @@ const TTFT_BUCKETS: &[f64] = &[
|
||||
400.0,
|
||||
];
|
||||
|
||||
/// Histogram bucket upper bounds (blocks) for
|
||||
/// `sgl_router_diverted_overlap_blocks`. Powers of two up to 8192 blocks;
|
||||
/// block size is engine-configured (commonly 16–64 tokens), so the ladder
|
||||
/// spans ~16 tokens to ~512K tokens of forfeited prefix.
|
||||
const OVERLAP_BLOCK_BUCKETS: &[f64] = &[
|
||||
1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0, 128.0, 256.0, 512.0, 1024.0, 2048.0, 4096.0, 8192.0,
|
||||
];
|
||||
|
||||
/// Recordable outcome for a request — narrowed to a handful of variants so
|
||||
/// the label cardinality stays bounded.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
@@ -223,6 +265,27 @@ pub(crate) enum PolicySelectionFailureReason {
|
||||
ProposalEmpty,
|
||||
}
|
||||
|
||||
/// Final cache-aware routing decision, one per prefill selection. See the
|
||||
/// module doc for how the labels read against each other.
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub enum CacheAwareDecision {
|
||||
CacheHit,
|
||||
CacheMiss,
|
||||
CacheWorkerQueued,
|
||||
AllQueued,
|
||||
}
|
||||
|
||||
impl CacheAwareDecision {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::CacheHit => "cache_hit",
|
||||
Self::CacheMiss => "cache_miss",
|
||||
Self::CacheWorkerQueued => "cache_worker_queued",
|
||||
Self::AllQueued => "all_queued",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PolicySelectionFailureReason {
|
||||
pub(crate) fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
@@ -280,6 +343,8 @@ pub struct MetricsRegistry {
|
||||
cache_pressure_guard_compared_total: AtomicU64,
|
||||
cache_pressure_guard_override_total: AtomicU64,
|
||||
cache_monitor_decisions_total: Mutex<HashMap<&'static str, Arc<AtomicU64>>>,
|
||||
cache_aware_decisions_total: Mutex<HashMap<CacheAwareDecisionKey, Arc<AtomicU64>>>,
|
||||
diverted_overlap_blocks: Mutex<HashMap<String, Histogram>>,
|
||||
ingress_tokenize_errors_total: Mutex<HashMap<String, Arc<AtomicU64>>>,
|
||||
}
|
||||
|
||||
@@ -344,6 +409,12 @@ struct PolicyDecisionKey {
|
||||
reason: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
|
||||
struct CacheAwareDecisionKey {
|
||||
model_id: String,
|
||||
decision: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Histogram {
|
||||
/// Bucket upper bounds this histogram observes against. Held per-instance
|
||||
@@ -621,6 +692,38 @@ impl MetricsRegistry {
|
||||
counter.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Record the final cache-aware routing decision for one prefill
|
||||
/// selection — exactly one call per cache-aware request, so the labels
|
||||
/// sum to the cache-aware request rate.
|
||||
pub fn record_cache_aware_decision(&self, model_id: &str, decision: CacheAwareDecision) {
|
||||
let key = CacheAwareDecisionKey {
|
||||
model_id: model_id.to_owned(),
|
||||
decision: decision.as_str(),
|
||||
};
|
||||
let mut guard = self.cache_aware_decisions_total.lock();
|
||||
let counter = guard
|
||||
.entry(key)
|
||||
.or_insert_with(|| Arc::new(AtomicU64::new(0)))
|
||||
.clone();
|
||||
drop(guard);
|
||||
counter.fetch_add(1, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
/// Observe the matched-prefix depth (blocks) a queue-gate diversion gave
|
||||
/// up, for `sgl_router_diverted_overlap_blocks`. Recorded ONLY when the
|
||||
/// gate emptied the candidate set (`cache_worker_queued`), so the
|
||||
/// histogram measures sacrifice rather than traffic. A PARTIAL diversion
|
||||
/// — the gate removed the deepest owner but a shallower one still won —
|
||||
/// is therefore not represented here even though some locality was given
|
||||
/// up; it books as `cache_hit`.
|
||||
pub fn observe_diverted_overlap_blocks(&self, model_id: &str, blocks: u64) {
|
||||
let mut guard = self.diverted_overlap_blocks.lock();
|
||||
let hist = guard
|
||||
.entry(model_id.to_owned())
|
||||
.or_insert_with(|| Histogram::new(OVERLAP_BLOCK_BUCKETS));
|
||||
hist.observe(blocks as f64);
|
||||
}
|
||||
|
||||
/// Bump `sgl_router_ingress_tokenize_errors_total{model_id}`.
|
||||
///
|
||||
/// Recorded ONLY when the tokenization offload SHOULD have fired but the
|
||||
@@ -1024,6 +1127,47 @@ impl MetricsRegistry {
|
||||
}
|
||||
drop(guard);
|
||||
|
||||
// cache_aware_decisions_total
|
||||
out.push_str(
|
||||
"# HELP sgl_router_cache_aware_decisions_total Final Cache-Aware routing decisions, one per selection that resolved a worker: cache_hit = prefix owner won; cache_miss = no usable owner (a tree miss books here even under saturation, because the gate never fired); cache_worker_queued = queue gate diverted the request off its prefix; all_queued = queue gate fired but every worker is queueing (fleet-saturation signal, not a cache hit).\n",
|
||||
);
|
||||
out.push_str("# TYPE sgl_router_cache_aware_decisions_total counter\n");
|
||||
let guard = self.cache_aware_decisions_total.lock();
|
||||
let mut entries: Vec<(&CacheAwareDecisionKey, u64)> = guard
|
||||
.iter()
|
||||
.map(|(k, v)| (k, v.load(Ordering::Relaxed)))
|
||||
.collect();
|
||||
entries.sort_by(|a, b| (&a.0.model_id, a.0.decision).cmp(&(&b.0.model_id, b.0.decision)));
|
||||
for (key, value) in entries {
|
||||
out.push_str(&format!(
|
||||
"sgl_router_cache_aware_decisions_total{{model_id=\"{}\",decision=\"{}\"}} {}\n",
|
||||
escape_label(&key.model_id),
|
||||
key.decision,
|
||||
value,
|
||||
));
|
||||
}
|
||||
drop(guard);
|
||||
|
||||
// diverted_overlap_blocks histogram
|
||||
out.push_str(
|
||||
"# HELP sgl_router_diverted_overlap_blocks Matched-prefix depth (blocks) given up by queue-gate diversions (decision=cache_worker_queued). Read against the overlap of all selections: a curve skewing high means the gate is trading large cached prefixes for short waits.\n",
|
||||
);
|
||||
out.push_str("# TYPE sgl_router_diverted_overlap_blocks histogram\n");
|
||||
let guard = self.diverted_overlap_blocks.lock();
|
||||
let mut models: Vec<&String> = guard.keys().collect();
|
||||
models.sort();
|
||||
for model_id in models {
|
||||
let hist = guard.get(model_id).unwrap();
|
||||
let label_body = format!("model_id=\"{}\"", escape_label(model_id));
|
||||
render_histogram(
|
||||
&mut out,
|
||||
"sgl_router_diverted_overlap_blocks",
|
||||
&label_body,
|
||||
hist,
|
||||
);
|
||||
}
|
||||
drop(guard);
|
||||
|
||||
// ingress_tokenize_errors_total
|
||||
out.push_str(
|
||||
"# HELP sgl_router_ingress_tokenize_errors_total Chat requests on a chat-encoder model whose ingress tokenization failed, silently falling back to engine-side tokenization (the input_ids offload was defeated).\n",
|
||||
@@ -1511,6 +1655,41 @@ mod tests {
|
||||
assert!(out.contains("sgl_router_cache_pressure_guard_override_total 1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_aware_decisions_and_diverted_overlap_render() {
|
||||
let reg = MetricsRegistry::new();
|
||||
reg.record_cache_aware_decision("tiny", CacheAwareDecision::CacheHit);
|
||||
reg.record_cache_aware_decision("tiny", CacheAwareDecision::CacheWorkerQueued);
|
||||
reg.record_cache_aware_decision("tiny", CacheAwareDecision::CacheWorkerQueued);
|
||||
reg.record_cache_aware_decision("tiny", CacheAwareDecision::AllQueued);
|
||||
reg.observe_diverted_overlap_blocks("tiny", 40);
|
||||
|
||||
let out = reg.render();
|
||||
assert!(out.contains(
|
||||
r#"sgl_router_cache_aware_decisions_total{model_id="tiny",decision="cache_hit"} 1"#
|
||||
));
|
||||
assert!(out.contains(
|
||||
r#"sgl_router_cache_aware_decisions_total{model_id="tiny",decision="cache_worker_queued"} 2"#
|
||||
));
|
||||
assert!(out.contains(
|
||||
r#"sgl_router_cache_aware_decisions_total{model_id="tiny",decision="all_queued"} 1"#
|
||||
));
|
||||
// The saturation label must not be absorbed by a `cache_hit.*`
|
||||
// hit-rate query.
|
||||
assert!(!out.contains(r#"decision="cache_hit_all_queued""#));
|
||||
assert!(
|
||||
out.contains(r#"sgl_router_diverted_overlap_blocks_count{model_id="tiny"} 1"#),
|
||||
"expected one diverted observation; got:\n{out}"
|
||||
);
|
||||
// 40 blocks lands in the le=64 bucket, not le=32.
|
||||
assert!(
|
||||
out.contains(r#"sgl_router_diverted_overlap_blocks_bucket{model_id="tiny",le="64"} 1"#)
|
||||
);
|
||||
assert!(
|
||||
out.contains(r#"sgl_router_diverted_overlap_blocks_bucket{model_id="tiny",le="32"} 0"#)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ingress_tokenize_error_counter_increments_per_model() {
|
||||
let reg = MetricsRegistry::new();
|
||||
|
||||
@@ -289,6 +289,18 @@ pub async fn chat_completions(
|
||||
.as_ref()
|
||||
.map(|config| config.session_affinity_mode)
|
||||
.unwrap_or(SessionAffinityMode::Bucket);
|
||||
// The queue gate (`--worker-queue-limit`) applies to the cache-aware
|
||||
// candidate resolution and, beneath it, to primary/backup admission and
|
||||
// the min-load range fallback. It does NOT reach the
|
||||
// `CapacityFallbackPowerOfTwo` last resort: by the time that fires no
|
||||
// worker in the domain is capacity-admitted, so there is no unqueued
|
||||
// destination left to prefer.
|
||||
let worker_queue_limit = ctx
|
||||
.config
|
||||
.model
|
||||
.affinity
|
||||
.as_ref()
|
||||
.and_then(|config| config.worker_queue_limit);
|
||||
// Each Bucket retry rebuilds the proposal and reruns Admission/Guard.
|
||||
let worker = select_prefill_worker(&PrefillSelectionInputs {
|
||||
policy: policy.as_ref(),
|
||||
@@ -307,6 +319,7 @@ pub async fn chat_completions(
|
||||
ttft_slo_ms,
|
||||
tps_slo,
|
||||
session_affinity_mode,
|
||||
worker_queue_limit,
|
||||
})
|
||||
.map_err(|reason| policy_selection_failed(&ctx, &model_str, reason))?;
|
||||
|
||||
|
||||
@@ -154,6 +154,7 @@ fn cache_candidate_uses_bucket_metadata_without_extend_range_filtering() {
|
||||
worker: cached,
|
||||
matched_prefix_tokens: 128,
|
||||
uncached_tokens: 128,
|
||||
matched_prefix_blocks: 8,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
};
|
||||
|
||||
@@ -124,6 +124,7 @@ impl Policy for CacheCandidatesPolicy {
|
||||
worker: Arc::clone(&self.worker),
|
||||
matched_prefix_tokens: 1,
|
||||
uncached_tokens: 1,
|
||||
matched_prefix_blocks: 1,
|
||||
candidate_range_id: "global".into(),
|
||||
max_pending_prefill_tokens: None,
|
||||
}],
|
||||
|
||||
Reference in New Issue
Block a user