i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(pool,scheduling): 号池调度维度、配额冷却机制与管理后台重构

Browse Source 
- 新增 scheduling_dimensions 模块，为每个 Key 计算多维调度状态(手动/冷却/熔断/成本/健康)
- 新增 quota_cooldown 模块，统一判定 Key 的有效冷却原因
- Pool 管理后台 API 扩展 Key 详情字段(调度状态/维度/配额/OAuth 信息)
- 前端 Pool 管理页面重写，支持调度状态展示、批量清理封禁 Key
- Handler 基类增加请求调度元数据采集，stream telemetry 增强
- 请求时间线组件增强，支持 attempted 候选展示
- Kiro OAuth 凭证导入解析改进
- 新增 usage 表 provider_key 索引迁移
- 补充调度维度、配额冷却、候选枚举等单元测试

Closes #197

Co-authored-by: AAEE86 <33052466+AAEE86@users.noreply.github.com>
This commit is contained in:
fawney19
2026-03-03 09:22:20 +08:00
parent f787b1b02a
commit 11997c024e
40 changed files with 4419 additions and 823 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
src/services/orchestration/candidate_resolver.py
View File

@@ -214,6 +214,11 @@ class CandidateResolver:
provider = candidate.provider
endpoint = candidate.endpoint
key = candidate.key
pool_extra = (
getattr(candidate, "_pool_extra_data", None)
if isinstance(getattr(candidate, "_pool_extra_data", None), dict)
else {}
)
if candidate.is_skipped:
record_id = str(uuid.uuid4())
@@ -235,6 +240,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),
@@ -269,6 +275,7 @@ class CandidateResolver:
"needs_conversion": candidate.needs_conversion,
"provider_api_format": candidate.provider_api_format or None,
"mapping_matched_model": candidate.mapping_matched_model or None,
**pool_extra,
},
"required_capabilities": active_capabilities,
"created_at": datetime.now(timezone.utc),

48
src/services/provider/adapters/kiro/models/credentials.py
View File

