diff --git a/DEPLOY.md b/DEPLOY.md index 02e135b..194cac1 100644 --- a/DEPLOY.md +++ b/DEPLOY.md @@ -156,7 +156,7 @@ hf download ZhengPeng7/BiRefNet_HR-matting --local-dir models/BiRefNet_HR- Override the base directory with `BGFILTER_WEIGHTS_DIR`. `models/` is gitignored. -### Cross-check veto (optional, off by default in the service) +### Cross-check veto (on by default) The pipeline runs a second, high-resolution matting model (`ZhengPeng7/BiRefNet_HR-matting`, ~425 MB) as a "second opinion" that clears @@ -291,6 +291,21 @@ service has **no authentication**. To reach it from other hosts: balancer (each with `--workers 1`). - Memory: one worker holds all three models (~1 GB weights + runtime). `MemoryMax=8G` is a safe ceiling for a single CPU worker. +- **Keeping RSS flat (glibc).** After each image the service calls glibc + `malloc_trim(0)` (`bgfilter/memtune.py`) to hand a large BiRefNet@2048 forward's + freed activations back to the OS — without it, ptmalloc keeps them in the arena + and RSS ratchets up across requests. For lower memory *and* better CPU + throughput, preload **jemalloc** (PyTorch's recommendation for CPU inference) + and let it decay dirty pages aggressively: + ```ini + # in the [Service] section of the systemd unit + Environment=LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 + Environment=MALLOC_CONF=background_thread:true,dirty_decay_ms:0,muzzy_decay_ms:0 + ``` + (`apt install libjemalloc2`; tcmalloc via `libtcmalloc_minimal.so.4` works too.) + jemalloc/tcmalloc manage their own release, so `malloc_trim` simply no-ops under + them. If you cannot preload an allocator, `MALLOC_ARENA_MAX=2` and + `MALLOC_TRIM_THRESHOLD_=0` in the unit environment curb ptmalloc retention. ## 12. Troubleshooting diff --git a/DEPLOY_ZH.md b/DEPLOY_ZH.md index f93cfd0..7cae7d4 100644 --- a/DEPLOY_ZH.md +++ b/DEPLOY_ZH.md @@ -272,6 +272,18 @@ curl -sS -F "file=@input.png" -F "screen_color=#CFEFFF" \ 和/或多实例 + 负载均衡(每个都 `--workers 1`)。 - 内存:单个 worker 持有全部三个模型(约 1 GB 权重 + 运行时)。单 CPU worker 用 `MemoryMax=8G` 作上限比较稳妥。 +- **让 RSS 保持平稳(glibc)。** 每处理完一张图,服务会调用 glibc 的 `malloc_trim(0)` + (`bgfilter/memtune.py`)把 BiRefNet@2048 前向释放的大块激活还给 OS——否则 ptmalloc + 会把它们留在 arena 里,RSS 会随请求逐步抬高。想进一步降内存**并**提升 CPU 吞吐,可以 + 预载 **jemalloc**(PyTorch 对 CPU 推理的推荐),并让它激进归还脏页: + ```ini + # 放进 systemd unit 的 [Service] 段 + Environment=LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 + Environment=MALLOC_CONF=background_thread:true,dirty_decay_ms:0,muzzy_decay_ms:0 + ``` + (`apt install libjemalloc2`;tcmalloc 用 `libtcmalloc_minimal.so.4` 同理。) + jemalloc/tcmalloc 自己管理归还,此时 `malloc_trim` 会自动空转。若无法预载分配器, + 在 unit 环境里设 `MALLOC_ARENA_MAX=2`、`MALLOC_TRIM_THRESHOLD_=0` 也能抑制 ptmalloc 滞留。 ## 12. 故障排查 diff --git a/bgfilter/__init__.py b/bgfilter/__init__.py index e002b31..16a274e 100644 --- a/bgfilter/__init__.py +++ b/bgfilter/__init__.py @@ -7,6 +7,8 @@ import os # BiRefNet@2048 pass, stacking under ViTMatte's attention peak (measured # 2048x2048 bf16: peak 25.2 -> 21.4 GB with purge 0, and slightly *faster*). # Must be set before the torch DLL loads; harmless where mimalloc is not used. +# The Linux/glibc equivalent (ptmalloc arena retention) is handled at runtime by +# bgfilter.memtune.release_freed_memory(); see DEPLOY for jemalloc as an alt. os.environ.setdefault("MIMALLOC_PURGE_DELAY", "0") __version__ = "0.1.0" diff --git a/bgfilter/cli.py b/bgfilter/cli.py index bcf6ee9..d326f41 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -6,6 +6,7 @@ import typer from rich import print from .config import load_settings, override_settings +from .memtune import release_freed_memory from .pipeline import MattingPipeline, run_image app = typer.Typer(help="Offline flat-background character matting.") @@ -69,6 +70,9 @@ def main( dbg = None if debug_dir is None else debug_dir / image_path.stem result = pipeline.run_image(image_path, out_path, dbg) print(f"[green]wrote[/green] {result['output']}") + # Free this image's activations back to the OS before the next one + # so batch RSS tracks a single image, not the whole run. + release_freed_memory() return if input is None or output is None: diff --git a/bgfilter/memtune.py b/bgfilter/memtune.py new file mode 100644 index 0000000..6a5d06c --- /dev/null +++ b/bgfilter/memtune.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +"""Cross-platform release of freed allocator memory after a heavy inference. + +Both of the process allocators this pipeline runs under retain large freed +blocks instead of returning them to the OS, so a BiRefNet@2048 forward's +activations keep inflating RSS long after the tensors are gone: + +- Windows PyTorch bundles mimalloc; its lazy purge is defused at import time by + ``MIMALLOC_PURGE_DELAY=0`` (set in ``bgfilter/__init__.py`` before torch loads). +- Linux/glibc PyTorch uses ptmalloc, which keeps freed arenas committed; + ``malloc_trim(0)`` hands those pages back and is called at runtime via + :func:`release_freed_memory` after each image. + +Neither hook exists on the other's platform, so this stays a no-op there (and on +musl, which has no ``malloc_trim``, or under a preloaded jemalloc/tcmalloc that +manages its own release). See DEPLOY for the recommended jemalloc LD_PRELOAD. +""" + +import ctypes +import ctypes.util +import sys + +_malloc_trim = None +if sys.platform.startswith("linux"): + try: + _libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6", use_errno=True) + _malloc_trim = _libc.malloc_trim # glibc only; AttributeError on musl + _malloc_trim.argtypes = [ctypes.c_size_t] + _malloc_trim.restype = ctypes.c_int + except (OSError, AttributeError): + _malloc_trim = None + + +def release_freed_memory() -> bool: + """Return freed heap pages to the OS; True if a release was actually issued. + + Cheap (sub-millisecond) and safe to call after every processed image. No-op + where no supported hook is available (Windows, musl, non-glibc allocators). + """ + if _malloc_trim is None: + return False + _malloc_trim(0) + return True diff --git a/bgfilter/service.py b/bgfilter/service.py index 1ce15df..29bb967 100644 --- a/bgfilter/service.py +++ b/bgfilter/service.py @@ -7,6 +7,7 @@ from dataclasses import replace import numpy as np from .config import _SEG_MODELS, override_settings +from .memtune import release_freed_memory from .pipeline import MattingPipeline, MattingResult from .segmentation import BiRefNetSegmenter, make_segmenter from .settings import PipelineSettings, SegmentationSettings @@ -151,4 +152,9 @@ class PipelineManager: pipeline = MattingPipeline( settings, runner=runner, segmenter=segmenter, cross_checker=cross_checker ) - return pipeline.run_rgb(rgb) + try: + return pipeline.run_rgb(rgb) + finally: + # Return this request's large freed activations to the OS so RSS does + # not ratchet up across requests in the long-lived service process. + release_freed_memory()