fix: reduce SAM2 video backend resource usage - #924
Conversation
There was a problem hiding this comment.
AI review below - I think this PR is larger than it needs to be
Blocking findings
-
Label Studio credentials can leak across redirects.
model.pypassesAuthorizationheaders to ffmpeg/ffprobe for uploaded-file streaming. FFmpeg forwards custom headers across redirects, including cross-origin redirects. I reproduced this with an LS-like endpoint redirecting to another server; the target received the originalAuthorization: Token ...header.Please keep authenticated streaming disabled/out of scope until redirects are constrained and credentials are stripped whenever the origin changes.
-
Concurrent prewarms can still duplicate decoding/network reads.
The flow is currently:
FRAME_CACHE.missing()→_prefetch_window()→FRAME_CACHE.submit().missing()does not reserve frames, so concurrent requests can both observe the same missing set and prefetch it before either request records the frames as pending. I reproduced two concurrent callers both prefetching the same frame set.Reservation must happen atomically before any decode or network I/O.
-
GPU concurrency is not globally bounded.
Tracking and frame encoding use separate executors, while synchronous predictions run in request threads. Setting both executors to one worker still permits a tracking job, frame encoder, and interactive prediction to execute against SAM2 concurrently.
Use one process-wide accelerator semaphore around all SAM2 operations. During tracking, acquire it per propagation step so long tracks do not monopolize interactive inference.
-
TTL and cleanup do not cover all task resources safely.
Frame-cache TTL eviction only runs after another encode, and
VideoRegistryhas no periodic expiry. Idle resources can therefore remain indefinitely.Conversely, tracking completion can drop frame/video state while an unrelated predict or prewarm request for the same task is still using it. Removing keyed resolve locks while they are held can also recreate the duplicate-resolution race.
Resource cleanup needs lease/reference-count semantics and one periodic maintenance loop.
-
The custom SAM2 integration depends on private, unpinned internals.
_init_tracking_inference_state()reconstructs SAM2’s private state dictionary and calls_get_image_feature(), while the Dockerfile installs mutable SAM2main. An upstream change can therefore break a previously working image without any change in this repository.Pin SAM2 to a tested commit and avoid copying its complete private state contract.
Recommended reduced design
-
Put single-flight resolution inside
VideoRegistry.Expose something like:
with VIDEOS.acquire(task_id, raw_url, resolve_fn) as video: ...Internally, store one
Future[VideoHandle]per(task_id, raw_url). The first caller resolves/downloads via the existing Label Studio SDK; concurrent callers await the same future. Track active leases sodrop_task()closes a handle only after current users finish.This removes
_video_resolve_locksand safely guarantees one SDK download. Authenticated ffmpeg streaming should be a separate follow-up. -
Replace
missing()plussubmit()with an atomic batch API.FRAME_CACHE.schedule_missing(task_id, indices, encode_batch)should reserve every newly missing frame under the task lock before submitting work. The batch worker should split sparse indices into contiguous runs and decode only those runs. -
Use one accelerator semaphore.
A process-wide
BoundedSemaphore(1)should guard image prediction, frame embedding, tracking initialization, prompt insertion, and each tracking propagation step. -
Use a bounded producer-consumer tracking session.
Each session should have one executor future and a standard
queue.Queue(maxsize=N)for generated masks. This bounds RAM and naturally applies backpressure when polling is slow. Limit active/queued sessions before executor submission, and process bidirectional ranges sequentially on the single GPU. -
Use one maintenance loop.
Periodically expire abandoned sessions, idle frame caches, and video handles. Cleanup should go through lease-safe registry methods and should not automatically destroy task state merely because tracking completed.
-
Minimize the SAM2 workaround.
Pin SAM2, initialize through its public
init_state(), stop its retaining async loader, and replace onlyinference_state["images"]with a non-retaining lazy sequence. Avoid model unloading,malloc_trim, copied private state, and pruning until measurements show they are necessary.
matt-bernstein
left a comment
There was a problem hiding this comment.
Blocking resource-control issues
-
Tracking workers are bounded, but submissions are not.
ThreadPoolExecutor(max_workers=TRACKING_WORKERS)has an unbounded internal queue. Everytrackrequest adds a session and one or two futures, retaining prompts, video leases, and session state until execution or max-age cleanup.Add an explicit maximum for active + queued tracking sessions and return a busy response when capacity is exhausted.
-
Client-controlled work sizes are not hard-capped.
window,max_frames, andmax_duration_msare accepted directly.MAX_FRAMES_TO_TRACKacts only as a default, not a maximum. A caller can therefore request the entire video, while_encode_frame_batch()materializes each contiguous run of decoded raw frames before encoding it.Clamp these values to server-side maxima, reject negative/invalid values, and consider limiting batch decode size independently of the requested prewarm window.
-
The guarantees only hold with one Gunicorn worker.
VideoRegistry,FrameCache,_ACCELERATOR, and the tracking executor are process-local. SettingWORKERS > 1restores duplicate downloads, concurrent GPU inference, duplicate caches, and one SAM2 model per process.Since this backend targets one GPU, enforce
WORKERS=1at startup or implement cross-process coordination. Documentation alone is easy to bypass accidentally.
Correctness issues
-
Progress responses can exceed
TRACK_PROGRESS_MAX_BATCH.drain_new()can consume a full batch each time. The micro-batching loop checks the current total, then appends another independently full drain. For example, with a maximum of 32, an initial batch of 10 can be followed by 32 more.Change this to
drain_new(limit)and passTRACK_PROGRESS_MAX_BATCH - len(new_frames)on subsequent drains. -
TTL eviction ignores empty task caches.
_evict_expired_tasks()requireslen(t.embeddings) > 0. Requests with no prompts and failed encodes can create zero-embedding, zero-pending_TaskCacheentries that are never removed.Expire every idle task with no pending work, regardless of whether it contains embeddings.
Removable or consolidatable code
The diff remains larger than necessary because old and new APIs coexist:
FrameCache.missing()is now used only by tests.FrameCache.ensure_encoded()is now used only by tests.SamVideoInteractive._encode_frame()is unused.VideoRegistry.get_or_create()is unused.FrameCache.submit()andschedule_missing()overlap and could become one scheduling API.- Bidirectional tracking could run as one executor future that processes both ranges sequentially, removing producer-count and multi-future bookkeeping.
- Safe defaults such as one tracking worker and on-demand frames could be fixed for this single-GPU example instead of exposing additional tuning flags.
Reserve tracking capacity before media resolution and count completed sessions with undrained results. Apply CUDA autocast in each inference thread, preserve MPS-safe initialization, bind sessions to tasks, deduplicate bidirectional prompt frames, and align cancellation and TTL cleanup semantics.
Sweep expired sessions before admission, preserve the track_busy payload contract, and bootstrap MPS inference from one frame so on-demand loading remains memory bounded.
Initialize an empty repository, fetch only the tested SAM2 commit, install it in the same Docker layer, and remove Git metadata to avoid cloning full history into the image.
Problem
The SAM2 video interactive backend degrades over time and uses excessive resources, especially RAM. DevOps also observed that the backend can download the same task video multiple times.
Root causes found in this branch:
VideoRegistryreuse, so concurrent prewarm/predict/track requests could all miss the cache and download the same large video at once.Solution
VideoRegistryhandles before invoking Label Studio downloads.VideoHandleso the resolved local path can be safely reused across events for the same task asset.TRACKING_WORKERS, default1).FRAME_ENCODER_WORKERS, default1) to reduce GPU/RAM spikes.Validation
.venv/bin/python -m py_compile \ label_studio_ml/examples/segment_anything_video_interactive/model.py \ label_studio_ml/examples/segment_anything_video_interactive/video_state.py \ label_studio_ml/examples/segment_anything_video_interactive/frame_cache.py cd label_studio_ml/examples/segment_anything_video_interactive ../../../.venv/bin/python -m pytest -q --ignore=test_api.pyResult:
35 passed.Note:
test_api.pyis a live-backend smoke test and was excluded from local validation.