Skip to content

fix: reduce SAM2 video backend resource usage - #924

Merged
nick-skriabin merged 9 commits into
masterfrom
fix/sam2-video-resource-usage
Aug 4, 2026
Merged

fix: reduce SAM2 video backend resource usage#924
nick-skriabin merged 9 commits into
masterfrom
fix/sam2-video-resource-usage

Conversation

@nick-skriabin

Copy link
Copy Markdown
Member

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:

  • Video source resolution/download happened before VideoRegistry reuse, so concurrent prewarm/predict/track requests could all miss the cache and download the same large video at once.
  • Tracking sessions retained already-drained base64 mask payloads until the final progress poll.
  • Completed, cancelled, or abandoned tracking sessions could remain in process memory indefinitely.
  • Prewarm decoded a full frame window before checking whether those frames were already cached or pending.
  • Frame embedding and tracking work were too concurrent by default for a single GPU-backed backend, causing memory spikes and contention.
  • Frame cache byte accounting could drift upward when an embedding was replaced.

Solution

  • Serialize video resolve/download/register per task + raw URL, and reuse VideoRegistry handles before invoking Label Studio downloads.
  • Track the original raw URL on VideoHandle so the resolved local path can be safely reused across events for the same task asset.
  • Make tracking progress draining destructive so already-sent mask payloads are freed immediately.
  • Add idle/max-age cleanup for tracking sessions and remove cancelled sessions from the registry.
  • Replace unbounded tracking daemon threads with a bounded tracking executor (TRACKING_WORKERS, default 1).
  • Make frame embedding concurrency configurable (FRAME_ENCODER_WORKERS, default 1) to reduce GPU/RAM spikes.
  • Prewarm only decodes frames that are missing from cache and not already queued.
  • Correct frame cache byte accounting when an existing frame embedding is replaced.
  • Document the new resource-control env vars in README and docker-compose.
  • Add unit coverage for frame-cache missing-frame detection and byte accounting.

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.py

Result: 35 passed.

Note: test_api.py is a live-backend smoke test and was excluded from local validation.

@github-actions github-actions Bot added the fix label Jul 29, 2026

@matt-bernstein matt-bernstein left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI review below - I think this PR is larger than it needs to be


Blocking findings

  1. Label Studio credentials can leak across redirects.

    model.py passes Authorization headers 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 original Authorization: Token ... header.

    Please keep authenticated streaming disabled/out of scope until redirects are constrained and credentials are stripped whenever the origin changes.

  2. 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.

  3. 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.

  4. TTL and cleanup do not cover all task resources safely.

    Frame-cache TTL eviction only runs after another encode, and VideoRegistry has 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.

  5. 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 SAM2 main. 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

  1. 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 so drop_task() closes a handle only after current users finish.

    This removes _video_resolve_locks and safely guarantees one SDK download. Authenticated ffmpeg streaming should be a separate follow-up.

  2. Replace missing() plus submit() 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.

  3. Use one accelerator semaphore.

    A process-wide BoundedSemaphore(1) should guard image prediction, frame embedding, tracking initialization, prompt insertion, and each tracking propagation step.

  4. 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.

  5. 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.

  6. Minimize the SAM2 workaround.

    Pin SAM2, initialize through its public init_state(), stop its retaining async loader, and replace only inference_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 matt-bernstein left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking resource-control issues

  1. Tracking workers are bounded, but submissions are not.

    ThreadPoolExecutor(max_workers=TRACKING_WORKERS) has an unbounded internal queue. Every track request 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.

  2. Client-controlled work sizes are not hard-capped.

    window, max_frames, and max_duration_ms are accepted directly. MAX_FRAMES_TO_TRACK acts 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.

  3. The guarantees only hold with one Gunicorn worker.

    VideoRegistry, FrameCache, _ACCELERATOR, and the tracking executor are process-local. Setting WORKERS > 1 restores duplicate downloads, concurrent GPU inference, duplicate caches, and one SAM2 model per process.

    Since this backend targets one GPU, enforce WORKERS=1 at startup or implement cross-process coordination. Documentation alone is easy to bypass accidentally.

Correctness issues

  1. 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 pass TRACK_PROGRESS_MAX_BATCH - len(new_frames) on subsequent drains.

  2. TTL eviction ignores empty task caches.

    _evict_expired_tasks() requires len(t.embeddings) > 0. Requests with no prompts and failed encodes can create zero-embedding, zero-pending _TaskCache entries 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() and schedule_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.

nick-skriabin and others added 4 commits August 3, 2026 17:19
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.
@nick-skriabin
nick-skriabin merged commit d89f3b2 into master Aug 4, 2026
9 checks passed
@nick-skriabin
nick-skriabin deleted the fix/sam2-video-resource-usage branch August 4, 2026 13:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants