mirror of
https://github.com/Mesh-LLM/mesh-llm.git
synced 2026-08-08 22:23:19 -04:00
router: weight 'auto' selection by locally observed tok/s
Before, 'auto' picked uniformly at random within the multi-digit-B
tier. On the public mesh this meant a fast MiniMax on a 4090 and a
slow 35B-A3B on an M2 Air were equally likely to be chosen, even
though we'd already measured the throughput gap in routing_metrics
and were just not reading it.
Now: each big-tier candidate is weighted by its locally observed
avg_tokens_per_second (clamped to [5, 100] tok/s so nothing fully
starves and no outlier monopolizes). Models without enough samples
(< 3) get a neutral weight so they compete fairly until data
accumulates. A 15% exploration probability ignores weights and
picks uniformly, which keeps the system from locking onto stale
rankings and guarantees cold peers see traffic.
Plumbing:
- RoutingMetrics::tps_for_model(name) -> Option<(f64, u64)>: cheap
per-model lookup that locks only the relevant shard, avoiding the
per-call HashMap allocation model_snapshots() does in the hot path.
- Node::routing_metrics() public accessor (Arc-backed, cheap).
- RoutingCandidate { name, caps, tps_hint, throughput_samples }
replaces the anonymous (&str, f64, ModelCapabilities) tuple whose
middle slot was literally always 0.0 at every populated call site.
The struct makes the tps hint a real, typed concept rather than a
dangling hook.
Behaviour preserved:
- Single-digit-B partition (smalls stay last-resort) unchanged.
- All-cold candidate pool falls back to ~uniform pick (regression
test confirms no model is starved when there's no data yet).
- Capability filtering for tools / reasoning / vision unchanged.
Plumbing per call site:
- ingress.rs + transport.rs: live routing path, look up tps_hint
from the local RoutingMetrics handle for each candidate.
- discovery.rs + integrations.rs: pre-startup paths with no live
metrics; build candidates with RoutingCandidate::unscored() so
they get the cold-neutral weight.
Tests:
- weighted_pick_all_cold_is_roughly_uniform — regression safety.
- weighted_pick_fast_wins_majority_but_slow_still_gets_some —
fast wins by >=1.5x but slow still gets >30/600 picks.
- weighted_pick_cold_model_competes_with_hot_fast — newcomer gets
>100/600 picks against an established fast peer (so it can
actually accumulate samples and earn its score).
- weighted_pick_low_sample_count_treated_as_cold — 1-sample
measurements don't dominate routing.
- candidate_weight_clamps_extremes — weight stays in [5, 100],
cold = 25.
Removed:
- shuffle_in_place (replaced by SplitMix64 + pick_weighted).
- The dishonest 0.0 f64 slot in the candidate tuple, everywhere.
Validation:
cargo fmt --all -- --check # clean
cargo check -p mesh-llm-host-runtime # clean
cargo clippy -p mesh-llm-host-runtime --lib # clean
cargo clippy -p mesh-mixture-of-agents --all-targets -- -D warnings # clean
cargo test -p mesh-llm-host-runtime --lib # 1398 passed (17 in router)
cargo test -p mesh-mixture-of-agents --lib # 59 passed (no regression)
This commit is contained in:
Michael Neale
parent
ebb167dc4c
commit
2524840991
7 changed files with 448 additions and 92 deletions
|
|
@ -255,11 +255,13 @@ async fn fetch_mesh_models(
|
|||
}
|
||||
model.clone()
|
||||
} else {
|
||||
let available: Vec<(&str, f64, crate::models::ModelCapabilities)> = models
|
||||
// Pre-startup path: no live routing metrics yet, so candidates
|
||||
// are scored as cold (uniform weight).
|
||||
let available: Vec<crate::network::router::RoutingCandidate<'_>> = models
|
||||
.iter()
|
||||
.map(|name| {
|
||||
let caps = crate::models::installed_model_capabilities(name);
|
||||
(name.as_str(), 0.0, caps)
|
||||
crate::network::router::RoutingCandidate::unscored(name.as_str(), caps)
|
||||
})
|
||||
.collect();
|
||||
let agentic = crate::network::router::Classification {
|
||||
|
|
|
|||
|
|
@ -1851,6 +1851,12 @@ impl Node {
|
|||
.load(std::sync::atomic::Ordering::Relaxed) as u64
|
||||
}
|
||||
|
||||
/// Locally observed routing metrics, used by the auto-router to score
|
||||
/// models by their measured throughput from this node's perspective.
|
||||
pub fn routing_metrics(&self) -> &crate::network::metrics::RoutingMetrics {
|
||||
&self.routing_metrics
|
||||
}
|
||||
|
||||
pub fn inflight_change_rx(&self) -> watch::Receiver<u64> {
|
||||
self.inflight_change_tx.subscribe()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -422,6 +422,26 @@ impl RoutingMetrics {
|
|||
}
|
||||
}
|
||||
|
||||
/// Cheap per-model throughput lookup for routing decisions.
|
||||
///
|
||||
/// Returns `(avg_tokens_per_second, throughput_samples)` if the model has
|
||||
/// observed throughput, `None` if the model is unknown or has never
|
||||
/// recorded a token-bearing attempt. Avoids the per-call HashMap
|
||||
/// allocation that [`model_snapshots`](Self::model_snapshots) does —
|
||||
/// callers in the routing hot path can poll this once per candidate
|
||||
/// without rebuilding every model's full snapshot.
|
||||
pub fn tps_for_model(&self, model: &str) -> Option<(f64, u64)> {
|
||||
let shard_index = self.shard_index(model);
|
||||
let shard = self.shards[shard_index].lock().unwrap();
|
||||
let metrics = shard.models.get(model)?;
|
||||
let samples = metrics.throughput_samples;
|
||||
if samples == 0 {
|
||||
return None;
|
||||
}
|
||||
let tps = average_milli(metrics.throughput_tps_milli_sum, samples)?;
|
||||
Some((tps, samples))
|
||||
}
|
||||
|
||||
fn shard_index(&self, model: &str) -> usize {
|
||||
let mut hasher = DefaultHasher::new();
|
||||
model.hash(&mut hasher);
|
||||
|
|
|
|||
|
|
@ -183,15 +183,23 @@ pub(crate) async fn api_proxy(
|
|||
}
|
||||
}
|
||||
}
|
||||
let available: Vec<(&str, f64, crate::models::ModelCapabilities)> =
|
||||
available_models
|
||||
.iter()
|
||||
.map(|name| {
|
||||
let caps =
|
||||
proxy::capabilities_for_model(name, &descriptors);
|
||||
(name.as_str(), 0.0, caps)
|
||||
})
|
||||
.collect();
|
||||
let routing_metrics = node.routing_metrics();
|
||||
let available: Vec<router::RoutingCandidate<'_>> = available_models
|
||||
.iter()
|
||||
.map(|name| {
|
||||
let caps = proxy::capabilities_for_model(name, &descriptors);
|
||||
let (tps_hint, throughput_samples) = routing_metrics
|
||||
.tps_for_model(name)
|
||||
.map(|(tps, samples)| (Some(tps), samples))
|
||||
.unwrap_or((None, 0));
|
||||
router::RoutingCandidate {
|
||||
name: name.as_str(),
|
||||
caps,
|
||||
tps_hint,
|
||||
throughput_samples,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let Some(available) =
|
||||
router::filter_media_compatible_candidates(&available, &media)
|
||||
else {
|
||||
|
|
|
|||
|
|
@ -2386,11 +2386,21 @@ pub async fn handle_mesh_request(
|
|||
Some(name)
|
||||
} else {
|
||||
let cl = router::classify(body_json);
|
||||
let with_caps: Vec<(&str, f64, crate::models::ModelCapabilities)> = served
|
||||
let routing_metrics = node.routing_metrics();
|
||||
let with_caps: Vec<router::RoutingCandidate<'_>> = served
|
||||
.iter()
|
||||
.map(|name| {
|
||||
let caps = capabilities_for_model(name, &descriptors);
|
||||
(name.as_str(), 0.0, caps)
|
||||
let (tps_hint, throughput_samples) = routing_metrics
|
||||
.tps_for_model(name)
|
||||
.map(|(tps, samples)| (Some(tps), samples))
|
||||
.unwrap_or((None, 0));
|
||||
router::RoutingCandidate {
|
||||
name: name.as_str(),
|
||||
caps,
|
||||
tps_hint,
|
||||
throughput_samples,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
let Some(available) =
|
||||
|
|
|
|||
|
|
@ -381,12 +381,12 @@ pub(crate) fn model_satisfies_media_requirements(
|
|||
}
|
||||
|
||||
pub(crate) fn filter_media_compatible_candidates<'a>(
|
||||
candidates: &[(&'a str, f64, crate::models::ModelCapabilities)],
|
||||
candidates: &[RoutingCandidate<'a>],
|
||||
media: &MediaRequirements,
|
||||
) -> Option<Vec<(&'a str, f64, crate::models::ModelCapabilities)>> {
|
||||
) -> Option<Vec<RoutingCandidate<'a>>> {
|
||||
let media_available: Vec<_> = candidates
|
||||
.iter()
|
||||
.filter(|(_, _, caps)| model_satisfies_media_requirements(caps, media))
|
||||
.filter(|c| model_satisfies_media_requirements(&c.caps, media))
|
||||
.cloned()
|
||||
.collect();
|
||||
if media_available.is_empty() && media.requires_runtime_modality() {
|
||||
|
|
@ -449,8 +449,64 @@ fn message_text(msg: &Value) -> String {
|
|||
|
||||
// ── Model selection ─────────────────────────────────────────────────
|
||||
|
||||
/// Pick the best model using full classification (category + complexity + tools).
|
||||
/// Pick the best model for a classified request using gossiped capabilities.
|
||||
/// A candidate model in the auto-routing pool.
|
||||
///
|
||||
/// `tps_hint` and `throughput_samples` come from the node's locally
|
||||
/// observed `RoutingMetrics`. They are `None` / `0` when we've never
|
||||
/// successfully completed a token-bearing request against this model
|
||||
/// (cold start, brand-new peer) — such candidates get a neutral weight
|
||||
/// so they still participate in routing while accumulating data.
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RoutingCandidate<'a> {
|
||||
pub name: &'a str,
|
||||
pub caps: crate::models::ModelCapabilities,
|
||||
/// Locally observed throughput in tokens/sec. `None` if no
|
||||
/// throughput-bearing attempts have completed for this model yet.
|
||||
pub tps_hint: Option<f64>,
|
||||
/// How many throughput samples back `tps_hint`. Used to decide
|
||||
/// whether we trust the hint or treat the model as "cold".
|
||||
pub throughput_samples: u64,
|
||||
}
|
||||
|
||||
impl<'a> RoutingCandidate<'a> {
|
||||
/// Build a candidate without any throughput hint. Useful for
|
||||
/// pre-startup paths or test fixtures.
|
||||
pub fn unscored(name: &'a str, caps: crate::models::ModelCapabilities) -> Self {
|
||||
Self {
|
||||
name,
|
||||
caps,
|
||||
tps_hint: None,
|
||||
throughput_samples: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Minimum number of throughput samples before `tps_hint` is allowed
|
||||
/// to influence weighting. Below this, the candidate is treated as
|
||||
/// cold (neutral weight).
|
||||
const TPS_MIN_SAMPLES: u64 = 3;
|
||||
|
||||
/// Lower clamp on tps used as a weight. Prevents catastrophic peers
|
||||
/// (~1 tok/s) from being completely starved — they still get the
|
||||
/// occasional request so they can re-prove themselves.
|
||||
const TPS_WEIGHT_MIN: f64 = 5.0;
|
||||
|
||||
/// Upper clamp on tps used as a weight. Prevents a single very fast
|
||||
/// outlier from monopolizing routing.
|
||||
const TPS_WEIGHT_MAX: f64 = 100.0;
|
||||
|
||||
/// Neutral weight assigned to cold / unscored candidates so they get
|
||||
/// a fair shot at picking up data. Set to the midpoint of the clamp
|
||||
/// range so a cold model is treated as "average" against scored peers.
|
||||
const TPS_NEUTRAL_WEIGHT: f64 = 25.0;
|
||||
|
||||
/// Probability of ignoring weights and picking uniformly. Keeps the
|
||||
/// system from locking onto stale rankings and gives cold models a
|
||||
/// guaranteed share of traffic so they accumulate data.
|
||||
const EXPLORATION_PROBABILITY: f64 = 0.15;
|
||||
|
||||
/// Pick the best model for a classified request using gossiped capabilities
|
||||
/// and locally observed throughput.
|
||||
///
|
||||
/// Filtering:
|
||||
/// - `needs_tools` → prefer models with `tool_use != None`
|
||||
|
|
@ -458,11 +514,14 @@ fn message_text(msg: &Value) -> String {
|
|||
/// - `Image` → prefer models with `vision != None`
|
||||
/// - anything else → no capability filter
|
||||
///
|
||||
/// Falls back to all models if the filter matches nothing.
|
||||
/// Among candidates, pick randomly to spread load.
|
||||
/// Falls back to all models if the filter matches nothing. Then biases
|
||||
/// toward larger models by partitioning single-digit-B names to the
|
||||
/// bottom tier. Within the chosen tier, picks weighted by observed
|
||||
/// tok/s (with cold models treated as average), with a configurable
|
||||
/// exploration probability that ignores weights and picks uniformly.
|
||||
pub fn pick_model_classified<'a>(
|
||||
classification: &Classification,
|
||||
available_models: &[(&'a str, f64, crate::models::ModelCapabilities)],
|
||||
available_models: &[RoutingCandidate<'a>],
|
||||
) -> Option<&'a str> {
|
||||
use crate::models::CapabilityLevel;
|
||||
|
||||
|
|
@ -470,29 +529,28 @@ pub fn pick_model_classified<'a>(
|
|||
return None;
|
||||
}
|
||||
if available_models.len() == 1 {
|
||||
return Some(available_models[0].0);
|
||||
return Some(available_models[0].name);
|
||||
}
|
||||
|
||||
// Capability filter based on what the request needs
|
||||
let filtered: Vec<&(&str, f64, crate::models::ModelCapabilities)> =
|
||||
match classification.category {
|
||||
_ if classification.needs_tools => available_models
|
||||
.iter()
|
||||
.filter(|(_, _, caps)| caps.tool_use != CapabilityLevel::None)
|
||||
.collect(),
|
||||
Category::Reasoning => available_models
|
||||
.iter()
|
||||
.filter(|(_, _, caps)| caps.reasoning != CapabilityLevel::None)
|
||||
.collect(),
|
||||
Category::Image => available_models
|
||||
.iter()
|
||||
.filter(|(_, _, caps)| caps.vision != CapabilityLevel::None)
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
// Capability filter based on what the request needs.
|
||||
let filtered: Vec<&RoutingCandidate<'a>> = match classification.category {
|
||||
_ if classification.needs_tools => available_models
|
||||
.iter()
|
||||
.filter(|c| c.caps.tool_use != CapabilityLevel::None)
|
||||
.collect(),
|
||||
Category::Reasoning => available_models
|
||||
.iter()
|
||||
.filter(|c| c.caps.reasoning != CapabilityLevel::None)
|
||||
.collect(),
|
||||
Category::Image => available_models
|
||||
.iter()
|
||||
.filter(|c| c.caps.vision != CapabilityLevel::None)
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
|
||||
// Fall back to all models if filter matched nothing
|
||||
let candidates: Vec<&(&str, f64, crate::models::ModelCapabilities)> = if filtered.is_empty() {
|
||||
// Fall back to all models if the filter matched nothing.
|
||||
let candidates: Vec<&RoutingCandidate<'a>> = if filtered.is_empty() {
|
||||
available_models.iter().collect()
|
||||
} else {
|
||||
filtered
|
||||
|
|
@ -502,21 +560,103 @@ pub fn pick_model_classified<'a>(
|
|||
// parameter count (e.g. "2B", "9B") go to the bottom. Everything
|
||||
// else — multi-digit billions (31B, 70B) or names that don't encode
|
||||
// a size at all (MiniMax, Coder-Next, fine-tune tags) — stays on
|
||||
// top. Each tier is shuffled independently so sessions organically
|
||||
// spread across the strong-tier models over time while smalls still
|
||||
// act as a fallback when nothing stronger is around.
|
||||
let (mut big, mut small): (Vec<_>, Vec<_>) = candidates
|
||||
// top. The big tier is sampled by tok/s-weighted draw; the small
|
||||
// tier acts only as a fallback when the big tier is empty.
|
||||
let (big, small): (Vec<_>, Vec<_>) = candidates
|
||||
.into_iter()
|
||||
.partition(|(name, _, _)| !is_single_digit_b_name(name));
|
||||
.partition(|c| !is_single_digit_b_name(c.name));
|
||||
|
||||
let nanos = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.subsec_nanos() as u64;
|
||||
shuffle_in_place(&mut big, nanos);
|
||||
shuffle_in_place(&mut small, nanos.wrapping_add(0x9E37_79B9_7F4A_7C15));
|
||||
|
||||
big.into_iter().chain(small).next().map(|&(n, _, _)| n)
|
||||
if !big.is_empty() {
|
||||
return Some(pick_weighted(&big, nanos));
|
||||
}
|
||||
if !small.is_empty() {
|
||||
return Some(pick_weighted(
|
||||
&small,
|
||||
nanos.wrapping_add(0x9E37_79B9_7F4A_7C15),
|
||||
));
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Compute the routing weight for a single candidate. See the module-level
|
||||
/// `TPS_*` constants for the rationale on each clamp.
|
||||
fn candidate_weight(candidate: &RoutingCandidate<'_>) -> f64 {
|
||||
if candidate.throughput_samples >= TPS_MIN_SAMPLES {
|
||||
candidate
|
||||
.tps_hint
|
||||
.unwrap_or(TPS_NEUTRAL_WEIGHT)
|
||||
.clamp(TPS_WEIGHT_MIN, TPS_WEIGHT_MAX)
|
||||
} else {
|
||||
TPS_NEUTRAL_WEIGHT
|
||||
}
|
||||
}
|
||||
|
||||
/// Pick one candidate from a non-empty slice using tok/s-weighted draw,
|
||||
/// with `EXPLORATION_PROBABILITY` chance of a uniform pick.
|
||||
fn pick_weighted<'a>(candidates: &[&RoutingCandidate<'a>], seed: u64) -> &'a str {
|
||||
debug_assert!(!candidates.is_empty(), "pick_weighted requires non-empty");
|
||||
|
||||
let mut rng = SplitMix64::new(seed);
|
||||
|
||||
// Exploration branch: ignore weights, pick uniformly. Keeps the
|
||||
// system from locking onto stale rankings.
|
||||
if rng.next_f64() < EXPLORATION_PROBABILITY {
|
||||
let idx = (rng.next_u64() as usize) % candidates.len();
|
||||
return candidates[idx].name;
|
||||
}
|
||||
|
||||
let total_weight: f64 = candidates.iter().map(|c| candidate_weight(c)).sum();
|
||||
// Defensive: if all weights are somehow zero (shouldn't happen given
|
||||
// TPS_WEIGHT_MIN > 0), fall back to a uniform pick.
|
||||
if total_weight <= 0.0 {
|
||||
let idx = (rng.next_u64() as usize) % candidates.len();
|
||||
return candidates[idx].name;
|
||||
}
|
||||
|
||||
let pick = rng.next_f64() * total_weight;
|
||||
let mut acc = 0.0;
|
||||
for c in candidates {
|
||||
acc += candidate_weight(c);
|
||||
if pick < acc {
|
||||
return c.name;
|
||||
}
|
||||
}
|
||||
// Numerical tail: pick the last candidate.
|
||||
candidates[candidates.len() - 1].name
|
||||
}
|
||||
|
||||
/// Small deterministic PRNG so a single seed drives both the
|
||||
/// exploration coin flip and the weighted draw. Avoids pulling in a
|
||||
/// rand dependency just for routing.
|
||||
struct SplitMix64 {
|
||||
state: u64,
|
||||
}
|
||||
|
||||
impl SplitMix64 {
|
||||
fn new(seed: u64) -> Self {
|
||||
// Avoid the zero state which gives a degenerate sequence.
|
||||
Self {
|
||||
state: seed.wrapping_add(0x9E37_79B9_7F4A_7C15),
|
||||
}
|
||||
}
|
||||
|
||||
fn next_u64(&mut self) -> u64 {
|
||||
self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
|
||||
let mut z = self.state;
|
||||
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
|
||||
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
|
||||
z ^ (z >> 31)
|
||||
}
|
||||
|
||||
fn next_f64(&mut self) -> f64 {
|
||||
// Use the top 53 bits for a uniform float in [0, 1).
|
||||
(self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
|
||||
}
|
||||
}
|
||||
|
||||
/// Return true if `name` advertises a single-digit billion-parameter
|
||||
|
|
@ -570,21 +710,6 @@ fn is_single_digit_b_name(name: &str) -> bool {
|
|||
false
|
||||
}
|
||||
|
||||
/// In-place Fisher-Yates shuffle seeded from `seed`.
|
||||
fn shuffle_in_place<T>(items: &mut [T], seed: u64) {
|
||||
if items.len() < 2 {
|
||||
return;
|
||||
}
|
||||
let mut state = seed.wrapping_mul(0x2545_F491_4F6C_DD1D).wrapping_add(1);
|
||||
for i in (1..items.len()).rev() {
|
||||
state ^= state << 13;
|
||||
state ^= state >> 7;
|
||||
state ^= state << 17;
|
||||
let j = (state as usize) % (i + 1);
|
||||
items.swap(i, j);
|
||||
}
|
||||
}
|
||||
|
||||
// ── Tests ───────────────────────────────────────────────────────────
|
||||
|
||||
#[cfg(test)]
|
||||
|
|
@ -873,18 +998,23 @@ mod tests {
|
|||
};
|
||||
let text_only = MediaRequirements::default();
|
||||
|
||||
let text_candidates = vec![("text", 0.0, text_caps)];
|
||||
let text_candidates = vec![RoutingCandidate::unscored("text", text_caps)];
|
||||
assert!(filter_media_compatible_candidates(&text_candidates, &image).is_none());
|
||||
|
||||
let mixed_candidates = vec![("text", 0.0, text_caps), ("vision", 0.0, vision_caps)];
|
||||
let mixed_candidates = vec![
|
||||
RoutingCandidate::unscored("text", text_caps),
|
||||
RoutingCandidate::unscored("vision", vision_caps),
|
||||
];
|
||||
let filtered = filter_media_compatible_candidates(&mixed_candidates, &image)
|
||||
.expect("vision candidate should satisfy image media request");
|
||||
assert_eq!(filtered.len(), 1);
|
||||
assert_eq!(filtered[0].0, "vision");
|
||||
assert_eq!(filtered[0].name, "vision");
|
||||
|
||||
let unfiltered = filter_media_compatible_candidates(&text_candidates, &text_only)
|
||||
.expect("text-only requests should keep normal router fallback behavior");
|
||||
assert_eq!(unfiltered, text_candidates);
|
||||
// Text-only requests with text-only candidates pass through unfiltered.
|
||||
assert_eq!(unfiltered.len(), text_candidates.len());
|
||||
assert_eq!(unfiltered[0].name, text_candidates[0].name);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
@ -898,8 +1028,8 @@ mod tests {
|
|||
let no_caps = ModelCapabilities::default();
|
||||
|
||||
let available = vec![
|
||||
("reasoning-model", 10.0, no_caps),
|
||||
("tool-model", 10.0, tool_caps),
|
||||
RoutingCandidate::unscored("reasoning-model", no_caps),
|
||||
RoutingCandidate::unscored("tool-model", tool_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Code,
|
||||
|
|
@ -922,8 +1052,8 @@ mod tests {
|
|||
let no_caps = ModelCapabilities::default();
|
||||
|
||||
let available = vec![
|
||||
("chat-model", 10.0, no_caps),
|
||||
("reasoning-model", 10.0, reasoning_caps),
|
||||
RoutingCandidate::unscored("chat-model", no_caps),
|
||||
RoutingCandidate::unscored("reasoning-model", reasoning_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Reasoning,
|
||||
|
|
@ -946,8 +1076,8 @@ mod tests {
|
|||
let no_caps = ModelCapabilities::default();
|
||||
|
||||
let available = vec![
|
||||
("text-model", 10.0, no_caps),
|
||||
("vision-model", 10.0, vision_caps),
|
||||
RoutingCandidate::unscored("text-model", no_caps),
|
||||
RoutingCandidate::unscored("vision-model", vision_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Image,
|
||||
|
|
@ -964,7 +1094,10 @@ mod tests {
|
|||
use crate::models::ModelCapabilities;
|
||||
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![("model-a", 10.0, no_caps), ("model-b", 10.0, no_caps)];
|
||||
let available = vec![
|
||||
RoutingCandidate::unscored("model-a", no_caps),
|
||||
RoutingCandidate::unscored("model-b", no_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Code,
|
||||
complexity: Complexity::Moderate,
|
||||
|
|
@ -978,7 +1111,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn test_pick_empty_returns_none() {
|
||||
let available: Vec<(&str, f64, crate::models::ModelCapabilities)> = vec![];
|
||||
let available: Vec<RoutingCandidate<'_>> = vec![];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
complexity: Complexity::Moderate,
|
||||
|
|
@ -992,7 +1125,10 @@ mod tests {
|
|||
fn test_pick_single_model() {
|
||||
use crate::models::ModelCapabilities;
|
||||
|
||||
let available = vec![("only-model", 10.0, ModelCapabilities::default())];
|
||||
let available = vec![RoutingCandidate::unscored(
|
||||
"only-model",
|
||||
ModelCapabilities::default(),
|
||||
)];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
complexity: Complexity::Moderate,
|
||||
|
|
@ -1007,7 +1143,10 @@ mod tests {
|
|||
use crate::models::ModelCapabilities;
|
||||
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![("model-a", 10.0, no_caps), ("model-b", 10.0, no_caps)];
|
||||
let available = vec![
|
||||
RoutingCandidate::unscored("model-a", no_caps),
|
||||
RoutingCandidate::unscored("model-b", no_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
complexity: Complexity::Moderate,
|
||||
|
|
@ -1074,12 +1213,12 @@ mod tests {
|
|||
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
("Qwen3.5-2B-Q4_K_M", 0.0, no_caps),
|
||||
("Qwen3.5-9B-Q4_K_M", 0.0, no_caps),
|
||||
("gemma-4-31B-it-Q8_0", 0.0, no_caps),
|
||||
("Qwen3.6-35B-A3B-BF16", 0.0, no_caps),
|
||||
("MiniMax-M2.5-Q4_K_M", 0.0, no_caps),
|
||||
("Qwen3-Coder-Next-Q4_K_M", 0.0, no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.5-2B-Q4_K_M", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.5-9B-Q4_K_M", no_caps),
|
||||
RoutingCandidate::unscored("gemma-4-31B-it-Q8_0", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.6-35B-A3B-BF16", no_caps),
|
||||
RoutingCandidate::unscored("MiniMax-M2.5-Q4_K_M", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3-Coder-Next-Q4_K_M", no_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
|
|
@ -1105,8 +1244,8 @@ mod tests {
|
|||
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
("Qwen3.5-2B-Q4_K_M", 0.0, no_caps),
|
||||
("Qwen3.5-9B-Q4_K_M", 0.0, no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.5-2B-Q4_K_M", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.5-9B-Q4_K_M", no_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
|
|
@ -1126,10 +1265,10 @@ mod tests {
|
|||
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
("gemma-4-31B-it-Q8_0", 0.0, no_caps),
|
||||
("Qwen3.6-35B-A3B-BF16", 0.0, no_caps),
|
||||
("MiniMax-M2.5-Q4_K_M", 0.0, no_caps),
|
||||
("Qwen3-Coder-Next-Q4_K_M", 0.0, no_caps),
|
||||
RoutingCandidate::unscored("gemma-4-31B-it-Q8_0", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3.6-35B-A3B-BF16", no_caps),
|
||||
RoutingCandidate::unscored("MiniMax-M2.5-Q4_K_M", no_caps),
|
||||
RoutingCandidate::unscored("Qwen3-Coder-Next-Q4_K_M", no_caps),
|
||||
];
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
|
|
@ -1154,4 +1293,173 @@ mod tests {
|
|||
"expected spread across big-tier models, only saw {seen:?}"
|
||||
);
|
||||
}
|
||||
|
||||
// ── tok/s-aware weighting ────────────────────────────────────────
|
||||
|
||||
/// Helper: build a scored candidate.
|
||||
fn scored<'a>(
|
||||
name: &'a str,
|
||||
caps: crate::models::ModelCapabilities,
|
||||
tps: f64,
|
||||
samples: u64,
|
||||
) -> RoutingCandidate<'a> {
|
||||
RoutingCandidate {
|
||||
name,
|
||||
caps,
|
||||
tps_hint: Some(tps),
|
||||
throughput_samples: samples,
|
||||
}
|
||||
}
|
||||
|
||||
fn count_picks(available: &[RoutingCandidate<'_>], iterations: usize) -> HashMapCounts {
|
||||
use std::collections::HashMap;
|
||||
let cl = Classification {
|
||||
category: Category::Chat,
|
||||
complexity: Complexity::Moderate,
|
||||
needs_tools: false,
|
||||
has_media_inputs: false,
|
||||
};
|
||||
let mut counts: HashMap<String, usize> = HashMap::new();
|
||||
for _ in 0..iterations {
|
||||
if let Some(name) = pick_model_classified(&cl, available) {
|
||||
*counts.entry(name.to_string()).or_insert(0) += 1;
|
||||
}
|
||||
// Bump the nanosecond seed between iterations.
|
||||
std::thread::sleep(std::time::Duration::from_nanos(1));
|
||||
}
|
||||
HashMapCounts(counts)
|
||||
}
|
||||
|
||||
struct HashMapCounts(std::collections::HashMap<String, usize>);
|
||||
|
||||
impl HashMapCounts {
|
||||
fn get(&self, name: &str) -> usize {
|
||||
self.0.get(name).copied().unwrap_or(0)
|
||||
}
|
||||
fn total(&self) -> usize {
|
||||
self.0.values().sum()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weighted_pick_all_cold_is_roughly_uniform() {
|
||||
// Backwards-compat: when nothing has tps data, picks should be
|
||||
// roughly uniform — same effective shape as the old random shuffle.
|
||||
use crate::models::ModelCapabilities;
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
RoutingCandidate::unscored("alpha-31B", no_caps),
|
||||
RoutingCandidate::unscored("beta-31B", no_caps),
|
||||
RoutingCandidate::unscored("gamma-31B", no_caps),
|
||||
];
|
||||
let counts = count_picks(&available, 600);
|
||||
let expected = counts.total() / 3;
|
||||
// Allow ±50% (300 picks across 3 models is loose statistical ground,
|
||||
// but we just need to see no model is starved).
|
||||
for name in ["alpha-31B", "beta-31B", "gamma-31B"] {
|
||||
let got = counts.get(name);
|
||||
assert!(
|
||||
got > expected / 2,
|
||||
"cold model {name} was starved: {got}/{expected} expected"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weighted_pick_fast_wins_majority_but_slow_still_gets_some() {
|
||||
// Core design claim: fast tok/s tilts routing without starving the slow.
|
||||
use crate::models::ModelCapabilities;
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
scored("fast-31B", no_caps, 80.0, 50),
|
||||
scored("slow-31B", no_caps, 6.0, 50),
|
||||
];
|
||||
let counts = count_picks(&available, 600);
|
||||
let fast = counts.get("fast-31B");
|
||||
let slow = counts.get("slow-31B");
|
||||
// Fast should win clearly more often.
|
||||
assert!(
|
||||
fast > slow,
|
||||
"fast tok/s model should win majority: fast={fast} slow={slow}"
|
||||
);
|
||||
// Fast wins by a meaningful margin (≥ 1.5x).
|
||||
assert!(
|
||||
fast as f64 > 1.5 * slow as f64,
|
||||
"fast model should win by at least 1.5x: fast={fast} slow={slow}",
|
||||
);
|
||||
// Slow model still gets meaningful traffic (exploration + clamp keep it alive).
|
||||
assert!(
|
||||
slow > 30,
|
||||
"slow model must not be starved (exploration keeps it alive): got {slow}",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weighted_pick_cold_model_competes_with_hot_fast() {
|
||||
// A brand-new peer (no samples) must still get meaningful traffic
|
||||
// against an established fast peer — otherwise it can never
|
||||
// accumulate the data it needs to be scored.
|
||||
use crate::models::ModelCapabilities;
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
scored("hot-fast-31B", no_caps, 80.0, 50),
|
||||
RoutingCandidate::unscored("cold-newcomer-31B", no_caps),
|
||||
];
|
||||
let counts = count_picks(&available, 600);
|
||||
let cold = counts.get("cold-newcomer-31B");
|
||||
// Cold gets NEUTRAL_WEIGHT (25) vs hot's clamped 80 — so cold
|
||||
// should still see at least a healthy minority of traffic.
|
||||
assert!(
|
||||
cold > 100,
|
||||
"cold newcomer must get fair traffic to accumulate samples: got {cold}/600"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn weighted_pick_low_sample_count_treated_as_cold() {
|
||||
// A model with only 1-2 samples shouldn't have those samples
|
||||
// dominate routing — we want a few real measurements before tps
|
||||
// participates.
|
||||
use crate::models::ModelCapabilities;
|
||||
let no_caps = ModelCapabilities::default();
|
||||
let available = vec![
|
||||
// Both are 31B "big-tier" names — single-digit-B partition
|
||||
// doesn't separate them.
|
||||
scored("alpha-31B", no_caps, 100.0, 1), // 1 sample of "fast" — should be ignored
|
||||
scored("beta-31B", no_caps, 100.0, 1),
|
||||
scored("gamma-31B", no_caps, 100.0, 1),
|
||||
];
|
||||
let counts = count_picks(&available, 600);
|
||||
// All three should land near uniform since none has enough samples.
|
||||
let expected = counts.total() / 3;
|
||||
for name in ["alpha-31B", "beta-31B", "gamma-31B"] {
|
||||
let got = counts.get(name);
|
||||
assert!(
|
||||
got > expected / 2,
|
||||
"low-sample model {name} was treated as scored instead of cold: {got}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn candidate_weight_clamps_extremes() {
|
||||
// Sanity: weight stays bounded so no peer can fully starve or monopolize.
|
||||
use crate::models::ModelCapabilities;
|
||||
let no_caps = ModelCapabilities::default();
|
||||
|
||||
let glacial = scored("glacial", no_caps, 0.5, 100);
|
||||
let blazing = scored("blazing", no_caps, 500.0, 100);
|
||||
let cold = RoutingCandidate::unscored("cold", no_caps);
|
||||
|
||||
let wg = candidate_weight(&glacial);
|
||||
let wb = candidate_weight(&blazing);
|
||||
let wc = candidate_weight(&cold);
|
||||
|
||||
assert!(wg >= TPS_WEIGHT_MIN, "glacial weight floored: {wg}");
|
||||
assert!(wb <= TPS_WEIGHT_MAX, "blazing weight capped: {wb}");
|
||||
assert!(
|
||||
(wc - TPS_NEUTRAL_WEIGHT).abs() < f64::EPSILON,
|
||||
"cold weight should be neutral: {wc}",
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -272,11 +272,13 @@ pub(crate) async fn check_mesh(
|
|||
}
|
||||
m.clone()
|
||||
} else {
|
||||
let available: Vec<(&str, f64, crate::models::ModelCapabilities)> = models
|
||||
// Pre-startup path: no live routing metrics yet, so candidates
|
||||
// are scored as cold (uniform weight).
|
||||
let available: Vec<router::RoutingCandidate<'_>> = models
|
||||
.iter()
|
||||
.map(|n| {
|
||||
let caps = crate::models::installed_model_capabilities(n);
|
||||
(n.as_str(), 0.0, caps)
|
||||
router::RoutingCandidate::unscored(n.as_str(), caps)
|
||||
})
|
||||
.collect();
|
||||
let agentic = router::Classification {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue