mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +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)
266 lines
10 KiB
Python
266 lines
10 KiB
Python
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
from src.core.logger import logger
|
||
|
||
|
||
class TaskPoolOperationsService:
|
||
"""任务池化相关操作（重排、展开、健康回写）。"""
|
||
|
||
def extract_session_uuid(
|
||
self,
|
||
provider_type: str,
|
||
request_body: dict[str, Any] | None,
|
||
) -> str | None:
|
||
"""Extract a session UUID from the request body (provider-type aware)."""
|
||
if not isinstance(request_body, dict):
|
||
return None
|
||
from src.services.provider.pool.hooks import get_pool_hook
|
||
|
||
hook = get_pool_hook(provider_type)
|
||
if hook is not None:
|
||
return hook.extract_session_uuid(request_body)
|
||
return None
|
||
|
||
async def apply_pool_reorder(
|
||
self,
|
||
candidates: list[Any],
|
||
request_body: dict[str, Any] | None,
|
||
) -> tuple[list[Any], list[Any]]:
|
||
"""Apply pool key ordering for PoolCandidate objects."""
|
||
if not candidates:
|
||
return candidates, []
|
||
|
||
pool_traces: list[Any] = []
|
||
|
||
try:
|
||
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
|
||
|
||
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 = self.extract_session_uuid(provider_type, request_body)
|
||
manager = PoolManager(provider_id, pool_cfg)
|
||
|
||
candidate_keys = list(candidate.pool_keys or [])
|
||
if not candidate_keys and getattr(candidate, "key", None) is not None:
|
||
candidate_keys = [candidate.key]
|
||
|
||
# 构造延迟可用性检查回调（从 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(candidate.pool_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:
|
||
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
|
||
|
||
async def pool_on_success(
|
||
self,
|
||
candidate: Any,
|
||
request_body: dict[str, Any] | None,
|
||
) -> None:
|
||
"""Notify the pool manager about a successful request (sticky + LRU)."""
|
||
try:
|
||
from src.services.provider.pool.config import parse_pool_config
|
||
from src.services.provider.pool.manager import PoolManager
|
||
|
||
provider = candidate.provider
|
||
provider_config = getattr(provider, "config", None)
|
||
pool_cfg = parse_pool_config(provider_config)
|
||
if pool_cfg is None:
|
||
return
|
||
|
||
provider_id = str(getattr(provider, "id", "") or "")
|
||
key_id = str(getattr(candidate.key, "id", "") or "")
|
||
if not provider_id or not key_id:
|
||
return
|
||
|
||
provider_type = str(getattr(provider, "provider_type", "") or "")
|
||
session_uuid = self.extract_session_uuid(provider_type, request_body)
|
||
|
||
mgr = PoolManager(provider_id, pool_cfg)
|
||
await mgr.on_request_success(
|
||
session_uuid=session_uuid,
|
||
key_id=key_id,
|
||
)
|
||
except Exception:
|
||
logger.opt(exception=True).debug("Pool on_request_success failed (non-blocking)")
|
||
|
||
@staticmethod
|
||
async def pool_on_error(
|
||
provider: Any,
|
||
key: Any,
|
||
status_code: int,
|
||
cause: Any,
|
||
) -> None:
|
||
"""Notify the pool manager about an upstream error (health policy)."""
|
||
try:
|
||
from src.services.provider.pool.config import parse_pool_config
|
||
from src.services.provider.pool.health_policy import apply_health_policy
|
||
|
||
pool_cfg = parse_pool_config(getattr(provider, "config", None))
|
||
if pool_cfg is None:
|
||
return
|
||
|
||
error_text = ""
|
||
resp_headers: dict[str, str] = {}
|
||
if getattr(cause, "response", None) is not None:
|
||
try:
|
||
error_text = (cause.response.text or "")[:4000]
|
||
except Exception:
|
||
pass
|
||
try:
|
||
resp_headers = dict(cause.response.headers)
|
||
except Exception:
|
||
pass
|
||
|
||
await apply_health_policy(
|
||
provider_id=str(provider.id),
|
||
key_id=str(key.id),
|
||
status_code=status_code,
|
||
error_body=error_text,
|
||
response_headers=resp_headers,
|
||
config=pool_cfg,
|
||
)
|
||
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
|