mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
refactor: defer ProviderAPIKey 大 JSON 字段,减少调度器和模型获取的内存占用
- candidate_builder: defer adjustment_history/utilization_samples/upstream_metadata, 号池路径在释放 DB 连接前预计算 _pool_account_state 避免后续 N+1 查询 - pool/manager: 优先读取预计算的 _pool_account_state,fallback 到实时解析 - fetch_scheduler: 三处 ProviderAPIKey 查询改用 defer/load_only 排除无关字段
This commit is contained in:
@@ -19,7 +19,7 @@ from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy.orm import joinedload
|
||||
from sqlalchemy.orm import defer, joinedload, load_only
|
||||
|
||||
from src.core.cache_service import CacheService
|
||||
from src.core.crypto import crypto_service
|
||||
@@ -284,7 +284,19 @@ class ModelFetchScheduler:
|
||||
"""更新 Key 的错误信息(独立事务)"""
|
||||
try:
|
||||
with create_session() as db:
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
defer(ProviderAPIKey.adjustment_history),
|
||||
defer(ProviderAPIKey.utilization_samples),
|
||||
defer(ProviderAPIKey.health_by_format),
|
||||
defer(ProviderAPIKey.circuit_breaker_by_format),
|
||||
defer(ProviderAPIKey.upstream_metadata),
|
||||
defer(ProviderAPIKey.allowed_models),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
if key:
|
||||
key.last_models_fetch_at = datetime.now(timezone.utc)
|
||||
key.last_models_fetch_error = error_msg
|
||||
@@ -384,7 +396,18 @@ class ModelFetchScheduler:
|
||||
with create_session() as db:
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(joinedload(ProviderAPIKey.provider))
|
||||
.options(
|
||||
load_only(
|
||||
ProviderAPIKey.id,
|
||||
ProviderAPIKey.is_active,
|
||||
ProviderAPIKey.auto_fetch_models,
|
||||
ProviderAPIKey.api_key,
|
||||
ProviderAPIKey.auth_type,
|
||||
ProviderAPIKey.auth_config,
|
||||
ProviderAPIKey.provider_id,
|
||||
ProviderAPIKey.proxy,
|
||||
),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
@@ -482,7 +505,18 @@ class ModelFetchScheduler:
|
||||
|
||||
with create_session() as db:
|
||||
# 重新获取 Key(因为之前的连接已关闭)
|
||||
key = db.query(ProviderAPIKey).filter(ProviderAPIKey.id == key_id).first()
|
||||
# defer 不需要的大 JSON 字段,减少内存占用
|
||||
key = (
|
||||
db.query(ProviderAPIKey)
|
||||
.options(
|
||||
defer(ProviderAPIKey.adjustment_history),
|
||||
defer(ProviderAPIKey.utilization_samples),
|
||||
defer(ProviderAPIKey.health_by_format),
|
||||
defer(ProviderAPIKey.circuit_breaker_by_format),
|
||||
)
|
||||
.filter(ProviderAPIKey.id == key_id)
|
||||
.first()
|
||||
)
|
||||
if not key:
|
||||
logger.warning(f"Key {key_id} 在更新时不存在")
|
||||
return "error"
|
||||
|
||||
@@ -221,16 +221,22 @@ class PoolManager:
|
||||
pass
|
||||
|
||||
# --- 3. Classify candidates -----------------------------------
|
||||
# Pre-compute account states to avoid repeated dict parsing inside loop.
|
||||
# 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:
|
||||
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),
|
||||
)
|
||||
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] = []
|
||||
@@ -582,16 +588,20 @@ class PoolManager:
|
||||
pass
|
||||
|
||||
# --- 3. Classify keys -------------------------------------------------
|
||||
# Pre-compute account states to avoid repeated dict parsing inside loop.
|
||||
# Use precomputed account states when available.
|
||||
account_states_sk: dict[str, Any] = {}
|
||||
for k in keys:
|
||||
kid = str(k.id)
|
||||
if kid not in account_states_sk:
|
||||
account_states_sk[kid] = resolve_pool_account_state(
|
||||
provider_type=self.provider_type,
|
||||
upstream_metadata=getattr(k, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
|
||||
)
|
||||
precomputed = getattr(k, "_pool_account_state", None)
|
||||
if precomputed is not None:
|
||||
account_states_sk[kid] = precomputed
|
||||
else:
|
||||
account_states_sk[kid] = resolve_pool_account_state(
|
||||
provider_type=self.provider_type,
|
||||
upstream_metadata=getattr(k, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(k, "oauth_invalid_reason", None),
|
||||
)
|
||||
|
||||
sticky_key: ProviderAPIKey | None = None
|
||||
available: list[ProviderAPIKey] = []
|
||||
|
||||
@@ -35,6 +35,9 @@ from src.models.database import (
|
||||
)
|
||||
from src.services.health.monitor import health_monitor
|
||||
from src.services.provider.format import normalize_endpoint_signature
|
||||
from src.services.provider.pool.account_state import (
|
||||
resolve_pool_account_state as _resolve_pool_account_state,
|
||||
)
|
||||
from src.services.scheduling.quota_skipper import is_key_quota_exhausted
|
||||
from src.services.scheduling.utils import release_db_connection_before_await
|
||||
|
||||
@@ -97,10 +100,13 @@ class CandidateBuilder:
|
||||
db.query(Provider)
|
||||
.options(
|
||||
# 预加载 Provider 级别的 api_keys
|
||||
# defer 排除调度热路径不需要的字段,减少内存占用
|
||||
# 凭证类字段(api_key/auth_config)在执行阶段才由 get_provider_auth() 按需加载
|
||||
# 注意:adjustment_history/utilization_samples 在并发检查阶段
|
||||
# (AdaptiveReservationManager) 被访问,不可 defer
|
||||
# defer 排除调度热路径不需要的大 JSON 字段,减少号池场景内存占用
|
||||
# - 凭证类: api_key/auth_config 在执行阶段由 get_provider_auth() 按需加载
|
||||
# - adjustment_history/utilization_samples: 仅 AdaptiveReservationManager
|
||||
# 在并发检查时读取单个 key,号池/非号池均可 lazy load
|
||||
# - upstream_metadata: 号池模式在 _build_candidates 中预计算账号封禁
|
||||
# 状态并挂到 key._pool_account_state,排序阶段不再需要原始 JSON;
|
||||
# 非号池模式 key 少,lazy load 可忽略
|
||||
selectinload(Provider.api_keys)
|
||||
.defer(ProviderAPIKey.api_key)
|
||||
.defer(ProviderAPIKey.auth_config)
|
||||
@@ -113,7 +119,10 @@ class CandidateBuilder:
|
||||
.defer(ProviderAPIKey.last_models_fetch_at)
|
||||
.defer(ProviderAPIKey.last_models_fetch_error)
|
||||
.defer(ProviderAPIKey.max_probe_interval_minutes)
|
||||
.defer(ProviderAPIKey.expires_at),
|
||||
.defer(ProviderAPIKey.expires_at)
|
||||
.defer(ProviderAPIKey.adjustment_history)
|
||||
.defer(ProviderAPIKey.utilization_samples)
|
||||
.defer(ProviderAPIKey.upstream_metadata),
|
||||
# 预加载 endpoints(用于按 api_format 选择请求配置)
|
||||
selectinload(Provider.endpoints),
|
||||
# 同时加载 models 和 global_model 关系
|
||||
@@ -570,15 +579,33 @@ class CandidateBuilder:
|
||||
if pool_cfg is not None:
|
||||
# 号池优化:跳过逐 key 的 _check_key_availability 检查,
|
||||
# 直接收集全部 active key,将检查推迟到 PoolManager 排序后分页执行。
|
||||
# 在此释放 DB 连接,因为后续的 PoolManager 排序涉及大量 Redis 操作,
|
||||
# 避免在 Redis I/O 期间长时间占用 DB 连接池。
|
||||
release_db_connection_before_await(db)
|
||||
|
||||
pool_keys = list(keys_to_check)
|
||||
|
||||
if not pool_keys:
|
||||
continue
|
||||
|
||||
# 在释放 DB 连接前预计算账号封禁状态并挂到 Key 对象上。
|
||||
# upstream_metadata 是 deferred 字段,逐条 lazy load 会产生 N+1 查询;
|
||||
# 这里集中触发后,PoolManager 排序时直接读取 _pool_account_state 即可。
|
||||
provider_type_str = (
|
||||
str(getattr(provider, "provider_type", "") or "").strip().lower() or None
|
||||
)
|
||||
for pk in pool_keys:
|
||||
setattr(
|
||||
pk,
|
||||
"_pool_account_state",
|
||||
_resolve_pool_account_state(
|
||||
provider_type=provider_type_str,
|
||||
upstream_metadata=getattr(pk, "upstream_metadata", None),
|
||||
oauth_invalid_reason=getattr(pk, "oauth_invalid_reason", None),
|
||||
),
|
||||
)
|
||||
|
||||
# 释放 DB 连接,因为后续的 PoolManager 排序涉及大量 Redis 操作,
|
||||
# 避免在 Redis I/O 期间长时间占用 DB 连接池。
|
||||
release_db_connection_before_await(db)
|
||||
|
||||
provider_priority_raw = getattr(provider, "provider_priority", None)
|
||||
try:
|
||||
provider_priority = (
|
||||
|
||||
Reference in New Issue
Block a user