mirror of
https://github.com/vectorize-io/hindsight.git
synced 2026-09-14 19:31:49 +08:00
perf(local): skip full GC for CPU reranker providers (#3858)
CPU local reranker providers release short-lived tensor and ONNX buffers through normal reference counting, while Python's cyclic collector runs on its threshold schedule. A full process-wide collection after every batch adds avoidable latency to both LocalST and FlashRank CPU inference. Keep full collection and allocator cleanup for CUDA, XPU, and MPS inference, where releasing Python wrappers precedes accelerator cache cleanup. Unknown device types retain the conservative full-collection behavior. Update cleanup tests to verify that CPU inference still trims the heap without calling gc.collect(). Tests: pytest -q hindsight-api-slim/tests/test_local_device.py
This commit is contained in:
@@ -163,10 +163,14 @@ def _empty_gpu_cache(device_type: str | None) -> None:
|
||||
def release_local_inference_memory(device_type: str | None = None) -> None:
|
||||
"""Release transient heap (and GPU allocator) memory after a local inference batch.
|
||||
|
||||
Frees Python objects, returns freed native pages to the OS, and empties the GPU
|
||||
allocator pool when the model ran on a GPU. Safe to call on every platform and
|
||||
device; the pieces that don't apply are cheap no-ops.
|
||||
Returns freed native pages to the OS, and empties the GPU allocator pool when
|
||||
the model ran on a GPU. Python's normal reference counting already releases
|
||||
the short-lived CPU inference buffers; a full cyclic-GC scan on every batch is
|
||||
needlessly expensive on that hot path, so CPU callers leave cyclic GC to its
|
||||
normal threshold-based schedule. GPU callers retain the existing full cleanup
|
||||
because allocator release is part of the opt-in accelerator memory policy.
|
||||
"""
|
||||
gc.collect()
|
||||
if device_type != "cpu":
|
||||
gc.collect()
|
||||
_heap_trim()
|
||||
_empty_gpu_cache(device_type)
|
||||
|
||||
@@ -161,8 +161,14 @@ class TestReleaseLocalInferenceMemory:
|
||||
assert calls == ["gc", "trim"]
|
||||
assert log == ["cuda"]
|
||||
|
||||
def test_release_cpu_skips_empty_cache(self):
|
||||
def test_release_cpu_skips_gc_and_empty_cache_but_trims_heap(self):
|
||||
calls = []
|
||||
log = []
|
||||
with patch.dict(sys.modules, {"torch": _fake_torch(empty_cache_log=log)}):
|
||||
with (
|
||||
patch.dict(sys.modules, {"torch": _fake_torch(empty_cache_log=log)}),
|
||||
patch.object(local_device.gc, "collect", lambda: calls.append("gc")),
|
||||
patch.object(local_device, "_heap_trim", lambda: calls.append("trim")),
|
||||
):
|
||||
release_local_inference_memory("cpu")
|
||||
assert calls == ["trim"]
|
||||
assert log == []
|
||||
|
||||
Reference in New Issue
Block a user