From f1b8b7bd82a342d40d626877d81e877babd94ddb Mon Sep 17 00:00:00 2001 From: xsl Date: Tue, 22 Sep 2026 23:48:21 +0800 Subject: [PATCH] =?UTF-8?q?feat(gateway):=20worker=20=E6=B1=A0=E4=B8=BB?= =?UTF-8?q?=E5=A4=87=E5=88=86=E7=BB=84=E6=B4=BE=E5=8F=91=EF=BC=8C=E4=B8=BB?= =?UTF-8?q?=E5=85=A8=E6=8C=82=E6=89=8D=E7=94=A8=E5=A4=87=E6=9C=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 配置说明，改配置需重启网关生效 --- docs/实现说明.md | 3 ++ gateway/app.py | 2 + gateway/config.example.json | 1 + gateway/config.py | 10 +++++ gateway/pool.py | 35 ++++++++++++---- tests/test_gateway_pool.py | 82 +++++++++++++++++++++++++++++++++++++ 6 files changed, 125 insertions(+), 8 deletions(-) diff --git a/docs/实现说明.md b/docs/实现说明.md index 74beff5..c8cf9e0 100644 --- a/docs/实现说明.md +++ b/docs/实现说明.md @@ -147,6 +147,9 @@ - 很轻：FastAPI/uvicorn/httpx + **接口4 的 `volcengine-python-sdk[ark]`**（或直接 httpx 调，OpenAI 兼容）。 - 不装 torch/mediapipe/opencv。配置 `gateway/config.json`(不入 git)：`workers` 列表、`shared_password`、 `ark` 的 api_key/base_url/model、`public_base_url`、超时(**生发接口慢，`request_timeout_seconds` 调大 ≥120s**)。 + 可选 `primary_workers`（`workers` 的子集）启用主备分组：主 worker 全部下线时才派发备用 worker； + 主组并行度由 `dispatch.max_global_concurrency` 控制（N 台主 worker 就设 N，依赖单进程 uvicorn）。 + 修改配置需重启 `hair-gateway.service` 生效（配置启动时加载一次）。 - 托管 `/static/annotations/`(落盘的图)，定期清理。 --- diff --git a/gateway/app.py b/gateway/app.py index f3c2e32..60a0e81 100644 --- a/gateway/app.py +++ b/gateway/app.py @@ -123,6 +123,8 @@ async def gateway_health(): "workers_total": status["total"], "workers_healthy": status["healthy"], "workers_busy": status["busy"], + "workers_primary_total": status.get("primary_total", 0), + "workers_primary_healthy": status.get("primary_healthy", 0), "global_max": status.get("global_max", 1), "global_busy": status.get("global_busy", 0), "global_waiting": status.get("global_waiting", 0), diff --git a/gateway/config.example.json b/gateway/config.example.json index d4a3430..aa757af 100644 --- a/gateway/config.example.json +++ b/gateway/config.example.json @@ -3,6 +3,7 @@ "http://127.0.0.1:28187", "http://127.0.0.1:28188" ], + "primary_workers": [], "shared_password": "REPLACE_ME_ROTATE_PERIODICALLY", "public_base_url": "https://hair.xiangsilian.com", "static_dir": "static/annotations", diff --git a/gateway/config.py b/gateway/config.py index 5c58bc5..c533b08 100644 --- a/gateway/config.py +++ b/gateway/config.py @@ -15,6 +15,7 @@ CONFIG_PATH = Path(__file__).resolve().parent / "config.json" DEFAULTS = { "public_base_url": "https://hair.xiangsilian.com", "static_dir": "static/annotations", + "primary_workers": [], # 主 worker 地址列表（须为 workers 的子集）；空 = 全部视为主 worker "health_check": { "path": "/health", "interval_seconds": 8, @@ -83,6 +84,15 @@ def load_config() -> dict: if not isinstance(w, str) or not w: raise ValueError(f"配置错误: workers[{i}] 必须是非空字符串") + primary_workers: List[str] = cfg.get("primary_workers", []) + for i, w in enumerate(primary_workers): + if not isinstance(w, str) or not w: + raise ValueError(f"配置错误: primary_workers[{i}] 必须是非空字符串") + for w in primary_workers: + if w not in workers: + logger.warning("⚠ primary_workers 中的 %s 不在 workers 列表里，将被忽略", w) + cfg["primary_workers"] = [w for w in primary_workers if w in workers] + password: str = cfg.get("shared_password", "") if not password or password == "REPLACE_ME_ROTATE_PERIODICALLY": logger.warning( diff --git a/gateway/pool.py b/gateway/pool.py index ab9e3fa..b538cb0 100644 --- a/gateway/pool.py +++ b/gateway/pool.py @@ -46,6 +46,7 @@ class WorkerState: 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() 基准） @@ -149,7 +150,12 @@ async def init_pool(cfg: dict) -> None: """初始化 worker 池并启动健康检查后台任务。""" global _workers, _pool_condition, _health_task, _shutdown_event - _workers = {url: WorkerState(url=url) for url in cfg["workers"]} + # 主 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() @@ -179,7 +185,8 @@ async def init_pool(cfg: dict) -> None: # 对于 healthy_threshold == 1 的情况，已在上面的检查中上线 # 标记初始在线状态 online_count = sum(1 for w in _workers.values() if w.online) - logger.info("Worker 池初始化完成 | 总数=%d | 在线=%d", len(_workers), online_count) + logger.info("Worker 池初始化完成 | 总数=%d | 在线=%d | 主worker=%d", + len(_workers), online_count, n_primary) # 启动后台健康检查 _health_task = asyncio.create_task(_health_check_loop(cfg)) @@ -273,6 +280,8 @@ async def mark_worker_unhealthy(w: WorkerState) -> None: async def acquire_worker(cfg: dict) -> WorkerState: """从在线池中获取一个空闲 worker。 + 主备分组：只要还有主 worker（primary）在线，就只在主组内挑空闲 + （全忙则排队等待，不用备组）；只有主组全部下线时才用备用 worker。 全忙则排队等待，最多 queue_wait_seconds；池空立即抛异常。 """ queue_wait = cfg["dispatch"]["queue_wait_seconds"] @@ -280,11 +289,13 @@ async def acquire_worker(cfg: dict) -> WorkerState: async with _pool_condition: while True: - # 在线 + 空闲 - idle = [w for w in _workers.values() if w.online and not w.busy] + 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: - online = [w for w in _workers.values() if w.online] if not online: raise NoWorkerAvailable("后端服务暂不可用，请稍后重试") @@ -293,7 +304,8 @@ async def acquire_worker(cfg: dict) -> WorkerState: if remaining <= 0: raise NoWorkerAvailable("后端服务繁忙，请稍后重试") - logger.info("所有 worker 全忙（%d 在线），等待 %.1fs...", len(online), remaining) + logger.info("所有可用 worker 全忙（%d 在线），等待 %.1fs...", + len(candidates), remaining) try: await asyncio.wait_for(_pool_condition.wait(), timeout=remaining) except asyncio.TimeoutError: @@ -303,6 +315,8 @@ async def acquire_worker(cfg: dict) -> WorkerState: # 取第一个空闲 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 @@ -327,8 +341,13 @@ def get_pool_status() -> dict: "global_waiting": _global_waiting, } if not _workers: - return {"total": 0, "healthy": 0, "busy": 0, **global_info} + 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) - return {"total": total, "healthy": healthy, "busy": busy, **global_info} + 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} diff --git a/tests/test_gateway_pool.py b/tests/test_gateway_pool.py index 4490528..a4e0b61 100644 --- a/tests/test_gateway_pool.py +++ b/tests/test_gateway_pool.py @@ -152,3 +152,85 @@ 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"