xsl
f1b8b7bd82
- pool.py: WorkerState 增加 primary 标记；acquire_worker 分级派发—— 主组有在线只在主组内挑空闲（全忙排队，不退备组），主组全部下线才用备机 - config.py: 新增可选 primary_workers 配置（须为 workers 子集，自动过滤+告警）； 不配置/空列表 = 全部视为主 worker，兼容旧行为 - config.json(不入 git): 3 台 GPU worker 设为主，max_global_concurrency 1→3 - app.py: /health 增加 workers_primary_total / workers_primary_healthy - tests: 新增 4 条主备派发单测（主优先/主全挂用备/主全忙不走备/无配置兼容） - docs(实现说明): 补充 primary_workers 配置说明，改配置需重启网关生效
237 lines
8.0 KiB
Python
237 lines
8.0 KiB
Python
"""gateway/pool.py 隔离期逻辑单测（不依赖真实网络）。
|
||
|
||
覆盖《docs/网关摘除隔离期实施方案.md》§3.3 的三个用例：
|
||
- 被动摘除设置隔离期（offline_until > now）
|
||
- 隔离期内探测成功不复活
|
||
- 隔离期满后探测成功恢复上线
|
||
外加：_try_revive 的阈值/在线短路、mark_worker_unhealthy 对已下线 worker 的幂等、
|
||
_health_check_loop 整链路验证、_cooldown_seconds 配置读取。
|
||
|
||
与 test_gateway_serialization.py 同款：asyncio.run 套同步测试，不依赖 pytest-asyncio。
|
||
"""
|
||
|
||
import asyncio
|
||
import time
|
||
from unittest.mock import patch
|
||
|
||
from gateway import pool
|
||
from gateway.pool import WorkerState
|
||
|
||
_HC_CFG = {
|
||
"health_check": {
|
||
"path": "/health",
|
||
"interval_seconds": 0.05, # 测试用：快速轮询
|
||
"timeout_seconds": 3,
|
||
"unhealthy_threshold": 2,
|
||
"healthy_threshold": 1,
|
||
},
|
||
"shared_password": "test-token",
|
||
}
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 被动摘除 → 隔离期
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def test_passive_ejection_sets_cooldown():
|
||
w = WorkerState(url="http://127.0.0.1:9000", online=True)
|
||
|
||
async def run():
|
||
pool._pool_condition = asyncio.Condition()
|
||
with patch.object(pool, "_cooldown_seconds", return_value=60.0):
|
||
await pool.mark_worker_unhealthy(w)
|
||
|
||
asyncio.run(run())
|
||
assert w.online is False
|
||
assert w.offline_until > time.monotonic(), "被动摘除必须设置未来的隔离截止时刻"
|
||
assert w.consecutive_failures == 1
|
||
|
||
|
||
def test_passive_ejection_idempotent_when_offline():
|
||
"""已下线的 worker 再次被标记：不重复计失败、不续期隔离。"""
|
||
w = WorkerState(url="http://127.0.0.1:9000", online=False)
|
||
w.offline_until = time.monotonic() + 30
|
||
before = w.offline_until
|
||
|
||
async def run():
|
||
pool._pool_condition = asyncio.Condition()
|
||
with patch.object(pool, "_cooldown_seconds", return_value=60.0):
|
||
await pool.mark_worker_unhealthy(w)
|
||
|
||
asyncio.run(run())
|
||
assert w.online is False
|
||
assert w.consecutive_failures == 0
|
||
assert w.offline_until == before
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _try_revive：探测成功后是否允许复活
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def test_probe_success_blocked_during_cooldown():
|
||
w = WorkerState(url="http://127.0.0.1:9000") # offline
|
||
w.consecutive_successes = 1
|
||
w.offline_until = time.monotonic() + 60
|
||
assert pool._try_revive(w, healthy_threshold=1) is False
|
||
assert w.online is False
|
||
|
||
|
||
def test_probe_success_revives_after_cooldown():
|
||
w = WorkerState(url="http://127.0.0.1:9000")
|
||
w.consecutive_successes = 1
|
||
w.offline_until = time.monotonic() - 1 # 隔离期已过
|
||
assert pool._try_revive(w, healthy_threshold=1) is True
|
||
|
||
|
||
def test_try_revive_requires_success_threshold():
|
||
w = WorkerState(url="http://127.0.0.1:9000")
|
||
w.consecutive_successes = 1
|
||
w.offline_until = 0.0
|
||
assert pool._try_revive(w, healthy_threshold=2) is False
|
||
|
||
|
||
def test_try_revive_noop_when_already_online():
|
||
w = WorkerState(url="http://127.0.0.1:9000", online=True)
|
||
w.consecutive_successes = 1
|
||
assert pool._try_revive(w, healthy_threshold=1) is False
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _health_check_loop 整链路：隔离期内探测 200 不上线，期满后上线
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def test_health_check_loop_respects_cooldown(monkeypatch):
|
||
w = WorkerState(url="http://127.0.0.1:9000", online=False)
|
||
w.offline_until = time.monotonic() + 60 # 处于隔离期
|
||
|
||
class _FakeResp:
|
||
status_code = 200
|
||
|
||
class _FakeClient:
|
||
async def __aenter__(self):
|
||
return self
|
||
|
||
async def __aexit__(self, *args):
|
||
return False
|
||
|
||
async def get(self, url, headers=None, timeout=None):
|
||
return _FakeResp()
|
||
|
||
async def run():
|
||
monkeypatch.setattr(pool, "_workers", {"w1": w})
|
||
monkeypatch.setattr(pool, "_shutdown_event", asyncio.Event())
|
||
monkeypatch.setattr(pool.httpx, "AsyncClient", _FakeClient)
|
||
task = asyncio.create_task(pool._health_check_loop(_HC_CFG))
|
||
try:
|
||
await asyncio.sleep(0.2) # 跑若干轮探测
|
||
assert w.consecutive_successes >= 1, "探测应已成功过"
|
||
assert w.online is False, "隔离期内即使探测成功也不得复活"
|
||
|
||
w.offline_until = time.monotonic() - 1 # 隔离期结束
|
||
await asyncio.sleep(0.2)
|
||
assert w.online is True, "隔离期满后下一次探测成功应上线"
|
||
finally:
|
||
pool._shutdown_event.set()
|
||
await asyncio.wait_for(task, timeout=2)
|
||
|
||
asyncio.run(run())
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# _cooldown_seconds 配置读取
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def test_cooldown_seconds_reads_config(monkeypatch):
|
||
from gateway import config as gw_config
|
||
monkeypatch.setattr(gw_config, "_config_cache",
|
||
{"health_check": {"cooldown_seconds": 120}})
|
||
assert pool._cooldown_seconds() == 120.0
|
||
|
||
|
||
def test_cooldown_seconds_falls_back_to_60(monkeypatch):
|
||
from gateway import config as gw_config
|
||
monkeypatch.setattr(gw_config, "_config_cache", None) # get_config 会抛 RuntimeError
|
||
assert pool._cooldown_seconds() == 60.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 主备分组派发（primary_workers）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_DISPATCH_CFG = {
|
||
"dispatch": {
|
||
"queue_wait_seconds": 0.05, # 测试用：快速超时
|
||
}
|
||
}
|
||
|
||
|
||
def _make_pool(specs):
|
||
"""specs: [(url, online, busy, primary), ...] → 构造 pool._workers。"""
|
||
return {
|
||
url: WorkerState(url=url, online=online, busy=busy, primary=primary)
|
||
for url, online, busy, primary in specs
|
||
}
|
||
|
||
|
||
def _run_acquire(workers):
|
||
async def run():
|
||
pool._pool_condition = asyncio.Condition()
|
||
orig = pool._workers
|
||
pool._workers = workers
|
||
try:
|
||
return await pool.acquire_worker(_DISPATCH_CFG)
|
||
finally:
|
||
pool._workers = orig
|
||
return asyncio.run(run())
|
||
|
||
|
||
def test_primary_preferred_over_backup():
|
||
"""主 worker 在线时空闲主优先于备 worker。"""
|
||
workers = _make_pool([
|
||
("http://backup:1", True, False, False),
|
||
("http://primary:1", True, False, True),
|
||
])
|
||
w = _run_acquire(workers)
|
||
assert w.url == "http://primary:1"
|
||
assert w.busy is True
|
||
|
||
|
||
def test_backup_used_only_when_all_primary_down():
|
||
"""主 worker 全部下线时才用备 worker。"""
|
||
workers = _make_pool([
|
||
("http://backup:1", True, False, False),
|
||
("http://primary:1", False, False, True),
|
||
("http://primary:2", False, False, True),
|
||
])
|
||
w = _run_acquire(workers)
|
||
assert w.url == "http://backup:1"
|
||
|
||
|
||
def test_backup_not_used_when_primary_busy():
|
||
"""主 worker 在线但全忙时不退到备 worker（等待后超时拒绝）。"""
|
||
workers = _make_pool([
|
||
("http://backup:1", True, False, False),
|
||
("http://primary:1", True, True, True),
|
||
])
|
||
|
||
async def run():
|
||
pool._pool_condition = asyncio.Condition()
|
||
try:
|
||
await pool.acquire_worker(_DISPATCH_CFG)
|
||
except pool.NoWorkerAvailable:
|
||
return True
|
||
return False
|
||
|
||
assert asyncio.run(run()) is True, "主组全忙时不得使用备 worker"
|
||
assert workers["http://backup:1"].busy is False
|
||
|
||
|
||
def test_no_primary_config_all_are_primary():
|
||
"""未配置 primary_workers 时全部视为主 worker（兼容旧行为）。"""
|
||
workers = _make_pool([
|
||
("http://w1:1", True, False, True),
|
||
("http://w2:1", True, False, True),
|
||
])
|
||
w = _run_acquire(workers)
|
||
assert w.url == "http://w1:1"
|