Return freed activations to the OS on Linux; recommend jemalloc (A+B)
MIMALLOC_PURGE_DELAY=0 only bites on Windows (bundled mimalloc). On a Linux server PyTorch uses glibc ptmalloc, which keeps a BiRefNet@2048 forward's freed activations in the arena, so RSS ratchets up across requests and the env var is a no-op there. A. bgfilter/memtune.py: release_freed_memory() calls glibc malloc_trim(0) at runtime (no-op on Windows/musl/other allocators). Wired into service.process (after each request) and the CLI batch loop (after each image), the two long-lived paths where RSS accumulates. Single-image CLI exits, so it is left alone. B. DEPLOY(_ZH): document preloading jemalloc via LD_PRELOAD + MALLOC_CONF (also improves CPU throughput) as the production alternative, with MALLOC_ARENA_MAX/ MALLOC_TRIM_THRESHOLD_ as an allocator-free fallback. Under jemalloc/tcmalloc malloc_trim simply no-ops. Also fixes the stale off-by-default cross-check heading in DEPLOY.md. Verified: modules import; release_freed_memory() returns False (no-op) on Windows. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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. 故障排查
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
+7
-1
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user