diff --git a/.env.example b/.env.example index 4a90c50c..deb6cdd2 100644 --- a/.env.example +++ b/.env.example @@ -137,6 +137,12 @@ WINDUP_SSE_REDIS_CHANNEL=windup:pubsub:generation-task-events # WINDUP_MATTE_REFINE=0 # 浅肤色脸颊/小腿可能漏检。ORT CPU arena 已在代码里关闭,不必再配。 WINDUP_MATTE_REFINE=1 + +# 抠图 provider。u2net(默认)/ birefnet。 +# birefnet 在同一帧上把主体内部非实心从 5541px 降到 106px,但**单帧峰值 6.85GB** —— +# 生产 worker 容器上限 5GiB、宿主 7.7GB,**服务器上不要开**,开了会 OOM。 +# 它是给组员本机(16GB+)做高质量素材用的。取值拼错会回落 u2net 并留一条 WARNING。 +WINDUP_MATTE_PROVIDER=u2net # web / worker 各有一份连接池。生成已短 session，不必再按 handler 并发配满。 POSTGRES_POOL_SIZE=5 POSTGRES_MAX_OVERFLOW=10 diff --git a/backend/packages/app/src/windup_app/bootstrap/worker.py b/backend/packages/app/src/windup_app/bootstrap/worker.py index 578050b1..5544538f 100644 --- a/backend/packages/app/src/windup_app/bootstrap/worker.py +++ b/backend/packages/app/src/windup_app/bootstrap/worker.py @@ -143,9 +143,11 @@ def _warmup_local_inference() -> None: except Exception: logger.warning("抽帧后端预热失败", exc_info=True) try: - from windup_framework.providers import OnnxU2NetMatteProvider + from windup_framework.providers import get_matte_provider - matte = OnnxU2NetMatteProvider() + # 取进程里那唯一一份 —— 三个 executor 的惰性兜底取到的是同一个对象, + # 所以就算下面 warmup() 抛了、bind_matte 没跑成,也不会各建一份。 + matte = get_matte_provider() matte.warmup() bind_matte(matte) logger.info("ONNX 抠图会话已预热") diff --git a/backend/packages/app/src/windup_app/server/orchestrator/executor.py b/backend/packages/app/src/windup_app/server/orchestrator/executor.py index f70a8f11..cbb876b7 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/executor.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/executor.py @@ -899,12 +899,12 @@ def _assemble(self, directions: int) -> CharacterGeneratorPort: from windup_common.models import GenRoute from windup_framework.gateway import build_image_gateway, build_video_gateway from windup_framework.gateway.image import _CIRCUIT - from windup_framework.providers import OnnxU2NetMatteProvider + from windup_framework.providers import get_matte_provider from windup_app.server.media.first_frame import MediaFirstFrameUploader if self._matte is None: - self._matte = OnnxU2NetMatteProvider() + self._matte = get_matte_provider() if self._image is None: self._image = build_image_gateway(circuit=_CIRCUIT) # uploader 在这里注入而不是让 provider 自己去拿:framework 不认识 app 的对象存储。 @@ -1025,7 +1025,7 @@ def __init__( self, *, image=None, # None → 懒加载 ImageGateway - matte: MatteProvider | None = None, # None → 懒加载 OnnxU2NetMatteProvider + matte: MatteProvider | None = None, # None → 懒加载,按 WINDUP_MATTE_PROVIDER 选 upload: Callable[[bytes], str] | None = None, # None → 真实对象存储上传 fetch_ref: Callable[[str], bytes] | None = None, # None → 下载 reference_image_url @@ -1236,9 +1236,9 @@ def _get_matte(self): return self._matte with self._assembly_lock: if self._matte is None: - from windup_framework.providers import OnnxU2NetMatteProvider + from windup_framework.providers import get_matte_provider - self._matte = OnnxU2NetMatteProvider() + self._matte = get_matte_provider() return self._matte def _download(self, url: str) -> bytes: diff --git a/backend/packages/app/src/windup_app/server/orchestrator/view_sheet_executor.py b/backend/packages/app/src/windup_app/server/orchestrator/view_sheet_executor.py index 332cf302..9790b565 100644 --- a/backend/packages/app/src/windup_app/server/orchestrator/view_sheet_executor.py +++ b/backend/packages/app/src/windup_app/server/orchestrator/view_sheet_executor.py @@ -338,9 +338,9 @@ def _get_matte(self): return self._matte with self._assembly_lock: if self._matte is None: - from windup_framework.providers import OnnxU2NetMatteProvider + from windup_framework.providers import get_matte_provider - self._matte = OnnxU2NetMatteProvider() + self._matte = get_matte_provider() return self._matte def _download(self, url: str) -> bytes: diff --git a/backend/packages/framework/src/windup_framework/providers/__init__.py b/backend/packages/framework/src/windup_framework/providers/__init__.py index cf6f2de7..37c91f89 100644 --- a/backend/packages/framework/src/windup_framework/providers/__init__.py +++ b/backend/packages/framework/src/windup_framework/providers/__init__.py @@ -13,6 +13,11 @@ ) from windup_framework.providers.matte import OnnxU2NetMatteProvider from windup_framework.providers.matte_birefnet import BiRefNetMatteProvider +from windup_framework.providers.matte_factory import ( + get_matte_provider, + make_matte_provider, + reset_matte_provider, +) from windup_framework.providers.sufy import ( SufyImageProvider, SufyVideoProvider, @@ -34,6 +39,9 @@ # FAL 队列面的 i2v(现役接口形态);首帧要公网 URL,故与 uploader 成对出现 "SufyImageProvider", "BiRefNetMatteProvider", + "get_matte_provider", + "make_matte_provider", + "reset_matte_provider", "OnnxU2NetMatteProvider", # Gateway 工厂(executor 从 windup_framework.gateway 取;此处再导出方便装配) "bind_call_context", diff --git a/backend/packages/framework/src/windup_framework/providers/matte_birefnet.py b/backend/packages/framework/src/windup_framework/providers/matte_birefnet.py index f7650cdd..0cdf028c 100644 --- a/backend/packages/framework/src/windup_framework/providers/matte_birefnet.py +++ b/backend/packages/framework/src/windup_framework/providers/matte_birefnet.py @@ -25,6 +25,7 @@ from __future__ import annotations import io +import threading import logging from pathlib import Path @@ -32,6 +33,7 @@ from PIL import Image from .interfaces import MatteProvider +from .matte import _RUN_LOCK from .matte import _download_atomic, _fill_enclosed_holes, _flat_bg_penalty logger = logging.getLogger("windup.matte.birefnet") @@ -66,6 +68,7 @@ def __init__(self, model_path: str | Path | None = None, union_with_u2net: bool self._path = Path(model_path) if model_path else _CACHE self._session = None self._union = union_with_u2net + self._init_lock = threading.Lock() self._u2net = None def _ensure_model(self) -> Path: @@ -83,13 +86,57 @@ def _get_session(self): ) return self._session + def warmup(self) -> None: + """把会话装进内存。**必须有这个方法** —— 没有它整条接线会静默失效。 + + ``bootstrap.worker`` 是 ``matte.warmup()`` 然后 ``bind_matte(matte)``,两句包在 + 同一个 ``except Exception`` 里。缺 ``warmup`` 时第一句抛 AttributeError, + **``bind_matte`` 就到不了** —— 于是三个 executor 各自惰性 new 一份 provider, + 而本类默认 ``union_with_u2net=True``,每份内部再 new 一个 u2net,进程里就是 + 6 个 ONNX 会话。生产 worker 容器上限 5GiB,而本模型单帧峰值 6.85GB。 + 表面上只有一条 "ONNX 预热失败" 的 WARNING,开发机上完全跑得通。 + + 并集那一路的 u2net 也一并预热:它是每帧都要跑的,留到首帧再装等于把两次冷启动 + 叠在一起。 + """ + self._get_session() + if self._union: + self._ensure_u2net().warmup() + + def _ensure_u2net(self): + """并集那一路的 u2net,**带锁的惰性初始化**。 + + 原先是 ``if self._u2net is None: self._u2net = ...`` 的检查后赋值,两个线程能 + 同时通过检查、各建一份 —— 而每份自带一套 ONNX 会话。worker 的生成并发是 + IMAGE=4 / ACTION=2,这个竞态在生产上是够得着的。 + """ + with self._init_lock: + if self._u2net is None: + from .matte import OnnxU2NetMatteProvider + + self._u2net = OnnxU2NetMatteProvider() + return self._u2net + def _predict_mask(self, img: Image.Image) -> Image.Image: session = self._get_session() arr = np.asarray( img.convert("RGB").resize((_SIDE, _SIDE), Image.LANCZOS), dtype=np.float32 ) tensor = (((arr / 255.0) - _MEAN) / _STD).transpose(2, 0, 1)[None] - outputs = session.run(None, {session.get_inputs()[0].name: tensor.astype(np.float32)}) + # **与 u2net 共用同一把锁**,让进程里同时只有一个 ONNX 前向在跑。 + # + # 这一路的单帧峰值是 6.85GB(u2net 是 1.10GB),而 worker 的生成并发是 + # IMAGE=4 / ACTION=2 —— 不串行的话四路并发前向就是四份激活值同时在内存里, + # 谁都装不下。u2net 那边早就这么做了(见 matte._RUN_LOCK 上方那段注释: + # "进程内串行 Run,吞吐靠多 worker 进程而不是同一进程里叠会话"),本类漏了。 + # + # 共用而不是各持一把:并集模式下两个模型每帧都要各跑一次,各持一把锁只能保证 + # "同类不并发",跨类照样叠 6.85+1.10。锁的作用域只包住 run 本身, + # 下面调 u2net 时这把锁已经释放,不会自锁。 + with _RUN_LOCK: + outputs = session.run( + None, {session.get_inputs()[0].name: tensor.astype(np.float32)} + ) # 多尺度监督的导出会给一串输出,最后一个是最高分辨率的那张;取 [0] 会拿到 # 1/32 尺度的粗图,放大回来就是一团模糊。 raw = np.asarray(outputs[-1] if isinstance(outputs, list) else outputs).squeeze() @@ -105,12 +152,8 @@ def cutout(self, frame: bytes) -> bytes: rgb = np.asarray(img.convert("RGB"), dtype=np.float32) alpha = np.asarray(self._predict_mask(img), dtype=np.float32) / 255.0 if self._union: - if self._u2net is None: - from .matte import OnnxU2NetMatteProvider - - self._u2net = OnnxU2NetMatteProvider() other = np.asarray( - Image.open(io.BytesIO(self._u2net.cutout(frame))).convert("RGBA") + Image.open(io.BytesIO(self._ensure_u2net().cutout(frame))).convert("RGBA") )[:, :, 3].astype(np.float32) / 255.0 alpha = np.maximum(alpha, other) alpha = alpha * _flat_bg_penalty(rgb) diff --git a/backend/packages/framework/src/windup_framework/providers/matte_factory.py b/backend/packages/framework/src/windup_framework/providers/matte_factory.py new file mode 100644 index 00000000..dd00fef5 --- /dev/null +++ b/backend/packages/framework/src/windup_framework/providers/matte_factory.py @@ -0,0 +1,146 @@ +"""按配置选抠图 provider。 + +为什么要这个开关:BiRefNet 在同一帧上把主体内部非实心从 5,541 px 降到 106 px(-98%), +但单帧峰值 6.85GB —— 生产 worker 容器上限 5GiB、宿主总共 7.7GB,**跑不了**,而组员 +本机 16GB 跑得很轻松。同一份代码两种装配,好过为它开一条分支:分支一定会漂,而漂出来 +的"更好的管线"产出的素材,产品复现不出来。 + +默认必须是 u2net:忘配等于用得起的那个,而不是忘配就把生产打 OOM。 +""" + +from __future__ import annotations + +import logging +import os +import pathlib +import threading + +from .interfaces import MatteProvider + +logger = logging.getLogger("windup.matte.factory") + +#: 环境变量名。取值 ``u2net``(默认) / ``birefnet``。 +ENV = "WINDUP_MATTE_PROVIDER" +_U2NET = "u2net" +_BIREFNET = "birefnet" + +#: BiRefNet 跑得起来所需的内存下限。 +#: +#: 实测(部署机 4 核 / 7.7GiB / 无 GPU / ORT 1.23.2,输入形状与实现一致):单帧前向峰值 +#: **6.85GB**,而 worker 进程自身稳态就占 1.7–3.9GB,且本 provider 默认与 u2net 取并集 +#: (两份 ONNX 会话同时在)。cgroup 2g / 3g / 4g 实测全部 rc=137。 +#: 8GiB 是"6.85 峰值 + worker 自身"取整,不是拍的。 +_BIREFNET_MIN_BYTES = 8 * 1024**3 + +#: cgroup v2 / v1 的内存上限文件。**必须读它,不能只读 /proc/meminfo** —— +#: OOM killer 按 cgroup 判,而容器里 ``/proc/meminfo`` 报的是宿主的总量: +#: 生产宿主 7.7GiB、worker 容器上限 5GiB,只看前者会得出"够用"的错结论。 +_CGROUP_MAX = ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes") + + +def _read_int(path: str) -> int | None: + try: + raw = pathlib.Path(path).read_text(encoding="utf-8").strip() + except OSError: + return None + if raw == "max": # cgroup v2 的"不限" + return None + try: + value = int(raw) + except ValueError: + return None + # cgroup v1 用一个接近 2^63 的哨兵表示"不限",别把它当成真上限。 + return None if value <= 0 or value >= 2**62 else value + + +def memory_budget_bytes() -> int | None: + """这个进程实际能用到的内存上限;判不出来时给 None。 + + 取 cgroup 上限与宿主总量的较小者。``None`` 表示两处都读不到(非 Linux 的开发机), + 这时**不拦** —— 拦一台判不出内存的机器,等于因为量不到就把功能关掉。 + """ + limits = [v for p in _CGROUP_MAX if (v := _read_int(p)) is not None] + try: + for line in pathlib.Path("/proc/meminfo").read_text(encoding="utf-8").splitlines(): + if line.startswith("MemTotal:"): + limits.append(int(line.split()[1]) * 1024) + break + except (OSError, ValueError, IndexError): + pass + return min(limits) if limits else None + + +def make_matte_provider(name: str | None = None) -> MatteProvider: + """按名字造 provider;不认识的名字回落 u2net 并留一条 WARNING。 + + 不认识就抛错的话,一个拼错的环境变量会让整个 worker 起不来;而回落是安全方向 —— + u2net 在任何机器上都跑得起来,坏处只是抠图差一点,且这条 WARNING 说明了原因。 + """ + choice = (name or os.environ.get(ENV) or _U2NET).strip().lower() + if choice == _BIREFNET: + budget = memory_budget_bytes() + if budget is not None and budget < _BIREFNET_MIN_BYTES: + # 回落而不是抛错,理由与下面那条不认识的值一样:抛错会让一次配置失误变成 + # 整个 worker 起不来。但级别是 ERROR 不是 WARNING —— 它是被明确要求了却 + # 没能给上,而不是没配。 + # + # 这道闸拦的坏例:在 5GiB 的生产 worker 上把这个变量设成 birefnet。 + # 没有它时,进程会一路跑到第一帧推理才被 OOM killer 杀掉,而现场看到的是 + # worker 无声重启、任务卡在 RUNNING —— 没有任何一行提到内存。 + logger.error( + "%s=%s 但本进程内存上限只有 %.1fGiB(需要 ≥%.0fGiB),回落 u2net。" + "BiRefNet 单帧峰值约 6.85GB,在这里开会被 OOM kill 而不是抠得差一点。" + "要用它请换内存更大的机器或调高容器上限。", + ENV, _BIREFNET, budget / 1024**3, _BIREFNET_MIN_BYTES / 1024**3, + ) + else: + from .matte_birefnet import BiRefNetMatteProvider + + logger.info( + "抠图用 BiRefNet(与 u2net 取并集);单帧峰值约 6.85GB,内存上限 %s", + f"{budget / 1024**3:.1f}GiB" if budget else "判不出(非 Linux?)", + ) + return BiRefNetMatteProvider() + if choice != _U2NET: + logger.warning("%s=%r 不认识,回落 u2net;可选:%s / %s", ENV, choice, _U2NET, _BIREFNET) + from .matte import OnnxU2NetMatteProvider + + return OnnxU2NetMatteProvider() + + +#: 进程级唯一实例。**这是硬保证,不是优化。** +#: +#: 每份 provider 自带一套 ONNX 会话:u2net 常驻 ~0.53GB、BiRefNet ~0.45GB,而并集模式 +#: 下 BiRefNet 内部再挂一个 u2net。生产 worker 容器上限 5GiB —— 多几份就不是"稍微费点 +#: 内存",是起不来。 +#: +#: 此前唯一性靠的是运气:``bootstrap.worker`` 建一份 → ``warmup()`` → ``bind_matte()`` +#: 把它塞给三个 executor。但那三个 executor 各自还有一条惰性兜底 +#: (``if self._matte is None: self._matte = make_matte_provider()``),而 ``warmup()`` +#: 一抛异常 ``bind_matte`` 就不执行 —— 于是三个 executor 各建一份,日志里只有一条 +#: "ONNX 预热失败,首个抠图任务会再加载" 的 WARNING,没有任何一处说"你现在有三份"。 +#: 那条惰性路径本身也不是线程安全的(检查后赋值),并发下还能更多。 +_INSTANCE: MatteProvider | None = None +_INSTANCE_LOCK = threading.Lock() + + +def get_matte_provider() -> MatteProvider: + """进程里那**唯一**一份抠图 provider。所有生产调用点都走这里。 + + 与 :func:`make_matte_provider` 的分工:那个是**工厂**(每调一次造一份,给测试和 + 本地脚本用),这个是**取唯一那份**。生产代码一律用这个 —— 工厂直接暴露给调用点, + 唯一性就只能靠每个调用点自己守规矩,而实测已经证明守不住。 + """ + global _INSTANCE + if _INSTANCE is None: + with _INSTANCE_LOCK: + if _INSTANCE is None: # 双检:锁外读、锁内再确认 + _INSTANCE = make_matte_provider() + return _INSTANCE + + +def reset_matte_provider() -> None: + """丢掉那份单例。**只给测试用** —— 生产里换 provider 要重启进程。""" + global _INSTANCE + with _INSTANCE_LOCK: + _INSTANCE = None diff --git a/backend/tests/test_custom_action.py b/backend/tests/test_custom_action.py index 662cc6c3..88e56aee 100644 --- a/backend/tests/test_custom_action.py +++ b/backend/tests/test_custom_action.py @@ -264,7 +264,6 @@ def test_concurrent_first_requests_build_one_shared_provider_set(monkeypatch): 选哪个 kling 是 Gateway 的事,不同 video_model 仍共用同一个 generator。 """ import windup_framework.gateway as gateway - from windup_framework import providers from windup_app.server.orchestrator.executor import ActionTaskExecutor @@ -279,7 +278,13 @@ def _factory(*_args, **_kwargs): return object() return _factory - monkeypatch.setattr(providers, "OnnxU2NetMatteProvider", _counting("matte")) + # 桩在**工厂**上并清掉单例:生产已改走 get_matte_provider()(它内部调工厂)。 + # 桩 get_matte_provider 本身会把"单例只建一份"那层绕过去,而那正是这条要守的东西 —— + # 现在这条断言的是"连工厂都只被调一次"，比之前更强。 + from windup_framework.providers import matte_factory as _mf + + monkeypatch.setattr(_mf, "make_matte_provider", _counting("matte")) + _mf.reset_matte_provider() monkeypatch.setattr(gateway, "build_image_gateway", _counting("image")) monkeypatch.setattr(gateway, "build_video_gateway", _counting("video")) @@ -319,9 +324,9 @@ def test_injected_matte_is_not_replaced_on_assemble(monkeypatch): sent = object() def _boom(*_a, **_k): - raise AssertionError("不该再构造 OnnxU2NetMatteProvider") + raise AssertionError("不该再构造抠图 provider") - monkeypatch.setattr(providers, "OnnxU2NetMatteProvider", _boom) + monkeypatch.setattr(providers, "make_matte_provider", _boom) monkeypatch.setattr(gateway, "build_image_gateway", lambda **_k: object()) monkeypatch.setattr(gateway, "build_video_gateway", lambda **_k: object()) diff --git a/backend/tests/test_env_example_keys_are_live.py b/backend/tests/test_env_example_keys_are_live.py index 4659e480..ce6b6ae8 100644 --- a/backend/tests/test_env_example_keys_are_live.py +++ b/backend/tests/test_env_example_keys_are_live.py @@ -34,6 +34,10 @@ "WINDUP_SSE_REDIS_CHANNEL", # windup_framework.providers.matte._refine_enabled:os.environ,不走 BaseSettings "WINDUP_MATTE_REFINE", + # providers.matte_factory.make_matte_provider:同上,os.environ。 + # 走 BaseSettings 的话 framework 的配置层就要认识 provider 的名字,而选哪个 provider + # 是装配决定,不是配置数据。 + "WINDUP_MATTE_PROVIDER", # render3d._tc3.TencentCredentials.resolve:环境变量 → 加锁文件,不走 BaseSettings "TENCENT_SECRET_ID", "TENCENT_SECRET_KEY", "TENCENT_REGION", # orchestrator.client_bake / docker-compose 的构建目标,都是 os.getenv diff --git a/backend/tests/test_matte_provider_choice.py b/backend/tests/test_matte_provider_choice.py new file mode 100644 index 00000000..e26479e8 --- /dev/null +++ b/backend/tests/test_matte_provider_choice.py @@ -0,0 +1,378 @@ +"""抠图 provider 的选择开关。 + +这个开关存在的理由:BiRefNet 单帧峰值 6.85GB,生产 worker 上限 5GiB —— 开错方向的代价 +不是"抠图差一点",是把 worker 打 OOM。所以默认值与回落方向都要有用例钉住。 +""" +from __future__ import annotations + +import io +import time + +from PIL import Image + +import pytest + +from windup_framework.providers import make_matte_provider +from windup_framework.providers.matte import OnnxU2NetMatteProvider +from windup_framework.providers.matte_factory import ENV + + +def test_default_is_u2net_so_a_missing_env_cannot_oom_the_worker(monkeypatch): + """拦的坏例:默认值给成 BiRefNet。 + + 忘配就该拿到跑得起来的那个。反过来的话,一台没设这个变量的机器会在第一帧抠图时 + 被 OOM kill,而表现是 worker 无声重启、任务卡在 RUNNING。 + """ + monkeypatch.delenv(ENV, raising=False) + assert isinstance(make_matte_provider(), OnnxU2NetMatteProvider) + + +def test_an_unknown_value_falls_back_instead_of_killing_the_worker(monkeypatch): + """拦的坏例:不认识的值直接抛错。 + + 一个拼错的环境变量(``bierfnet``)会让整个 worker 起不来,而回落只是抠图差一点。 + 两个方向的代价不对称。 + """ + monkeypatch.setenv(ENV, "bierfnet") + assert isinstance(make_matte_provider(), OnnxU2NetMatteProvider) + + +def test_explicit_argument_beats_the_environment(monkeypatch): + """显式传参优先于环境变量 —— 否则测试与本地脚本没法覆盖部署的设置。""" + monkeypatch.setenv(ENV, "birefnet") + assert isinstance(make_matte_provider("u2net"), OnnxU2NetMatteProvider) + + +def test_birefnet_is_reachable_by_name_not_dead_code(monkeypatch): + """拦的坏例:provider 合进仓里却没有任何路径能选到它(#686 就是这么变成死代码的)。 + + 只断言"造出来的是那个类",不真跑推理:权重 224MB,CI 上不该下载,而这条要证明的是 + **接线通了**,不是模型好不好。 + """ + pytest.importorskip("onnxruntime") + from windup_framework.providers.matte_birefnet import BiRefNetMatteProvider + + monkeypatch.setenv(ENV, "birefnet") + assert isinstance(make_matte_provider(), BiRefNetMatteProvider) + + +def test_every_selectable_provider_survives_the_bootstrap_warmup_call(monkeypatch): + """拦的坏例:某个 provider 缺 ``warmup``,整条共享接线静默失效。 + + ``bootstrap.worker`` 是 ``matte.warmup()`` 然后 ``bind_matte(matte)``,两句包在同一个 + ``except Exception`` 里。缺 ``warmup`` 时第一句抛 AttributeError,``bind_matte`` + **就到不了** —— 三个 executor 各自惰性 new 一份,而 BiRefNet 默认与 u2net 取并集, + 每份内部再 new 一个 u2net,进程里 6 个 ONNX 会话。生产 worker 上限 5GiB, + BiRefNet 单帧峰值 6.85GB。表面上只有一条 "ONNX 预热失败" 的 WARNING。 + (FennoAI 式审查在 #823 上指出;本用例把它钉住。) + + 断言的是**协议齐全**,不真跑推理:权重 224MB + 176MB,CI 上不该下载。 + """ + from windup_framework.providers import make_matte_provider + from windup_framework.providers.matte_factory import ENV, _BIREFNET, _U2NET + + for choice in (_U2NET, _BIREFNET): + monkeypatch.setenv(ENV, choice) + provider = make_matte_provider() + assert callable(getattr(provider, "warmup", None)), ( + f"{type(provider).__name__} 缺 warmup —— bind_matte 会被跳过," + "共享实例失效,进程里会装多份 ONNX 会话" + ) + assert callable(getattr(provider, "cutout", None)) + + +# ── 内存闸 ─────────────────────────────────────────────────────────────── +# +# 上面几条钉的是"选哪个",这几条钉的是"选了但机器扛不住"。两者的失败长相完全不同: +# 前者拿到的是跑得起来的 provider,后者拿到的是一个会在第一帧推理时被 OOM killer +# 杀掉的进程 —— 而现场只看得到 worker 无声重启、任务卡在 RUNNING。 + + +def _no_cgroup(monkeypatch, host_gib: float | None): + """把内存探测换成一台指定大小的机器;None = 判不出(非 Linux 开发机)。""" + from windup_framework.providers import matte_factory as f + + monkeypatch.setattr( + f, "memory_budget_bytes", + lambda: None if host_gib is None else int(host_gib * 1024**3), + ) + + +def test_birefnet_on_the_production_worker_falls_back_instead_of_being_oom_killed( + monkeypatch, caplog +): + """拦的坏例:在生产 worker 上把这个变量设成 birefnet。 + + 5GiB 是线上实测值(``/sys/fs/cgroup/memory.max`` = 5368709120),而 BiRefNet 单帧 + 峰值 6.85GB。没有这道闸时,进程一路跑到第一帧推理才被杀,没有任何一行提到内存。 + """ + monkeypatch.setenv(ENV, "birefnet") + _no_cgroup(monkeypatch, 5.0) + with caplog.at_level("ERROR"): + provider = make_matte_provider() + assert isinstance(provider, OnnxU2NetMatteProvider) + # 光回落不够 —— 得说清是内存不够,否则运维只会看到"抠图怎么还是旧的"。 + assert any("内存上限" in r.getMessage() for r in caplog.records), caplog.text + + +def test_a_developer_laptop_still_gets_birefnet(monkeypatch): + """反方向:16GiB 的本机必须拿得到它。 + + 闸门只该拦扛不住的机器。把组员本机也一起拦掉,这个 provider 就又变回死代码 —— + 而 #823 存在的全部理由就是让它别是死代码。 + """ + pytest.importorskip("onnxruntime") + from windup_framework.providers.matte_birefnet import BiRefNetMatteProvider + + monkeypatch.setenv(ENV, "birefnet") + _no_cgroup(monkeypatch, 16.0) + assert isinstance(make_matte_provider(), BiRefNetMatteProvider) + + +def test_an_unmeasurable_machine_is_not_blocked(monkeypatch): + """判不出内存时放行。 + + macOS 上没有 cgroup 也没有 /proc/meminfo。因为量不到就把功能关掉,是拿"我不知道" + 当"不行"用。 + """ + pytest.importorskip("onnxruntime") + from windup_framework.providers.matte_birefnet import BiRefNetMatteProvider + + monkeypatch.setenv(ENV, "birefnet") + _no_cgroup(monkeypatch, None) + assert isinstance(make_matte_provider(), BiRefNetMatteProvider) + + +def test_the_gate_reads_the_cgroup_limit_not_just_host_memory(monkeypatch, tmp_path): + """拦的坏例:只读 ``/proc/meminfo`` 就下结论。 + + 容器里那个文件报的是**宿主**的总量:生产宿主 7.7GiB、worker 容器上限 5GiB。 + 只看宿主会得出"7.7 也不够、正好也拦住了"——碰巧对,但换一台 32GiB 宿主 + 5GiB + 容器的机器就会放行,然后被 OOM kill。取两者较小值才是对的。 + """ + from windup_framework.providers import matte_factory as f + + cg = tmp_path / "memory.max" + cg.write_text("5368709120") # 容器 5GiB —— 线上实测值 + meminfo = tmp_path / "meminfo" + meminfo.write_text("MemTotal: 33554432 kB\n") # 宿主 32GiB + monkeypatch.setattr(f, "_CGROUP_MAX", (str(cg),)) + monkeypatch.setattr(f.pathlib, "Path", f.pathlib.Path) # 保持真实实现 + orig = f.pathlib.Path + + def _fake(arg): + return orig(str(meminfo)) if str(arg) == "/proc/meminfo" else orig(arg) + + monkeypatch.setattr(f.pathlib, "Path", _fake) + assert f.memory_budget_bytes() == 5368709120, "取的不是较小值" + + +@pytest.mark.parametrize("sentinel", ["max", "9223372036854771712"]) +def test_an_unlimited_cgroup_is_not_read_as_a_tiny_limit(monkeypatch, tmp_path, sentinel): + """拦的坏例:把"不限"的哨兵值当成真上限。 + + cgroup v2 用字面量 ``max``,v1 用一个接近 2^63 的数。前者解析失败会拿到 None + (碰巧安全),后者会得出一个天文数字的"上限"然后放行 —— 而真正该看的是宿主总量。 + """ + from windup_framework.providers import matte_factory as f + + cg = tmp_path / "memory.max" + cg.write_text(sentinel) + assert f._read_int(str(cg)) is None + + +# ── 并发安全 ───────────────────────────────────────────────────────────── + + +def test_birefnet_forward_passes_never_overlap(monkeypatch): + """拦的坏例:并发前向叠内存。 + + BiRefNet 单帧峰值 6.85GB,而 worker 的生成并发是 IMAGE=4 / ACTION=2。不串行的话 + 四路并发就是四份激活值同时在内存里 —— 而单份就已经超过容器上限。u2net 那边一直 + 有这把锁(``matte._RUN_LOCK``),本类此前漏了。 + + 判据是**实测重叠数**,不是"代码里有没有 Lock 字样":后者在把锁挪错位置时照样绿。 + """ + import threading + + from windup_framework.providers import matte_birefnet as mb + + live = 0 + peak = 0 + guard = threading.Lock() + + class _FakeSession: + def get_inputs(self): + return [type("I", (), {"name": "x"})()] + + def run(self, _out, _feed): + nonlocal live, peak + with guard: + live += 1 + peak = max(peak, live) + time.sleep(0.02) # 让并发真的有机会重叠 + with guard: + live -= 1 + import numpy as np + return [np.zeros((1, 1, 8, 8), dtype="float32")] + + prov = mb.BiRefNetMatteProvider(union_with_u2net=False) + monkeypatch.setattr(prov, "_get_session", lambda: _FakeSession()) + + img = Image.new("RGB", (16, 16), (120, 120, 120)) + buf = io.BytesIO() + img.save(buf, format="PNG") + png = buf.getvalue() + + threads = [threading.Thread(target=prov.cutout, args=(png,)) for _ in range(4)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert peak == 1, f"实测同时有 {peak} 个前向在跑;每个峰值 6.85GB,叠起来必 OOM" + + +def test_the_union_u2net_is_created_once_under_concurrency(monkeypatch): + """拦的坏例:并集那一路的惰性初始化是检查后赋值,两个线程各建一份。 + + 每份自带一套 ONNX 会话。worker 并发 4,这个竞态在生产上够得着。 + """ + import threading + + from windup_framework.providers import matte_birefnet as mb + + made = [] + + class _FakeU2Net: + def __init__(self): + made.append(1) + time.sleep(0.01) # 放大竞态窗口 + + def warmup(self): + pass + + monkeypatch.setattr( + "windup_framework.providers.matte.OnnxU2NetMatteProvider", _FakeU2Net + ) + prov = mb.BiRefNetMatteProvider(union_with_u2net=True) + threads = [threading.Thread(target=prov._ensure_u2net) for _ in range(6)] + for t in threads: + t.start() + for t in threads: + t.join() + assert len(made) == 1, f"并发下建了 {len(made)} 份 u2net,每份一套 ONNX 会话" + + +# ── 进程里只能有一份 ───────────────────────────────────────────────────── + + +def test_every_caller_gets_the_same_single_instance(monkeypatch): + """拦的坏例:进程里存在第二份 provider。 + + 每份自带一套 ONNX 会话(u2net 常驻 ~0.53GB,并集模式下 BiRefNet 内部再挂一个 + u2net)。生产 worker 容器上限 5GiB —— 多几份不是"稍微费点内存",是起不来。 + """ + from windup_framework.providers import get_matte_provider, reset_matte_provider + + monkeypatch.delenv(ENV, raising=False) + reset_matte_provider() + try: + a, b, c = get_matte_provider(), get_matte_provider(), get_matte_provider() + assert a is b is c + finally: + reset_matte_provider() + + +def test_a_failed_warmup_does_not_triple_the_instances(monkeypatch): + """拦的坏例:``warmup()`` 抛异常 → ``bind_matte`` 不执行 → 三个 executor 各建一份。 + + 这正是此前唯一性靠运气的那条路径:``bootstrap.worker`` 把 warmup + bind_matte 包在 + 同一个 ``except Exception`` 里,warmup 一抛,bind_matte 就到不了,而三个 executor + 各自还有惰性兜底。日志里只有一条"ONNX 预热失败,首个抠图任务会再加载"的 WARNING, + 没有任何一处说"你现在有三份"。 + """ + from windup_framework.providers import get_matte_provider, reset_matte_provider + from windup_framework.providers import matte_factory as mf + + built = [] + + class _Boom: + def __init__(self): + built.append(1) + + def warmup(self): + raise RuntimeError("模型下载失败") + + def cutout(self, frame): + return frame + + monkeypatch.setattr(mf, "make_matte_provider", lambda *a, **k: _Boom()) + reset_matte_provider() + try: + # 模拟 bootstrap:拿一份、预热炸了 + p0 = get_matte_provider() + with pytest.raises(RuntimeError): + p0.warmup() + # 三个 executor 各自走惰性兜底 + got = [get_matte_provider() for _ in range(3)] + assert all(g is p0 for g in got), "预热失败后又建了新的实例" + assert len(built) == 1, f"进程里建了 {len(built)} 份 provider" + finally: + reset_matte_provider() + + +def test_concurrent_first_calls_build_exactly_one(monkeypatch): + """拦的坏例:惰性初始化是检查后赋值,并发下多建几份。 + + worker 的生成并发默认 IMAGE=4 / ACTION=2,首个抠图任务并发到达是常态。 + """ + import threading + + from windup_framework.providers import get_matte_provider, reset_matte_provider + from windup_framework.providers import matte_factory as mf + + built = [] + + class _Slow: + def __init__(self): + built.append(1) + time.sleep(0.02) # 放大竞态窗口 + + monkeypatch.setattr(mf, "make_matte_provider", lambda *a, **k: _Slow()) + reset_matte_provider() + try: + out = [] + threads = [threading.Thread(target=lambda: out.append(get_matte_provider())) + for _ in range(8)] + for t in threads: + t.start() + for t in threads: + t.join() + assert len(built) == 1, f"并发下建了 {len(built)} 份" + assert len({id(o) for o in out}) == 1, "拿到了不同的实例" + finally: + reset_matte_provider() + + +def test_no_production_call_site_bypasses_the_singleton(): + """拦的坏例:新增一个调用点直接用工厂,唯一性就又只能靠人守规矩。 + + 实测:此前四个调用点全部直接调工厂,唯一性完全靠 ``bind_matte`` 恰好跑成。 + """ + import pathlib + + app_src = pathlib.Path(__file__).resolve().parents[1] / "packages/app/src" + # 两个口子都要堵:调工厂,和直接 new 具体类。只堵前者的话,一句 + # ``OnnxU2NetMatteProvider()`` 照样多出一份,而它连工厂都没经过。 + banned = ("make_matte_provider", "OnnxU2NetMatteProvider(", "BiRefNetMatteProvider(") + offenders = [] + for f in app_src.rglob("*.py"): + text = f.read_text(encoding="utf-8") + hit = [b for b in banned if b in text] + if hit: + offenders.append(f"{f.relative_to(app_src)} → {hit}") + assert not offenders, ( + f"这些生产文件绕过了单例:{offenders};" + "工厂与具体类只给测试和本地脚本用,生产一律走 get_matte_provider()" + ) diff --git a/backend/tests/test_mq_worker.py b/backend/tests/test_mq_worker.py index c461067b..c3d5517d 100644 --- a/backend/tests/test_mq_worker.py +++ b/backend/tests/test_mq_worker.py @@ -1464,9 +1464,12 @@ def warmup(self): self.warmed = True fake = _Fake() - monkeypatch.setattr( - "windup_framework.providers.OnnxU2NetMatteProvider", lambda *a, **k: fake - ) + # 桩在**工厂**上,并清掉单例 —— 生产已改走 get_matte_provider(),它内部调工厂。 + # 桩 get_matte_provider 本身会把"单例真的只建一份"这层一起绕过去,那正是要守的东西。 + from windup_framework.providers import matte_factory as _mf + + monkeypatch.setattr(_mf, "make_matte_provider", lambda *a, **k: fake) + _mf.reset_matte_provider() prev_a, prev_i = ex.executor._matte, ex.image_executor._matte try: w._warmup_local_inference() @@ -1476,6 +1479,7 @@ def warmup(self): finally: ex.executor._matte = prev_a ex.image_executor._matte = prev_i + _mf.reset_matte_provider() # 单例是进程级的,不清会污染同进程后面的用例 def test_consumer_trims_heap_after_image_not_email(engine, worker_session, monkeypatch):