@@ -36,6 +36,28 @@ def _nonempty(s: str | None) -> str | None:
return None
def _normalize_auth_method(value: str | None) -> str:
method = (value or "").strip().lower()
if not method:
return "social"
# 历史/别名兼容：统一映射到 idc
if method in {
"idc",
"builder-id",
"builder_id",
"builderid",
"identity-center",
"identity_center",
"identitycenter",
"iam",
"device",
"device_authorization",
"device-auth",
}:
return "idc"
return method
def _parse_iso_to_epoch_seconds(value: object) -> int | None:
if not isinstance(value, str) or not value.strip():
return None
@@ -111,6 +133,11 @@ class KiroAuthConfig:
- 包含 clientId + clientSecret -> IdC
- 仅含 refreshToken -> Social
"""
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
normalized_explicit = _normalize_auth_method(explicit_method)
if normalized_explicit != "social":
return normalized_explicit
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -136,7 +163,12 @@ class KiroAuthConfig:
return False, "refreshToken 不完整（含有 ...），请导出完整的 Token"
# IdC 类型需要 clientId 和 clientSecret
auth_method = KiroAuthConfig.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else KiroAuthConfig.infer_auth_method(raw)
)
if auth_method == "idc":
client_id = raw.get("client_id") or raw.get("clientId")
client_secret = raw.get("client_secret") or raw.get("clientSecret")
@@ -155,8 +187,12 @@ class KiroAuthConfig:
provider_type = _get_str(raw, "provider_type", "providerType") or "kiro"
# 自动推断 auth_method（如果未显式指定）
explicit_method = _get_str(raw, "auth_method", "authMethod")
auth_method = explicit_method.lower() if explicit_method else cls.infer_auth_method(raw)
explicit_method = _get_str(raw, "auth_method", "authMethod", "auth_type", "authType")
auth_method = (
_normalize_auth_method(explicit_method)
if explicit_method
else cls.infer_auth_method(raw)
)
refresh_token = (_get_str(raw, "refresh_token", "refreshToken") or "").strip()
@@ -168,7 +204,7 @@ class KiroAuthConfig:
cfg = cls(
provider_type=provider_type,
auth_method=(auth_method or "social").lower(),
auth_method=_normalize_auth_method(auth_method),
refresh_token=refresh_token,
expires_at=int(expires_at),
profile_arn=_get_str(raw, "profile_arn", "profileArn"),
@@ -185,10 +221,6 @@ class KiroAuthConfig:
access_token=_get_str(raw, "access_token", "accessToken"),
)
# Normalize auth_method aliases.
if cfg.auth_method in {"builder-id", "builder_id", "iam"}:
cfg.auth_method = "idc"
return cfg
def to_dict(self) -> dict[str, Any]:

44
src/services/provider/pool/redis_ops.py
View File

@@ -440,6 +440,50 @@ async def get_key_sticky_count(provider_id: str, key_id: str) -> int:
return 0
async def batch_get_key_sticky_counts(
provider_id: str,
key_ids: list[str],
) -> dict[str, int]:
"""Count sticky sessions for multiple keys in a single scan (admin only)."""
if not key_ids:
return {}
redis = await _get_redis()
if redis is None:
return {kid: 0 for kid in key_ids}
target_ids = set(key_ids)
counts: dict[str, int] = {kid: 0 for kid in key_ids}
try:
pattern = f"{PREFIX}:{provider_id}:sticky:*"
batch: list[bytes | str] = []
async for key in redis.scan_iter(match=pattern, count=200):
batch.append(key)
if len(batch) >= 200:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
batch.clear()
if batch:
vals = await redis.mget(batch)
for val in vals:
if not val:
continue
bound_id = val.decode() if isinstance(val, bytes) else str(val)
if bound_id in target_ids:
counts[bound_id] = counts.get(bound_id, 0) + 1
return counts
except Exception:
return {kid: 0 for kid in key_ids}
async def get_cooldown_ttl(provider_id: str, key_id: str) -> int | None:
"""Get remaining cooldown TTL in seconds. None = no cooldown."""
redis = await _get_redis()

373
src/services/provider/pool/scheduling_dimensions.py Normal file
View File

@@ -0,0 +1,373 @@
"""Pool scheduling dimension registry and evaluation helpers.
This module keeps pool scheduling scoring isolated from API layer code.
Callers build a :class:`PoolSchedulingSnapshot` and evaluate it against
registered dimensions to obtain a normalized summary.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Protocol
PoolDimensionStatus = str # ok / degraded / blocked
@dataclass(frozen=True, slots=True)
class PoolSchedulingSnapshot:
"""Point-in-time scheduling inputs for one key."""
is_active: bool
cooldown_reason: str | None
cooldown_ttl_seconds: int | None
circuit_breaker_open: bool
cost_window_usage: int
cost_limit: int | None
cost_soft_threshold_percent: int = 80
health_score: float = 1.0
@dataclass(frozen=True, slots=True)
class PoolSchedulingDimensionResult:
"""Evaluation output for one scheduling dimension."""
code: str
label: str
status: PoolDimensionStatus = "ok"
blocking: bool = False
source: str = "pool"
weight: int = 1
score: float = 1.0
detail: str | None = None
ttl_seconds: int | None = None
@dataclass(frozen=True, slots=True)
class PoolSchedulingSummary:
"""Merged scheduling state across all dimensions."""
status: str # available / degraded / blocked
reason: str
label: str
score: float
candidate_eligible: bool
blocked_count: int
degraded_count: int
class PoolSchedulingDimension(Protocol):
"""Dimension evaluator protocol."""
code: str
label: str
source: str
weight: int
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
"""Evaluate one dimension from snapshot."""
@dataclass(frozen=True, slots=True)
class _ManualEnableDimension:
code: str = "manual_disabled"
label: str = "已禁用"
source: str = "manual"
weight: int = 8
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if snapshot.is_active:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail="账号被手动禁用",
)
@dataclass(frozen=True, slots=True)
class _CooldownDimension:
code: str = "cooldown"
label: str = "冷却中"
source: str = "pool"
weight: int = 7
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.cooldown_reason:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=snapshot.cooldown_reason,
ttl_seconds=snapshot.cooldown_ttl_seconds,
)
@dataclass(frozen=True, slots=True)
class _CircuitBreakerDimension:
code: str = "circuit_open"
label: str = "熔断中"
source: str = "health"
weight: int = 6
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
if not snapshot.circuit_breaker_open:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
)
@dataclass(frozen=True, slots=True)
class _CostDimension:
code: str = "cost"
label: str = "成本"
source: str = "pool"
weight: int = 5
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
limit = snapshot.cost_limit
usage = max(snapshot.cost_window_usage, 0)
if limit is None or limit <= 0:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=f"{usage}/-",
)
ratio = usage / limit
detail = f"{usage}/{limit}"
if ratio >= 1.0:
return PoolSchedulingDimensionResult(
code="cost_exhausted",
label="成本超限",
source=self.source,
weight=self.weight,
status="blocked",
blocking=True,
score=0.0,
detail=detail,
)
soft_threshold = max(1, min(snapshot.cost_soft_threshold_percent, 100))
if ratio * 100 >= soft_threshold:
return PoolSchedulingDimensionResult(
code="cost_soft",
label="成本接近上限",
source=self.source,
weight=self.weight,
status="degraded",
score=0.45,
detail=detail,
)
if ratio >= 0.6:
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="degraded",
score=0.72,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
@dataclass(frozen=True, slots=True)
class _HealthDimension:
code: str = "health"
label: str = "健康度"
source: str = "health"
weight: int = 4
def evaluate(self, snapshot: PoolSchedulingSnapshot) -> PoolSchedulingDimensionResult:
score = max(0.0, min(snapshot.health_score, 1.0))
detail = f"{score:.2f}"
if score < 0.5:
return PoolSchedulingDimensionResult(
code="health_low",
label="健康度过低",
source=self.source,
weight=self.weight,
status="degraded",
score=0.3,
detail=detail,
)
if score < 0.8:
return PoolSchedulingDimensionResult(
code="health_degraded",
label="健康度下降",
source=self.source,
weight=self.weight,
status="degraded",
score=0.65,
detail=detail,
)
return PoolSchedulingDimensionResult(
code=self.code,
label=self.label,
source=self.source,
weight=self.weight,
status="ok",
score=1.0,
detail=detail,
)
_POOL_DIMENSION_REGISTRY: dict[str, PoolSchedulingDimension] = {}
_POOL_DIMENSION_ORDER: list[str] = []
def register_pool_scheduling_dimension(name: str, dimension: PoolSchedulingDimension) -> None:
"""Register a dimension evaluator by name."""
normalized = name.strip()
if not normalized:
return
if normalized not in _POOL_DIMENSION_ORDER:
_POOL_DIMENSION_ORDER.append(normalized)
_POOL_DIMENSION_REGISTRY[normalized] = dimension
def get_pool_scheduling_dimension(name: str) -> PoolSchedulingDimension | None:
"""Fetch a registered dimension evaluator."""
return _POOL_DIMENSION_REGISTRY.get(name.strip())
def list_pool_scheduling_dimensions() -> tuple[str, ...]:
"""List registered dimension names in evaluation order."""
return tuple(_POOL_DIMENSION_ORDER)
def evaluate_pool_scheduling_dimensions(
snapshot: PoolSchedulingSnapshot,
*,
dimension_names: tuple[str, ...] | None = None,
) -> list[PoolSchedulingDimensionResult]:
"""Evaluate snapshot across all registered dimensions."""
names = dimension_names or list_pool_scheduling_dimensions()
results: list[PoolSchedulingDimensionResult] = []
for name in names:
dimension = get_pool_scheduling_dimension(name)
if dimension is None:
continue
results.append(dimension.evaluate(snapshot))
return results
def summarize_pool_scheduling_dimensions(
dimensions: list[PoolSchedulingDimensionResult],
) -> PoolSchedulingSummary:
"""Summarize dimension outputs into a unified scheduling state."""
if not dimensions:
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=100.0,
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
blocked = [item for item in dimensions if item.status == "blocked" or item.blocking]
degraded = [item for item in dimensions if item.status == "degraded"]
total_weight = sum(max(item.weight, 1) for item in dimensions)
weighted_score = sum(
max(item.weight, 1) * max(min(item.score, 1.0), 0.0) for item in dimensions
) / max(total_weight, 1)
if blocked:
primary = blocked[0]
return PoolSchedulingSummary(
status="blocked",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=False,
blocked_count=len(blocked),
degraded_count=len(degraded),
)
if degraded:
primary = degraded[0]
return PoolSchedulingSummary(
status="degraded",
reason=primary.code,
label=primary.label,
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=len(degraded),
)
return PoolSchedulingSummary(
status="available",
reason="available",
label="可用",
score=round(weighted_score * 100, 1),
candidate_eligible=True,
blocked_count=0,
degraded_count=0,
)
def _register_default_dimensions() -> None:
register_pool_scheduling_dimension("manual", _ManualEnableDimension())
register_pool_scheduling_dimension("cooldown", _CooldownDimension())
register_pool_scheduling_dimension("circuit", _CircuitBreakerDimension())
register_pool_scheduling_dimension("cost", _CostDimension())
register_pool_scheduling_dimension("health", _HealthDimension())
_register_default_dimensions()

