[wayrise]

Summary

LeRobot v3 datasets store their clips in AV1, which FFmpeg decodes with libdav1d. For every clip PyAV opens, libdav1d starts a thread pool sized to the host, because FFmpeg passes its automatic thread count (about one thread per CPU core). With DataLoader workers × ranks decoding concurrently, these pools oversubscribe the machine and the GPUs wait for data. On LIBERO with 7×H200 and 8 workers per rank, the load average was about 150 on 160 cores.

decode_video_frames, the shared reader of every LeRobotV3Reader, now caps the libdav1d pool at 2 threads per clip.

• OPENWAM_VIDEO_DECODE_THREADS overrides the count; 0 keeps FFmpeg's default.
• Only libdav1d streams are touched, and only their thread count. Other decoders (H.264, HEVC, …) keep PyAV's defaults. Several H.264 datasets in the repo are slice-encoded and already decode efficiently with the default slice threading.
• Decoded frames are identical.
• The README documents the variable next to the other training-acceleration options.

Measurements

• End to end. Same config back to back: LIBERO, dual-system joint_self_attn from Wan2.2, 7×H200, 32 samples/GPU, 8 workers/rank, OPENWAM_SDPA_CUDNN=0.

  • Step time: main 1.940 s/step → this PR 1.380 s/step (−29 %), measured over steps 50–150.
  • Mean GPU utilisation: 82 % → 88 %.
  • Losses: step 1 is 1.9742 on both; step 50 is 0.6002 vs 0.5999.

• Decode throughput. 56 concurrent processes; each sample is 2 cameras × 9 frames of a 33-frame window; 3 interleaved repetitions.

  Setting	Throughput	CPU per sample
  main	77–88 samples/s	1.03–1.31 s
  this PR	162–213 samples/s	0.46–0.63 s

  Frame threading with 2 threads measured the same within noise, so the thread type is left unchanged. Decoded pixels are identical across all 2,240 clip decodes per setting (56 workers × 20 samples × 2 cameras).

• Corrupted input. LIBERO AV1 clips with 200 scattered byte flips or 60 % truncation, decoded over 48 random windows. The outcomes are identical to main: 15 decoded, 33 InvalidDataError, no hangs.

Why only libdav1d and only the thread count. An earlier revision enabled frame threading for all decoders. During review, a corrupted H.264 input then hung, where main raises InvalidDataError: the av.logging callback that this module installs takes the GIL while frame-thread workers are being joined. The slice-encoded H.264 datasets also gained nothing from frame threading. libdav1d runs its own thread pool regardless of the FFmpeg thread type, so capping its size changes nothing but the pool size.

Also fixed (found in review). The seek in decode_video_frames was guarded by except av.AVError, a name PyAV removed in 14.0. With the pinned PyAV 17.1, a failing seek therefore raised AttributeError from the except clause itself, and the intended fallback to a sequential decode never ran. The guard now catches av.FFmpegError, which exists in PyAV 12–17 and is the base of every PyAV error. The new test of this path fails on main.

Not addressed here, as a possible follow-up: in PyAV 17, frame.to_image() builds a swscale context with an automatic thread count for every frame. This setting does not control those threads.

Tests

• New tests/dataloader/test_video_io_threads.py:
  • Identical frames for 0/1/2/4 threads on synthetic AV1 (libsvtav1 → libdav1d), H.264 and MPEG-4 clips. Clips use an 8-frame GOP, and both a window starting at frame 0 (sequential decode) and a late window (seek to a mid-stream keyframe) are decoded. Each codec is skipped when its encoder is unavailable.
  • Parsing of the default, blank, override and invalid values of the environment variable. Values must be ASCII digits and at most 256 (libdav1d's maximum).
  • A failing seek falls back to a sequential decode with identical frames.
  • Only libdav1d streams are capped.
  • The cap is applied to the stream decode_video_frames opens for AV1, and not for MPEG-4. The AV1 case is skipped when libdav1d is not the available AV1 decoder.
  • 13 tests in this file pass.
• CPU suite (-m "not gpu"): 2147 passed, 1 failed; the failure, test_cosmos_predict25_build_with_bad_path_raises_clear_error, fails identically on main.

🤖 Generated with Claude Code