diff --git a/app.py b/app.py index d4a4755..164fd4b 100644 --- a/app.py +++ b/app.py @@ -2,15 +2,20 @@ 接口 1（四庭七眼测量）已为**真实算法实现**（见 face_analysis 包）；接口 2~5 仍为 Mock。 拆分架构：worker 跑算法、返回 `annotated_image_base64`（不落盘、不拼 URL，由网关完成）。 -worker 对 `/api/*` 校验内网鉴权头 `X-Internal-Token`，`/health` 供网关探测不校验。 +worker 对 `/api/*` 校验内网鉴权头 `X-Internal-Token`，`/health` 供网关探测不校验 +（深检：磁盘水位 + 下游依赖 TCP 探测，异常返回 503 degraded，见 _load_health_cfg）。 """ +import asyncio import base64 import json import logging import os +import shutil +import time from contextlib import asynccontextmanager from io import BytesIO from typing import Any, List, Optional +from urllib.parse import urlsplit import cv2 import numpy as np @@ -1911,12 +1916,127 @@ async def download_hairline_log(rid: Optional[str] = None, tail: int = 500): # 健康检查 # --------------------------------------------------------------------------- +def _load_health_cfg() -> dict: + """深健康检查配置：环境变量优先，回退 worker_config.json 的 "health" 段。 + + min_free_gb=0 关闭磁盘检查；check_deps=false 关闭依赖探测（均退化为旧行为， + 即只看模型就绪标志）。 + """ + cfg_path = os.path.join(os.path.dirname(__file__), "worker_config.json") + file_cfg: dict = {} + if os.path.isfile(cfg_path): + try: + with open(cfg_path, encoding="utf-8") as f: + file_cfg = json.load(f).get("health") or {} + except Exception as e: # noqa: BLE001 + logger.warning("读取 worker_config.json health 段失败：%s", e) + + def _get(env_key, file_key, default): + raw = os.getenv(env_key) + return raw if raw not in (None, "") else file_cfg.get(file_key, default) + + def _as_float(value, default): + try: + return float(value) + except (TypeError, ValueError): + return default + + return { + "min_free_gb": _as_float(_get("HEALTH_MIN_FREE_GB", "min_free_gb", 5.0), 5.0), + "check_deps": str(_get("HEALTH_CHECK_DEPS", "check_deps", True)).strip().lower() + not in ("0", "false", "no", "off"), + "dep_comfyui": str(_get("HEALTH_DEP_COMFYUI", "dep_comfyui", "127.0.0.1:8188")), + "deps_cache_seconds": max(0.0, _as_float( + _get("HEALTH_DEPS_CACHE_SECONDS", "deps_cache_seconds", 5.0), 5.0)), + } + + +_HEALTH_CFG = _load_health_cfg() + + +def _disk_free_gb() -> float: + """worker 所在文件系统的剩余空间（GB）。""" + st = shutil.disk_usage(os.path.dirname(os.path.abspath(__file__))) + return st.free / (1 << 30) + + +def _parse_host_port(addr: str, default_port: int): + """解析 "host:port"（仅支持 IPv4/主机名形式）。""" + host, _, port = str(addr).partition(":") + try: + return host or "127.0.0.1", int(port) if port else default_port + except ValueError: + return host or "127.0.0.1", default_port + + +def _change_hair_addr(): + """从 _CHANGE_HAIR_BASE 解析 (host, port)，失败返回 None。""" + try: + u = urlsplit(_CHANGE_HAIR_BASE) + if u.hostname: + return u.hostname, u.port or 80 + except ValueError: + pass + return None + + +async def _tcp_probe(host: str, port: int, timeout: float = 0.8) -> bool: + """TCP 连通性探测：只判下游进程死活，不过 HTTP 认证、不受业务繁忙影响。""" + try: + _, writer = await asyncio.wait_for(asyncio.open_connection(host, port), timeout) + writer.close() + return True + except Exception: # noqa: BLE001 + return False + + +_DEP_CACHE = {"ts": -1.0, "detail": None} + + +async def _check_deps(): + """并发探测下游依赖（ComfyUI、change_hair），结果缓存 deps_cache_seconds 秒。""" + now = time.monotonic() + detail = _DEP_CACHE["detail"] + if detail is None or now - _DEP_CACHE["ts"] >= _HEALTH_CFG["deps_cache_seconds"]: + probes = {"comfyui": _tcp_probe(*_parse_host_port(_HEALTH_CFG["dep_comfyui"], 8188))} + ch = _change_hair_addr() + if ch: + probes["change_hair"] = _tcp_probe(*ch) + names = list(probes) + results = await asyncio.gather(*probes.values()) + detail = dict(zip(names, results)) + _DEP_CACHE["ts"], _DEP_CACHE["detail"] = now, detail + return all(detail.values()), detail + + @app.get("/health", include_in_schema=False) async def health(): - # 模型就绪后才返回 200，供网关探测；未就绪返回 503。 + """网关探测入口。 + + 模型未就绪 → 503 starting；磁盘剩余低于阈值 / 下游依赖 TCP 不通 → 503 degraded； + 正常 → 200（附 disk_free_gb 与 deps 明细，供人工排查）。深检自身异常时按就绪返回 + （健康检查不能自己成为故障点）。 + """ if not _STATE["ready"]: return JSONResponse(status_code=503, content={"status": "starting"}) - return {"status": "ok"} + try: + free_gb = _disk_free_gb() + body = {"status": "ok", "disk_free_gb": round(free_gb, 1)} + min_free = _HEALTH_CFG["min_free_gb"] + if min_free > 0 and free_gb < min_free: + body.update(status="degraded", reason=f"disk_free_below_{min_free}g") + return JSONResponse(status_code=503, content=body) + if _HEALTH_CFG["check_deps"]: + deps_ok, detail = await _check_deps() + body["deps"] = detail + if not deps_ok: + down = ",".join(k for k, v in detail.items() if not v) + body.update(status="degraded", reason=f"dep_down:{down}") + return JSONResponse(status_code=503, content=body) + return body + except Exception: # noqa: BLE001 + logger.exception("深健康检查异常，按就绪返回") + return {"status": "ok"} @app.get("/", include_in_schema=False) diff --git a/tests/test_api.py b/tests/test_api.py index 2e71dd5..f531688 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -36,12 +36,76 @@ def test_auth_missing_token_401(client): assert r.status_code == 401 -def test_health_no_token_200(client): +def _patch_health_ok(monkeypatch): + """让 /health 深检确定性通过：不依赖本机磁盘水位与 8188/8801 服务存活。""" + monkeypatch.setattr(app_module, "_HEALTH_CFG", { + **app_module._HEALTH_CFG, "min_free_gb": 0.0, "check_deps": True}) + + async def _deps_ok(): + return True, {"comfyui": True, "change_hair": True} + monkeypatch.setattr(app_module, "_check_deps", _deps_ok) + + +def test_health_no_token_200(client, monkeypatch): + _patch_health_ok(monkeypatch) r = client.get("/health") assert r.status_code == 200 assert r.json()["status"] == "ok" +def test_health_disk_low_503(client, monkeypatch): + # 阈值调到不可能的 99999G → 必然触发磁盘水位 503 + monkeypatch.setattr(app_module, "_HEALTH_CFG", { + **app_module._HEALTH_CFG, "min_free_gb": 99999.0, "check_deps": False}) + r = client.get("/health") + assert r.status_code == 503 + body = r.json() + assert body["status"] == "degraded" + assert body["reason"].startswith("disk_free_below") + + +def test_health_dep_down_503(client, monkeypatch): + monkeypatch.setattr(app_module, "_HEALTH_CFG", { + **app_module._HEALTH_CFG, "min_free_gb": 0.0, "check_deps": True}) + + async def _deps_down(): + return False, {"comfyui": False, "change_hair": True} + monkeypatch.setattr(app_module, "_check_deps", _deps_down) + r = client.get("/health") + assert r.status_code == 503 + body = r.json() + assert body["status"] == "degraded" + assert "comfyui" in body["reason"] + + +def test_health_ok_carries_detail(client, monkeypatch): + _patch_health_ok(monkeypatch) + r = client.get("/health") + assert r.status_code == 200 + body = r.json() + assert "disk_free_gb" in body + assert body["deps"] == {"comfyui": True, "change_hair": True} + + +def test_tcp_probe_real_socket(): + # 用进程内真实 socket 验证探测本身（曾因 open_connection 返回元组误判全部不可达） + import asyncio + + async def scenario(): + server = await asyncio.start_server(lambda r, w: None, host="127.0.0.1", port=0) + port = server.sockets[0].getsockname()[1] + up = await app_module._tcp_probe("127.0.0.1", port) + server.close() + await server.wait_closed() + # 端口 1 无监听，本机 connect 立即被拒 + down = await app_module._tcp_probe("127.0.0.1", 1) + return up, down + + up, down = asyncio.run(scenario()) + assert up is True + assert down is False + + def test_param_none_provided_1007(client): r = client.post(URL, headers=H) assert r.json()["code"] == 1007 diff --git a/worker_config.example.json b/worker_config.example.json index 483fc13..9a13911 100644 --- a/worker_config.example.json +++ b/worker_config.example.json @@ -1,5 +1,11 @@ { - "_comment": "复制为 worker_config.json（不入 git）。accept_passwords：worker 内网鉴权密码列表，网关带 X-Internal-Token: <其中之一>（也可用环境变量 WORKER_ACCEPT_PASSWORDS 覆盖）。comfyui_password：本机 ComfyUI(8188) Basic Auth 密码（user 默认 admin；也可放 password.txt 或环境变量 COMFYUI_PASSWORD）。注：接口4 火山方舟 ark_api_key 已迁到网关配置。", + "_comment": "复制为 worker_config.json（不入 git）。accept_passwords：worker 内网鉴权密码列表，网关带 X-Internal-Token: <其中之一>（也可用环境变量 WORKER_ACCEPT_PASSWORDS 覆盖）。comfyui_password：本机 ComfyUI(8188) Basic Auth 密码（user 默认 admin；也可放 password.txt 或环境变量 COMFYUI_PASSWORD）。health：深健康检查（/health 磁盘水位+下游依赖 TCP 探测，均可用环境变量 HEALTH_MIN_FREE_GB / HEALTH_CHECK_DEPS / HEALTH_DEP_COMFYUI / HEALTH_DEPS_CACHE_SECONDS 覆盖）。注：接口4 火山方舟 ark_api_key 已迁到网关配置。", "accept_passwords": ["change-me-to-a-strong-secret"], - "comfyui_password": "your-comfyui-basic-auth-password" + "comfyui_password": "your-comfyui-basic-auth-password", + "health": { + "min_free_gb": 5, + "check_deps": true, + "dep_comfyui": "127.0.0.1:8188", + "deps_cache_seconds": 5 + } }