mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
feat: 配置导入导出支持 ProxyNode,号池 key 可用性检查延迟到排序后分页执行
配置导入导出: - 导出/导入新增 ProxyNode(代理节点)数据 - 导入时自动建立 old_id -> new_id 映射,重映射 Provider/Endpoint/Key 中的 node_id - 前端预览和结果展示新增代理节点统计 号池调度优化: - CandidateBuilder 不再逐 key 调用 _check_key_availability,直接收集全部 active key - 将可用性检查参数打包到 PoolCandidate._deferred_check_params - PoolManager.select_pool_keys 排序后分页调用 availability_checker,找到足够可用 key 即停止 - 减少大号池场景下不必要的可用性检查开销
This commit is contained in:
@@ -362,12 +362,23 @@ class PoolManager:
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
keys: list[ProviderAPIKey],
|
||||
*,
|
||||
availability_checker: (
|
||||
Callable[[ProviderAPIKey], tuple[bool, str | None, str | None]] | None
|
||||
) = None,
|
||||
page_size: int = 50,
|
||||
) -> tuple[list[ProviderAPIKey], PoolSchedulingTrace]:
|
||||
"""Select and order pool keys with trace output.
|
||||
|
||||
Reuses :meth:`reorder_candidates` logic by adapting keys to lightweight
|
||||
candidate-like wrappers, then propagates skip/trace metadata back onto
|
||||
each key object for downstream execution/recording.
|
||||
|
||||
When *availability_checker* is provided, post-sort availability checks
|
||||
are performed lazily: only the top *page_size* non-skipped keys are
|
||||
checked at a time; if all fail, the next page is checked, and so on.
|
||||
Keys beyond the last checked page are marked as ``deferred`` (skipped
|
||||
without checking) to avoid unnecessary CPU work on large pools.
|
||||
"""
|
||||
if not keys:
|
||||
return (
|
||||
@@ -424,6 +435,30 @@ class PoolManager:
|
||||
)
|
||||
ordered_keys.append(key)
|
||||
|
||||
# -- 分页可用性检查 --
|
||||
# 排序后对非 skipped key 分页调用 availability_checker,
|
||||
# 找到 page_size 个可用 key 后停止检查,剩余标记 deferred。
|
||||
if availability_checker is not None:
|
||||
available_count = 0
|
||||
found_enough = False
|
||||
for key in ordered_keys:
|
||||
if getattr(key, "_pool_skipped", False):
|
||||
continue
|
||||
if found_enough:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", "deferred")
|
||||
continue
|
||||
is_available, skip_reason_check, mapping_model = availability_checker(key)
|
||||
if not is_available:
|
||||
setattr(key, "_pool_skipped", True)
|
||||
setattr(key, "_pool_skip_reason", skip_reason_check)
|
||||
else:
|
||||
if mapping_model:
|
||||
setattr(key, "_pool_mapping_matched_model", mapping_model)
|
||||
available_count += 1
|
||||
if available_count >= page_size:
|
||||
found_enough = True
|
||||
|
||||
return ordered_keys, trace
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -563,33 +563,9 @@ class CandidateBuilder:
|
||||
)
|
||||
|
||||
if pool_cfg is not None:
|
||||
# 号池 Provider 仅构建一个 PoolCandidate,内部 key 选择延迟到执行阶段。
|
||||
pool_keys: list[ProviderAPIKey] = []
|
||||
pool_miss_counts: list[int] = []
|
||||
pool_mapping: dict[str, str | None] = {}
|
||||
for key in keys_to_check:
|
||||
is_available, _key_skip_reason, mapping_matched_model = (
|
||||
self._check_key_availability(
|
||||
key,
|
||||
endpoint_format_str,
|
||||
model_name,
|
||||
capability_requirements,
|
||||
model_mappings=model_mappings,
|
||||
candidate_models=provider_model_names,
|
||||
provider_type=getattr(provider, "provider_type", None),
|
||||
)
|
||||
)
|
||||
if not is_available:
|
||||
continue
|
||||
|
||||
pool_keys.append(key)
|
||||
pool_miss_counts.append(
|
||||
compute_capability_score(
|
||||
key.capabilities or {},
|
||||
capability_requirements,
|
||||
)
|
||||
)
|
||||
pool_mapping[str(key.id)] = mapping_matched_model
|
||||
# 号池优化:跳过逐 key 的 _check_key_availability 检查,
|
||||
# 直接收集全部 active key,将检查推迟到 PoolManager 排序后分页执行。
|
||||
pool_keys = list(keys_to_check)
|
||||
|
||||
if not pool_keys:
|
||||
continue
|
||||
@@ -619,20 +595,21 @@ class CandidateBuilder:
|
||||
pool_keys=pool_keys,
|
||||
pool_config=pool_cfg,
|
||||
pool_priority=pool_priority,
|
||||
mapping_matched_model=pool_mapping.get(str(pool_keys[0].id)),
|
||||
needs_conversion=needs_conversion,
|
||||
provider_api_format=str(endpoint_format_str or ""),
|
||||
output_limit=output_limit,
|
||||
capability_miss_count=min(pool_miss_counts) if pool_miss_counts else 0,
|
||||
capability_miss_count=0,
|
||||
)
|
||||
|
||||
# 在 key 对象上附加映射结果,供 PoolCandidate 运行时切 key 后同步模型名。
|
||||
for pool_key in pool_keys:
|
||||
setattr(
|
||||
pool_key,
|
||||
"_pool_mapping_matched_model",
|
||||
pool_mapping.get(str(pool_key.id)),
|
||||
)
|
||||
# 打包延迟检查参数,供 PoolManager 排序后分页调用
|
||||
pool_candidate._deferred_check_params = {
|
||||
"endpoint_format": endpoint_format_str,
|
||||
"model_name": model_name,
|
||||
"capability_requirements": capability_requirements,
|
||||
"model_mappings": model_mappings,
|
||||
"candidate_models": provider_model_names,
|
||||
"provider_type": getattr(provider, "provider_type", None),
|
||||
}
|
||||
|
||||
if needs_conversion:
|
||||
convertible_candidates.append(pool_candidate)
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from src.models.database import (
|
||||
Provider,
|
||||
@@ -84,6 +84,8 @@ class PoolCandidate(ProviderCandidate):
|
||||
pool_config: PoolConfig | None = None
|
||||
pool_priority: int = 999999
|
||||
_pool_key_index: int = 0
|
||||
# 延迟可用性检查参数(号池优化:先排序再分页检查)
|
||||
_deferred_check_params: dict[str, Any] | None = field(default=None, repr=False)
|
||||
|
||||
|
||||
@dataclass
|
||||
|
||||
@@ -63,12 +63,25 @@ class TaskPoolOperationsService:
|
||||
if not candidate_keys and getattr(candidate, "key", None) is not None:
|
||||
candidate_keys = [candidate.key]
|
||||
|
||||
ordered_keys, trace = await manager.select_pool_keys(session_uuid, candidate_keys)
|
||||
candidate.pool_keys = ordered_keys
|
||||
# 构造延迟可用性检查回调(从 CandidateBuilder 打包的参数)
|
||||
checker = None
|
||||
deferred_params = candidate._deferred_check_params
|
||||
if deferred_params is not None:
|
||||
checker = self._build_availability_checker(deferred_params)
|
||||
|
||||
ordered_keys, trace = await manager.select_pool_keys(
|
||||
session_uuid,
|
||||
candidate_keys,
|
||||
availability_checker=checker,
|
||||
)
|
||||
# 移除 deferred key(未检查的),避免为其创建 DB 记录
|
||||
candidate.pool_keys = [
|
||||
k for k in ordered_keys if getattr(k, "_pool_skip_reason", None) != "deferred"
|
||||
]
|
||||
|
||||
selected_key_index = 0
|
||||
selected_key = None
|
||||
for idx, pool_key in enumerate(ordered_keys):
|
||||
for idx, pool_key in enumerate(candidate.pool_keys):
|
||||
if not bool(getattr(pool_key, "_pool_skipped", False)):
|
||||
selected_key = pool_key
|
||||
selected_key_index = idx
|
||||
@@ -220,3 +233,33 @@ class TaskPoolOperationsService:
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _build_availability_checker(
|
||||
params: dict[str, Any],
|
||||
) -> Any:
|
||||
"""Construct a key availability checker from deferred check params."""
|
||||
from src.services.scheduling.candidate_builder import CandidateBuilder
|
||||
|
||||
endpoint_format = params.get("endpoint_format")
|
||||
model_name = params.get("model_name", "")
|
||||
capability_requirements = params.get("capability_requirements")
|
||||
model_mappings = params.get("model_mappings")
|
||||
candidate_models = params.get("candidate_models")
|
||||
provider_type = params.get("provider_type")
|
||||
|
||||
# _check_key_availability 不依赖 _sorter,传 None 安全
|
||||
builder = CandidateBuilder(candidate_sorter=None) # type: ignore[arg-type]
|
||||
|
||||
def _checker(key: Any) -> tuple[bool, str | None, str | None]:
|
||||
return builder._check_key_availability(
|
||||
key,
|
||||
endpoint_format,
|
||||
model_name,
|
||||
capability_requirements,
|
||||
model_mappings=model_mappings,
|
||||
candidate_models=candidate_models,
|
||||
provider_type=provider_type,
|
||||
)
|
||||
|
||||
return _checker
|
||||
|
||||
Reference in New Issue
Block a user