mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
1d9c77522a
- 删除全部 Python 源码 (src/) 及 Alembic 迁移脚本，归档至 _deprecated_py_src/ - 重构 Rust gateway ai_pipeline: 拆分 planner/finalize 模块，新增 contracts/adaptation 层 - 重组 handlers 模块为 admin/public/proxy/internal/shared 子模块结构 - 新增 executor 模块，引入 Rust 原生数据库迁移 (aether-data/migrations) - 简化 CI/Docker 构建流程，移除 base image 二级构建，统一为单一 app image - 移除 Python 相关基础设施文件 (entrypoint.sh, gunicorn_conf.py, Dockerfile.base)
342 lines
13 KiB
Python
342 lines
13 KiB
Python
"""Account Pool configuration (provider-agnostic)."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
from src.core.logger import logger
|
||
from src.services.provider.pool.dimensions import get_preset_dimension, get_preset_names
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class ScoringWeights:
|
||
"""Weights used by multi-score scheduling."""
|
||
|
||
lru: float = 0.3
|
||
latency: float = 0.25
|
||
health: float = 0.2
|
||
cost_remaining: float = 0.25
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class SchedulingPreset:
|
||
"""Single scheduling preset item with enable/disable and optional sub-config."""
|
||
|
||
preset: str
|
||
enabled: bool = True
|
||
mode: str | None = None
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class UnschedulableRule:
|
||
"""Keyword-based temporary unschedule rule."""
|
||
|
||
keyword: str
|
||
duration_minutes: int = 5
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class PoolConfig:
|
||
"""Parsed pool configuration for any Provider.
|
||
|
||
All transient state lives in Redis; this dataclass only holds
|
||
the *configuration* that controls pool behaviour.
|
||
"""
|
||
|
||
# -- Sticky Session -------------------------------------------------------
|
||
sticky_session_ttl_seconds: int = 3600 # 1 hour
|
||
# Key 优先模式下号池整体优先级（None 时回退 provider_priority）
|
||
global_priority: int | None = None
|
||
|
||
# -- Load-Aware Selection -------------------------------------------------
|
||
load_threshold_percent: int = 80
|
||
|
||
# -- Scheduling (unified preset list) -------------------------------------
|
||
scheduling_presets: tuple[SchedulingPreset, ...] = (
|
||
SchedulingPreset(preset="cache_affinity", enabled=True),
|
||
)
|
||
# Derived from scheduling_presets at parse time (backward compat for consumers)
|
||
lru_enabled: bool = True
|
||
scheduling_mode: str = "lru" # lru | multi_score
|
||
|
||
scoring_weights: ScoringWeights = field(default_factory=ScoringWeights)
|
||
latency_window_seconds: int = 3600
|
||
latency_sample_limit: int = 50
|
||
|
||
# -- Rolling-Window Cost Tracking -----------------------------------------
|
||
cost_window_seconds: int = 18000 # 5 hours
|
||
cost_limit_per_key_tokens: int | None = None # None = unlimited
|
||
cost_soft_threshold_percent: int = 80
|
||
|
||
# -- Cooldown Defaults ----------------------------------------------------
|
||
rate_limit_cooldown_seconds: int = 300 # 429
|
||
overload_cooldown_seconds: int = 30 # 529
|
||
|
||
# -- OAuth Proactive Refresh ----------------------------------------------
|
||
proactive_refresh_seconds: int = 180 # 3 minutes before expiry
|
||
|
||
# -- Health Policy --------------------------------------------------------
|
||
health_policy_enabled: bool = True
|
||
|
||
# -- Temporary Unschedulable Rules ----------------------------------------
|
||
unschedulable_rules: list[UnschedulableRule] = field(default_factory=list)
|
||
|
||
# -- Batch Operations -----------------------------------------------------
|
||
batch_concurrency: int = 8
|
||
|
||
# -- Quota Probing --------------------------------------------------------
|
||
probing_enabled: bool = False
|
||
probing_interval_minutes: int = 10
|
||
auto_remove_banned_keys: bool = False
|
||
|
||
# -- Stream Timeout Auto-Pause --------------------------------------------
|
||
stream_timeout_threshold: int = 3 # N timeouts within window trigger cooldown
|
||
stream_timeout_window_seconds: int = 1800 # 30 min counting window
|
||
stream_timeout_cooldown_seconds: int = 300 # 5 min cooldown
|
||
|
||
# -- Pluggable Strategies -------------------------------------------------
|
||
strategies: tuple[str, ...] = ()
|
||
|
||
|
||
def parse_pool_config(provider_config: Any) -> PoolConfig | None:
|
||
"""Parse PoolConfig from ``Provider.config``.
|
||
|
||
Only looks for the explicit ``pool_advanced`` key. Returns ``None``
|
||
when the provider has no pool section configured, meaning the caller
|
||
should use the normal (non-pool) scheduling path.
|
||
"""
|
||
config_dict = provider_config if isinstance(provider_config, dict) else {}
|
||
|
||
raw_advanced = config_dict.get("pool_advanced")
|
||
if raw_advanced is None:
|
||
return None
|
||
|
||
if not isinstance(raw_advanced, dict):
|
||
# Could be a pre-validated Pydantic model; grab its dict.
|
||
try:
|
||
raw_advanced = raw_advanced.model_dump() # type: ignore[union-attr]
|
||
except Exception:
|
||
logger.warning(
|
||
"PoolConfig: advanced config type invalid ({}), falling back to defaults",
|
||
type(raw_advanced).__name__,
|
||
)
|
||
return PoolConfig()
|
||
|
||
rules: list[UnschedulableRule] = []
|
||
raw_rules = raw_advanced.get("unschedulable_rules")
|
||
if isinstance(raw_rules, list):
|
||
for r in raw_rules:
|
||
if isinstance(r, dict) and isinstance(r.get("keyword"), str):
|
||
rules.append(
|
||
UnschedulableRule(
|
||
keyword=r["keyword"],
|
||
duration_minutes=int(r.get("duration_minutes", 5)),
|
||
)
|
||
)
|
||
|
||
def _int_or(key: str, default: int) -> int:
|
||
v = raw_advanced.get(key)
|
||
if v is None:
|
||
return default
|
||
try:
|
||
return int(v)
|
||
except (TypeError, ValueError):
|
||
return default
|
||
|
||
def _bool_or(key: str, default: bool) -> bool:
|
||
v = raw_advanced.get(key)
|
||
if v is None:
|
||
return default
|
||
return bool(v)
|
||
|
||
def _opt_int(key: str) -> int | None:
|
||
v = raw_advanced.get(key)
|
||
if v is None:
|
||
return None
|
||
try:
|
||
return int(v)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
scoring_weights = _parse_scoring_weights(raw_advanced.get("scoring_weights"))
|
||
|
||
# Parse scheduling presets (new object-list format or legacy string-list)
|
||
presets = _parse_scheduling_presets_v2(
|
||
raw_advanced.get("scheduling_presets"),
|
||
legacy_mode=raw_advanced.get("scheduling_mode"),
|
||
legacy_lru=raw_advanced.get("lru_enabled"),
|
||
)
|
||
|
||
# Derive scheduling_mode and lru_enabled from the presets list
|
||
enabled = [p for p in presets if p.enabled]
|
||
lru_enabled = any(p.preset == "lru" for p in enabled)
|
||
non_lru_enabled = [p for p in enabled if p.preset != "lru"]
|
||
scheduling_mode = "multi_score" if non_lru_enabled else "lru"
|
||
|
||
strategies = list(_parse_strategies(raw_advanced.get("strategies")))
|
||
if scheduling_mode == "multi_score" and "multi_score" not in strategies:
|
||
strategies.append("multi_score")
|
||
|
||
return PoolConfig(
|
||
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
|
||
global_priority=_opt_int("global_priority"),
|
||
load_threshold_percent=_int_or("load_threshold_percent", 80),
|
||
scheduling_presets=presets,
|
||
lru_enabled=lru_enabled,
|
||
scheduling_mode=scheduling_mode,
|
||
scoring_weights=scoring_weights,
|
||
latency_window_seconds=_int_or("latency_window_seconds", 3600),
|
||
latency_sample_limit=_int_or("latency_sample_limit", 50),
|
||
cost_window_seconds=_int_or("cost_window_seconds", 18000),
|
||
cost_limit_per_key_tokens=_opt_int("cost_limit_per_key_tokens"),
|
||
cost_soft_threshold_percent=_int_or("cost_soft_threshold_percent", 80),
|
||
rate_limit_cooldown_seconds=_int_or("rate_limit_cooldown_seconds", 300),
|
||
overload_cooldown_seconds=_int_or("overload_cooldown_seconds", 30),
|
||
proactive_refresh_seconds=_int_or("proactive_refresh_seconds", 180),
|
||
health_policy_enabled=_bool_or("health_policy_enabled", True),
|
||
unschedulable_rules=rules,
|
||
batch_concurrency=max(1, min(_int_or("batch_concurrency", 8), 32)),
|
||
probing_enabled=_bool_or("probing_enabled", False),
|
||
probing_interval_minutes=max(1, min(_int_or("probing_interval_minutes", 10), 1440)),
|
||
auto_remove_banned_keys=_bool_or("auto_remove_banned_keys", False),
|
||
stream_timeout_threshold=_int_or("stream_timeout_threshold", 3),
|
||
stream_timeout_window_seconds=_int_or("stream_timeout_window_seconds", 1800),
|
||
stream_timeout_cooldown_seconds=_int_or("stream_timeout_cooldown_seconds", 300),
|
||
strategies=tuple(strategies),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Internal parsers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _allowed_preset_names() -> set[str]:
|
||
return get_preset_names() | {"lru"}
|
||
|
||
|
||
def _get_preset_mode_meta(preset_name: str) -> tuple[tuple[str, ...], str | None]:
|
||
dim = get_preset_dimension(preset_name)
|
||
if dim is None or not dim.modes:
|
||
return (), None
|
||
|
||
modes = tuple(str(mode).strip().lower() for mode in dim.modes if str(mode).strip())
|
||
if not modes:
|
||
return (), None
|
||
|
||
raw_default = str(dim.default_mode or "").strip().lower()
|
||
default_mode = raw_default if raw_default in modes else modes[0]
|
||
return modes, default_mode
|
||
|
||
|
||
def _parse_strategies(raw: Any) -> tuple[str, ...]:
|
||
"""Parse strategy names from config (list[str] -> tuple[str, ...])."""
|
||
if not isinstance(raw, list):
|
||
return ()
|
||
return tuple(str(s) for s in raw if isinstance(s, str) and s)
|
||
|
||
|
||
def _parse_scoring_weights(raw: Any) -> ScoringWeights:
|
||
"""Parse scoring weights with graceful fallback."""
|
||
if not isinstance(raw, dict):
|
||
return ScoringWeights()
|
||
|
||
def _float_or(value: Any, default: float) -> float:
|
||
try:
|
||
parsed = float(value)
|
||
except (TypeError, ValueError):
|
||
return default
|
||
return max(0.0, min(parsed, 1.0))
|
||
|
||
return ScoringWeights(
|
||
lru=_float_or(raw.get("lru"), 0.3),
|
||
latency=_float_or(raw.get("latency"), 0.25),
|
||
health=_float_or(raw.get("health"), 0.2),
|
||
cost_remaining=_float_or(raw.get("cost_remaining"), 0.25),
|
||
)
|
||
|
||
|
||
def _parse_scheduling_presets_v2(
|
||
raw: Any,
|
||
*,
|
||
legacy_mode: Any = None,
|
||
legacy_lru: Any = None,
|
||
) -> tuple[SchedulingPreset, ...]:
|
||
"""Parse scheduling presets, supporting both new and legacy formats.
|
||
|
||
New format::
|
||
|
||
[{"preset": "lru", "enabled": true},
|
||
{"preset": "free_team_first", "enabled": true, "mode": "free_only"},
|
||
...]
|
||
|
||
Legacy format::
|
||
|
||
["free_team_first", "recent_refresh"] (with separate scheduling_mode / lru_enabled)
|
||
"""
|
||
if isinstance(raw, list) and raw:
|
||
first = raw[0]
|
||
if isinstance(first, dict):
|
||
return _parse_preset_object_list(raw)
|
||
if isinstance(first, str):
|
||
return _convert_legacy_string_list(raw, legacy_mode, legacy_lru)
|
||
|
||
# No presets at all: derive from legacy fields
|
||
return _build_from_legacy_fields(legacy_mode, legacy_lru)
|
||
|
||
|
||
def _parse_preset_object_list(raw: list[Any]) -> tuple[SchedulingPreset, ...]:
|
||
"""Parse new-format object list into SchedulingPreset tuple."""
|
||
allowed = _allowed_preset_names()
|
||
ordered: list[SchedulingPreset] = []
|
||
seen: set[str] = set()
|
||
for item in raw:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
name = str(item.get("preset", "")).strip().lower()
|
||
if name not in allowed or name in seen:
|
||
continue
|
||
seen.add(name)
|
||
enabled = bool(item.get("enabled", True))
|
||
mode: str | None = None
|
||
modes, default_mode = _get_preset_mode_meta(name)
|
||
if modes:
|
||
raw_mode = str(item.get("mode", default_mode) or "").strip().lower()
|
||
mode = raw_mode if raw_mode in modes else default_mode
|
||
ordered.append(SchedulingPreset(preset=name, enabled=enabled, mode=mode))
|
||
return tuple(ordered) if ordered else (SchedulingPreset(preset="lru", enabled=True),)
|
||
|
||
|
||
def _convert_legacy_string_list(
|
||
raw: list[Any],
|
||
legacy_mode: Any,
|
||
legacy_lru: Any,
|
||
) -> tuple[SchedulingPreset, ...]:
|
||
"""Convert legacy string list + mode/lru fields to new format."""
|
||
lru_enabled = legacy_lru if isinstance(legacy_lru, bool) else True
|
||
|
||
allowed_non_lru = _allowed_preset_names() - {"lru"}
|
||
items: list[SchedulingPreset] = [SchedulingPreset(preset="lru", enabled=lru_enabled)]
|
||
seen: set[str] = {"lru"}
|
||
for p in raw:
|
||
if not isinstance(p, str):
|
||
continue
|
||
name = p.strip().lower()
|
||
if name not in allowed_non_lru or name in seen:
|
||
continue
|
||
seen.add(name)
|
||
items.append(SchedulingPreset(preset=name, enabled=True))
|
||
return tuple(items)
|
||
|
||
|
||
def _build_from_legacy_fields(legacy_mode: Any, legacy_lru: Any) -> tuple[SchedulingPreset, ...]:
|
||
"""Build presets from legacy scheduling_mode / lru_enabled only."""
|
||
# Explicit lru_enabled=True -> LRU; explicit lru_enabled=False -> cache_affinity.
|
||
# No legacy fields at all -> default to cache_affinity.
|
||
if isinstance(legacy_lru, bool):
|
||
if legacy_lru:
|
||
return (SchedulingPreset(preset="lru", enabled=True),)
|
||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|
||
return (SchedulingPreset(preset="cache_affinity", enabled=True),)
|