Pipeline concurrent requests over a shared GPU with a narrowed lock
The single asyncio lock serialized the whole request (chroma -> seg ->
trimap -> ViTMatte -> cross-check -> foreground -> despill), so the GPU
sat idle ~70-80% of each request while the CPU stages ran. Replace it
with a thread pool that runs several requests concurrently, and move the
mutex down to just the GPU forwards: when one request steps off the GPU
to do CPU work, another steps on, keeping the device fed.
Two opt-in env knobs, both defaulting to the old strictly-serial
behaviour:
* BGFILTER_CPU_WORKERS (pool size) attacks the CPU-bound regime -- with
cross-check off, throughput went 0.91 -> 1.67 req/s (1.84x) at pool=8
on a 4080S, vs a mere +6% from a second worker process, at half the
VRAM (one model copy).
* BGFILTER_GPU_CONCURRENCY caps how many forwards may co-schedule on the
SMs. batch-1 underfills the GPU, so =2 recovers the SM overlap a lock
forgoes: cross-check on went 0.83 -> 1.01 req/s, beating two worker
processes (0.92) with one model copy (~19 GB vs 24 GB).
The CLI/single-thread path passes no lock and uses a nullcontext, so it
stays allocation- and contention-free. Concurrent output verified
identical (10x same-frame -> one hash). A per-worker admission semaphore
bounds in-flight decoded images to the pool size.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
+18
-2
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import replace
|
||||
|
||||
@@ -37,6 +38,14 @@ class PipelineManager:
|
||||
self._segmenters: dict[str, object] = {}
|
||||
self._cross_checker: BiRefNetSegmenter | None = None
|
||||
self._load_lock = threading.Lock()
|
||||
# Caps how many concurrent requests may be inside a GPU forward at once
|
||||
# (the app may run several process() calls on a thread pool). Shared by every
|
||||
# pipeline this manager builds. =1 (default) makes the device strictly
|
||||
# single-stream; >1 lets that many forwards overlap so batch-1 kernels that
|
||||
# underfill the SMs can co-schedule (higher GPU-bound throughput, but the
|
||||
# concurrent activations multiply GPU memory). Uncontended when CPU_WORKERS=1.
|
||||
gpu_conc = max(1, int(os.environ.get("BGFILTER_GPU_CONCURRENCY", "1")))
|
||||
self._gpu_lock = threading.Lock() if gpu_conc == 1 else threading.BoundedSemaphore(gpu_conc)
|
||||
if preload:
|
||||
if base_settings.model.matting_method == "vitmatte":
|
||||
self._get_runner()
|
||||
@@ -169,11 +178,18 @@ class PipelineManager:
|
||||
self._get_cross_checker() if settings.cross_check.enabled else None
|
||||
)
|
||||
pipeline = MattingPipeline(
|
||||
settings, runner=runner, segmenter=segmenter, cross_checker=cross_checker
|
||||
settings,
|
||||
runner=runner,
|
||||
segmenter=segmenter,
|
||||
cross_checker=cross_checker,
|
||||
gpu_lock=self._gpu_lock,
|
||||
)
|
||||
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()
|
||||
# Hold the GPU lock: empty_cache must not race a concurrent request's
|
||||
# in-flight allocations on the same device.
|
||||
with self._gpu_lock:
|
||||
release_freed_memory()
|
||||
|
||||
Reference in New Issue
Block a user