20
src/services/provider/pool/trace.py
View File

@@ -63,7 +63,12 @@ class PoolSchedulingTrace:
session_uuid: str | None = None
candidate_traces: dict[str, PoolCandidateTrace] = field(default_factory=dict)
def build_summary(self, success_key_id: str | None = None) -> dict[str, Any]:
def build_summary(
self,
success_key_id: str | None = None,
*,
attempted_key_ids: set[str] | None = None,
) -> dict[str, Any]:
"""Build compact dict for ``Usage.request_metadata["pool_summary"]``."""
skipped_cooldown = 0
skipped_cost = 0
@@ -74,8 +79,17 @@ class PoolSchedulingTrace:
skipped_cooldown += 1
elif t.skip_type == "cost_exhausted":
skipped_cost += 1
else:
attempted += 1
if attempted_key_ids is None:
# Backward-compatible behavior: count all schedulable keys.
attempted = sum(1 for t in self.candidate_traces.values() if not t.skipped)
else:
# Preferred behavior: count only keys that were actually executed.
attempted = sum(
1
for kid in attempted_key_ids
if kid in self.candidate_traces and not self.candidate_traces[kid].skipped
)
success_reason: str | None = None
if success_key_id and success_key_id in self.candidate_traces:

35
src/services/provider_keys/quota_cooldown.py Normal file
View File

