"""Multi-endpoint load balancing: one logical model, N local ports. from evalharness.model.pool import PooledAdapter from evalharness.model.adapter import resolve_adapter base = resolve_adapter('openai/http://127.0.0.1:8123/v1?Qwen3-8B') pool = PooledAdapter([resolve_adapter(f'openai/http://127.0.0.1:{p}/v1?Qwen3-8B') for p in range(8123, 8131)]) out = await pool.generate(...) # round-robin over instances Traffic management: round-robin keeps per-endpoint traffic even; backends that fail repeatedly enter a cool-down window and are skipped until it expires, so one sick endpoint cannot absorb its share of the load. """ import asyncio import itertools import os import statistics import time from typing import Dict, List, Optional from ..data.sample import ChatMessage from .adapter import ModelAdapter from .output import ModelOutput, Usage class PooledAdapter(ModelAdapter): """Round-robin over N equivalent backend instances with health cooling.""" name = 'pool' COOLDOWN_S = 60.0 # a backend that failed EVERY attempt rests this long COOLDOWN_AFTER = 2 # consecutive full-pass failures before cooling def __init__(self, adapters: List[ModelAdapter]): if not adapters: raise ValueError('PooledAdapter needs at least one backend') super().__init__(model=adapters[0].model, api_base=adapters[0].api_base) self.adapters = adapters self._cycle = itertools.cycle(range(len(adapters))) self.usage = Usage() # request outcome counters (success rate accounting) self.stats = {'requests': 0, 'ok': 0, 'failed': 0, 'retried': 0} # per-backend health: consecutive_failures, cooling_until, per-endpoint counts self._health = [{ 'fails': 0, 'until': 0.0, 'ok': 0, 'req': 0} for _ in adapters] # adaptive per-endpoint concurrency gates (AIMD over /metrics signals) self._gates = [AdaptiveGate(a) for a in adapters] def _next(self) -> ModelAdapter: """Round-robin, skipping endpoints inside their cool-down window.""" n = len(self.adapters) now = time.time() for _ in range(n): i = next(self._cycle) h = self._health[i] if h['until'] <= now or all(x['until'] <= now for x in self._health): self._health[i]['req'] += 1 return self.adapters[i] # everything cooling: take the next anyway (better to try than stall) i = next(self._cycle) self._health[i]['req'] += 1 return self.adapters[i] def _mark(self, adapter: ModelAdapter, ok: bool) -> None: try: i = self.adapters.index(adapter) except ValueError: return h = self._health[i] if ok: h['fails'] = 0 h['until'] = 0.0 h['ok'] += 1 else: h['fails'] += 1 if h['fails'] >= self.COOLDOWN_AFTER: h['until'] = time.time() + self.COOLDOWN_S h['fails'] = 0 def request_stats(self) -> Dict[str, float]: """Success-rate + per-endpoint traffic view (load-balance audit).""" n = self.stats['requests'] out = { 'requests': n, 'success_rate': self.stats['ok'] / n if n else 0.0, 'retry_rate': self.stats['retried'] / n if n else 0.0, 'failure_rate': self.stats['failed'] / n if n else 0.0, } for i, (a, h) in enumerate(zip(self.adapters, self._health)): tag = a.api_base.rsplit('//', 1)[-1].replace('/', '_') out[f'ep{i}_{tag}_reqs'] = h['req'] out[f'ep{i}_{tag}_ok'] = h['ok'] if i < len(self._gates): for k, v in self._gates[i].report().items(): out[f'ep{i}_{tag}_{k}'] = v return out async def generate(self, messages: List[ChatMessage], tools: Optional[list] = None, **kw) -> ModelOutput: last_exc = None self.stats['requests'] += 1 for _ in range(len(self.adapters)): # try each instance once adapter = self._next() try: i = self.adapters.index(adapter) stamp = await self._gates[i].acquire() # NB: service time ONLY -- starting the clock before acquire # bills the gate's own QUEUE wait as service slowness, and # the shrinking medians convinced the gate the endpoint # was degrading under load it was not under t0 = time.monotonic() ok = False try: out = await adapter.generate(messages, tools=tools, **kw) ok = True finally: # ONE release per acquire: the old code released True in # this finally AND False in the except handler, double- # decrementing _inflight (gate over-admits) and applying # the x0.7 backoff twice per failure self._gates[i].release( ok, stamp=stamp, dur_s=time.monotonic() - t0, tokens=(out.usage.output_tokens if ok else 0)) self.usage = self.usage + out.usage self.stats['ok'] += 1 self._mark(adapter, True) if out.usage.retries: self.stats['retried'] += 1 return out except Exception as e: # dead/overloaded instance -> next last_exc = e self._mark(adapter, False) # 4xx (e.g. 400 overloaded) still worth trying ANOTHER instance: # one backend's state can differ from the rest continue self.stats['failed'] += 1 raise last_exc async def close(self) -> None: for a in self.adapters: await a.close() for g in self._gates: g.stop() class AdaptiveGate: """Per-endpoint adaptive concurrency limiter (Netflix Gradient2 style). One formula, no thresholds, no server telemetry, no mode state machine: when a measurement window fills: gradient = median per-stream speed this window / previous window limit = limit * gradient + sqrt(limit) - healthy (gradient ~1): grows by sqrt(limit) -- big steps at low limits (fast escape from 1), small relative steps at high limits (conservative where overshoot is fatal) - degraded (gradient <1): multiplicative shrink proportional to how much each stream actually slowed - window-vs-window comparison self-normalizes context growth: swe sessions decode slower as they mature, and an all-time-best baseline would wrongly shrink the gate forever for that alone - windows are COUNT driven (min(10, 2*limit) same-stamp samples): sparse workloads judge slowly by construction -- TCP's trick of using RTT as the clock, inherited for free. Time-based windows judge on thin evidence (a 30s window at limit 1 sees zero swe turns) - every sample is stamped with the limit at admission; only current-stamp samples close a window. Early completions at a new level are survivorship-biased (short requests finish first) and always read optimistic - per-stream speed = completion_tokens/duration; gateways that report no usage fall back to 1/duration (completion frequency) - a failed request: limit x0.7 + a 30s admission pause so the server can drain (anti death-spiral). Fully dead endpoints are the runner-level circuit breaker's job. """ LO = 1 HI = 96 INITIAL = 1.0 # '--concurrency auto' rebinds this SILENCE_S = 30.0 # post-failure drain pause before admitting again WINDOW_MAX = 10 # samples to close a window (2*limit when smaller) def __init__(self, adapter: ModelAdapter): self.adapter = adapter self.limit = self.INITIAL self._inflight = 0 self._cond: Optional[asyncio.Condition] = None self._loop = None self._stopped = False self._quiet_until = 0.0 # post-failure drain window self._window: List[tuple] = [] # (stamp, speed) since last close self._prev_median: Optional[float] = None self.stats = {'window': 0, 'grow': 0, 'shrink': 0, 'backoff_fail': 0} # ---- bar plumbing ---- def _push_limit(self) -> None: rep = (self.adapter.extra or {}).get('progress_reporter') fn = getattr(rep, 'set_gate', None) if fn is not None: try: fn(max(1, int(self.limit))) except Exception: pass def push_inflight(self) -> None: rep = (self.adapter.extra or {}).get('progress_reporter') fn = getattr(rep, 'set_admitted', None) if fn is not None: try: fn(self._inflight) except Exception: pass # ---- gate semantics ---- async def acquire(self) -> int: """Admit one request; the return value is its admission stamp (the limit it was admitted under) -- pass it back to release().""" loop = asyncio.get_running_loop() if self._cond is None or self._loop is not loop or self._stopped: # lazy init OR LOOP CHANGE: pools survive across asyncio.run() # boundaries and a Condition is loop-bound -- rebuild on the # new loop (nothing is in flight there by construction) self._loop = loop self._cond = asyncio.Condition() self._inflight = 0 self._stopped = False self._push_limit() while (self._inflight >= max(1, int(self.limit)) or time.monotonic() < self._quiet_until): # timed wait: silence expiry must re-open admission even if no # request completes to wake us await self._cond.acquire() try: await asyncio.wait_for(self._cond.wait(), timeout=5.0) except asyncio.TimeoutError: pass finally: self._cond.release() self._inflight += 1 self.push_inflight() return max(1, int(self.limit)) def release(self, ok: bool, stamp: int = 0, dur_s: float = 0.0, tokens: float = 0.0) -> None: self._inflight = max(0, self._inflight - 1) if not ok: self.stats['backoff_fail'] += 1 self.limit = max(self.LO, self.limit * 0.7) self._quiet_until = time.monotonic() + self.SILENCE_S self._push_limit() else: speed = (tokens / dur_s if tokens > 0 else 1.0 / dur_s if dur_s > 0 else 0.0) if speed > 0: self._window.append((stamp, speed)) self._maybe_close_window() self.push_inflight() self._wake() def _maybe_close_window(self) -> None: cur = max(1, int(self.limit)) fresh = [s for st, s in self._window if st == cur] if len(fresh) < min(self.WINDOW_MAX, 2 * cur): return # not enough evidence yet: keep waiting med = statistics.median(fresh) if self._prev_median: # dampers (Netflix/Envoy both do this): small-window medians # are noisy -- a bare ratio swings the limit 5 -> 1 -> 9 -> 2 # in simulation. Clamp the gradient band and smooth the # reference so one jittery window cannot whipsaw the gate. g = max(0.5, min(1.5, med / self._prev_median)) if g < 0.95: # clearly degraded: PURE multiplicative shrink -- adding # sqrt headroom here let a halved speed at low limits # still RAISE the limit (sim: 3.4 x 0.5 + sqrt(3.4) = 3.6) new = self.limit * g self.stats['shrink'] += 1 else: # healthy (or within noise): gentle sqrt probe upward new = self.limit * g + self.limit ** 0.5 if new > self.limit: self.stats['grow'] += 1 self.limit = max(float(self.LO), min(float(self.HI), new)) self._push_limit() self._prev_median = (med if self._prev_median is None else 0.6 * self._prev_median + 0.4 * med) self._window = [] self.stats['window'] += 1 def _wake(self) -> None: if self._cond is None: return try: asyncio.get_running_loop().create_task(self._notify_all()) except RuntimeError: pass # no running loop (shutdown) async def _notify_all(self) -> None: async with self._cond: self._cond.notify_all() def stop(self) -> None: self._stopped = True def report(self) -> Dict[str, float]: return {'limit': max(1, int(self.limit)), 'inflight': self._inflight, **{f'gate_{k}': v for k, v in self.stats.items()}} def pooled(specs: List[str], api_key: str = '') -> PooledAdapter: """['openai/http://127.0.0.1:8123/v1?M', ...] -> PooledAdapter. api_key: explicit key applied to EVERY member (two-key setups should build two pools, or use env resolution per host).""" from .adapter import resolve_adapter members = [resolve_adapter(s) for s in specs] if api_key: for m in members: m.api_key = api_key return PooledAdapter(members)