Fix wan2.2 VBench accuracy pipeline: drain timeouts, .mp4 staging, torch>=2.6, wget preflight

[wu6u3tw]

Fix wan2.2 VBench accuracy pipeline: configurable drain timeouts, .mp4 staging, torch>=2.6 and wget cold-start failures

Summary

Running the wan2.2 (videogen) benchmark end to end on GB300-NVL72 (aarch64, SLURM/enroot) surfaced four independent client-side failures. Each one only manifests after a full generation pass (30-50 min on 18 nodes), and together they made VBench accuracy runs impossible to complete. This PR fixes all four, with regression tests for the two library-level changes.

All fixes were validated on hardware: with them applied (initially as runtime monkeypatches), a full Offline run completed with 144/144 performance samples, a clean warmed-up perf phase, and a VBench score of 0.69998 over 248 samples.

Fix 1: configurable per-phase drain timeouts (settings.drain)

BenchmarkSession._drain_inflight hard-coded a 240 s wait for in-flight responses after every phase, then abandoned whatever was left and moved on.

For workloads at minutes per sample this is fatal, and silently so:

• A performance phase can end at min_duration with every sample still in flight. Observed: Drain timed out after 240 s with 144 responses still in flight followed by Total samples issued: 144 / Total samples completed: 0, and the run exits green.
• Accuracy phases issue every sample up front, so they can never fit in 240 s; the scorer then fails on missing videos.
• A large warmup (e.g. 10 requests x 72 single-GPU servers = ~18 min of backlog) leaks into the measured performance phase and corrupts its numbers.

New settings.drain (DrainConfig), wired through PhaseConfig.drain_timeout_s:

Field	Default	Rationale
warmup_timeout_s	240	unchanged; raise for slow-sample workloads so warmup drains before measurement
performance_timeout_s	240	unchanged
accuracy_timeout_s	None (wait for all)	abandoning in-flight accuracy samples silently invalidates the score

None waits indefinitely. The only behavior change with default configs is the accuracy drain, which previously produced wrong results whenever it triggered.

Example for slow generative workloads:

settings:
  drain:
    warmup_timeout_s: 3600
    performance_timeout_s: null   # wait for all in-flight videos
    # accuracy_timeout_s defaults to null

Fix 2: stage VBench videos as .mp4 regardless of source suffix

VBenchScorer._stage_videos kept the source extension (src.suffix or '.mp4'). VBench's load_video dispatches purely on the extension and raises bare NotImplementedError for anything but .mp4/.gif/frame dirs. trtllm-serve emits MJPEG .avi when ffmpeg is not installed server-side, so a full accuracy pass generated all 248 videos and then died at scoring.

decord (libav) detects the container by content, not extension; an MJPEG-AVI symlinked under an .mp4 name decodes correctly (verified on GB300: 81 frames, 720x1280). Staged symlinks are now always named {prompt}-{idx}.mp4.

Fix 3: vbench_runner.py fails on torch >= 2.6 checkpoint loading

The wan2.2 accuracy subproject resolves torch 2.12, where torch.load defaults to weights_only=True (changed in 2.6). VBench's reference checkpoints (motion_smoothness AMT, RAFT) are full pickles and fail with UnpicklingError. The runner now defaults torch.load to weights_only=False within the scorer subprocess only; callers passing weights_only explicitly (e.g. torch.hub) are unaffected, and the parent benchmark process keeps stock semantics.

Fix 4: opaque failure when wget/unzip are missing

vbench.utils.init_submodules downloads per-dimension weights via literal wget/unzip subprocesses on a cold cache. In a minimal client container this dies mid-evaluation with FileNotFoundError: 'wget', after the videos were already generated. The runner now preflights both tools and exits with a structured error naming what is missing and how to fix it.

Testing

• New: test_drain_respects_per_phase_timeout (session honors PhaseConfig.drain_timeout_s).
• New: test_stage_videos_renames_non_mp4_sources_to_mp4.
• tests/unit (minus two modules needing optional deps not in the default group): 1024 passed; the 5 pre-existing test_profiler.py failures are identical on the unmodified base commit.
• Hardware validation on GB300-NVL72 x72 (18 nodes): full Offline perf+accuracy run completed; warmup (720 requests) drained fully under a 3600 s warmup timeout before the measured phase; VBench scored 248/248 staged .avi videos (0.69998).

Notes for reviewers

• The accuracy_timeout_s: None default is a deliberate behavior change: the old fixed bound could only ever truncate accuracy runs into invalid scores. A hung endpoint now blocks until the job walltime, which is loud rather than silently wrong.
• Not addressed here (can file separately): accuracy-only configs are still rejected ("At least one performance dataset required"), which forces regenerating the performance dataset on every accuracy retry.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[gemini-code-assist]

Code Review

This pull request introduces configurable per-phase drain timeouts to prevent slow generative workloads from timing out prematurely, patches torch.load to support older VBench reference checkpoints, and ensures staged videos are named with an .mp4 extension to satisfy VBench's loader. The review feedback is highly constructive: it suggests checking if the VBench cache is already populated before failing on missing system tools (wget/unzip), and points out a potential TypeError when formatting a None timeout in the drain exception handler.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If the VBench cache directory (typically ~/.cache/vbench) is already pre-populated with the required model weights, wget and unzip are not needed. However, this preflight check unconditionally fails and blocks execution if either tool is missing, even if the cache is fully populated.

To support offline or air-gapped environments where the cache is pre-populated but these tools are not installed, we should check if the cache is already populated before failing.

    from pathlib import Path
    cache_dir = Path.home() / ".cache" / "vbench"
    cache_populated = cache_dir.exists() and any(cache_dir.iterdir())

    missing_tools = [t for t in ("wget", "unzip") if shutil.which(t) is None]
    if missing_tools and not cache_populated:

[gemini-code-assist]

If timeout_s is None (which is the default for accuracy phases), asyncio.wait_for will wait indefinitely and won't raise a TimeoutError. However, if a TimeoutError is raised from within the try block (e.g., propagated from another task or manually raised), formatting timeout_s with %.0f in the exception handler will raise a TypeError: must be real number, not NoneType.

To prevent the exception handler from crashing and masking the original error, we should format timeout_s safely.

Suggested change

	except TimeoutError:
	logger.error(
	"Drain timed out after 240 s with %d responses still in flight; "
	"Drain timed out after %.0f s with %d responses still in flight; "
	"proceeding to next phase.",
	timeout_s,
	phase_issuer.inflight,
	)
	except TimeoutError:
	timeout_str = f"{timeout_s:.0f}" if timeout_s is not None else "infinite"
	logger.error(
	"Drain timed out after %s s with %d responses still in flight; "
	"proceeding to next phase.",
	timeout_str,
	phase_issuer.inflight,
	)

[wu6u3tw]

Split into two focused PRs per review scoping: drain timeout configurability in #402, VBench scorer cold-start fixes (.mp4 staging, torch>=2.6, wget preflight) in #403. Closing this combined version.