mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +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)
145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
"""
|
||
并发控制检查器 (ConcurrencyChecker)
|
||
|
||
从 CacheAwareScheduler 提取的 RPM 限流和动态预留逻辑。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
from typing import Any
|
||
|
||
from src.core.logger import logger
|
||
from src.models.database import ProviderAPIKey
|
||
from src.services.rate_limit.adaptive_reservation import AdaptiveReservationManager
|
||
from src.services.rate_limit.adaptive_rpm import get_adaptive_rpm_manager
|
||
from src.services.scheduling.schemas import ConcurrencySnapshot
|
||
|
||
|
||
class ConcurrencyChecker:
|
||
"""并发控制检查器，封装 RPM 限流和动态预留逻辑。"""
|
||
|
||
def __init__(
|
||
self,
|
||
concurrency_manager: Any,
|
||
reservation_manager: AdaptiveReservationManager,
|
||
) -> None:
|
||
self._concurrency_manager = concurrency_manager
|
||
self._reservation_manager = reservation_manager
|
||
|
||
@staticmethod
|
||
def get_effective_rpm_limit(key: ProviderAPIKey) -> int | None:
|
||
"""获取有效的 RPM 限制（委托给 AdaptiveRPMManager 统一逻辑）"""
|
||
return get_adaptive_rpm_manager().get_effective_limit(key)
|
||
|
||
async def check_available(
|
||
self,
|
||
key: ProviderAPIKey,
|
||
is_cached_user: bool = False,
|
||
) -> tuple[bool, ConcurrencySnapshot]:
|
||
"""
|
||
检查 RPM 限制是否可用（使用动态预留机制）
|
||
|
||
核心逻辑 - 动态缓存预留机制:
|
||
- 总槽位: 有效 RPM 限制（固定值或学习到的值）
|
||
- 预留比例: 由 AdaptiveReservationManager 根据置信度和负载动态计算
|
||
- 缓存用户可用: 全部槽位
|
||
- 新用户可用: 总槽位 x (1 - 动态预留比例)
|
||
|
||
Args:
|
||
key: ProviderAPIKey对象
|
||
is_cached_user: 是否是缓存用户
|
||
|
||
Returns:
|
||
(是否可用, 并发快照)
|
||
"""
|
||
# 获取有效的并发限制
|
||
effective_key_limit = self.get_effective_rpm_limit(key)
|
||
|
||
logger.debug(
|
||
" -> 并发检查: _concurrency_manager={}, "
|
||
"is_cached_user={}, effective_limit={}",
|
||
self._concurrency_manager is not None,
|
||
is_cached_user,
|
||
effective_key_limit,
|
||
)
|
||
|
||
if not self._concurrency_manager:
|
||
# 并发管理器不可用，直接返回True
|
||
logger.debug(" -> 无并发管理器，直接通过")
|
||
snapshot = ConcurrencySnapshot(
|
||
key_current=0,
|
||
key_limit=effective_key_limit,
|
||
is_cached_user=is_cached_user,
|
||
)
|
||
return True, snapshot
|
||
|
||
# 获取当前 RPM 计数
|
||
key_count = await self._concurrency_manager.get_key_rpm_count(
|
||
key_id=str(key.id),
|
||
)
|
||
|
||
can_use = True
|
||
|
||
# 计算动态预留比例
|
||
reservation_result = self._reservation_manager.calculate_reservation(
|
||
key=key,
|
||
current_usage=key_count,
|
||
effective_limit=effective_key_limit,
|
||
)
|
||
|
||
available_for_new = None
|
||
reservation_ratio = reservation_result.ratio
|
||
|
||
# 检查Key级别限制（使用动态预留比例）
|
||
if effective_key_limit is not None:
|
||
if is_cached_user:
|
||
# 缓存用户: 可以使用全部槽位
|
||
if key_count >= effective_key_limit:
|
||
can_use = False
|
||
else:
|
||
# 新用户: 只能使用 (1 - 动态预留比例) 的槽位
|
||
# 使用 max 确保至少有 1 个槽位可用
|
||
|
||
# 与 ConcurrencyManager 的 Lua 脚本保持一致：使用 floor 计算新用户可用槽位
|
||
available_for_new = max(
|
||
1, math.floor(effective_key_limit * (1 - reservation_ratio))
|
||
)
|
||
if key_count >= available_for_new:
|
||
logger.debug(
|
||
"Key {}... 新用户配额已满 " "({}/{}, 总{}, 预留{:.0%}[{}])",
|
||
key.id[:8],
|
||
key_count,
|
||
available_for_new,
|
||
effective_key_limit,
|
||
reservation_ratio,
|
||
reservation_result.phase,
|
||
)
|
||
can_use = False
|
||
|
||
key_limit_for_snapshot: int | None
|
||
if is_cached_user:
|
||
key_limit_for_snapshot = effective_key_limit
|
||
elif effective_key_limit is not None:
|
||
key_limit_for_snapshot = (
|
||
available_for_new if available_for_new is not None else effective_key_limit
|
||
)
|
||
else:
|
||
key_limit_for_snapshot = None
|
||
|
||
snapshot = ConcurrencySnapshot(
|
||
key_current=key_count,
|
||
key_limit=key_limit_for_snapshot,
|
||
is_cached_user=is_cached_user,
|
||
reservation_ratio=reservation_ratio,
|
||
reservation_phase=reservation_result.phase,
|
||
reservation_confidence=reservation_result.confidence,
|
||
load_factor=reservation_result.load_factor,
|
||
)
|
||
|
||
return can_use, snapshot
|
||
|
||
def get_reservation_stats(self) -> dict[str, Any]:
|
||
"""获取动态预留管理器的统计信息"""
|
||
return self._reservation_manager.get_stats()
|