mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
7c580e843f
- 移除 token/latency Prometheus 指标的 model 标签，避免 provider x model 笛卡尔积 - HealthMonitor 滑动窗口从 DB JSON 迁移至进程内存，减少写放大 - ModelCostService 三层缓存增加 500 条上限，超限时清空 - StickyPriority 粘性缓存和健康状态字典增加容量淘汰 - AffinityManager 请求锁字典增加 500 条上限，淘汰空闲锁 - 配额刷新/探测查询使用 defer/load_only 避免加载大 JSON 列 - Alembic 迁移清理 DB 中遗留的 request_results_window 数据 - 同步更新测试适配 batch_get_cooldowns 返回值和批量删除异步化
241 lines
7.6 KiB
Python
241 lines
7.6 KiB
Python
"""Tests for pool manager trace data collection and attachment."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import pytest
|
|
|
|
from src.services.provider.pool.config import PoolConfig
|
|
from src.services.provider.pool.manager import PoolManager
|
|
from src.services.provider.pool.trace import PoolSchedulingTrace
|
|
|
|
|
|
def _make_candidate(key_id: str, *, is_skipped: bool = False) -> SimpleNamespace:
|
|
return SimpleNamespace(
|
|
key=SimpleNamespace(id=key_id),
|
|
is_skipped=is_skipped,
|
|
skip_reason=None,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Trace attachment to candidates
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_attached_to_first_candidate() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig())
|
|
c1 = _make_candidate("key-1")
|
|
c2 = _make_candidate("key-2")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-1": None, "key-2": None},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_lru_scores",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-1": 100.0, "key-2": 50.0},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates("sess-1", [c1, c2])
|
|
|
|
# First candidate should have _pool_scheduling_trace
|
|
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
|
assert trace is not None
|
|
assert isinstance(trace, PoolSchedulingTrace)
|
|
assert trace.total_keys == 2
|
|
assert trace.provider_id == "prov-1"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pool_extra_data_on_selected_candidate() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig())
|
|
c1 = _make_candidate("key-1")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-1": None},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_lru_scores",
|
|
new_callable=AsyncMock,
|
|
return_value={},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates(None, [c1])
|
|
|
|
extra = getattr(result[0], "_pool_extra_data", None)
|
|
assert extra is not None
|
|
assert "pool_selection" in extra
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pool_extra_data_on_skipped_candidate() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig())
|
|
c1 = _make_candidate("key-1")
|
|
c2 = _make_candidate("key-2")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={
|
|
"key-1": "rate_limited_429",
|
|
"key-2": None,
|
|
},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_lru_scores",
|
|
new_callable=AsyncMock,
|
|
return_value={},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates(None, [c1, c2])
|
|
|
|
# key-1 is skipped with pool_skip extra data
|
|
skipped = [r for r in result if r.is_skipped]
|
|
assert len(skipped) == 1
|
|
extra = getattr(skipped[0], "_pool_extra_data", None)
|
|
assert extra is not None
|
|
assert extra["pool_skip"]["type"] == "cooldown"
|
|
assert extra["pool_skip"]["cooldown_reason"] == "rate_limited_429"
|
|
# TTL is not fetched on the scheduling hot path (include_ttl=False)
|
|
assert "cooldown_ttl" not in extra["pool_skip"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_build_summary_matches() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig(cost_limit_per_key_tokens=1000))
|
|
c1 = _make_candidate("key-1")
|
|
c2 = _make_candidate("key-2")
|
|
c3 = _make_candidate("key-3")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={
|
|
"key-1": "overloaded_529",
|
|
"key-2": None,
|
|
"key-3": None,
|
|
},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cost_totals",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-2": 1500, "key-3": 200},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates(None, [c1, c2, c3])
|
|
|
|
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
|
assert trace is not None
|
|
|
|
summary = trace.build_summary(success_key_id="key-3")
|
|
assert summary["total_keys"] == 3
|
|
assert summary["skipped_cooldown"] == 1 # key-1
|
|
assert summary["skipped_cost"] == 1 # key-2
|
|
assert summary["attempted"] == 1 # key-3
|
|
assert summary["success_key_id"] == "key-3"[:8]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_trace_build_summary_uses_attempted_key_ids_when_provided() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig())
|
|
c1 = _make_candidate("key-1")
|
|
c2 = _make_candidate("key-2")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value=None,
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-1": None, "key-2": None},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_lru_scores",
|
|
new_callable=AsyncMock,
|
|
return_value={},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates(None, [c1, c2])
|
|
|
|
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
|
assert trace is not None
|
|
|
|
summary = trace.build_summary(
|
|
success_key_id="key-1",
|
|
attempted_key_ids={"key-1"},
|
|
)
|
|
assert summary["total_keys"] == 2
|
|
assert summary["attempted"] == 1
|
|
assert summary["success_key_id"] == "key-1"[:8]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_sticky_trace_info() -> None:
|
|
pool = PoolManager("prov-1", PoolConfig(sticky_session_ttl_seconds=3600))
|
|
c1 = _make_candidate("key-1")
|
|
c2 = _make_candidate("key-2")
|
|
|
|
with (
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_sticky_binding",
|
|
new_callable=AsyncMock,
|
|
return_value="key-2",
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.batch_get_cooldowns",
|
|
new_callable=AsyncMock,
|
|
return_value={"key-1": None, "key-2": None},
|
|
),
|
|
patch(
|
|
"src.services.provider.pool.redis_ops.get_lru_scores",
|
|
new_callable=AsyncMock,
|
|
return_value={},
|
|
),
|
|
):
|
|
result = await pool.reorder_candidates("sess-1", [c1, c2])
|
|
|
|
trace = getattr(result[0], "_pool_scheduling_trace", None)
|
|
assert trace is not None
|
|
assert trace.sticky_session_used is True
|
|
|
|
# key-2 (sticky hit) should be first
|
|
assert result[0].key.id == "key-2"
|
|
extra = getattr(result[0], "_pool_extra_data", None)
|
|
assert extra is not None
|
|
assert extra["pool_selection"]["sticky_hit"] is True
|