"""Account Pool Manager (provider-agnostic). Stateless facade that coordinates pool operations for any Provider with pool configuration enabled. All state lives in Redis via :mod:`redis_ops`. Usage:: mgr = PoolManager(provider_id, pool_config) reordered = await mgr.reorder_candidates(session_uuid, candidates) # ... execute request ... await mgr.on_request_success(session_uuid=..., key_id=..., tokens_used=...) """ from __future__ import annotations import asyncio import random from collections.abc import Callable from typing import TYPE_CHECKING, Any, TypeVar from src.core.logger import logger from src.services.provider.pool import redis_ops from src.services.provider.pool.account_state import resolve_pool_account_state from src.services.provider.pool.config import PoolConfig from src.services.provider.pool.health_cache import get_health_scores from src.services.provider.pool.trace import PoolCandidateTrace, PoolSchedulingTrace if TYPE_CHECKING: from src.models.database import ProviderAPIKey from src.services.scheduling.schemas import ProviderCandidate class PoolManager: """Coordinate pool-level scheduling for a single Provider.""" __slots__ = ("provider_id", "config", "provider_type") def __init__( self, provider_id: str, config: PoolConfig, provider_type: str | None = None, ) -> None: self.provider_id = provider_id self.config = config self.provider_type = str(provider_type or "").strip().lower() or None # ------------------------------------------------------------------ # Core scheduling: reorder candidate list for pool-aware selection # ------------------------------------------------------------------ async def reorder_candidates( self, session_uuid: str | None, candidates: list[ProviderCandidate], ) -> list[ProviderCandidate]: """Reorder *candidates* according to pool rules. The returned list keeps the same elements but in a new order: 1. **Sticky session hit** -- if the session is already bound to a key and that key appears in *candidates* and is not in cooldown, move it to position 0. 2. **Filter** out keys in account-blocked / cooldown / cost-exhausted state (mark ``is_skipped``). 3. **LRU sort** -- among remaining candidates at the same priority level, sort by least-recently-used. 4. **Random tiebreak** -- among candidates with identical LRU score. Also builds a :class:`PoolSchedulingTrace` and attaches per-candidate trace data via ``_pool_extra_data`` / ``_pool_scheduling_trace`` attributes on candidate objects. """ if not candidates: return candidates pid = self.provider_id # Build trace trace = PoolSchedulingTrace( provider_id=pid, total_keys=len(candidates), session_uuid=session_uuid[:8] if session_uuid else None, ) # --- Strategy: before_select ---------------------------------- strategies = _get_active_strategies(self.config) key_ids = [str(c.key.id) for c in candidates] strategy_context: dict[str, Any] = {"session_uuid": session_uuid} for strategy in strategies: if hasattr(strategy, "on_before_select"): try: filtered = strategy.on_before_select( provider_id=pid, key_ids=key_ids, config=self.config, context=strategy_context, ) if filtered is not None: key_ids = filtered except Exception: logger.opt(exception=True).debug( "Pool[{}]: strategy before_select failed", pid[:8] ) # --- 1. Sticky session ---------------------------------------- sticky_key_id: str | None = None if session_uuid and self.config.sticky_session_ttl_seconds > 0: sticky_key_id = await redis_ops.get_sticky_binding( pid, session_uuid, self.config.sticky_session_ttl_seconds ) provider_type = self.provider_type if provider_type is None and candidates: first_provider = getattr(candidates[0], "provider", None) provider_type = str(getattr(first_provider, "provider_type", "") or "").strip().lower() if not provider_type: provider_type = None # --- 2. Batch fetch pool state (parallel) --------------------- all_key_ids = [str(c.key.id) for c in candidates] # Fire independent Redis queries concurrently. # Only fetch reason (no TTL) on the scheduling hot path -- TTL is only # used for trace display and costs an extra pipeline command per key. _cooldown_coro = redis_ops.batch_get_cooldowns(pid, all_key_ids, include_ttl=False) _cost_coro = ( redis_ops.batch_get_cost_totals(pid, all_key_ids, self.config.cost_window_seconds) if ( self.config.cost_limit_per_key_tokens is not None or self.config.scheduling_mode == "multi_score" ) else None ) # LRU scores are needed both for plain LRU sorting and for multi_score # dimensions (e.g. cache_affinity / single_account) that rely on # lru_scores data. _need_lru = self.config.lru_enabled or self.config.scheduling_mode == "multi_score" _lru_coro = redis_ops.get_lru_scores(pid, all_key_ids) if _need_lru else None _latency_coro = ( redis_ops.batch_get_latency_avgs(pid, all_key_ids, self.config.latency_window_seconds) if self.config.scheduling_mode == "multi_score" else None ) # Gather all non-None coroutines in parallel. coros: list[Any] = [_cooldown_coro] _cost_idx = -1 _lru_idx = -1 _latency_idx = -1 if _cost_coro is not None: _cost_idx = len(coros) coros.append(_cost_coro) if _lru_coro is not None: _lru_idx = len(coros) coros.append(_lru_coro) if _latency_coro is not None: _latency_idx = len(coros) coros.append(_latency_coro) gathered = await asyncio.gather(*coros) cooldowns: dict[str, str | None] = gathered[0] # Cost check cost_exhausted: set[str] = set() cost_soft: set[str] = set() cost_totals: dict[str, int] = {} if _cost_idx >= 0: cost_totals = gathered[_cost_idx] limit = self.config.cost_limit_per_key_tokens if limit is not None: for kid, total in cost_totals.items(): if total >= limit: cost_exhausted.add(kid) elif total >= limit * self.config.cost_soft_threshold_percent / 100: cost_soft.add(kid) # LRU scores lru_scores: dict[str, float] = {} if _lru_idx >= 0: lru_scores = gathered[_lru_idx] # Latency averages latency_avgs: dict[str, float] = {} if _latency_idx >= 0: latency_avgs = gathered[_latency_idx] # Health scores (TTL cached, no Redis round-trip) -- only needed for multi_score health_scores: dict[str, float] = {} if self.config.scheduling_mode == "multi_score": health_scores = get_health_scores(pid, [c.key for c in candidates]) strategy_context.update( { "provider_type": provider_type, "all_key_ids": all_key_ids, "lru_scores": lru_scores, "cost_totals": cost_totals, "cost_limit_per_key_tokens": self.config.cost_limit_per_key_tokens, "latency_avgs": latency_avgs, "health_scores": health_scores, "keys_by_id": {str(c.key.id): c.key for c in candidates}, } ) # --- Strategy: compute_score ---------------------------------- custom_scores: dict[str, float] = {} for strategy in strategies: if hasattr(strategy, "compute_score"): for kid in all_key_ids: try: custom = strategy.compute_score( key_id=kid, config=self.config, context=strategy_context, ) if custom is not None: custom_scores[kid] = float(custom) lru_scores[kid] = float(custom) except Exception: pass # --- 3. Classify candidates ----------------------------------- # Use precomputed account states from CandidateBuilder when available # (upstream_metadata is deferred on pool keys to save memory). # Fall back to on-the-fly resolution for non-pool candidates. account_states: dict[str, Any] = {} for c in candidates: kid = str(c.key.id) if kid not in account_states: precomputed = getattr(c.key, "_pool_account_state", None) if precomputed is not None: account_states[kid] = precomputed else: account_states[kid] = resolve_pool_account_state( provider_type=provider_type, upstream_metadata=getattr(c.key, "upstream_metadata", None), oauth_invalid_reason=getattr(c.key, "oauth_invalid_reason", None), ) sticky_candidate: ProviderCandidate | None = None available: list[ProviderCandidate] = [] skipped: list[ProviderCandidate] = [] for c in candidates: kid = str(c.key.id) ct = PoolCandidateTrace(key_id=kid) ct.scoring_mode = self.config.scheduling_mode ct.latency_avg_ms = float(latency_avgs.get(kid, 0.0) or 0.0) ct.health_score = float(health_scores.get(kid, 1.0) or 1.0) if kid in custom_scores: ct.composite_score = float(custom_scores[kid]) # Already skipped upstream? if c.is_skipped: skipped.append(c) ct.skipped = True ct.skip_type = "upstream" trace.candidate_traces[kid] = ct continue # Account blocked? account_state = account_states[kid] if account_state.blocked: c.is_skipped = True skip_reason = account_state.reason or account_state.label or "account blocked" c.skip_reason = f"pool account blocked: {skip_reason}" skipped.append(c) ct.skipped = True ct.skip_type = "account_blocked" ct.account_block_code = account_state.code ct.account_block_label = account_state.label ct.account_block_reason = account_state.reason _attach_pool_extra(c, ct) trace.candidate_traces[kid] = ct continue cd_reason = cooldowns.get(kid) if cd_reason is not None: c.is_skipped = True c.skip_reason = f"pool cooldown: {cd_reason}" skipped.append(c) ct.skipped = True ct.skip_type = "cooldown" ct.cooldown_reason = cd_reason ct.cooldown_ttl = None # TTL skipped on hot path for perf _attach_pool_extra(c, ct) trace.candidate_traces[kid] = ct continue # Cost exhausted? if kid in cost_exhausted: c.is_skipped = True c.skip_reason = "pool cost limit reached" skipped.append(c) ct.skipped = True ct.skip_type = "cost_exhausted" ct.cost_window_usage = cost_totals.get(kid, 0) ct.cost_limit = self.config.cost_limit_per_key_tokens _attach_pool_extra(c, ct) trace.candidate_traces[kid] = ct continue # Sticky hit? if sticky_key_id and kid == sticky_key_id: sticky_candidate = c ct.reason = "sticky" ct.sticky_hit = True trace.sticky_session_used = True else: available.append(c) if kid in custom_scores and self.config.scheduling_mode == "multi_score": ct.reason = "multi_score" else: ct.reason = "lru" if lru_scores.get(kid, 0) > 0 else "random" ct.lru_score = lru_scores.get(kid, 0.0) ct.cost_window_usage = cost_totals.get(kid, 0) ct.cost_limit = self.config.cost_limit_per_key_tokens if kid in cost_soft: ct.cost_soft_threshold = True _attach_pool_extra(c, ct) trace.candidate_traces[kid] = ct # --- 4. Sort available by LRU --------------------------------- if lru_scores and available: available.sort(key=lambda c: lru_scores.get(str(c.key.id), 0.0)) # Random tiebreak among candidates with the same LRU score if len(available) > 1 and lru_scores: _shuffle_same_score_groups(available, lru_scores) # --- 5. Assemble final order ---------------------------------- result: list[ProviderCandidate] = [] if sticky_candidate is not None: result.append(sticky_candidate) result.extend(available) result.extend(skipped) if sticky_candidate: logger.debug( "Pool[{}]: sticky hit key={}", pid[:8], sticky_key_id and sticky_key_id[:8], ) # --- Strategy: after_select ----------------------------------- if result: first_kid = str(result[0].key.id) first_trace = trace.candidate_traces.get(first_kid) for strategy in strategies: if hasattr(strategy, "on_after_select") and first_trace: try: strategy.on_after_select( provider_id=pid, selected_key_id=first_kid, trace=first_trace, config=self.config, context=strategy_context, ) except Exception: pass # Attach the full trace to the first candidate for downstream use. if result: setattr(result[0], "_pool_scheduling_trace", trace) return result async def select_pool_keys( self, session_uuid: str | None, keys: list[ProviderAPIKey], *, availability_checker: ( Callable[[ProviderAPIKey], tuple[bool, str | None, str | None]] | None ) = None, page_size: int = 50, ) -> tuple[list[ProviderAPIKey], PoolSchedulingTrace]: """Select and order pool keys with trace output. Reuses :meth:`reorder_candidates` logic by adapting keys to lightweight candidate-like wrappers, then propagates skip/trace metadata back onto each key object for downstream execution/recording. When *availability_checker* is provided, post-sort availability checks are performed lazily: only the top *page_size* non-skipped keys are checked at a time; if all fail, the next page is checked, and so on. Keys beyond the last checked page are marked as ``deferred`` (skipped without checking) to avoid unnecessary CPU work on large pools. """ if not keys: return ( [], PoolSchedulingTrace( provider_id=self.provider_id, total_keys=0, session_uuid=session_uuid[:8] if session_uuid else None, ), ) class _KeyCandidate: __slots__ = ( "key", "is_skipped", "skip_reason", "_pool_extra_data", "_pool_scheduling_trace", ) def __init__(self, key: ProviderAPIKey) -> None: self.key = key self.is_skipped = False self.skip_reason: str | None = None self._pool_extra_data: dict | None = None self._pool_scheduling_trace: PoolSchedulingTrace | None = None wrappers = [_KeyCandidate(k) for k in keys] reordered_wrappers = await self.reorder_candidates(session_uuid, wrappers) # type: ignore[arg-type] trace: PoolSchedulingTrace | None = None if reordered_wrappers: maybe_trace = getattr(reordered_wrappers[0], "_pool_scheduling_trace", None) if isinstance(maybe_trace, PoolSchedulingTrace): trace = maybe_trace if trace is None: trace = PoolSchedulingTrace( provider_id=self.provider_id, total_keys=len(keys), session_uuid=session_uuid[:8] if session_uuid else None, ) ordered_keys: list[ProviderAPIKey] = [] for order_idx, wrapped in enumerate(reordered_wrappers): key = wrapped.key is_skipped = bool(getattr(wrapped, "is_skipped", False)) skip_reason = str(getattr(wrapped, "skip_reason", "") or "") setattr(key, "_pool_skipped", is_skipped) setattr(key, "_pool_skip_reason", skip_reason if skip_reason else None) setattr(key, "_pool_order_index", order_idx) pool_extra = getattr(wrapped, "_pool_extra_data", None) setattr( key, "_pool_extra_data", dict(pool_extra) if isinstance(pool_extra, dict) else {} ) ordered_keys.append(key) # -- 分页可用性检查 -- # 排序后对非 skipped key 分页调用 availability_checker， # 找到 page_size 个可用 key 后停止检查，剩余标记 deferred。 if availability_checker is not None: available_count = 0 found_enough = False for key in ordered_keys: if getattr(key, "_pool_skipped", False): continue if found_enough: setattr(key, "_pool_skipped", True) setattr(key, "_pool_skip_reason", "deferred") continue is_available, skip_reason_check, mapping_model = availability_checker(key) if not is_available: setattr(key, "_pool_skipped", True) setattr(key, "_pool_skip_reason", skip_reason_check) else: if mapping_model: setattr(key, "_pool_mapping_matched_model", mapping_model) available_count += 1 if available_count >= page_size: found_enough = True return ordered_keys, trace # ------------------------------------------------------------------ # Post-request hooks # ------------------------------------------------------------------ async def on_request_success( self, *, session_uuid: str | None, key_id: str, tokens_used: int = 0, ttfb_ms: int | None = None, ) -> None: """Called after a successful upstream request.""" pid = self.provider_id # Bind sticky session if session_uuid and self.config.sticky_session_ttl_seconds > 0: await redis_ops.set_sticky_binding( pid, session_uuid, key_id, self.config.sticky_session_ttl_seconds ) # Touch LRU -- needed for both plain LRU mode and multi_score dimensions # (e.g. cache_affinity) that rely on LRU timestamps. if self.config.lru_enabled or self.config.scheduling_mode == "multi_score": await redis_ops.touch_lru(pid, key_id) # Record cost if tokens_used > 0 and self.config.cost_limit_per_key_tokens is not None: await redis_ops.add_cost_entry( pid, key_id, tokens_used, self.config.cost_window_seconds ) # Record latency sample for multi-score scheduling. if self.config.scheduling_mode == "multi_score" and ttfb_ms is not None and ttfb_ms >= 0: await redis_ops.record_latency( pid, key_id, ttfb_ms, self.config.latency_window_seconds, self.config.latency_sample_limit, ) async def on_request_error( self, *, key_id: str, status_code: int, error_body: str | None = None, response_headers: dict[str, str] | None = None, ) -> None: """Called after an upstream error. Delegates to health policy.""" # Import lazily to avoid circular deps from src.services.provider.pool.health_policy import apply_health_policy await apply_health_policy( provider_id=self.provider_id, key_id=key_id, status_code=status_code, error_body=error_body, response_headers=response_headers, config=self.config, ) # ------------------------------------------------------------------ # Key schedulability check (used by candidate_builder) # ------------------------------------------------------------------ async def is_key_schedulable(self, key_id: str) -> tuple[bool, str | None]: """Check if *key_id* is currently schedulable (not in cooldown, not cost-exhausted). Returns ``(True, None)`` or ``(False, reason)``. """ pid = self.provider_id # Cooldown check cd = await redis_ops.get_cooldown(pid, key_id) if cd is not None: return False, f"pool cooldown: {cd}" # Cost check if self.config.cost_limit_per_key_tokens is not None: total = await redis_ops.get_cost_window_total( pid, key_id, self.config.cost_window_seconds ) if total >= self.config.cost_limit_per_key_tokens: return False, "pool cost limit reached" return True, None # Backward-compatible alias ClaudeCodePoolManager = PoolManager # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- _T = TypeVar("_T") def _attach_pool_extra(candidate: Any, ct: PoolCandidateTrace) -> None: """Attach pool trace extra_data onto a candidate object.""" existing = getattr(candidate, "_pool_extra_data", None) or {} existing.update(ct.to_extra_data()) setattr(candidate, "_pool_extra_data", existing) def _get_active_strategies(config: PoolConfig) -> list[Any]: """Get active strategies for the given config (lazy import).""" if not config.strategies: return [] try: # Import triggers built-in strategy registration via module-level side effects. from src.services.provider.pool import strategies as _builtin_strategies # noqa: F401 from src.services.provider.pool.strategy import get_active_strategies return get_active_strategies(config.strategies) except Exception: return [] def _shuffle_same_score( items: list[_T], lru_scores: dict[str, float], key_fn: Callable[[_T], str], ) -> None: """In-place random shuffle within groups that share the same LRU score.""" if len(items) <= 1: return i = 0 while i < len(items): score_i = lru_scores.get(key_fn(items[i]), 0.0) j = i + 1 while j < len(items) and lru_scores.get(key_fn(items[j]), 0.0) == score_i: j += 1 if j - i > 1: group = items[i:j] random.shuffle(group) items[i:j] = group i = j def _shuffle_same_score_groups( candidates: list[ProviderCandidate], lru_scores: dict[str, float], ) -> None: _shuffle_same_score(candidates, lru_scores, lambda c: str(c.key.id))