@@ -0,0 +1,35 @@
"""配额冷却判定工具。"""
from __future__ import annotations
from typing import Any
from src.core.logger import logger
from src.services.scheduling.quota_skipper import is_key_quota_exhausted
def resolve_effective_cooldown_reason(
*,
provider_type: str | None,
key: Any,
redis_reason: str | None,
) -> str | None:
"""返回 Key 的有效冷却原因。
规则：
- Redis 冷却存在时，优先返回 Redis 原因（429/403/quota_exhausted 等）。
- Redis 冷却不存在时，回退到 upstream_metadata 配额判断：
若账号级配额耗尽（Codex/Kiro），返回 ``quota_exhausted``。
"""
if redis_reason:
return redis_reason
try:
exhausted, _ = is_key_quota_exhausted(provider_type, key, model_name="")
except Exception:
logger.opt(exception=True).debug(
"quota_cooldown: is_key_quota_exhausted failed for key={}",
getattr(key, "id", "?"),
)
return None
return "quota_exhausted" if exhausted else None

56
src/services/scheduling/candidate_builder.py
View File

@@ -576,32 +576,21 @@ class CandidateBuilder:
if not active_keys:
continue
# --- Pool branch: select a single key internally ------
# Pool provider should still expose all key candidates here.
# Runtime pool scheduling/failover is handled later by TaskService._apply_pool_reorder.
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if pool_cfg is not None:
selected_key = await self._pool_select_key(
db, provider, pool_cfg, active_keys, request_body
)
if selected_key is None:
logger.debug(
"Pool[{}]: no schedulable key for endpoint {}",
str(provider.id)[:8],
endpoint_format_str,
)
continue
keys_to_check: list[ProviderAPIKey] = [selected_key]
else:
# --- Normal branch: check all keys ----
use_random = all((key.cache_ttl_minutes or 0) == 0 for key in active_keys)
if use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
use_random = False
elif use_random and len(active_keys) > 1:
logger.debug(
" Provider {} 启用 Key 轮换模式 (endpoint_format={}, {} keys)",
provider.name,
endpoint_format_str,
len(active_keys),
)
keys_to_check = self._sorter.shuffle_keys_by_internal_priority(
active_keys, affinity_key, use_random
)
for key in keys_to_check:
# Key 级别检查（健康度/熔断按 provider_format bucket）
@@ -660,22 +649,3 @@ class CandidateBuilder:
candidates = candidates[:max_candidates]
return candidates
async def _pool_select_key(
self,
db: Session,
provider: Provider,
pool_cfg: "PoolConfig",
active_keys: list[ProviderAPIKey],
request_body: dict | None,
) -> ProviderAPIKey | None:
"""Select a single key via pool scheduling (sticky -> cooldown/cost -> LRU)."""
from src.services.provider.pool.hooks import get_pool_hook
from src.services.provider.pool.manager import PoolManager
provider_type = str(getattr(provider, "provider_type", "") or "")
hook = get_pool_hook(provider_type)
session_uuid = hook.extract_session_uuid(request_body) if hook and request_body else None
mgr = PoolManager(str(provider.id), pool_cfg)
release_db_connection_before_await(db)
return await mgr.select_key(session_uuid, active_keys)

