mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(pool): 号池候选重构为 PoolCandidate 单候选模式与池内 key 故障转移
- 新增 PoolCandidate 子类，排序阶段作为单候选参与，执行阶段在 pool_keys 内部选择/切换 key - FailoverEngine 新增 _execute_pool_candidate 方法，支持池内 key 级别故障转移与重试 - 提取 _execute_attempt / _attach_attempt_context / _classify_attempt_error 公共方法 - CandidateBuilder 对号池 Provider 构建单个 PoolCandidate（包含所有可用 key） - CandidateSorter 支持 PoolCandidate 独立优先级分组（global_priority / pool_priority） - CandidateResolver PRE_EXPAND 模式按 pool_keys 展开预创建记录，附加 pool_group_id - TaskService._apply_pool_reorder 改为对 PoolCandidate 调用 select_pool_keys - PoolManager 新增 select_pool_keys 方法，复用 reorder_candidates 逻辑 - 新增 global_priority 号池配置字段（前后端同步） - 前端 Timeline 支持按 pool_group_id 分组显示多号池尝试 - 异步提交路径新增 _expand_pool_candidates_for_async_submit 展开逻辑
This commit is contained in:
fawney19
@@ -130,6 +130,12 @@ class FailoverRulesConfig(BaseModel):
|
||||
class PoolAdvancedConfig(BaseModel):
|
||||
"""通用号池配置（适用于所有 Provider 类型）。"""
|
||||
|
||||
global_priority: int | None = Field(
|
||||
None,
|
||||
ge=0,
|
||||
le=999999,
|
||||
description="global_key 模式下号池整体优先级（数字越小越优先）",
|
||||
)
|
||||
sticky_session_ttl_seconds: int | None = Field(
|
||||
None,
|
||||
ge=60,
|
||||
|
||||
@@ -15,7 +15,7 @@ from src.core.logger import logger
|
||||
from src.models.database import RequestCandidate
|
||||
from src.services.orchestration.error_classifier import ErrorAction, ErrorClassifier
|
||||
from src.services.request.candidate import RequestCandidateService
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||||
from src.services.task.exceptions import StreamProbeError
|
||||
from src.services.task.protocol import AttemptFunc, AttemptKind, AttemptResult
|
||||
from src.services.task.schema import ExecutionResult
|
||||
@@ -141,6 +141,26 @@ class FailoverEngine:
|
||||
)
|
||||
continue
|
||||
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
pool_result, attempt_count, last_status_code = await self._execute_pool_candidate(
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
attempt_func=attempt_func,
|
||||
retry_policy=retry_policy,
|
||||
request_id=request_id,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidate_keys_fallback=candidate_keys_fallback,
|
||||
candidates=candidates,
|
||||
attempt_count=attempt_count,
|
||||
max_attempts=max_attempts,
|
||||
execution_error_handler=execution_error_handler,
|
||||
)
|
||||
if pool_result is not None:
|
||||
return pool_result
|
||||
continue
|
||||
|
||||
max_retries = self._get_max_retries(candidate, retry_policy)
|
||||
retry_index = 0
|
||||
while retry_index < max_retries:
|
||||
@@ -163,16 +183,9 @@ class FailoverEngine:
|
||||
api_key_id=api_key_id,
|
||||
)
|
||||
|
||||
# Attach per-attempt context onto candidate for attempt_func (keeps AttemptFunc signature stable).
|
||||
try:
|
||||
setattr(candidate, "_utf_candidate_index", candidate_index)
|
||||
setattr(candidate, "_utf_retry_index", retry_index)
|
||||
setattr(candidate, "_utf_candidate_record_id", record_id)
|
||||
setattr(candidate, "_utf_attempt_count", attempt_count)
|
||||
setattr(candidate, "_utf_max_attempts", max_attempts)
|
||||
except Exception:
|
||||
# Best-effort only; attempt_func may not rely on these attributes.
|
||||
pass
|
||||
self._attach_attempt_context(
|
||||
candidate, candidate_index, retry_index, record_id, attempt_count, max_attempts
|
||||
)
|
||||
|
||||
# Mark pending
|
||||
now = datetime.now(timezone.utc)
|
||||
@@ -187,43 +200,13 @@ class FailoverEngine:
|
||||
self._commit_before_await()
|
||||
|
||||
try:
|
||||
attempt_result = await attempt_func(candidate)
|
||||
attempt_result = await self._execute_attempt(
|
||||
candidate=candidate,
|
||||
record_id=record_id,
|
||||
attempt_func=attempt_func,
|
||||
)
|
||||
last_status_code = int(getattr(attempt_result, "http_status", 0) or 0)
|
||||
|
||||
# Stream: probe first chunk, failover only before first chunk
|
||||
if attempt_result.kind == AttemptKind.STREAM:
|
||||
attempt_result = await self._probe_stream_first_chunk(
|
||||
attempt_result=attempt_result,
|
||||
record_id=record_id,
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
# Sync: check success_failover_patterns on response body
|
||||
if attempt_result.kind == AttemptKind.SYNC_RESPONSE:
|
||||
body = getattr(attempt_result, "response_body", None)
|
||||
if body:
|
||||
if isinstance(body, bytes):
|
||||
body_text = body.decode("utf-8", errors="replace")
|
||||
elif isinstance(body, (dict, list)):
|
||||
body_text = json.dumps(body, ensure_ascii=False)
|
||||
else:
|
||||
body_text = str(body)
|
||||
rule_action = self._check_provider_failover_rules(
|
||||
candidate, is_success=True, response_text=body_text
|
||||
)
|
||||
if rule_action == FailoverAction.CONTINUE:
|
||||
self._record_attempt_failure(
|
||||
record_id,
|
||||
Exception("success_failover_pattern matched"),
|
||||
200,
|
||||
)
|
||||
raise StreamProbeError(
|
||||
"Success failover pattern matched",
|
||||
http_status=200,
|
||||
)
|
||||
|
||||
self._record_attempt_success(record_id, attempt_result)
|
||||
|
||||
# PRE_EXPAND: mark unused slots after request ends (success)
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_remaining_slots_unused(
|
||||
@@ -316,6 +299,307 @@ class FailoverEngine:
|
||||
attempt_count=attempt_count,
|
||||
)
|
||||
|
||||
async def _execute_pool_candidate(
|
||||
self,
|
||||
*,
|
||||
candidate: PoolCandidate,
|
||||
candidate_index: int,
|
||||
attempt_func: AttemptFunc,
|
||||
retry_policy: RetryPolicy,
|
||||
request_id: str | None,
|
||||
user_id: str | None,
|
||||
api_key_id: str | None,
|
||||
candidate_record_map: dict[tuple[int, int], str] | None,
|
||||
candidate_keys_fallback: list[CandidateKey],
|
||||
candidates: list[ProviderCandidate],
|
||||
attempt_count: int,
|
||||
max_attempts: int | None,
|
||||
execution_error_handler: Any,
|
||||
) -> tuple[ExecutionResult | None, int, int | None]:
|
||||
"""Execute a PoolCandidate with in-pool key failover."""
|
||||
last_status_code: int | None = None
|
||||
retry_slots_per_key = self._get_pool_key_max_retries(candidate, retry_policy)
|
||||
|
||||
for key_index, pool_key in enumerate(candidate.pool_keys or []):
|
||||
base_retry_index = key_index * retry_slots_per_key
|
||||
candidate.key = pool_key
|
||||
candidate._pool_key_index = key_index
|
||||
candidate.mapping_matched_model = getattr(pool_key, "_pool_mapping_matched_model", None)
|
||||
|
||||
if bool(getattr(pool_key, "_pool_skipped", False)):
|
||||
skip_reason = str(
|
||||
getattr(pool_key, "_pool_skip_reason", None)
|
||||
or getattr(candidate, "skip_reason", None)
|
||||
or "pool_skipped"
|
||||
)
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_retry_indices_status(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidate_idx=candidate_index,
|
||||
retry_indices=range(
|
||||
base_retry_index, base_retry_index + retry_slots_per_key
|
||||
),
|
||||
status="skipped",
|
||||
skip_reason=skip_reason,
|
||||
)
|
||||
elif request_id:
|
||||
await self._create_skipped_record(
|
||||
request_id=request_id,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=base_retry_index,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
skip_reason=skip_reason,
|
||||
)
|
||||
candidate_keys_fallback.append(
|
||||
self._make_candidate_key(
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=base_retry_index,
|
||||
status="skipped",
|
||||
skip_reason=skip_reason,
|
||||
)
|
||||
)
|
||||
continue
|
||||
|
||||
max_retries_for_key = retry_slots_per_key
|
||||
retry_index = 0
|
||||
while retry_index < max_retries_for_key:
|
||||
attempt_count += 1
|
||||
composite_retry_index = base_retry_index + retry_index
|
||||
|
||||
record_id = None
|
||||
if candidate_record_map:
|
||||
record_id = candidate_record_map.get((candidate_index, composite_retry_index))
|
||||
if record_id is None:
|
||||
# Rectify may extend retries beyond pre-created range.
|
||||
record_id = candidate_record_map.get((candidate_index, base_retry_index))
|
||||
if record_id is None and request_id and retry_policy.mode != RetryMode.PRE_EXPAND:
|
||||
record_id = await self._ensure_record_exists(
|
||||
request_id=request_id,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=composite_retry_index,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
)
|
||||
|
||||
self._attach_attempt_context(
|
||||
candidate,
|
||||
candidate_index,
|
||||
composite_retry_index,
|
||||
record_id,
|
||||
attempt_count,
|
||||
max_attempts,
|
||||
)
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
if record_id:
|
||||
self._update_record(
|
||||
record_id,
|
||||
status="pending",
|
||||
started_at=now,
|
||||
)
|
||||
|
||||
self._commit_before_await()
|
||||
|
||||
try:
|
||||
attempt_result = await self._execute_attempt(
|
||||
candidate=candidate,
|
||||
record_id=record_id,
|
||||
attempt_func=attempt_func,
|
||||
)
|
||||
last_status_code = int(getattr(attempt_result, "http_status", 0) or 0)
|
||||
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_remaining_slots_unused(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidates=candidates,
|
||||
success_candidate_idx=candidate_index,
|
||||
success_retry_idx=composite_retry_index,
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
|
||||
return (
|
||||
ExecutionResult(
|
||||
success=True,
|
||||
attempt_result=attempt_result,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=composite_retry_index,
|
||||
provider_id=str(candidate.provider.id),
|
||||
provider_name=str(candidate.provider.name),
|
||||
endpoint_id=str(candidate.endpoint.id),
|
||||
key_id=str(candidate.key.id),
|
||||
candidate_keys=self._get_candidate_keys(
|
||||
request_id=request_id,
|
||||
fallback=candidate_keys_fallback,
|
||||
candidates=candidates,
|
||||
),
|
||||
attempt_count=attempt_count,
|
||||
request_candidate_id=record_id,
|
||||
),
|
||||
attempt_count,
|
||||
last_status_code,
|
||||
)
|
||||
|
||||
except StreamProbeError as exc:
|
||||
last_status_code = exc.http_status
|
||||
self._record_attempt_failure(record_id, exc, exc.http_status)
|
||||
action = FailoverAction.CONTINUE
|
||||
|
||||
except Exception as exc:
|
||||
outcome = await self._handle_pool_attempt_error(
|
||||
exc,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
key_retry_index=retry_index,
|
||||
composite_retry_index=composite_retry_index,
|
||||
max_retries=max_retries_for_key,
|
||||
record_id=record_id,
|
||||
attempt_count=attempt_count,
|
||||
max_attempts=max_attempts,
|
||||
execution_error_handler=execution_error_handler,
|
||||
)
|
||||
action = outcome.action
|
||||
last_status_code = outcome.last_status_code
|
||||
max_retries_for_key = min(outcome.max_retries, retry_slots_per_key)
|
||||
|
||||
if action == FailoverAction.CONTINUE:
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
# max_retries_for_key may have been shrunk by error handler;
|
||||
# mark unused up to the *original* retry_slots_per_key to cover
|
||||
# all pre-created records.
|
||||
self._mark_retry_indices_status(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidate_idx=candidate_index,
|
||||
retry_indices=range(
|
||||
composite_retry_index + 1,
|
||||
base_retry_index + retry_slots_per_key,
|
||||
),
|
||||
status="unused",
|
||||
)
|
||||
break
|
||||
if action == FailoverAction.RETRY:
|
||||
retry_index += 1
|
||||
continue
|
||||
|
||||
# STOP: only stop this pool candidate; outer candidate traversal continues.
|
||||
# Rationale: pool-internal STOP (from error_stop_patterns on a non-ExecutionError)
|
||||
# should not terminate the entire request because other providers may still succeed.
|
||||
# When handler_used=True, TaskService raises directly for true STOP semantics.
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_candidate_remaining_retries_unused(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidate_idx=candidate_index,
|
||||
from_retry_idx=composite_retry_index + 1,
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
return None, attempt_count, last_status_code
|
||||
|
||||
return None, attempt_count, last_status_code
|
||||
|
||||
async def _handle_pool_attempt_error(
|
||||
self,
|
||||
exc: Exception,
|
||||
*,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
key_retry_index: int,
|
||||
composite_retry_index: int,
|
||||
max_retries: int,
|
||||
record_id: str | None,
|
||||
attempt_count: int,
|
||||
max_attempts: int | None,
|
||||
execution_error_handler: Any,
|
||||
) -> AttemptErrorOutcome:
|
||||
"""Handle pool attempt errors without forcing outer STOP semantics.
|
||||
|
||||
Args:
|
||||
key_retry_index: key 内部的重试索引 (用于判断 has_retry_left)
|
||||
composite_retry_index: 全局维度的重试索引 (传给 execution_error_handler,
|
||||
与 candidate_record_map 对齐)
|
||||
"""
|
||||
return await self._classify_attempt_error(
|
||||
exc,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=composite_retry_index,
|
||||
has_retry_left=key_retry_index + 1 < max_retries,
|
||||
max_retries=max_retries,
|
||||
record_id=record_id,
|
||||
attempt_count=attempt_count,
|
||||
max_attempts=max_attempts,
|
||||
execution_error_handler=execution_error_handler,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _attach_attempt_context(
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
retry_index: int,
|
||||
record_id: str | None,
|
||||
attempt_count: int,
|
||||
max_attempts: int | None,
|
||||
) -> None:
|
||||
"""Attach per-attempt context onto candidate for attempt_func (best-effort)."""
|
||||
try:
|
||||
setattr(candidate, "_utf_candidate_index", candidate_index)
|
||||
setattr(candidate, "_utf_retry_index", retry_index)
|
||||
setattr(candidate, "_utf_candidate_record_id", record_id)
|
||||
setattr(candidate, "_utf_attempt_count", attempt_count)
|
||||
setattr(candidate, "_utf_max_attempts", max_attempts)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _execute_attempt(
|
||||
self,
|
||||
*,
|
||||
candidate: ProviderCandidate,
|
||||
record_id: str | None,
|
||||
attempt_func: AttemptFunc,
|
||||
) -> AttemptResult:
|
||||
"""Run attempt_func with stream probe and sync failover-pattern checks.
|
||||
|
||||
On success records the attempt; raises StreamProbeError on failover-pattern
|
||||
match or stream probe failure so the caller can handle retries uniformly.
|
||||
"""
|
||||
attempt_result = await attempt_func(candidate)
|
||||
|
||||
if attempt_result.kind == AttemptKind.STREAM:
|
||||
attempt_result = await self._probe_stream_first_chunk(
|
||||
attempt_result=attempt_result,
|
||||
record_id=record_id,
|
||||
candidate=candidate,
|
||||
)
|
||||
|
||||
if attempt_result.kind == AttemptKind.SYNC_RESPONSE:
|
||||
body = getattr(attempt_result, "response_body", None)
|
||||
if body:
|
||||
if isinstance(body, bytes):
|
||||
body_text = body.decode("utf-8", errors="replace")
|
||||
elif isinstance(body, (dict, list)):
|
||||
body_text = json.dumps(body, ensure_ascii=False)
|
||||
else:
|
||||
body_text = str(body)
|
||||
rule_action = self._check_provider_failover_rules(
|
||||
candidate, is_success=True, response_text=body_text
|
||||
)
|
||||
if rule_action == FailoverAction.CONTINUE:
|
||||
self._record_attempt_failure(
|
||||
record_id,
|
||||
Exception("success_failover_pattern matched"),
|
||||
200,
|
||||
)
|
||||
raise StreamProbeError(
|
||||
"Success failover pattern matched",
|
||||
http_status=200,
|
||||
)
|
||||
|
||||
self._record_attempt_success(record_id, attempt_result)
|
||||
return attempt_result
|
||||
|
||||
def _record_attempt_success(self, record_id: str | None, attempt_result: AttemptResult) -> None:
|
||||
"""Mark attempt record as success/streaming."""
|
||||
if not record_id:
|
||||
@@ -375,9 +659,62 @@ class FailoverEngine:
|
||||
Returns:
|
||||
AttemptErrorOutcome; stop_result is non-None only when action==STOP.
|
||||
"""
|
||||
has_retry_left = retry_index + 1 < max_retries
|
||||
outcome = await self._classify_attempt_error(
|
||||
exc,
|
||||
candidate=candidate,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=retry_index,
|
||||
has_retry_left=retry_index + 1 < max_retries,
|
||||
max_retries=max_retries,
|
||||
record_id=record_id,
|
||||
attempt_count=attempt_count,
|
||||
max_attempts=max_attempts,
|
||||
execution_error_handler=execution_error_handler,
|
||||
)
|
||||
|
||||
# If caller provides an execution_error_handler, prefer it for ExecutionError.
|
||||
if outcome.action == FailoverAction.STOP:
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_remaining_slots_unused(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidates=candidates,
|
||||
success_candidate_idx=candidate_index,
|
||||
success_retry_idx=retry_index,
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
outcome.stop_result = ExecutionResult(
|
||||
success=False,
|
||||
error_type=type(exc).__name__,
|
||||
error_message=self._sanitize(str(exc)),
|
||||
last_status_code=outcome.last_status_code or None,
|
||||
candidate_keys=self._get_candidate_keys(
|
||||
request_id=request_id,
|
||||
fallback=candidate_keys_fallback,
|
||||
candidates=candidates,
|
||||
),
|
||||
attempt_count=attempt_count,
|
||||
)
|
||||
|
||||
return outcome
|
||||
|
||||
async def _classify_attempt_error(
|
||||
self,
|
||||
exc: Exception,
|
||||
*,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
retry_index: int,
|
||||
has_retry_left: bool,
|
||||
max_retries: int,
|
||||
record_id: str | None,
|
||||
attempt_count: int,
|
||||
max_attempts: int | None,
|
||||
execution_error_handler: Any,
|
||||
) -> AttemptErrorOutcome:
|
||||
"""Classify an attempt error: delegate to external handler or internal classifier.
|
||||
|
||||
Returns a base AttemptErrorOutcome (without stop_result). Callers add
|
||||
STOP-specific logic (e.g. PRE_EXPAND cleanup, stop_result construction) as needed.
|
||||
"""
|
||||
handler_used = False
|
||||
action = FailoverAction.CONTINUE
|
||||
if execution_error_handler is not None:
|
||||
@@ -408,40 +745,11 @@ class FailoverEngine:
|
||||
candidate=candidate,
|
||||
has_retry_left=has_retry_left,
|
||||
)
|
||||
|
||||
last_status_code = int(getattr(exc, "status_code", 0) or 0) or int(
|
||||
getattr(exc, "http_status", 0) or 0
|
||||
)
|
||||
|
||||
self._record_attempt_failure(record_id, exc, last_status_code or None)
|
||||
|
||||
if action == FailoverAction.STOP:
|
||||
if retry_policy.mode == RetryMode.PRE_EXPAND and candidate_record_map:
|
||||
self._mark_remaining_slots_unused(
|
||||
candidate_record_map=candidate_record_map,
|
||||
candidates=candidates,
|
||||
success_candidate_idx=candidate_index,
|
||||
success_retry_idx=retry_index,
|
||||
retry_policy=retry_policy,
|
||||
)
|
||||
return AttemptErrorOutcome(
|
||||
action=action,
|
||||
last_status_code=last_status_code,
|
||||
max_retries=max_retries,
|
||||
stop_result=ExecutionResult(
|
||||
success=False,
|
||||
error_type=type(exc).__name__,
|
||||
error_message=self._sanitize(str(exc)),
|
||||
last_status_code=last_status_code or None,
|
||||
candidate_keys=self._get_candidate_keys(
|
||||
request_id=request_id,
|
||||
fallback=candidate_keys_fallback,
|
||||
candidates=candidates,
|
||||
),
|
||||
attempt_count=attempt_count,
|
||||
),
|
||||
)
|
||||
|
||||
return AttemptErrorOutcome(
|
||||
action=action,
|
||||
last_status_code=last_status_code,
|
||||
@@ -538,10 +846,7 @@ class FailoverEngine:
|
||||
api_key_id: str | None,
|
||||
) -> str:
|
||||
# Create "available" record, then caller will mark pending.
|
||||
extra: dict = {}
|
||||
pool_extra = getattr(candidate, "_pool_extra_data", None)
|
||||
if pool_extra:
|
||||
extra.update(pool_extra)
|
||||
extra = self._build_pool_extra_data(candidate)
|
||||
row = RequestCandidateService.create_candidate(
|
||||
db=self.db,
|
||||
request_id=request_id,
|
||||
@@ -564,19 +869,17 @@ class FailoverEngine:
|
||||
request_id: str,
|
||||
candidate: ProviderCandidate,
|
||||
candidate_index: int,
|
||||
retry_index: int = 0,
|
||||
user_id: str | None,
|
||||
api_key_id: str | None,
|
||||
skip_reason: str | None,
|
||||
) -> str:
|
||||
extra: dict = {}
|
||||
pool_extra = getattr(candidate, "_pool_extra_data", None)
|
||||
if pool_extra:
|
||||
extra.update(pool_extra)
|
||||
extra = self._build_pool_extra_data(candidate)
|
||||
row = RequestCandidateService.create_candidate(
|
||||
db=self.db,
|
||||
request_id=request_id,
|
||||
candidate_index=candidate_index,
|
||||
retry_index=0,
|
||||
retry_index=retry_index,
|
||||
user_id=user_id,
|
||||
api_key_id=api_key_id,
|
||||
provider_id=str(candidate.provider.id),
|
||||
@@ -592,6 +895,20 @@ class FailoverEngine:
|
||||
self.db.commit()
|
||||
return str(row.id)
|
||||
|
||||
def _build_pool_extra_data(self, candidate: ProviderCandidate) -> dict[str, Any]:
|
||||
extra: dict[str, Any] = {}
|
||||
pool_extra = getattr(candidate, "_pool_extra_data", None)
|
||||
if isinstance(pool_extra, dict):
|
||||
extra.update(pool_extra)
|
||||
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
extra["pool_group_id"] = str(getattr(candidate.provider, "id", "") or "")
|
||||
extra["pool_key_index"] = int(getattr(candidate, "_pool_key_index", 0) or 0)
|
||||
key_extra = getattr(candidate.key, "_pool_extra_data", None)
|
||||
if isinstance(key_extra, dict):
|
||||
extra.update(key_extra)
|
||||
return extra
|
||||
|
||||
def _should_skip(
|
||||
self, candidate: ProviderCandidate, skip_policy: SkipPolicy
|
||||
) -> tuple[bool, str | None]:
|
||||
@@ -612,6 +929,15 @@ class FailoverEngine:
|
||||
return False, None
|
||||
|
||||
def _get_max_retries(self, candidate: ProviderCandidate, retry_policy: RetryPolicy) -> int:
|
||||
per_key_retries = self._get_pool_key_max_retries(candidate, retry_policy)
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
key_count = len(candidate.pool_keys or []) or 1
|
||||
return max(1, key_count * per_key_retries)
|
||||
return per_key_retries
|
||||
|
||||
def _get_pool_key_max_retries(
|
||||
self, candidate: ProviderCandidate, retry_policy: RetryPolicy
|
||||
) -> int:
|
||||
if retry_policy.mode == RetryMode.DISABLED:
|
||||
return 1
|
||||
if retry_policy.retry_on_cached_only and not bool(getattr(candidate, "is_cached", False)):
|
||||
@@ -941,6 +1267,26 @@ class FailoverEngine:
|
||||
self._update_record(record_id, status="unused", finished_at=now)
|
||||
self.db.commit()
|
||||
|
||||
def _mark_retry_indices_status(
|
||||
self,
|
||||
*,
|
||||
candidate_record_map: dict[tuple[int, int], str],
|
||||
candidate_idx: int,
|
||||
retry_indices: range,
|
||||
status: str,
|
||||
skip_reason: str | None = None,
|
||||
) -> None:
|
||||
now = datetime.now(timezone.utc)
|
||||
for retry_idx in retry_indices:
|
||||
record_id = candidate_record_map.get((candidate_idx, retry_idx))
|
||||
if not record_id:
|
||||
continue
|
||||
values: dict[str, Any] = {"status": status, "finished_at": now}
|
||||
if status == "skipped":
|
||||
values["skip_reason"] = skip_reason
|
||||
self._update_record(record_id, **values)
|
||||
self.db.commit()
|
||||
|
||||
def _mark_all_remaining_available_unused(
|
||||
self, candidate_record_map: dict[tuple[int, int], str]
|
||||
) -> None:
|
||||
|
||||
@@ -14,7 +14,8 @@ from src.core.exceptions import ProviderNotAvailableException
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey
|
||||
from src.services.provider.format import normalize_endpoint_signature
|
||||
from src.services.scheduling.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
||||
from src.services.scheduling.aware_scheduler import CacheAwareScheduler
|
||||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||||
|
||||
|
||||
class CandidateResolver:
|
||||
@@ -153,14 +154,27 @@ class CandidateResolver:
|
||||
if preferred_key_ids:
|
||||
preferred_set = {str(kid) for kid in preferred_key_ids if kid}
|
||||
if preferred_set:
|
||||
preferred_candidates = [
|
||||
c for c in all_candidates if c.key and str(c.key.id) in preferred_set
|
||||
]
|
||||
other_candidates = [
|
||||
c for c in all_candidates if not (c.key and str(c.key.id) in preferred_set)
|
||||
]
|
||||
|
||||
def _is_preferred_candidate(c: ProviderCandidate) -> bool:
|
||||
if c.key and str(c.key.id) in preferred_set:
|
||||
return True
|
||||
if isinstance(c, PoolCandidate):
|
||||
return any(str(pk.id) in preferred_set for pk in (c.pool_keys or []))
|
||||
return False
|
||||
|
||||
preferred_candidates = [c for c in all_candidates if _is_preferred_candidate(c)]
|
||||
other_candidates = [c for c in all_candidates if not _is_preferred_candidate(c)]
|
||||
if preferred_candidates:
|
||||
matched_key_ids = [str(c.key.id) for c in preferred_candidates if c.key]
|
||||
matched_key_ids: list[str] = []
|
||||
for candidate in preferred_candidates:
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
matched_key_ids.extend(
|
||||
str(pk.id)
|
||||
for pk in (candidate.pool_keys or [])
|
||||
if str(pk.id) in preferred_set
|
||||
)
|
||||
elif candidate.key:
|
||||
matched_key_ids.append(str(candidate.key.id))
|
||||
logger.debug(
|
||||
f" [{request_id}] 优先候选命中: {len(preferred_candidates)} 个 "
|
||||
f"(key_ids={matched_key_ids[:3]}{'...' if len(matched_key_ids) > 3 else ''})"
|
||||
@@ -210,6 +224,11 @@ class CandidateResolver:
|
||||
if not active_capabilities:
|
||||
active_capabilities = None
|
||||
|
||||
def _retry_slots_for_candidate(candidate: ProviderCandidate) -> int:
|
||||
if not expand_retries:
|
||||
return 1
|
||||
return int(candidate.provider.max_retries or 2) if candidate.is_cached else 1
|
||||
|
||||
for candidate_index, candidate in enumerate(all_candidates):
|
||||
provider = candidate.provider
|
||||
endpoint = candidate.endpoint
|
||||
@@ -219,6 +238,71 @@ class CandidateResolver:
|
||||
if isinstance(getattr(candidate, "_pool_extra_data", None), dict)
|
||||
else {}
|
||||
)
|
||||
base_extra = {
|
||||
"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,
|
||||
}
|
||||
|
||||
if isinstance(candidate, PoolCandidate) and candidate.pool_keys:
|
||||
retry_slots = _retry_slots_for_candidate(candidate)
|
||||
for key_idx, pool_key in enumerate(candidate.pool_keys):
|
||||
key_id = str(pool_key.id)
|
||||
key_pool_extra = (
|
||||
getattr(pool_key, "_pool_extra_data", None)
|
||||
if isinstance(getattr(pool_key, "_pool_extra_data", None), dict)
|
||||
else {}
|
||||
)
|
||||
key_skipped = candidate.is_skipped or bool(
|
||||
getattr(pool_key, "_pool_skipped", False)
|
||||
)
|
||||
key_skip_reason_raw = (
|
||||
getattr(pool_key, "_pool_skip_reason", None) if key_skipped else None
|
||||
)
|
||||
key_skip_reason = (
|
||||
str(key_skip_reason_raw)
|
||||
if key_skip_reason_raw
|
||||
else (candidate.skip_reason if key_skipped else None)
|
||||
)
|
||||
mapping_model = getattr(pool_key, "_pool_mapping_matched_model", None)
|
||||
extra_data = {
|
||||
**base_extra,
|
||||
"mapping_matched_model": (
|
||||
mapping_model
|
||||
if mapping_model
|
||||
else base_extra.get("mapping_matched_model")
|
||||
),
|
||||
"pool_group_id": str(provider.id),
|
||||
"pool_key_index": key_idx,
|
||||
**key_pool_extra,
|
||||
}
|
||||
|
||||
for retry_in_key in range(retry_slots):
|
||||
retry_index = key_idx * retry_slots + retry_in_key
|
||||
status = "skipped" if key_skipped else "available"
|
||||
record_id = str(uuid.uuid4())
|
||||
candidate_records_to_insert.append(
|
||||
{
|
||||
"id": record_id,
|
||||
"request_id": request_id,
|
||||
"candidate_index": candidate_index,
|
||||
"retry_index": retry_index,
|
||||
"user_id": user_id,
|
||||
"api_key_id": user_api_key.id if user_api_key else None,
|
||||
"provider_id": provider.id,
|
||||
"endpoint_id": endpoint.id,
|
||||
"key_id": key_id,
|
||||
"status": status,
|
||||
"skip_reason": key_skip_reason if key_skipped else None,
|
||||
"is_cached": candidate.is_cached,
|
||||
"extra_data": extra_data,
|
||||
"required_capabilities": active_capabilities,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
}
|
||||
)
|
||||
candidate_record_map[(candidate_index, retry_index)] = record_id
|
||||
continue
|
||||
|
||||
if candidate.is_skipped:
|
||||
record_id = str(uuid.uuid4())
|
||||
@@ -236,25 +320,14 @@ class CandidateResolver:
|
||||
"status": "skipped",
|
||||
"skip_reason": candidate.skip_reason,
|
||||
"is_cached": candidate.is_cached,
|
||||
"extra_data": {
|
||||
"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,
|
||||
},
|
||||
"extra_data": base_extra,
|
||||
"required_capabilities": active_capabilities,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
}
|
||||
)
|
||||
candidate_record_map[(candidate_index, 0)] = record_id
|
||||
else:
|
||||
# max_retries 已从 Endpoint 迁移到 Provider（Endpoint 仍可能保留旧字段用于兼容）
|
||||
if not expand_retries:
|
||||
max_retries_for_candidate = 1
|
||||
else:
|
||||
max_retries_for_candidate = (
|
||||
int(provider.max_retries or 2) if candidate.is_cached else 1
|
||||
)
|
||||
max_retries_for_candidate = _retry_slots_for_candidate(candidate)
|
||||
|
||||
for retry_index in range(max_retries_for_candidate):
|
||||
record_id = str(uuid.uuid4())
|
||||
@@ -271,12 +344,7 @@ class CandidateResolver:
|
||||
"key_id": key.id,
|
||||
"status": "available",
|
||||
"is_cached": candidate.is_cached,
|
||||
"extra_data": {
|
||||
"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,
|
||||
},
|
||||
"extra_data": base_extra,
|
||||
"required_capabilities": active_capabilities,
|
||||
"created_at": datetime.now(timezone.utc),
|
||||
}
|
||||
@@ -326,7 +394,22 @@ class CandidateResolver:
|
||||
total = 0
|
||||
for candidate in all_candidates:
|
||||
if not candidate.is_skipped:
|
||||
provider = candidate.provider
|
||||
max_retries = int(provider.max_retries or 2) if candidate.is_cached else 1
|
||||
total += max_retries
|
||||
retries_per_slot = (
|
||||
int(candidate.provider.max_retries or 2) if candidate.is_cached else 1
|
||||
)
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
schedulable_keys = [
|
||||
k
|
||||
for k in (candidate.pool_keys or [])
|
||||
if not bool(getattr(k, "_pool_skipped", False))
|
||||
]
|
||||
if schedulable_keys:
|
||||
total += len(schedulable_keys) * retries_per_slot
|
||||
elif candidate.pool_keys:
|
||||
# 兜底：尚未附加 _pool_skipped 标记时，按 key 数估算。
|
||||
total += len(candidate.pool_keys) * retries_per_slot
|
||||
else:
|
||||
total += retries_per_slot
|
||||
else:
|
||||
total += retries_per_slot
|
||||
return total
|
||||
|
||||
@@ -26,6 +26,8 @@ class PoolConfig:
|
||||
|
||||
# -- Sticky Session -------------------------------------------------------
|
||||
sticky_session_ttl_seconds: int = 3600 # 1 hour
|
||||
# Key 优先模式下号池整体优先级（None 时回退 provider_priority）
|
||||
global_priority: int | None = None
|
||||
|
||||
# -- Load-Aware Selection -------------------------------------------------
|
||||
load_threshold_percent: int = 80
|
||||
@@ -122,6 +124,7 @@ def parse_pool_config(provider_config: Any) -> PoolConfig | None:
|
||||
|
||||
return PoolConfig(
|
||||
sticky_session_ttl_seconds=_int_or("sticky_session_ttl_seconds", 3600),
|
||||
global_priority=_opt_int("global_priority"),
|
||||
load_threshold_percent=_int_or("load_threshold_percent", 80),
|
||||
lru_enabled=_bool_or("lru_enabled", True),
|
||||
cost_window_seconds=_int_or("cost_window_seconds", 18000),
|
||||
|
||||
@@ -280,6 +280,66 @@ class PoolManager:
|
||||
|
||||
return result
|
||||
|
||||
async def select_pool_keys(
|
||||
self,
|
||||
session_uuid: str | None,
|
||||
keys: list[ProviderAPIKey],
|
||||
) -> 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.
|
||||
"""
|
||||
if not keys:
|
||||
return (
|
||||
[],
|
||||
PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=0,
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
),
|
||||
)
|
||||
|
||||
class _KeyCandidate:
|
||||
__slots__ = ("key", "is_skipped", "skip_reason")
|
||||
|
||||
def __init__(self, key: ProviderAPIKey) -> None:
|
||||
self.key = key
|
||||
self.is_skipped = False
|
||||
self.skip_reason: str | None = None
|
||||
|
||||
wrappers = [_KeyCandidate(k) for k in keys]
|
||||
reordered_wrappers = await self.reorder_candidates(session_uuid, wrappers) # type: ignore[arg-type]
|
||||
|
||||
trace: PoolSchedulingTrace | None = None
|
||||
if reordered_wrappers:
|
||||
maybe_trace = getattr(reordered_wrappers[0], "_pool_scheduling_trace", None)
|
||||
if isinstance(maybe_trace, PoolSchedulingTrace):
|
||||
trace = maybe_trace
|
||||
if trace is None:
|
||||
trace = PoolSchedulingTrace(
|
||||
provider_id=self.provider_id,
|
||||
total_keys=len(keys),
|
||||
session_uuid=session_uuid[:8] if session_uuid else None,
|
||||
)
|
||||
|
||||
ordered_keys: list[ProviderAPIKey] = []
|
||||
for order_idx, wrapped in enumerate(reordered_wrappers):
|
||||
key = wrapped.key
|
||||
is_skipped = bool(getattr(wrapped, "is_skipped", False))
|
||||
skip_reason = str(getattr(wrapped, "skip_reason", "") or "")
|
||||
setattr(key, "_pool_skipped", is_skipped)
|
||||
setattr(key, "_pool_skip_reason", skip_reason if skip_reason else None)
|
||||
setattr(key, "_pool_order_index", order_idx)
|
||||
pool_extra = getattr(wrapped, "_pool_extra_data", None)
|
||||
setattr(
|
||||
key, "_pool_extra_data", dict(pool_extra) if isinstance(pool_extra, dict) else {}
|
||||
)
|
||||
ordered_keys.append(key)
|
||||
|
||||
return ordered_keys, trace
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Single-key selection (used by CandidateBuilder for pooled providers)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -76,6 +76,7 @@ from src.services.scheduling.concurrency_checker import ConcurrencyChecker
|
||||
from src.services.scheduling.restriction_checker import get_effective_restrictions
|
||||
from src.services.scheduling.scheduling_config import SchedulingConfig
|
||||
from src.services.scheduling.schemas import ConcurrencySnapshot as ConcurrencySnapshot # re-export
|
||||
from src.services.scheduling.schemas import PoolCandidate as PoolCandidate # re-export
|
||||
from src.services.scheduling.schemas import ProviderCandidate as ProviderCandidate # re-export
|
||||
from src.services.scheduling.utils import affinity_hash as _affinity_hash # re-export compat
|
||||
from src.services.scheduling.utils import (
|
||||
@@ -652,11 +653,20 @@ class CacheAwareScheduler:
|
||||
endpoint = candidate.endpoint
|
||||
key = candidate.key
|
||||
|
||||
if (
|
||||
provider.id == affinity.provider_id
|
||||
is_pool_candidate = isinstance(candidate, PoolCandidate)
|
||||
pool_matched = (
|
||||
is_pool_candidate
|
||||
and provider.id == affinity.provider_id
|
||||
and endpoint.id == affinity.endpoint_id
|
||||
)
|
||||
key_matched = (
|
||||
(not is_pool_candidate)
|
||||
and provider.id == affinity.provider_id
|
||||
and endpoint.id == affinity.endpoint_id
|
||||
and key.id == affinity.key_id
|
||||
):
|
||||
)
|
||||
|
||||
if pool_matched or key_matched:
|
||||
candidate.is_cached = True
|
||||
matched_candidate = candidate
|
||||
matched = True
|
||||
|
||||
@@ -398,7 +398,7 @@ class CandidateBuilder:
|
||||
Returns:
|
||||
候选列表
|
||||
"""
|
||||
from src.services.scheduling.schemas import ProviderCandidate
|
||||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||||
|
||||
candidates: list[ProviderCandidate] = []
|
||||
client_format_str = normalize_endpoint_signature(client_format)
|
||||
@@ -438,7 +438,6 @@ class CandidateBuilder:
|
||||
] = {}
|
||||
exact_candidates: list[ProviderCandidate] = []
|
||||
convertible_candidates: list[ProviderCandidate] = []
|
||||
pool_has_usable = False
|
||||
pool_cfg = _get_pool_config(provider)
|
||||
|
||||
# 使用新架构字段 (api_family, endpoint_kind) 进行预过滤与排序：
|
||||
@@ -576,8 +575,6 @@ class CandidateBuilder:
|
||||
if not active_keys:
|
||||
continue
|
||||
|
||||
# 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:
|
||||
use_random = False
|
||||
@@ -592,6 +589,84 @@ class CandidateBuilder:
|
||||
active_keys, affinity_key, use_random
|
||||
)
|
||||
|
||||
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
|
||||
|
||||
if not pool_keys:
|
||||
continue
|
||||
|
||||
provider_priority_raw = getattr(provider, "provider_priority", None)
|
||||
try:
|
||||
provider_priority = (
|
||||
int(provider_priority_raw)
|
||||
if provider_priority_raw is not None
|
||||
else 999999
|
||||
)
|
||||
except Exception:
|
||||
provider_priority = 999999
|
||||
try:
|
||||
pool_priority = (
|
||||
int(pool_cfg.global_priority)
|
||||
if pool_cfg.global_priority is not None
|
||||
else provider_priority
|
||||
)
|
||||
except Exception:
|
||||
pool_priority = provider_priority
|
||||
|
||||
pool_candidate = PoolCandidate(
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
key=pool_keys[0],
|
||||
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,
|
||||
)
|
||||
|
||||
# 在 key 对象上附加映射结果，供 PoolCandidate 运行时切 key 后同步模型名。
|
||||
for pool_key in pool_keys:
|
||||
setattr(
|
||||
pool_key,
|
||||
"_pool_mapping_matched_model",
|
||||
pool_mapping.get(str(pool_key.id)),
|
||||
)
|
||||
|
||||
if needs_conversion:
|
||||
convertible_candidates.append(pool_candidate)
|
||||
else:
|
||||
exact_candidates.append(pool_candidate)
|
||||
break
|
||||
|
||||
for key in keys_to_check:
|
||||
# Key 级别检查（健康度/熔断按 provider_format bucket）
|
||||
# 传入 provider_model_names 作为 candidate_models，
|
||||
@@ -634,13 +709,6 @@ class CandidateBuilder:
|
||||
else:
|
||||
exact_candidates.append(candidate)
|
||||
|
||||
if is_available:
|
||||
pool_has_usable = True
|
||||
|
||||
# Pool mode: stop after the first endpoint that produced a usable candidate.
|
||||
if pool_cfg is not None and pool_has_usable:
|
||||
break
|
||||
|
||||
candidates.extend(exact_candidates)
|
||||
candidates.extend(convertible_candidates)
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ from collections import defaultdict
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.services.scheduling.scheduling_config import SchedulingConfig
|
||||
from src.services.scheduling.schemas import PoolCandidate
|
||||
from src.services.scheduling.utils import affinity_hash
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
@@ -149,6 +150,8 @@ class CandidateSorter:
|
||||
|
||||
def get_priority(candidate: ProviderCandidate) -> int:
|
||||
"""获取候选的优先级"""
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
return int(getattr(candidate, "pool_priority", 999999) or 999999)
|
||||
if not candidate.key:
|
||||
return 999999
|
||||
priority_by_format = candidate.key.global_priority_by_format or {}
|
||||
@@ -170,8 +173,11 @@ class CandidateSorter:
|
||||
# 同优先级内哈希分散负载均衡
|
||||
scored_candidates = []
|
||||
for candidate in group:
|
||||
key_id = candidate.key.id if candidate.key else ""
|
||||
hash_value = affinity_hash(affinity_key, key_id)
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
hash_id = str(getattr(candidate.provider, "id", "") or "")
|
||||
else:
|
||||
hash_id = candidate.key.id if candidate.key else ""
|
||||
hash_value = affinity_hash(affinity_key, hash_id)
|
||||
scored_candidates.append((hash_value, candidate))
|
||||
|
||||
# 按哈希值排序
|
||||
@@ -181,11 +187,16 @@ class CandidateSorter:
|
||||
# 单个候选或没有 affinity_key，按次要排序条件排序
|
||||
def secondary_sort(c: ProviderCandidate) -> tuple[int, int, str]:
|
||||
pp = c.provider.provider_priority
|
||||
ip = c.key.internal_priority if c.key else None
|
||||
if isinstance(c, PoolCandidate):
|
||||
ip = int(getattr(c, "pool_priority", 999999) or 999999)
|
||||
key_id = str(getattr(c.provider, "id", "") or "")
|
||||
else:
|
||||
ip = c.key.internal_priority if c.key else None
|
||||
key_id = c.key.id if c.key else ""
|
||||
return (
|
||||
pp if pp is not None else 999999,
|
||||
ip if ip is not None else 999999,
|
||||
c.key.id if c.key else "",
|
||||
key_id,
|
||||
)
|
||||
|
||||
result.extend(sorted(group, key=secondary_sort))
|
||||
@@ -222,17 +233,27 @@ class CandidateSorter:
|
||||
if self._config.priority_mode == SchedulingConfig.PRIORITY_MODE_GLOBAL_KEY:
|
||||
# 全局 Key 优先模式：按格式特定优先级分组
|
||||
for candidate in candidates:
|
||||
priority = 999999
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
priority = int(getattr(candidate, "pool_priority", 999999) or 999999)
|
||||
# -1 使号池候选独立成组，不与普通 key 候选 (0) 混组打乱
|
||||
priority_groups[(priority, -1)].append(candidate)
|
||||
continue
|
||||
else:
|
||||
priority = 999999
|
||||
if candidate.key:
|
||||
priority_by_format = candidate.key.global_priority_by_format or {}
|
||||
if api_format and api_format in priority_by_format:
|
||||
priority = priority_by_format[api_format]
|
||||
priority_groups[(priority,)].append(candidate)
|
||||
priority_groups[(priority, 0)].append(candidate)
|
||||
else:
|
||||
# 提供商优先模式：按 (provider_priority, internal_priority) 分组
|
||||
for candidate in candidates:
|
||||
pp = candidate.provider.provider_priority
|
||||
ip = candidate.key.internal_priority if candidate.key else None
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
# 号池候选独立成组，不与普通 key 候选混组打乱。
|
||||
ip = -1
|
||||
else:
|
||||
ip = candidate.key.internal_priority if candidate.key else None
|
||||
key = (
|
||||
pp if pp is not None else 999999,
|
||||
ip if ip is not None else 999999,
|
||||
|
||||
@@ -6,7 +6,8 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.models.database import (
|
||||
Provider,
|
||||
@@ -14,6 +15,9 @@ from src.models.database import (
|
||||
ProviderEndpoint,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.services.provider.pool.config import PoolConfig
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProviderCandidate:
|
||||
@@ -69,6 +73,19 @@ class ProviderCandidate:
|
||||
return self._stable_order_key() < other._stable_order_key()
|
||||
|
||||
|
||||
@dataclass
|
||||
class PoolCandidate(ProviderCandidate):
|
||||
"""号池候选。
|
||||
|
||||
排序阶段作为单个候选参与；执行阶段再在 pool_keys 内部选择/切换 key。
|
||||
"""
|
||||
|
||||
pool_keys: list[ProviderAPIKey] = field(default_factory=list)
|
||||
pool_config: PoolConfig | None = None
|
||||
pool_priority: int = 999999
|
||||
_pool_key_index: int = 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConcurrencySnapshot:
|
||||
key_current: int
|
||||
|
||||
@@ -206,61 +206,130 @@ class TaskService:
|
||||
candidates: list[Any],
|
||||
request_body: dict[str, Any] | None,
|
||||
) -> tuple[list[Any], list[Any]]:
|
||||
"""Apply Account Pool reordering when applicable.
|
||||
|
||||
Groups candidates by provider_id and applies pool reordering
|
||||
independently per provider, then reassembles in original group order.
|
||||
Non-pool providers are left in their original order.
|
||||
|
||||
Returns:
|
||||
Tuple of (reordered_candidates, pool_traces) where pool_traces
|
||||
is a list of :class:`PoolSchedulingTrace` objects (one per
|
||||
pooled provider group, may be empty).
|
||||
"""
|
||||
"""Apply pool key ordering for PoolCandidate objects."""
|
||||
if not candidates:
|
||||
return candidates, []
|
||||
|
||||
pool_traces: list[Any] = []
|
||||
|
||||
try:
|
||||
from collections import OrderedDict
|
||||
|
||||
from src.services.provider.pool.config import parse_pool_config
|
||||
from src.services.provider.pool.manager import PoolManager
|
||||
from src.services.scheduling.schemas import PoolCandidate
|
||||
|
||||
# Group candidates by provider_id while preserving order.
|
||||
groups: OrderedDict[str, list[Any]] = OrderedDict()
|
||||
for c in candidates:
|
||||
pid = str(getattr(c.provider, "id", "") or "")
|
||||
groups.setdefault(pid, []).append(c)
|
||||
|
||||
result: list[Any] = []
|
||||
for pid, group in groups.items():
|
||||
provider = group[0].provider
|
||||
pool_cfg = parse_pool_config(getattr(provider, "config", None))
|
||||
if pool_cfg is None or not pid:
|
||||
result.extend(group)
|
||||
for candidate in candidates:
|
||||
if not isinstance(candidate, PoolCandidate):
|
||||
continue
|
||||
|
||||
provider = candidate.provider
|
||||
provider_id = str(getattr(provider, "id", "") or "")
|
||||
if not provider_id:
|
||||
continue
|
||||
|
||||
pool_cfg = candidate.pool_config or parse_pool_config(
|
||||
getattr(provider, "config", None)
|
||||
)
|
||||
if pool_cfg is None:
|
||||
continue
|
||||
candidate.pool_config = pool_cfg
|
||||
|
||||
provider_type = str(getattr(provider, "provider_type", "") or "")
|
||||
session_uuid = TaskService._extract_session_uuid(provider_type, request_body)
|
||||
mgr = PoolManager(pid, pool_cfg)
|
||||
reordered = await mgr.reorder_candidates(session_uuid, group)
|
||||
result.extend(reordered)
|
||||
manager = PoolManager(provider_id, pool_cfg)
|
||||
|
||||
# Extract trace attached by PoolManager.reorder_candidates
|
||||
if reordered:
|
||||
trace = getattr(reordered[0], "_pool_scheduling_trace", None)
|
||||
if trace is not None:
|
||||
pool_traces.append(trace)
|
||||
candidate_keys = list(candidate.pool_keys or [])
|
||||
if not candidate_keys and getattr(candidate, "key", None) is not None:
|
||||
candidate_keys = [candidate.key]
|
||||
|
||||
return result, pool_traces
|
||||
ordered_keys, trace = await manager.select_pool_keys(session_uuid, candidate_keys)
|
||||
candidate.pool_keys = ordered_keys
|
||||
|
||||
selected_key_index = 0
|
||||
selected_key = None
|
||||
for idx, pool_key in enumerate(ordered_keys):
|
||||
if not bool(getattr(pool_key, "_pool_skipped", False)):
|
||||
selected_key = pool_key
|
||||
selected_key_index = idx
|
||||
break
|
||||
|
||||
if selected_key is not None:
|
||||
candidate.key = selected_key
|
||||
candidate._pool_key_index = selected_key_index
|
||||
candidate.mapping_matched_model = getattr(
|
||||
selected_key, "_pool_mapping_matched_model", None
|
||||
)
|
||||
candidate.is_skipped = False
|
||||
candidate.skip_reason = None
|
||||
else:
|
||||
candidate.is_skipped = True
|
||||
candidate.skip_reason = "pool: all keys unavailable"
|
||||
|
||||
if trace is not None:
|
||||
pool_traces.append(trace)
|
||||
|
||||
return candidates, pool_traces
|
||||
except Exception:
|
||||
from src.core.logger import logger
|
||||
|
||||
logger.opt(exception=True).debug("Pool reorder failed, using original order")
|
||||
return candidates, []
|
||||
|
||||
@staticmethod
|
||||
def _expand_pool_candidates_for_async_submit(candidates: list[Any]) -> list[Any]:
|
||||
"""Expand PoolCandidate to key-level candidates for async submit traversal."""
|
||||
from src.services.scheduling.schemas import PoolCandidate, ProviderCandidate
|
||||
|
||||
expanded: list[Any] = []
|
||||
for candidate in candidates:
|
||||
if not isinstance(candidate, PoolCandidate):
|
||||
expanded.append(candidate)
|
||||
continue
|
||||
|
||||
pool_keys = list(candidate.pool_keys or [])
|
||||
if not pool_keys:
|
||||
expanded.append(candidate)
|
||||
continue
|
||||
|
||||
for key_index, pool_key in enumerate(pool_keys):
|
||||
key_skipped = bool(getattr(pool_key, "_pool_skipped", False))
|
||||
key_skip_reason = (
|
||||
str(getattr(pool_key, "_pool_skip_reason", "") or "") or candidate.skip_reason
|
||||
)
|
||||
key_extra = (
|
||||
getattr(pool_key, "_pool_extra_data", None)
|
||||
if isinstance(getattr(pool_key, "_pool_extra_data", None), dict)
|
||||
else {}
|
||||
)
|
||||
|
||||
key_candidate = ProviderCandidate(
|
||||
provider=candidate.provider,
|
||||
endpoint=candidate.endpoint,
|
||||
key=pool_key,
|
||||
is_cached=candidate.is_cached,
|
||||
is_skipped=bool(candidate.is_skipped) or key_skipped,
|
||||
skip_reason=(
|
||||
key_skip_reason if (bool(candidate.is_skipped) or key_skipped) else None
|
||||
),
|
||||
mapping_matched_model=getattr(pool_key, "_pool_mapping_matched_model", None)
|
||||
or candidate.mapping_matched_model,
|
||||
needs_conversion=candidate.needs_conversion,
|
||||
provider_api_format=candidate.provider_api_format,
|
||||
output_limit=candidate.output_limit,
|
||||
capability_miss_count=candidate.capability_miss_count,
|
||||
)
|
||||
setattr(
|
||||
key_candidate,
|
||||
"_pool_extra_data",
|
||||
{
|
||||
"pool_group_id": str(candidate.provider.id),
|
||||
"pool_key_index": key_index,
|
||||
**key_extra,
|
||||
},
|
||||
)
|
||||
expanded.append(key_candidate)
|
||||
|
||||
return expanded
|
||||
|
||||
@staticmethod
|
||||
async def _pool_on_success(
|
||||
candidate: Any,
|
||||
@@ -466,6 +535,25 @@ class TaskService:
|
||||
|
||||
# Safety net: if record_id missing, create an "available" record on-demand.
|
||||
if not candidate_record_id:
|
||||
from src.services.scheduling.schemas import PoolCandidate
|
||||
|
||||
pool_extra = (
|
||||
getattr(candidate.key, "_pool_extra_data", None)
|
||||
if isinstance(getattr(candidate.key, "_pool_extra_data", None), dict)
|
||||
else {}
|
||||
)
|
||||
extra_data: dict[str, Any] = {
|
||||
"needs_conversion": bool(getattr(candidate, "needs_conversion", False)),
|
||||
"provider_api_format": getattr(candidate, "provider_api_format", None) or None,
|
||||
"mapping_matched_model": getattr(candidate, "mapping_matched_model", None)
|
||||
or None,
|
||||
**pool_extra,
|
||||
}
|
||||
if isinstance(candidate, PoolCandidate):
|
||||
extra_data["pool_group_id"] = str(candidate.provider.id)
|
||||
extra_data["pool_key_index"] = int(
|
||||
getattr(candidate, "_pool_key_index", 0) or 0
|
||||
)
|
||||
created = RequestCandidateService.create_candidate(
|
||||
db=self.db,
|
||||
request_id=request_id,
|
||||
@@ -478,6 +566,7 @@ class TaskService:
|
||||
key_id=str(candidate.key.id),
|
||||
status="available",
|
||||
is_cached=bool(getattr(candidate, "is_cached", False)),
|
||||
extra_data=extra_data,
|
||||
)
|
||||
candidate_record_id = str(created.id)
|
||||
candidate_record_map[(candidate_index, retry_index)] = candidate_record_id
|
||||
@@ -1331,6 +1420,7 @@ class TaskService:
|
||||
candidates, _pool_traces = await self._apply_pool_reorder(
|
||||
candidates, request_body=request_body
|
||||
)
|
||||
candidates = self._expand_pool_candidates_for_async_submit(candidates)
|
||||
|
||||
if max_candidates is not None and max_candidates > 0:
|
||||
candidates = candidates[:max_candidates]
|
||||
|
||||
Reference in New Issue
Block a user