fix: TTL=0时启用Key随机轮换模式

- 当所有Key的cache_ttl_minutes都为0时,使用随机排序代替确定性哈希
- 将hashlib和random的import移到文件顶部
- 简化单Key场景的处理逻辑

Closes #57
This commit is contained in:
fawney19
2025-12-28 19:07:25 +08:00
parent 2b1d197047
commit d0ce798881

View File

@@ -30,6 +30,8 @@
from __future__ import annotations from __future__ import annotations
import hashlib
import random
import time import time
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
@@ -956,7 +958,16 @@ class CacheAwareScheduler:
# 获取活跃的 Key 并按 internal_priority + 负载均衡排序 # 获取活跃的 Key 并按 internal_priority + 负载均衡排序
active_keys = [key for key in endpoint.api_keys if key.is_active] active_keys = [key for key in endpoint.api_keys if key.is_active]
keys = self._shuffle_keys_by_internal_priority(active_keys, affinity_key) # 检查是否所有 Key 都是 TTL=0轮换模式
# 如果所有 Key 的 cache_ttl_minutes 都是 0 或 None则使用随机排序
use_random = all(
(key.cache_ttl_minutes or 0) == 0 for key in active_keys
) if active_keys else False
if use_random and len(active_keys) > 1:
logger.debug(
f" Endpoint {endpoint.id[:8]}... 启用 Key 轮换模式 (TTL=0, {len(active_keys)} keys)"
)
keys = self._shuffle_keys_by_internal_priority(active_keys, affinity_key, use_random)
for key in keys: for key in keys:
# Key 级别的能力检查(模型级别的能力检查已在上面完成) # Key 级别的能力检查(模型级别的能力检查已在上面完成)
@@ -1170,6 +1181,7 @@ class CacheAwareScheduler:
self, self,
keys: List[ProviderAPIKey], keys: List[ProviderAPIKey],
affinity_key: Optional[str] = None, affinity_key: Optional[str] = None,
use_random: bool = False,
) -> List[ProviderAPIKey]: ) -> List[ProviderAPIKey]:
""" """
对 API Key 按 internal_priority 分组,同优先级内部基于 affinity_key 进行确定性打乱 对 API Key 按 internal_priority 分组,同优先级内部基于 affinity_key 进行确定性打乱
@@ -1178,10 +1190,12 @@ class CacheAwareScheduler:
- 数字越小越优先使用 - 数字越小越优先使用
- 同优先级 Key 之间实现负载均衡 - 同优先级 Key 之间实现负载均衡
- 使用 affinity_key 哈希确保同一请求 Key 的请求稳定(避免破坏缓存亲和性) - 使用 affinity_key 哈希确保同一请求 Key 的请求稳定(避免破坏缓存亲和性)
- 当 use_random=True 时,使用随机排序实现轮换(用于 TTL=0 的场景)
Args: Args:
keys: API Key 列表 keys: API Key 列表
affinity_key: 亲和性标识符(通常为 API Key ID用于确定性打乱 affinity_key: 亲和性标识符(通常为 API Key ID用于确定性打乱
use_random: 是否使用随机排序TTL=0 时为 True
Returns: Returns:
排序后的 Key 列表 排序后的 Key 列表
@@ -1198,15 +1212,19 @@ class CacheAwareScheduler:
priority = key.internal_priority if key.internal_priority is not None else 999999 priority = key.internal_priority if key.internal_priority is not None else 999999
priority_groups[priority].append(key) priority_groups[priority].append(key)
# 对每个优先级组内的 Key 进行确定性打乱 # 对每个优先级组内的 Key 进行打乱
result = [] result = []
for priority in sorted(priority_groups.keys()): # 数字小的优先级高,排前面 for priority in sorted(priority_groups.keys()): # 数字小的优先级高,排前面
group_keys = priority_groups[priority] group_keys = priority_groups[priority]
if len(group_keys) > 1 and affinity_key: if len(group_keys) > 1:
# 改进的哈希策略:为每个 key 计算独立的哈希值 if use_random:
import hashlib # TTL=0 模式:使用随机排序实现 Key 轮换
shuffled = list(group_keys)
random.shuffle(shuffled)
result.extend(shuffled)
elif affinity_key:
# 正常模式:使用哈希确定性打乱(保持缓存亲和性)
key_scores = [] key_scores = []
for key in group_keys: for key in group_keys:
# 使用 affinity_key + key.id 的组合哈希 # 使用 affinity_key + key.id 的组合哈希
@@ -1218,8 +1236,11 @@ class CacheAwareScheduler:
sorted_group = [key for _, key in sorted(key_scores)] sorted_group = [key for _, key in sorted(key_scores)]
result.extend(sorted_group) result.extend(sorted_group)
else: else:
# 单个 Key 或没有 affinity_key 时保持原顺序 # 没有 affinity_key 时按 ID 排序保持稳定性
result.extend(sorted(group_keys, key=lambda k: k.id)) result.extend(sorted(group_keys, key=lambda k: k.id))
else:
# 单个 Key 直接添加
result.extend(group_keys)
return result return result