Trim process heap after document analysis

Author: medisean

mineru/backend/vlm/vlm_analyze.py

@@ -35,6 +35,7 @@
     open_pdfium_document,
 )
 from ...utils.models_download_utils import auto_download_and_get_model_root_path
+from ...utils.model_utils import clean_memory
 
 from mineru_vl_utils import MinerUClient
 from packaging import version
@@ -518,6 +519,7 @@ def doc_analyze(
     finally:
         if not doc_closed:
             close_pdfium_document(pdf_doc)
+        clean_memory(get_device())
 
 
 async def aio_doc_analyze(
@@ -617,3 +619,4 @@ async def aio_doc_analyze(
     finally:
         if not doc_closed:
             close_pdfium_document(pdf_doc)
+        clean_memory(get_device())

mineru/utils/model_utils.py

@@ -1,8 +1,13 @@
 # Copyright (c) Opendatalab. All rights reserved.
+import ctypes
+import gc
 import math
 import os
+import sys
 import time
-import gc
+from functools import lru_cache
+from typing import Callable
+
 from PIL import Image
 from loguru import logger
 import numpy as np
@@ -35,6 +40,39 @@
 }
 
 
+@lru_cache(maxsize=1)
+def _get_malloc_trim() -> Callable[[int], int] | None:
+    """Return glibc's heap trimming function when it is available."""
+    if sys.platform != "linux":
+        return None
+
+    try:
+        libc = ctypes.CDLL(None)
+        malloc_trim = getattr(libc, "malloc_trim", None)
+    except OSError:
+        return None
+
+    if malloc_trim is None:
+        return None
+
+    malloc_trim.argtypes = [ctypes.c_size_t]
+    malloc_trim.restype = ctypes.c_int
+    return malloc_trim
+
+
+def trim_process_memory() -> None:
+    """Return unused glibc heap pages to the operating system when possible."""
+    malloc_trim = _get_malloc_trim()
+    if malloc_trim is None:
+        return
+
+    try:
+        malloc_trim(0)
+    except Exception as exc:
+        # Heap trimming is an optional optimization and must never break parsing.
+        logger.debug(f"Unable to trim process heap: {exc}")
+
+
 def _get_bbox(item):
     bbox = item.get("bbox")
     if bbox is not None:
@@ -203,6 +241,7 @@ def clean_memory(device='cuda'):
         if torch.sdaa.is_available():
             torch.sdaa.empty_cache()  
     gc.collect()
+    trim_process_memory()
 
 
 def clean_vram(device, vram_threshold=8):

tests/unittest/test_model_utils.py

@@ -0,0 +1,54 @@
+from typing import Any
+
+from mineru.utils import model_utils
+
+
+def test_trim_process_memory_invokes_malloc_trim(monkeypatch: Any) -> None:
+    calls = []
+
+    def fake_malloc_trim(padding: int) -> int:
+        calls.append(padding)
+        return 1
+
+    monkeypatch.setattr(model_utils, "_get_malloc_trim", lambda: fake_malloc_trim)
+
+    model_utils.trim_process_memory()
+
+    assert calls == [0]
+
+
+def test_trim_process_memory_is_optional(monkeypatch: Any) -> None:
+    monkeypatch.setattr(model_utils, "_get_malloc_trim", lambda: None)
+
+    model_utils.trim_process_memory()
+
+
+def test_clean_memory_runs_process_heap_trim(monkeypatch: Any) -> None:
+    calls = []
+    monkeypatch.setattr(model_utils, "trim_process_memory", lambda: calls.append(True))
+
+    model_utils.clean_memory("cpu")
+
+    assert calls == [True]
+
+
+def test_get_malloc_trim_configures_available_symbol(monkeypatch: Any) -> None:
+    class FakeTrim:
+        argtypes = None
+        restype = None
+
+        def __call__(self, padding: int) -> int:
+            return padding
+
+    fake_trim = FakeTrim()
+    fake_libc = type("FakeLibc", (), {"malloc_trim": fake_trim})()
+    monkeypatch.setattr(model_utils.sys, "platform", "linux")
+    monkeypatch.setattr(model_utils.ctypes, "CDLL", lambda _: fake_libc)
+    model_utils._get_malloc_trim.cache_clear()
+
+    try:
+        assert model_utils._get_malloc_trim() is fake_trim
+        assert fake_trim.argtypes == [model_utils.ctypes.c_size_t]
+        assert fake_trim.restype is model_utils.ctypes.c_int
+    finally:
+        model_utils._get_malloc_trim.cache_clear()