25
src/services/task/service.py
View File

@@ -119,6 +119,7 @@ class TaskService:
allow_format_conversion=allow_format_conversion,
capability_requirements=capability_requirements,
max_candidates=max_candidates,
request_body=request_body,
)
candidate_keys = []
@@ -598,8 +599,22 @@ class TaskService:
# Build pool scheduling summary from traces collected during reorder.
if pool_traces and result.key_id:
try:
attempted_key_ids: set[str] = set()
for ck in result.candidate_keys or []:
status = str(getattr(ck, "status", "") or "").strip().lower()
if status in {"", "available", "pending", "skipped", "unused"}:
continue
kid = getattr(ck, "key_id", None)
if isinstance(kid, str) and kid:
attempted_key_ids.add(kid)
if not attempted_key_ids:
attempted_key_ids.add(str(result.key_id))
for pt in pool_traces:
summary = pt.build_summary(result.key_id)
summary = pt.build_summary(
result.key_id,
attempted_key_ids=attempted_key_ids,
)
if summary:
result.pool_summary = summary
break
@@ -1204,6 +1219,7 @@ class TaskService:
allow_format_conversion: bool = False,
capability_requirements: dict[str, bool] | None = None,
max_candidates: int | None = None,
request_body: dict[str, Any] | None = None,
) -> Any:
"""
Unified ASYNC submit entrypoint (Phase 3.2).
@@ -1300,6 +1316,7 @@ class TaskService:
request_id=request_id,
is_stream=False,
capability_requirements=capability_requirements,
request_body=request_body,
)
if not candidates:
@@ -1309,6 +1326,12 @@ class TaskService:
last_status_code=None,
)
# Account Pool: keep internal key failover order/skip behavior
# consistent with the SYNC path.
candidates, _pool_traces = await self._apply_pool_reorder(
candidates, request_body=request_body
)
if max_candidates is not None and max_candidates > 0:
candidates = candidates[:max_candidates]

1
src/services/usage/_recording_helpers.py
View File

@@ -39,6 +39,7 @@ METADATA_KEEP_KEYS: frozenset[str] = frozenset(
"billing_updated_at",
"perf",
"pool_summary",
"scheduling_audit",
"_metadata_truncated",
}
)