routing improvements - dynamic (wip)
This commit is contained in:
Dobromir Popov
257
packages/tracker/meshnet_tracker/routing_stats.py
Normal file
257
packages/tracker/meshnet_tracker/routing_stats.py
Normal file
@@ -0,0 +1,257 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user