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:
@@ -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
|
||||
Reference in New Issue
Block a user