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 配置说明，改配置需重启网关生效
354 lines
14 KiB
Python
354 lines
14 KiB
Python
"""Worker 健康池 + 空闲派发。
|
||
|
||
- 后台 asyncio 任务周期性探测每个 worker 的 /health
|
||
- 连续失败 N 次 → 下线；连续成功 N 次 → 上线
|
||
- 每个 worker 一个 busy 标志（per_worker_concurrency=1）
|
||
- acquire_worker() 从在线池挑空闲 worker；全忙排队；池空→NoWorkerAvailable
|
||
- release_worker() 释放 worker
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from typing import Dict, List, Optional
|
||
|
||
import httpx
|
||
|
||
logger = logging.getLogger("gateway.pool")
|
||
|
||
|
||
def _cooldown_seconds() -> float:
|
||
"""被动摘除后的隔离期时长（秒），配置 health_check.cooldown_seconds，默认 60。"""
|
||
try:
|
||
from gateway.config import get_config
|
||
return float(get_config()["health_check"].get("cooldown_seconds", 60))
|
||
except Exception: # noqa: BLE001
|
||
return 60.0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 异常
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class NoWorkerAvailable(Exception):
|
||
"""无可用 worker（池空或全忙排队超时）。"""
|
||
pass
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 数据结构
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@dataclass
|
||
class WorkerState:
|
||
"""单个 worker 的运行时状态。"""
|
||
url: str
|
||
online: bool = False # 初始 offline，等健康检查通过后上线
|
||
busy: bool = False
|
||
primary: bool = True # 主 worker（True 优先派发）；全部主 worker 下线才用备组
|
||
consecutive_failures: int = 0
|
||
consecutive_successes: int = 0
|
||
offline_until: float = 0.0 # 被动摘除后的隔离截止时刻（time.monotonic() 基准）
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 全局状态
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_workers: Dict[str, WorkerState] = {}
|
||
_pool_condition: Optional[asyncio.Condition] = None
|
||
_health_task: Optional[asyncio.Task] = None
|
||
_shutdown_event: Optional[asyncio.Event] = None
|
||
|
||
# 全局串行槽：同一时间最多 max_global_concurrency 个请求进入 worker 派发，其余排队。
|
||
# ⚠️ 依赖单进程 uvicorn 部署（hair-gateway.service 无 --workers）。
|
||
# 若改多 worker / gunicorn，进程内信号量会失效，需换成跨进程锁（文件锁/Redis）。
|
||
_global_sem: Optional[asyncio.Semaphore] = None
|
||
_global_max: int = 1
|
||
_global_busy: int = 0
|
||
_global_waiting: int = 0
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 健康检查后台任务
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _try_revive(w: WorkerState, healthy_threshold: int) -> bool:
|
||
"""探测成功后是否允许将该 worker 复活上线。
|
||
|
||
条件：未在线 + 连续成功达到阈值 + **不在被动摘除隔离期内**
|
||
（offline_until 之前即使 /health 探测成功也压着不上线，防横跳）。
|
||
"""
|
||
if w.online or w.consecutive_successes < healthy_threshold:
|
||
return False
|
||
if time.monotonic() < w.offline_until:
|
||
logger.debug("Worker %s 探测成功但处于隔离期，暂不复活", w.url)
|
||
return False
|
||
return True
|
||
|
||
|
||
async def _health_check_loop(cfg: dict) -> None:
|
||
"""后台循环：周期性探测所有 worker 健康状态。"""
|
||
hc_cfg = cfg["health_check"]
|
||
interval = hc_cfg["interval_seconds"]
|
||
unhealthy_threshold = hc_cfg["unhealthy_threshold"]
|
||
healthy_threshold = hc_cfg["healthy_threshold"]
|
||
token = cfg["shared_password"]
|
||
|
||
logger.info(
|
||
"健康检查循环启动 | 间隔=%ds | 下线阈值=%d | 上线阈值=%d | 隔离期=%.0fs | workers=%d",
|
||
interval, unhealthy_threshold, healthy_threshold, _cooldown_seconds(), len(_workers),
|
||
)
|
||
|
||
async with httpx.AsyncClient() as client:
|
||
while not _shutdown_event.is_set():
|
||
for w in _workers.values():
|
||
url = f"{w.url}{hc_cfg['path']}"
|
||
try:
|
||
resp = await client.get(
|
||
url,
|
||
headers={"X-Internal-Token": token},
|
||
timeout=hc_cfg["timeout_seconds"],
|
||
)
|
||
if resp.status_code == 200:
|
||
w.consecutive_failures = 0
|
||
w.consecutive_successes += 1
|
||
if _try_revive(w, healthy_threshold):
|
||
w.online = True
|
||
logger.info("✅ Worker 上线: %s", w.url)
|
||
else:
|
||
w.consecutive_successes = 0
|
||
w.consecutive_failures += 1
|
||
if w.consecutive_failures >= unhealthy_threshold and w.online:
|
||
w.online = False
|
||
logger.warning("⚠ Worker 下线: %s（HTTP %d，连续失败 %d 次）",
|
||
w.url, resp.status_code, w.consecutive_failures)
|
||
except Exception as exc:
|
||
w.consecutive_successes = 0
|
||
w.consecutive_failures += 1
|
||
if w.consecutive_failures >= unhealthy_threshold and w.online:
|
||
w.online = False
|
||
logger.warning("⚠ Worker 下线: %s（%s，连续失败 %d 次）",
|
||
w.url, exc, w.consecutive_failures)
|
||
|
||
# 等待下一次探测（支持快速关闭）
|
||
try:
|
||
await asyncio.wait_for(_shutdown_event.wait(), timeout=interval)
|
||
break # shutdown signaled
|
||
except asyncio.TimeoutError:
|
||
pass # 正常的 interval 到期
|
||
|
||
logger.info("健康检查循环已停止")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 初始化 / 关闭
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def init_pool(cfg: dict) -> None:
|
||
"""初始化 worker 池并启动健康检查后台任务。"""
|
||
global _workers, _pool_condition, _health_task, _shutdown_event
|
||
|
||
# 主 worker 分组：配置了 primary_workers 则按其标记，否则全部视为主 worker
|
||
primary_set = set(cfg.get("primary_workers") or cfg["workers"])
|
||
|
||
_workers = {url: WorkerState(url=url, primary=url in primary_set)
|
||
for url in cfg["workers"]}
|
||
n_primary = sum(1 for w in _workers.values() if w.primary)
|
||
_pool_condition = asyncio.Condition()
|
||
_shutdown_event = asyncio.Event()
|
||
|
||
# 全局串行槽（独立于 worker busy 标志）
|
||
init_global_slot(cfg)
|
||
|
||
# 立即做一轮健康检查以快速上线
|
||
hc_cfg = cfg["health_check"]
|
||
token = cfg["shared_password"]
|
||
async with httpx.AsyncClient() as client:
|
||
for w in _workers.values():
|
||
try:
|
||
resp = await client.get(
|
||
f"{w.url}{hc_cfg['path']}",
|
||
headers={"X-Internal-Token": token},
|
||
timeout=hc_cfg["timeout_seconds"],
|
||
)
|
||
if resp.status_code == 200:
|
||
w.consecutive_successes = 1
|
||
if hc_cfg["healthy_threshold"] <= 1:
|
||
w.online = True
|
||
else:
|
||
w.consecutive_failures = 1
|
||
except Exception:
|
||
w.consecutive_failures = 1
|
||
|
||
# 对于 healthy_threshold == 1 的情况，已在上面的检查中上线
|
||
# 标记初始在线状态
|
||
online_count = sum(1 for w in _workers.values() if w.online)
|
||
logger.info("Worker 池初始化完成 | 总数=%d | 在线=%d | 主worker=%d",
|
||
len(_workers), online_count, n_primary)
|
||
|
||
# 启动后台健康检查
|
||
_health_task = asyncio.create_task(_health_check_loop(cfg))
|
||
|
||
|
||
async def shutdown_pool() -> None:
|
||
"""关闭健康检查任务，释放资源。"""
|
||
global _shutdown_event, _health_task
|
||
if _shutdown_event:
|
||
_shutdown_event.set()
|
||
if _health_task:
|
||
_health_task.cancel()
|
||
try:
|
||
await _health_task
|
||
except asyncio.CancelledError:
|
||
pass
|
||
_health_task = None
|
||
logger.info("Worker 池已关闭")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 全局串行槽
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def init_global_slot(cfg: dict) -> None:
|
||
"""初始化全局串行信号量。
|
||
|
||
max_global_concurrency=1 → 同一时间只有一个请求进入 worker 派发，其余在信号量上排队。
|
||
依赖单进程 uvicorn 部署；多进程需换跨进程锁。
|
||
"""
|
||
global _global_sem, _global_max, _global_busy, _global_waiting
|
||
n = int(cfg.get("dispatch", {}).get("max_global_concurrency", 1))
|
||
if n < 1:
|
||
n = 1
|
||
_global_sem = asyncio.Semaphore(n)
|
||
_global_max = n
|
||
_global_busy = 0
|
||
_global_waiting = 0
|
||
logger.info("全局并发槽初始化: max=%d（单进程生效）", n)
|
||
|
||
|
||
async def acquire_global_slot(cfg: dict) -> bool:
|
||
"""获取一个全局处理槽。True=获得；False=排队超时。
|
||
|
||
超时由 dispatch.max_queue_wait_seconds 控制（需 < nginx proxy_read_timeout 600s）。
|
||
"""
|
||
global _global_busy, _global_waiting
|
||
if _global_sem is None:
|
||
return True # 未初始化，不限流
|
||
timeout = float(cfg.get("dispatch", {}).get("max_queue_wait_seconds", 590))
|
||
_global_waiting += 1
|
||
try:
|
||
await asyncio.wait_for(_global_sem.acquire(), timeout=timeout)
|
||
_global_busy += 1
|
||
return True
|
||
except asyncio.TimeoutError:
|
||
logger.warning("全局排队超时（%.1fs），拒绝请求 | waiting=%d", timeout, _global_waiting)
|
||
return False
|
||
finally:
|
||
_global_waiting -= 1
|
||
|
||
|
||
def release_global_slot() -> None:
|
||
"""释放全局处理槽。"""
|
||
global _global_busy
|
||
if _global_sem is not None:
|
||
_global_busy -= 1
|
||
_global_sem.release()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 派发
|
||
# ---------------------------------------------------------------------------
|
||
|
||
async def mark_worker_unhealthy(w: WorkerState) -> None:
|
||
"""被动标记：转发失败时立即将该 worker 下线，并进入隔离期。
|
||
|
||
隔离期内即使 /health 探测成功也不复活（由 _try_revive 守卫），
|
||
避免"摘除→8 秒复活→再失败"的横跳。
|
||
"""
|
||
async with _pool_condition:
|
||
if w.online:
|
||
w.online = False
|
||
w.consecutive_failures += 1
|
||
cooldown = _cooldown_seconds()
|
||
w.offline_until = time.monotonic() + cooldown
|
||
logger.warning("⚠ Worker 被动下线: %s（转发失败，隔离 %.0fs）",
|
||
w.url, cooldown)
|
||
|
||
|
||
async def acquire_worker(cfg: dict) -> WorkerState:
|
||
"""从在线池中获取一个空闲 worker。
|
||
|
||
主备分组：只要还有主 worker（primary）在线，就只在主组内挑空闲
|
||
（全忙则排队等待，不用备组）；只有主组全部下线时才用备用 worker。
|
||
全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。
|
||
"""
|
||
queue_wait = cfg["dispatch"]["queue_wait_seconds"]
|
||
deadline = time.monotonic() + queue_wait
|
||
|
||
async with _pool_condition:
|
||
while True:
|
||
online = [w for w in _workers.values() if w.online]
|
||
primary_online = [w for w in online if w.primary]
|
||
# 主组有在线 → 只在主组内挑；主组全挂 → 才轮到备组
|
||
candidates = primary_online if primary_online else online
|
||
idle = [w for w in candidates if not w.busy]
|
||
|
||
if not idle:
|
||
if not online:
|
||
raise NoWorkerAvailable("后端服务暂不可用，请稍后重试")
|
||
|
||
# 全忙，等待
|
||
remaining = deadline - time.monotonic()
|
||
if remaining <= 0:
|
||
raise NoWorkerAvailable("后端服务繁忙，请稍后重试")
|
||
|
||
logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...",
|
||
len(candidates), remaining)
|
||
try:
|
||
await asyncio.wait_for(_pool_condition.wait(), timeout=remaining)
|
||
except asyncio.TimeoutError:
|
||
raise NoWorkerAvailable("后端服务繁忙，请稍后重试")
|
||
continue # 重新检查
|
||
|
||
# 取第一个空闲 worker
|
||
w = idle[0]
|
||
w.busy = True
|
||
if not w.primary:
|
||
logger.info("主 worker 全部下线，使用备用 worker: %s", w.url)
|
||
logger.debug("Worker 分配: %s", w.url)
|
||
return w
|
||
|
||
|
||
async def release_worker(w: WorkerState) -> None:
|
||
"""释放 worker，标记为空闲并通知等待者。"""
|
||
async with _pool_condition:
|
||
w.busy = False
|
||
logger.debug("Worker 释放: %s", w.url)
|
||
_pool_condition.notify(1)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 状态查询
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def get_pool_status() -> dict:
|
||
"""返回当前池状态（供 /gateway-health 使用）。"""
|
||
global_info = {
|
||
"global_max": _global_max,
|
||
"global_busy": _global_busy,
|
||
"global_waiting": _global_waiting,
|
||
}
|
||
if not _workers:
|
||
return {"total": 0, "healthy": 0, "busy": 0,
|
||
"primary_total": 0, "primary_healthy": 0, **global_info}
|
||
total = len(_workers)
|
||
healthy = sum(1 for w in _workers.values() if w.online)
|
||
busy = sum(1 for w in _workers.values() if w.busy)
|
||
primary_total = sum(1 for w in _workers.values() if w.primary)
|
||
primary_healthy = sum(1 for w in _workers.values() if w.primary and w.online)
|
||
return {"total": total, "healthy": healthy, "busy": busy,
|
||
"primary_total": primary_total, "primary_healthy": primary_healthy,
|
||
**global_info}
|