258 lines
9.6 KiB
Python
258 lines
9.6 KiB
Python
"""Learned route statistics for dynamic bandit-style route selection (ADR-0021).
|
|
|
|
The tracker treats each viable route (ordered chain of node shards covering a
|
|
model) as a bandit arm. Observed end-to-end tokens/sec per route is kept as a
|
|
time-decayed EWMA. Selection splits traffic between:
|
|
|
|
- **exploit**: weighted-random among *proven* routes, weight ∝ tps ** alpha
|
|
(alpha=1.0 → a 1.5x-faster route gets 1.5x the traffic);
|
|
- **scout**: with probability `explore_share`, the least-measured unproven or
|
|
stale route is chosen so the tracker keeps learning as the network morphs.
|
|
|
|
Staleness has two mechanisms:
|
|
- continuous: sample mass decays with `stats_half_life_seconds`, so old
|
|
observations fade;
|
|
- abrupt: every node join/leave bumps the model's *topology epoch*; stats from
|
|
an older epoch keep their EWMA as a prior but drop back into the scout pool
|
|
until re-measured.
|
|
|
|
Route signatures embed node ids and shard ranges, so a node re-registering
|
|
with a different shard produces a new arm automatically.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import random
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Iterable
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RoutingConfig:
|
|
explore_share: float = 0.3
|
|
weight_alpha: float = 1.0
|
|
stats_half_life_seconds: float = 600.0
|
|
min_sample_tokens: int = 8
|
|
# One fresh sample has mass 1.0 and decays from there; 0.5 keeps a single
|
|
# observation "proven" for one half-life before demoting it to the scout pool.
|
|
min_proven_weight: float = 0.5
|
|
max_candidate_routes: int = 8
|
|
prune_after_seconds: float = 86400.0
|
|
|
|
|
|
@dataclass
|
|
class RouteStat:
|
|
ewma_tps: float = 0.0
|
|
weight: float = 0.0 # decayed effective sample mass
|
|
last_sample_ts: float = 0.0
|
|
epoch: int = 0
|
|
samples: int = 0 # lifetime raw sample count (display only)
|
|
|
|
def decayed_weight(self, now: float, half_life: float) -> float:
|
|
if self.weight <= 0.0:
|
|
return 0.0
|
|
age = max(0.0, now - self.last_sample_ts)
|
|
return self.weight * 0.5 ** (age / half_life)
|
|
|
|
|
|
@dataclass
|
|
class RouteCandidate:
|
|
nodes: list[Any]
|
|
signature: str
|
|
prior_tps: float = 0.0
|
|
|
|
|
|
def route_signature(model_key: str, nodes: Iterable[Any]) -> str:
|
|
hops = "->".join(
|
|
f"{getattr(n, 'node_id', '?')}[{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}]"
|
|
for n in nodes
|
|
)
|
|
return f"{model_key}|{hops}"
|
|
|
|
|
|
class RouteStatsStore:
|
|
"""Thread-safe per-route decayed throughput statistics."""
|
|
|
|
def __init__(self, config: RoutingConfig | None = None) -> None:
|
|
self.config = config or RoutingConfig()
|
|
self._lock = threading.Lock()
|
|
self._stats: dict[str, RouteStat] = {}
|
|
self._epochs: dict[str, int] = {}
|
|
|
|
def epoch(self, model_key: str) -> int:
|
|
with self._lock:
|
|
return self._epochs.get(model_key, 0)
|
|
|
|
def bump_epoch(self, model_keys: Iterable[str | None]) -> None:
|
|
"""Mark the topology changed for the given model keys (node join/leave)."""
|
|
with self._lock:
|
|
for key in model_keys:
|
|
if key:
|
|
self._epochs[key] = self._epochs.get(key, 0) + 1
|
|
|
|
def record_sample(
|
|
self,
|
|
model_key: str,
|
|
signature: str,
|
|
tokens: int,
|
|
elapsed_seconds: float,
|
|
now: float | None = None,
|
|
) -> bool:
|
|
"""Fold one completed request into the route's EWMA.
|
|
|
|
Returns False (and records nothing) for samples below
|
|
`min_sample_tokens` — near-empty completions come from broken routes
|
|
and would poison the arm with meaningless throughput values.
|
|
"""
|
|
cfg = self.config
|
|
if tokens < cfg.min_sample_tokens or elapsed_seconds <= 0.0:
|
|
return False
|
|
tps = tokens / elapsed_seconds
|
|
ts = time.time() if now is None else now
|
|
with self._lock:
|
|
stat = self._stats.get(signature)
|
|
if stat is None:
|
|
stat = RouteStat()
|
|
self._stats[signature] = stat
|
|
carried = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
|
|
total = carried + 1.0
|
|
stat.ewma_tps = (stat.ewma_tps * carried + tps) / total
|
|
stat.weight = total
|
|
stat.last_sample_ts = ts
|
|
stat.epoch = self._epochs.get(model_key, 0)
|
|
stat.samples += 1
|
|
return True
|
|
|
|
def snapshot(self, signature: str, model_key: str, now: float | None = None) -> dict:
|
|
"""Point-in-time view of one route's learned state."""
|
|
ts = time.time() if now is None else now
|
|
cfg = self.config
|
|
with self._lock:
|
|
stat = self._stats.get(signature)
|
|
current_epoch = self._epochs.get(model_key, 0)
|
|
if stat is None:
|
|
return {"tps": None, "weight": 0.0, "samples": 0, "status": "unsampled"}
|
|
weight = stat.decayed_weight(ts, cfg.stats_half_life_seconds)
|
|
if stat.epoch != current_epoch:
|
|
status = "stale"
|
|
elif weight < cfg.min_proven_weight:
|
|
status = "decayed" if stat.samples else "unsampled"
|
|
else:
|
|
status = "proven"
|
|
return {
|
|
"tps": round(stat.ewma_tps, 4) if stat.samples else None,
|
|
"weight": round(weight, 4),
|
|
"samples": stat.samples,
|
|
"status": status,
|
|
}
|
|
|
|
def prune(self, now: float | None = None) -> int:
|
|
"""Drop routes with no samples for `prune_after_seconds`."""
|
|
ts = time.time() if now is None else now
|
|
cutoff = ts - self.config.prune_after_seconds
|
|
with self._lock:
|
|
dead = [sig for sig, stat in self._stats.items() if stat.last_sample_ts < cutoff]
|
|
for sig in dead:
|
|
del self._stats[sig]
|
|
return len(dead)
|
|
|
|
|
|
def choose_route(
|
|
candidates: list[RouteCandidate],
|
|
store: RouteStatsStore,
|
|
model_key: str,
|
|
rng: random.Random | None = None,
|
|
now: float | None = None,
|
|
) -> tuple[RouteCandidate | None, dict]:
|
|
"""Pick a route: ε-scout among unproven arms, else weighted ∝ tps**alpha.
|
|
|
|
Returns (candidate, decision) where decision explains the pick for logs
|
|
and diagnostics: {"mode": "scout"|"exploit"|"prior", ...}.
|
|
"""
|
|
if not candidates:
|
|
return None, {"mode": "none"}
|
|
rng = rng or random
|
|
cfg = store.config
|
|
proven: list[tuple[RouteCandidate, float]] = []
|
|
scouts: list[tuple[RouteCandidate, float]] = []
|
|
for cand in candidates:
|
|
snap = store.snapshot(cand.signature, model_key, now=now)
|
|
if snap["status"] == "proven":
|
|
proven.append((cand, max(float(snap["tps"] or 0.0), 1e-6)))
|
|
else:
|
|
scouts.append((cand, float(snap["weight"])))
|
|
if scouts and (not proven or rng.random() < cfg.explore_share):
|
|
# Least-measured first so new/stale arms accumulate samples fastest;
|
|
# tiebreak on prior estimate so plausible routes get scouted first.
|
|
scouts.sort(key=lambda item: (item[1], -item[0].prior_tps))
|
|
pick = scouts[0][0]
|
|
return pick, {"mode": "scout", "signature": pick.signature}
|
|
if proven:
|
|
weights = [tps ** cfg.weight_alpha for _, tps in proven]
|
|
pick = rng.choices([cand for cand, _ in proven], weights=weights, k=1)[0]
|
|
return pick, {
|
|
"mode": "exploit",
|
|
"signature": pick.signature,
|
|
"candidates": len(proven),
|
|
}
|
|
# No stats anywhere yet — fall back to the prior (benchmark-derived) estimate.
|
|
weights = [max(cand.prior_tps, 1e-6) ** cfg.weight_alpha for cand in candidates]
|
|
pick = rng.choices(candidates, weights=weights, k=1)[0]
|
|
return pick, {"mode": "prior", "signature": pick.signature}
|
|
|
|
|
|
def route_table(
|
|
candidates: list[RouteCandidate],
|
|
store: RouteStatsStore,
|
|
model_key: str,
|
|
now: float | None = None,
|
|
) -> list[dict]:
|
|
"""Diagnostics rows: learned tps, coefficient vs best, expected traffic share."""
|
|
cfg = store.config
|
|
rows = []
|
|
for cand in candidates:
|
|
snap = store.snapshot(cand.signature, model_key, now=now)
|
|
rows.append({"candidate": cand, **snap})
|
|
proven = [r for r in rows if r["status"] == "proven"]
|
|
scouts = [r for r in rows if r["status"] != "proven"]
|
|
best_tps = max((float(r["tps"]) for r in proven), default=0.0)
|
|
exploit_budget = 1.0 - (cfg.explore_share if scouts and proven else 0.0)
|
|
if not proven:
|
|
exploit_budget = 0.0
|
|
weight_sum = sum(float(r["tps"]) ** cfg.weight_alpha for r in proven) or 1.0
|
|
out = []
|
|
for r in rows:
|
|
cand: RouteCandidate = r["candidate"]
|
|
if r["status"] == "proven":
|
|
share = exploit_budget * (float(r["tps"]) ** cfg.weight_alpha) / weight_sum
|
|
coefficient = round(float(r["tps"]) / best_tps, 3) if best_tps else None
|
|
else:
|
|
share = (
|
|
(cfg.explore_share if proven else 1.0) / len(scouts)
|
|
if scouts
|
|
else 0.0
|
|
)
|
|
coefficient = None
|
|
out.append({
|
|
"signature": cand.signature,
|
|
"hops": [
|
|
{
|
|
"node_id": getattr(n, "node_id", "?"),
|
|
"shard": f"{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}",
|
|
"endpoint": getattr(n, "endpoint", "?"),
|
|
}
|
|
for n in cand.nodes
|
|
],
|
|
"tps": r["tps"],
|
|
"coefficient": coefficient,
|
|
"expected_share": round(share, 4),
|
|
"samples": r["samples"],
|
|
"weight": r["weight"],
|
|
"status": r["status"],
|
|
"prior_tps": round(cand.prior_tps, 4),
|
|
})
|
|
out.sort(key=lambda r: (-(r["tps"] or 0.0), -r["prior_tps"]))
|
|
return out
|