diff --git a/app/control/account/quota_defaults.py b/app/control/account/quota_defaults.py index f5464806e..6ab8d861a 100644 --- a/app/control/account/quota_defaults.py +++ b/app/control/account/quota_defaults.py @@ -2,10 +2,10 @@ Canonical quota totals per pool type (from upstream rate-limits API): - auto fast expert heavy - basic 20 60 8 — window: 72000 / 36000 s - super 50 140 50 — window: 7200 s - heavy 150 400 150 20 window: 7200 s + auto fast expert heavy grok_4_3 + basic 20 60 8 — — window: 72000 / 36000 s + super 50 140 50 — 50 window: 7200 s + heavy 150 400 150 20 150 window: 7200 s Pool inference uses ``auto.total`` as the primary signal — the three values (20 / 50 / 150) are mutually exclusive across pool types. diff --git a/app/control/account/refresh.py b/app/control/account/refresh.py index c21238d45..203c27784 100644 --- a/app/control/account/refresh.py +++ b/app/control/account/refresh.py @@ -57,9 +57,9 @@ class AccountRefreshService: """Fetches real quota data from the upstream usage API and persists it. Triggers: - 1. Import — super accounts: fetch all 3 modes. - 2. Call — super accounts: fetch the called mode (async, non-blocking). - 3. Schedule — super: fetch all 3 modes; basic: static window reset check. + 1. Import — fetch all modes supported by the account's pool. + 2. Call — fetch the called mode only (async, non-blocking). + 3. Schedule — refresh one pool per loop using that pool's supported modes. """ def __init__(self, repository: "AccountRepository") -> None: @@ -75,7 +75,13 @@ def __init__(self, repository: "AccountRepository") -> None: async def _fetch_all_quotas( self, token: str, pool: str ) -> dict[int, QuotaWindow] | None: - """Fetch quota windows for modes supported by *pool*. Returns {mode_id: window}.""" + """Fetch quota windows for every mode supported by *pool*. + + Examples: + - basic -> auto / fast / expert + - super -> auto / fast / expert / grok_4_3 + - heavy -> auto / fast / expert / heavy / grok_4_3 + """ try: from app.dataplane.reverse.protocol.xai_usage import fetch_all_quotas @@ -142,7 +148,7 @@ async def refresh_on_import(self, tokens: list[str]) -> RefreshResult: return agg async def refresh_call_async(self, token: str, mode_id: int) -> None: - """Fire-and-forget quota sync after a successful call (all accounts).""" + """Fire-and-forget single-mode quota sync after a successful call.""" record = (await self._repo.get_accounts([token]) or [None])[0] if record is None or record.is_deleted(): return @@ -219,7 +225,7 @@ async def _refresh_one( *, apply_fallback: bool = False, ) -> RefreshResult: - """Fetch all 3 modes from the usage API and persist real quota data. + """Fetch all pool-supported modes from the usage API and persist them. apply_fallback=True — used by scheduled/import paths: when API fails, decrement REAL quotas or reset expired DEFAULT windows. diff --git a/app/control/account/runtime.py b/app/control/account/runtime.py index 00a5778c0..fbf5f7f3a 100644 --- a/app/control/account/runtime.py +++ b/app/control/account/runtime.py @@ -1,11 +1,19 @@ -"""Account runtime singletons exposed without importing app.main.""" +"""Account runtime singletons and hot-apply helpers. -from typing import TYPE_CHECKING +These helpers expose the process-local account refresh runtime without making +callers import ``app.main``. Admin handlers use them to reconcile strategy and +scheduler state after hot config updates. +""" + +from typing import TYPE_CHECKING, Literal if TYPE_CHECKING: from .refresh import AccountRefreshService + from .scheduler import AccountRefreshScheduler _refresh_service: "AccountRefreshService | None" = None +_refresh_scheduler: "AccountRefreshScheduler | None" = None +_refresh_scheduler_leader = False def set_refresh_service(service: "AccountRefreshService | None") -> None: @@ -19,4 +27,76 @@ def get_refresh_service() -> "AccountRefreshService | None": return _refresh_service -__all__ = ["get_refresh_service", "set_refresh_service"] +def set_refresh_scheduler(scheduler: "AccountRefreshScheduler | None") -> None: + """Register the process-global account refresh scheduler.""" + global _refresh_scheduler + _refresh_scheduler = scheduler + + +def get_refresh_scheduler() -> "AccountRefreshScheduler | None": + """Return the registered account refresh scheduler, if any.""" + return _refresh_scheduler + + +def set_refresh_scheduler_leader(is_leader: bool) -> None: + """Record whether this worker currently owns the refresh scheduler lock.""" + global _refresh_scheduler_leader + _refresh_scheduler_leader = bool(is_leader) + + +def is_refresh_scheduler_leader() -> bool: + """Return True when this worker is the active refresh-scheduler leader.""" + return _refresh_scheduler_leader + + +def reconcile_refresh_runtime( + enabled: bool | None = None, +) -> Literal["quota", "random"]: + """Hot-apply refresh strategy and scheduler state for the current worker.""" + from app.dataplane.account.selector import current_strategy, set_strategy + from app.platform.config.snapshot import config + from app.platform.logging.logger import logger + + refresh_enabled = ( + config.get_bool("account.refresh.enabled", False) + if enabled is None + else bool(enabled) + ) + target_strategy: Literal["quota", "random"] = ( + "quota" if refresh_enabled else "random" + ) + previous_strategy = current_strategy() + if previous_strategy != target_strategy: + set_strategy(target_strategy) + + scheduler_action = "unchanged" + scheduler = _refresh_scheduler + if scheduler is not None and _refresh_scheduler_leader: + if refresh_enabled: + if not scheduler.is_running(): + scheduler.start() + scheduler_action = "started" + elif scheduler.is_running(): + scheduler.stop() + scheduler_action = "stopped" + + if previous_strategy != target_strategy or scheduler_action != "unchanged": + logger.info( + "account refresh runtime reconciled: previous_strategy={} strategy={} leader={} scheduler_action={}", + previous_strategy, + target_strategy, + _refresh_scheduler_leader, + scheduler_action, + ) + return target_strategy + + +__all__ = [ + "get_refresh_service", + "set_refresh_service", + "get_refresh_scheduler", + "set_refresh_scheduler", + "is_refresh_scheduler_leader", + "set_refresh_scheduler_leader", + "reconcile_refresh_runtime", +] diff --git a/app/control/account/scheduler.py b/app/control/account/scheduler.py index 0fa265fae..8d384d219 100644 --- a/app/control/account/scheduler.py +++ b/app/control/account/scheduler.py @@ -39,8 +39,16 @@ def __init__(self, refresh_service: AccountRefreshService) -> None: self._tasks: list[asyncio.Task] = [] self._stop = asyncio.Event() + def bind_service(self, refresh_service: AccountRefreshService) -> None: + """Update the refresh service used by the singleton scheduler.""" + self._service = refresh_service + + def is_running(self) -> bool: + """Return True while any pool refresh loop is still active.""" + return any(not task.done() for task in self._tasks) + def start(self) -> None: - if self._tasks and not all(t.done() for t in self._tasks): + if self.is_running(): return self._stop.clear() self._tasks = [ @@ -54,11 +62,14 @@ def start(self) -> None: ) def stop(self) -> None: + was_running = self.is_running() self._stop.set() for t in self._tasks: if not t.done(): t.cancel() - logger.info("account refresh scheduler stopped") + self._tasks = [] + if was_running: + logger.info("account refresh scheduler stopped") async def _loop(self, pool: str) -> None: while not self._stop.is_set(): @@ -102,6 +113,8 @@ def get_account_refresh_scheduler( global _scheduler if _scheduler is None: _scheduler = AccountRefreshScheduler(refresh_service) + else: + _scheduler.bind_service(refresh_service) return _scheduler diff --git a/app/dataplane/account/__init__.py b/app/dataplane/account/__init__.py index b753edcb3..11c074f00 100644 --- a/app/dataplane/account/__init__.py +++ b/app/dataplane/account/__init__.py @@ -8,16 +8,17 @@ import asyncio from typing import TYPE_CHECKING +from app.platform.config.snapshot import get_config from app.platform.logging.logger import logger from app.platform.runtime.clock import now_s from app.control.account.repository import AccountRepository from app.control.account.enums import FeedbackKind from .table import AccountRuntimeTable from .lease import AccountLease, new_lease -from .selector import select, select_any +from .selector import current_strategy, select, select_any from .sync import bootstrap as _bootstrap, apply_changes from . import feedback as fb -from ..shared.enums import StatusId +from ..shared.enums import POOL_ID_TO_STR, StatusId if TYPE_CHECKING: pass @@ -250,12 +251,22 @@ async def feedback( ts = now_s_val if now_s_val is not None else now_s() + strategy = current_strategy() + async with self._lock: if kind == FeedbackKind.SUCCESS: - fb.apply_success(table, idx, mode_id) + if strategy == "random": + fb.apply_success_random(table, idx) + else: + fb.apply_success_quota(table, idx, mode_id) elif kind == FeedbackKind.RATE_LIMITED: - fb.apply_rate_limited(table, idx, mode_id) + if strategy == "random": + pool_id = int(table.pool_by_idx[idx]) + cooling_sec = _pool_cooling_sec(pool_id) + fb.apply_rate_limited_random(table, idx, cooling_sec=cooling_sec) + else: + fb.apply_rate_limited_quota(table, idx, mode_id) fb.update_last_fail(table, idx, ts) elif kind == FeedbackKind.UNAUTHORIZED: @@ -271,8 +282,9 @@ async def feedback( fb.apply_server_error(table, idx) fb.update_last_fail(table, idx, ts) - # Apply authoritative quota data from upstream response headers. - if remaining is not None and reset_at_ms is not None: + # Quota strategy may receive authoritative quota data from upstream + # response headers; the random strategy ignores this entirely. + if strategy == "quota" and remaining is not None and reset_at_ms is not None: reset_s = int(reset_at_ms // 1000) fb.apply_quota_update(table, idx, mode_id, remaining, reset_s) @@ -289,6 +301,27 @@ def revision(self) -> int: return self._table.revision if self._table else 0 +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +_POOL_INTERVAL_CONFIG: dict[str, tuple[str, int]] = { + "basic": ("account.refresh.basic_interval_sec", 36_000), + "super": ("account.refresh.super_interval_sec", 7_200), + "heavy": ("account.refresh.heavy_interval_sec", 7_200), +} + + +def _pool_cooling_sec(pool_id: int) -> int: + """Cooling seconds for a 429 on a given pool (random strategy only).""" + pool_str = POOL_ID_TO_STR.get(pool_id, "basic") + interval_key, default_interval = _POOL_INTERVAL_CONFIG.get( + pool_str, _POOL_INTERVAL_CONFIG["basic"] + ) + return max(0, int(get_config(interval_key, default_interval))) + + # --------------------------------------------------------------------------- # Module-level singleton # --------------------------------------------------------------------------- diff --git a/app/dataplane/account/feedback.py b/app/dataplane/account/feedback.py index b77942c1e..3372b56e3 100644 --- a/app/dataplane/account/feedback.py +++ b/app/dataplane/account/feedback.py @@ -1,39 +1,78 @@ """Apply feedback to the runtime table columns (in-place, lock-free inner ops). -The caller (AccountDirectory) is responsible for holding the state lock -before calling these functions. +The caller (AccountDirectory) is responsible for holding the state lock before +calling these functions. + +Strategy-split functions +------------------------ +``apply_success`` / ``apply_rate_limited`` are split into ``*_quota`` and +``*_random`` variants. The caller picks exactly one based on the active +selector strategy. No runtime ``if`` inside the apply helpers themselves. """ +from app.platform.runtime.clock import now_s from ..shared.enums import ALL_MODE_IDS, StatusId from .table import AccountRuntimeTable # Health adjustment constants. -_SUCCESS_STEP = 0.12 -_AUTH_FACTOR = 0.55 -_FORBIDDEN_FACTOR = 0.25 -_RATE_LIMIT_FACTOR = 0.45 +_SUCCESS_STEP = 0.12 +_AUTH_FACTOR = 0.55 +_FORBIDDEN_FACTOR = 0.25 +_RATE_LIMIT_FACTOR = 0.45 _SERVER_ERROR_FACTOR = 0.75 _MIN_HEALTH = 0.05 _MAX_HEALTH = 1.0 -def apply_success(table: AccountRuntimeTable, idx: int, mode_id: int) -> None: - """Record a successful call: decrement quota, improve health.""" +# --------------------------------------------------------------------------- +# apply_success — two independent implementations +# --------------------------------------------------------------------------- + + +def apply_success_quota(table: AccountRuntimeTable, idx: int, mode_id: int) -> None: + """Quota strategy: decrement per-mode quota and improve health.""" quota_col = table._quota_col(mode_id) - new_q = max(0, int(quota_col[idx]) - 1) - quota_col[idx] = new_q + quota_col[idx] = max(0, int(quota_col[idx]) - 1) + _bump_health(table, idx) + - # Recover health. - h = min(_MAX_HEALTH, float(table.health_by_idx[idx]) + _SUCCESS_STEP) - table.health_by_idx[idx] = h +def apply_success_random(table: AccountRuntimeTable, idx: int) -> None: + """Random strategy: only improve health; quota columns are ignored.""" + _bump_health(table, idx) -def apply_rate_limited(table: AccountRuntimeTable, idx: int, mode_id: int) -> None: - """Zero the mode quota and reduce health.""" +# --------------------------------------------------------------------------- +# apply_rate_limited — two independent implementations +# --------------------------------------------------------------------------- + + +def apply_rate_limited_quota( + table: AccountRuntimeTable, idx: int, mode_id: int +) -> None: + """Quota strategy: zero the mode quota and reduce health.""" table._quota_col(mode_id)[idx] = 0 _adjust_health(table, idx, _RATE_LIMIT_FACTOR) +def apply_rate_limited_random( + table: AccountRuntimeTable, idx: int, *, cooling_sec: int +) -> None: + """Random strategy: set per-account cooldown timestamp; do not touch quota. + + The account stays in ``mode_available``; ``_random_select`` filters it out + until ``cooling_until_s_by_idx`` has elapsed. + """ + ts = now_s() + max(0, cooling_sec) + cooling_col = table.cooling_until_s_by_idx + cooling_col[idx] = max(int(cooling_col[idx]), ts) + _adjust_health(table, idx, _RATE_LIMIT_FACTOR) + + +# --------------------------------------------------------------------------- +# Strategy-agnostic feedback helpers (shared) +# --------------------------------------------------------------------------- + + def apply_auth_failure(table: AccountRuntimeTable, idx: int) -> None: """Reduce health on 401; caller may mark account expired.""" _adjust_health(table, idx, _AUTH_FACTOR) @@ -62,13 +101,11 @@ def apply_status_change( table.status_by_idx[idx] = new_status_id if new_status_id != int(StatusId.ACTIVE): - # Remove from all mode availability buckets. for mode_id in ALL_MODE_IDS: bucket = table.mode_available.get((pool_id, mode_id)) if bucket: bucket.discard(idx) else: - # Re-add to mode buckets where quota > 0. for mode_id in ALL_MODE_IDS: if int(table._quota_col(mode_id)[idx]) > 0: table.mode_available.setdefault((pool_id, mode_id), set()).add(idx) @@ -81,7 +118,11 @@ def apply_quota_update( remaining: int, reset_s: int, ) -> None: - """Update quota and reset timestamp from upstream API data.""" + """Update quota and reset timestamp from upstream API data. + + Used exclusively by the quota strategy (refresh service). The random + strategy never calls this. + """ quota_col = table._quota_col(mode_id) reset_col = table._reset_col(mode_id) quota_col[idx] = max(0, min(remaining, 32767)) @@ -116,14 +157,23 @@ def update_last_fail(table: AccountRuntimeTable, idx: int, now_s: int) -> None: # --------------------------------------------------------------------------- +def _bump_health(table: AccountRuntimeTable, idx: int) -> None: + table.health_by_idx[idx] = min( + _MAX_HEALTH, float(table.health_by_idx[idx]) + _SUCCESS_STEP + ) + + def _adjust_health(table: AccountRuntimeTable, idx: int, factor: float) -> None: - h = max(_MIN_HEALTH, float(table.health_by_idx[idx]) * factor) - table.health_by_idx[idx] = h + table.health_by_idx[idx] = max( + _MIN_HEALTH, float(table.health_by_idx[idx]) * factor + ) __all__ = [ - "apply_success", - "apply_rate_limited", + "apply_success_quota", + "apply_success_random", + "apply_rate_limited_quota", + "apply_rate_limited_random", "apply_auth_failure", "apply_forbidden", "apply_server_error", diff --git a/app/dataplane/account/selector.py b/app/dataplane/account/selector.py index 814259049..c7d048cbe 100644 --- a/app/dataplane/account/selector.py +++ b/app/dataplane/account/selector.py @@ -1,22 +1,63 @@ -"""Hot-path account selector — pure function, no object allocation. +"""Hot-path account selector — pluggable strategies. -All scoring operates on typed ``array.array`` columns; no attribute lookup -overhead. The caller provides the frozen set of excluded indices and -pre-resolved tag index set. +Two fully independent strategies: + +* ``_quota_select`` — scores candidates by health / quota / inflight / fails. + Used when ``account.refresh.enabled=true``. Behaviour is the historical one, + kept unchanged. +* ``_random_select`` — uniform random choice among non-cooling candidates. + Used when ``account.refresh.enabled=false``. Ignores quota and health entirely. + +Strategy selection is process-global, registered once by the lifespan via +:func:`set_strategy`. Callers (``AccountDirectory``) invoke :func:`select` / +:func:`select_any` and are unaware of which strategy is active. """ +import array # noqa: F401 — used in forward-referenced type annotations +import random +from typing import Literal + +from app.platform.config.snapshot import get_config from ..shared.enums import PoolId from .table import AccountRuntimeTable -# Scoring weights — tuned for throughput/fairness balance. -_W_HEALTH = 100.0 -_W_QUOTA = 25.0 -_W_RECENT = 15.0 # penalty for recently used accounts +# Scoring weights used by the quota strategy. +_W_HEALTH = 100.0 +_W_QUOTA = 25.0 +_W_RECENT = 15.0 # penalty for recently used accounts _W_INFLIGHT = 20.0 -_W_FAIL = 4.0 +_W_FAIL = 4.0 _RECENT_WINDOW_S = 15 # seconds +# --------------------------------------------------------------------------- +# Strategy registry +# --------------------------------------------------------------------------- + +_StrategyName = Literal["quota", "random"] +_STRATEGY_NAME: _StrategyName = "random" + + +def set_strategy(name: _StrategyName) -> None: + """Register the process-global selection strategy. + + Called once by the lifespan after reading ``account.refresh.enabled``. + """ + global _STRATEGY_NAME + if name not in ("quota", "random"): + raise ValueError(f"unknown selection strategy: {name!r}") + _STRATEGY_NAME = name + + +def current_strategy() -> _StrategyName: + return _STRATEGY_NAME + + +# --------------------------------------------------------------------------- +# Public entry points — delegate to the active strategy +# --------------------------------------------------------------------------- + + def select( table: AccountRuntimeTable, pool_id: int, @@ -26,30 +67,79 @@ def select( prefer_tag_idxs: set[int] | None = None, now_s: int, ) -> int | None: - """Select the best available account slot index. + """Select an account slot for ``(pool_id, mode_id)``. + + Returns the slot index or ``None`` when no candidate is available. + Does not mutate the table — callers increment inflight separately. + """ + if _STRATEGY_NAME == "random": + return _random_select( + table, pool_id, + exclude_idxs=exclude_idxs, + prefer_tag_idxs=prefer_tag_idxs, + now_s=now_s, + ) + return _quota_select( + table, pool_id, mode_id, + exclude_idxs=exclude_idxs, + prefer_tag_idxs=prefer_tag_idxs, + now_s=now_s, + ) + + +def select_any( + table: AccountRuntimeTable, + pool_id: int, + *, + exclude_idxs: frozenset[int] | None = None, + prefer_tag_idxs: set[int] | None = None, + now_s: int, +) -> int | None: + """Select any active account in ``pool_id`` irrespective of per-mode quota. - Returns the slot index on success, ``None`` if no account is available. - Does **not** mutate any table state — callers must increment inflight. + Used by WebSocket-based products that manage their own rate limiting. """ + if _STRATEGY_NAME == "random": + return _random_select( + table, pool_id, + exclude_idxs=exclude_idxs, + prefer_tag_idxs=prefer_tag_idxs, + now_s=now_s, + ) + return _quota_select_any( + table, pool_id, + exclude_idxs=exclude_idxs, + prefer_tag_idxs=prefer_tag_idxs, + now_s=now_s, + ) + + +# --------------------------------------------------------------------------- +# Strategy: quota — score-based selection (unchanged behaviour) +# --------------------------------------------------------------------------- + + +def _quota_select( + table: AccountRuntimeTable, + pool_id: int, + mode_id: int, + *, + exclude_idxs: frozenset[int] | None, + prefer_tag_idxs: set[int] | None, + now_s: int, +) -> int | None: candidates: set[int] | None = table.mode_available.get((pool_id, mode_id)) if not candidates: return None - # Apply window-expiry resets for basic accounts inline. - reset_col = table._reset_col(mode_id) - quota_col = table._quota_col(mode_id) - total_col = table._total_col(mode_id) + reset_col = table._reset_col(mode_id) + quota_col = table._quota_col(mode_id) + total_col = table._total_col(mode_id) window_col = table._window_col(mode_id) _maybe_reset_windows( - table, - candidates, - mode_id, - reset_col, - quota_col, - total_col, - window_col, - pool_id, - now_s, + table, candidates, mode_id, + reset_col, quota_col, total_col, window_col, + pool_id, now_s, ) working: set[int] = candidates.copy() @@ -59,30 +149,50 @@ def select( if not working: return None - # Prefer tagged subset; fall back to full set if empty. if prefer_tag_idxs: preferred = working & prefer_tag_idxs working = preferred if preferred else working - return _best(table, working, mode_id, quota_col, now_s) + return _best(table, working, quota_col, now_s) + + +def _quota_select_any( + table: AccountRuntimeTable, + pool_id: int, + *, + exclude_idxs: frozenset[int] | None, + prefer_tag_idxs: set[int] | None, + now_s: int, +) -> int | None: + candidates: set[int] = _pool_union(table, pool_id) + if not candidates: + return None + + working = candidates.copy() + if exclude_idxs: + working -= exclude_idxs + if not working: + return None + + if prefer_tag_idxs: + preferred = working & prefer_tag_idxs + working = preferred if preferred else working + + return _best_no_quota(table, working, now_s) def _maybe_reset_windows( - table: AccountRuntimeTable, + table: AccountRuntimeTable, candidates: set[int], - mode_id: int, + mode_id: int, reset_col: "array.array", quota_col: "array.array", total_col: "array.array", window_col: "array.array", - pool_id: int, - now_s: int, + pool_id: int, + now_s: int, ) -> None: - """Reset expired windows for basic-pool accounts (no API call required). - - Only applies to the basic pool — super/heavy quotas are managed exclusively - by the periodic refresh service and must not be reset inline. - """ + """Reset expired windows for basic-pool accounts inline (no API call needed).""" if pool_id != int(PoolId.BASIC): return @@ -93,22 +203,19 @@ def _maybe_reset_windows( if int(table.pool_by_idx[idx]) != pool_id: continue new_total = int(total_col[idx]) - window_s = int(window_col[idx]) + window_s = int(window_col[idx]) if new_total <= 0 or window_s <= 0: continue - # Window expired — restore the last known total for this account/mode. quota_col[idx] = new_total reset_col[idx] = now_s + window_s def _best( - table: AccountRuntimeTable, - working: set[int], - mode_id: int, + table: AccountRuntimeTable, + working: set[int], quota_col: "array.array", - now_s: int, -) -> int: - """Return the index of the highest-scoring candidate.""" + now_s: int, +) -> int | None: best_idx = -1 best_score = -1e18 @@ -118,7 +225,7 @@ def _best( last_use_col = table.last_use_at_by_idx for idx in working: - quota = int(quota_col[idx]) + quota = int(quota_col[idx]) if quota <= 0: continue health = float(health_col[idx]) @@ -144,60 +251,23 @@ def _best( return best_idx if best_idx >= 0 else None -def select_any( - table: AccountRuntimeTable, - pool_id: int, - *, - exclude_idxs: frozenset[int] | None = None, - prefer_tag_idxs: set[int] | None = None, - now_s: int, -) -> int | None: - """Select any active account from *pool_id* without mode-specific quota checking. - - Used for WebSocket-based operations that manage their own upstream rate limiting - and do not consume tracked quota buckets. Candidates are the union of all - mode_available sets for the pool so that accounts are reachable even when one - particular mode bucket (e.g. AUTO) is exhausted. - """ - candidates: set[int] = set() - for (pid, _mid), accounts in table.mode_available.items(): - if pid == pool_id: - candidates |= accounts - - if not candidates: - return None - - working = candidates.copy() - if exclude_idxs: - working -= exclude_idxs - if not working: - return None - - if prefer_tag_idxs: - preferred = working & prefer_tag_idxs - working = preferred if preferred else working - - return _best_no_quota(table, working, now_s) - - def _best_no_quota( table: AccountRuntimeTable, working: set[int], now_s: int, ) -> int | None: - """Score candidates by health and inflight only — no quota weighting.""" - best_idx = -1 + best_idx = -1 best_score = -1e18 - health_col = table.health_by_idx + health_col = table.health_by_idx inflight_col = table.inflight_by_idx - fail_col = table.fail_count_by_idx + fail_col = table.fail_count_by_idx last_use_col = table.last_use_at_by_idx for idx in working: - health = float(health_col[idx]) + health = float(health_col[idx]) inflight = int(inflight_col[idx]) - fails = min(int(fail_col[idx]), 10) + fails = min(int(fail_col[idx]), 10) last_use = int(last_use_col[idx]) score = health * _W_HEALTH - inflight * _W_INFLIGHT - fails * _W_FAIL @@ -208,9 +278,62 @@ def _best_no_quota( if score > best_score: best_score = score - best_idx = idx + best_idx = idx return best_idx if best_idx >= 0 else None -__all__ = ["select", "select_any"] +# --------------------------------------------------------------------------- +# Strategy: random — uniform choice with cooling + inflight filter +# --------------------------------------------------------------------------- + + +def _random_select( + table: AccountRuntimeTable, + pool_id: int, + *, + exclude_idxs: frozenset[int] | None, + prefer_tag_idxs: set[int] | None, + now_s: int, +) -> int | None: + candidates: set[int] = _pool_union(table, pool_id) + if not candidates: + return None + + max_inflight = int(get_config("account.selection.max_inflight", 8)) + cooling_col = table.cooling_until_s_by_idx + inflight_col = table.inflight_by_idx + + working = candidates.copy() + if exclude_idxs: + working -= exclude_idxs + working = { + idx for idx in working + if int(cooling_col[idx]) <= now_s + and int(inflight_col[idx]) < max_inflight + } + if not working: + return None + + if prefer_tag_idxs: + preferred = working & prefer_tag_idxs + working = preferred if preferred else working + + return random.choice(tuple(working)) + + +# --------------------------------------------------------------------------- +# Shared helper +# --------------------------------------------------------------------------- + + +def _pool_union(table: AccountRuntimeTable, pool_id: int) -> set[int]: + """Union of all ``mode_available`` buckets for ``pool_id``.""" + out: set[int] = set() + for (pid, _mid), accounts in table.mode_available.items(): + if pid == pool_id: + out |= accounts + return out + + +__all__ = ["select", "select_any", "set_strategy", "current_strategy"] diff --git a/app/dataplane/account/table.py b/app/dataplane/account/table.py index 69ab143a8..16a5719e9 100644 --- a/app/dataplane/account/table.py +++ b/app/dataplane/account/table.py @@ -151,6 +151,11 @@ class AccountRuntimeTable: default_factory=lambda: array.array("L") ) + # --- Per-account cooldown (random strategy only; uint32 epoch-seconds; 0 = not cooling) --- + cooling_until_s_by_idx: "array.array[int]" = field( + default_factory=lambda: array.array("L") + ) + # --- Pre-computed selection indexes --- # (pool_id, mode_id) → set of idx with a supported quota window and status == ACTIVE mode_available: dict[tuple[int, int], set[int]] = field(default_factory=dict) @@ -300,6 +305,7 @@ def _append_slot( self.health_by_idx.append(health) self.last_use_at_by_idx.append(last_use_s) self.last_fail_at_by_idx.append(last_fail_s) + self.cooling_until_s_by_idx.append(0) self.size += 1 self._add_to_indexes(idx) self._add_to_tag_idx(idx, tags) diff --git a/app/dataplane/shared/enums.py b/app/dataplane/shared/enums.py index bdbe4a9ea..88b9ed66b 100644 --- a/app/dataplane/shared/enums.py +++ b/app/dataplane/shared/enums.py @@ -36,6 +36,8 @@ class StatusId(IntEnum): "heavy": int(PoolId.HEAVY), } +POOL_ID_TO_STR: dict[int, str] = {v: k for k, v in POOL_STR_TO_ID.items()} + STATUS_STR_TO_ID: dict[str, int] = { "active": int(StatusId.ACTIVE), "cooling": int(StatusId.COOLING), @@ -56,6 +58,7 @@ class StatusId(IntEnum): "PoolId", "StatusId", "POOL_STR_TO_ID", + "POOL_ID_TO_STR", "STATUS_STR_TO_ID", "ALL_MODE_IDS", ] diff --git a/app/main.py b/app/main.py index 9a5174928..e69498c8d 100644 --- a/app/main.py +++ b/app/main.py @@ -116,7 +116,12 @@ async def lifespan(app: FastAPI): create_repository, describe_repository_target, ) - from app.control.account.runtime import set_refresh_service + from app.control.account.runtime import ( + reconcile_refresh_runtime, + set_refresh_scheduler, + set_refresh_scheduler_leader, + set_refresh_service, + ) from app.control.account.scheduler import get_account_refresh_scheduler from app.dataplane.account import get_account_directory @@ -182,26 +187,47 @@ async def _sync_loop() -> None: # 4. Account refresh scheduler — only the leader worker. # Uses an advisory file lock so exactly one process runs the heavy # upstream quota-fetch loop regardless of worker count. + # + # ``account.refresh.enabled`` selects between two independent runtime + # strategies: + # - true → "quota" selector + scheduler running (default). + # - false → "random" selector + scheduler idle (no upstream probing). from app.control.account.refresh import AccountRefreshService + refresh_enabled = _config.get_bool("account.refresh.enabled", False) + refresh_svc = AccountRefreshService(repo) set_refresh_service(refresh_svc) app.state.refresh_service = refresh_svc is_leader = _try_acquire_scheduler_lock() scheduler = get_account_refresh_scheduler(refresh_svc) - if is_leader: - scheduler.start() + set_refresh_scheduler(scheduler) + set_refresh_scheduler_leader(is_leader) + app.state.account_refresh_scheduler = scheduler + app.state.account_refresh_is_leader = is_leader + + strategy_name = reconcile_refresh_runtime(refresh_enabled) + if is_leader and strategy_name == "quota": logger.info( - "scheduler leader: pid={} active_sync_s={} idle_sync_s={}", + "scheduler leader: pid={} strategy=quota active_sync_s={} idle_sync_s={}", + os.getpid(), + _SYNC_ACTIVE_INTERVAL, + _SYNC_IDLE_INTERVAL, + ) + elif is_leader: + logger.info( + "scheduler leader: pid={} strategy=random (scheduler idle) " + "active_sync_s={} idle_sync_s={}", os.getpid(), _SYNC_ACTIVE_INTERVAL, _SYNC_IDLE_INTERVAL, ) else: logger.info( - "scheduler follower: pid={} active_sync_s={} idle_sync_s={}", + "scheduler follower: pid={} strategy={} active_sync_s={} idle_sync_s={}", os.getpid(), + strategy_name, _SYNC_ACTIVE_INTERVAL, _SYNC_IDLE_INTERVAL, ) @@ -233,6 +259,8 @@ async def _sync_loop() -> None: proxy_scheduler.stop() _release_scheduler_lock() + set_refresh_scheduler(None) + set_refresh_scheduler_leader(False) set_refresh_service(None) await repo.close() logger.info("application shutdown completed") @@ -313,7 +341,10 @@ def create_app() -> FastAPI: # Ensure config is loaded on every request. @app.middleware("http") async def _ensure_config(request: Request, call_next): + from app.control.account.runtime import reconcile_refresh_runtime + await _config.load() + reconcile_refresh_runtime() return await call_next(request) # Global exception handler — converts AppError to JSON. diff --git a/app/products/_account_selection.py b/app/products/_account_selection.py index 056583e67..a2f9423d9 100644 --- a/app/products/_account_selection.py +++ b/app/products/_account_selection.py @@ -3,8 +3,26 @@ from app.control.model.enums import ModeId from app.control.model.spec import ModelSpec from app.control.account.runtime import get_refresh_service +from app.dataplane.account.selector import current_strategy from app.platform.config.snapshot import get_config +# Random strategy has no config key for retry count; it is pinned here so that +# every retry-driven call site (chat / images / video / anthropic) sees the same +# value without introducing scattered magic numbers. +_RANDOM_MAX_RETRIES = 5 + + +def selection_max_retries() -> int: + """Retry count for account-swap loops, aware of the active selection strategy. + + - ``random`` strategy: fixed at :data:`_RANDOM_MAX_RETRIES` (=5). + - ``quota`` strategy: reads ``retry.max_retries`` (default 1), preserving + the historical behaviour. + """ + if current_strategy() == "random": + return _RANDOM_MAX_RETRIES + return int(get_config("retry.max_retries", 1)) + def mode_candidates(spec: ModelSpec) -> tuple[int, ...]: """Return mode IDs to try for *spec* in priority order. @@ -32,7 +50,9 @@ async def reserve_account( ): """Reserve an account and return ``(lease, selected_mode_id)``. - Returns ``(None, original_mode_id)`` when no account is available. + Returns ``(None, original_mode_id)`` when no account is available. Under the + random strategy no on-demand refresh fallback is attempted — upstream quota + data is never probed. """ original_mode_id = int(spec.mode_id) @@ -52,6 +72,9 @@ async def _try_reserve(): if lease is not None: return lease, selected_mode_id + if current_strategy() == "random": + return None, original_mode_id + refresh_svc = get_refresh_service() if refresh_svc is not None: await refresh_svc.refresh_on_demand() diff --git a/app/products/anthropic/messages.py b/app/products/anthropic/messages.py index 31ef2f90e..42ba98a17 100644 --- a/app/products/anthropic/messages.py +++ b/app/products/anthropic/messages.py @@ -33,7 +33,7 @@ _quota_sync, _fail_sync, _parse_retry_codes, _feedback_kind, _log_task_exception, _configured_retry_codes, _should_retry_upstream, ) -from app.products._account_selection import reserve_account +from app.products._account_selection import reserve_account, selection_max_retries from app.products.openai._tool_sieve import ToolSieve @@ -310,7 +310,7 @@ async def create( raise RateLimitError("Account directory not initialised") directory = _acct_dir - max_retries = cfg.get_int("retry.max_retries", 1) + max_retries = selection_max_retries() retry_codes = _configured_retry_codes(cfg) timeout_s = cfg.get_float("chat.timeout", 120.0) msg_id = _make_msg_id() diff --git a/app/products/openai/chat.py b/app/products/openai/chat.py index 4eeab694a..0d45bb30c 100644 --- a/app/products/openai/chat.py +++ b/app/products/openai/chat.py @@ -22,6 +22,7 @@ from app.control.model.registry import resolve as resolve_model from app.control.model.enums import ModeId from app.control.account.enums import FeedbackKind +from app.dataplane.account.selector import current_strategy from app.dataplane.proxy.adapters.headers import build_http_headers from app.dataplane.proxy import get_proxy_runtime from app.dataplane.proxy.adapters.session import ( @@ -54,7 +55,7 @@ build_usage, ) from ._tool_sieve import ToolSieve -from app.products._account_selection import reserve_account +from app.products._account_selection import reserve_account, selection_max_retries def _to_chat_annotations(anns: list[dict]) -> list[dict]: @@ -106,6 +107,8 @@ def _transport_upstream_error(exc: BaseException, *, context: str) -> UpstreamEr async def _quota_sync(token: str, mode_id: int) -> None: """Fire-and-forget: fetch real quota after a successful call.""" try: + if current_strategy() != "quota": + return svc = get_refresh_service() if svc: await svc.refresh_call_async(token, mode_id) @@ -121,12 +124,20 @@ async def _quota_sync(token: str, mode_id: int) -> None: async def _fail_sync( token: str, mode_id: int, exc: BaseException | None = None ) -> None: - """Fire-and-forget: persist failure counter after a failed call.""" + """Fire-and-forget: persist failure metadata after a failed call. + + In random mode this helper must not trigger upstream quota probes. It still + records failures so 401 invalidation and local failure accounting continue + to work unchanged. + """ try: svc = get_refresh_service() if svc: await svc.record_failure_async(token, mode_id, exc) - if getattr(exc, "status", None) == 429: + if ( + current_strategy() == "quota" + and getattr(exc, "status", None) == 429 + ): result = await svc.refresh_on_demand() logger.info( "account on-demand refresh triggered: token={}... mode_id={} refreshed={} failed={} rate_limited={}", @@ -470,7 +481,7 @@ async def completions( raise RateLimitError("Account directory not initialised") directory = _acct_dir - max_retries = cfg.get_int("retry.max_retries", 1) + max_retries = selection_max_retries() retry_codes = _configured_retry_codes(cfg) response_id = make_response_id() timeout_s = cfg.get_float("chat.timeout", 120.0) diff --git a/app/products/openai/responses.py b/app/products/openai/responses.py index 95fdc3e7c..d816c7a92 100644 --- a/app/products/openai/responses.py +++ b/app/products/openai/responses.py @@ -18,7 +18,7 @@ from app.control.model.registry import resolve as resolve_model from app.control.account.enums import FeedbackKind from app.dataplane.reverse.protocol.xai_chat import classify_line, StreamAdapter -from app.products._account_selection import reserve_account +from app.products._account_selection import reserve_account, selection_max_retries from .chat import _stream_chat, _extract_message, _resolve_image, _quota_sync, _fail_sync, _parse_retry_codes, _feedback_kind, _log_task_exception, _upstream_body_excerpt from .chat import _configured_retry_codes, _should_retry_upstream @@ -247,7 +247,7 @@ async def create( raise RateLimitError("Account directory not initialised") directory = _acct_dir - max_retries = cfg.get_int("retry.max_retries", 1) + max_retries = selection_max_retries() retry_codes = _configured_retry_codes(cfg) response_id = make_resp_id("resp") reasoning_id = make_resp_id("rs") diff --git a/app/products/web/admin/__init__.py b/app/products/web/admin/__init__.py index 529452a4c..8b87c6d7d 100644 --- a/app/products/web/admin/__init__.py +++ b/app/products/web/admin/__init__.py @@ -184,6 +184,8 @@ async def get_config_endpoint(): @router.post("/config", tags=[_TAG_ADMIN_SYSTEM]) async def update_config(req: ConfigPatchRequest): + from app.control.account.runtime import reconcile_refresh_runtime + patch = _sanitize_proxy_config(req.root) _ensure_runtime_patch_allowed(patch) await config.update(patch) @@ -195,7 +197,12 @@ async def update_config(req: ConfigPatchRequest): file_level=config.get_str("logging.file_level", "") or None, max_files=config.get_int("logging.max_files", 7), ) - return {"status": "success", "message": "配置已更新"} + strategy_name = reconcile_refresh_runtime() + return { + "status": "success", + "message": "配置已更新", + "selection_strategy": strategy_name, + } @router.get("/storage", tags=[_TAG_ADMIN_SYSTEM]) @@ -205,6 +212,7 @@ async def get_storage_mode(): @router.get("/status", tags=[_TAG_ADMIN_SYSTEM]) async def runtime_status(): + from app.control.account.runtime import reconcile_refresh_runtime from app.dataplane.account import _directory if _directory is None: @@ -214,12 +222,14 @@ async def runtime_status(): code="directory_not_initialised", status=503, ) + strategy_name = reconcile_refresh_runtime() return Response( content=orjson.dumps( { "status": "ok", "size": _directory.size, "revision": _directory.revision, + "selection_strategy": strategy_name, } ), media_type="application/json", diff --git a/app/statics/admin/account.html b/app/statics/admin/account.html index f8983e3fa..d72e814c5 100644 --- a/app/statics/admin/account.html +++ b/app/statics/admin/account.html @@ -83,6 +83,9 @@ gap:12px; margin-bottom:20px; } + .stat-grid.mode-random { + grid-template-columns:repeat(4, minmax(0, 1fr)); + } .stat-cell { min-height:88px; padding:14px 16px; @@ -126,6 +129,16 @@ margin-top:auto; } .stat-muted { color:#a1a1aa !important; } + [data-random-only] { display:none } + .stat-grid.mode-random [data-random-only] { display:flex } + .stat-grid.mode-random [data-stat-card="active"] { order:1 } + .stat-grid.mode-random [data-stat-card="cooling"] { order:2 } + .stat-grid.mode-random [data-stat-card="invalid"] { order:3 } + .stat-grid.mode-random [data-stat-card="disabled"] { order:4 } + .stat-grid.mode-random [data-stat-card="total"] { order:5 } + .stat-grid.mode-random [data-stat-card="calls"] { order:6 } + .stat-grid.mode-random [data-stat-card="success"] { order:7 } + .stat-grid.mode-random [data-stat-card="rate"] { order:8 } /* Filters */ .filter-bar { @@ -726,9 +739,9 @@