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

examples/09_Wan22_VideoGen_Example/accuracy/vbench_runner.py

@@ -27,7 +27,9 @@
 """
 
 import argparse
+import functools
 import json
+import shutil
 import sys
 import traceback
 from importlib.resources import files as _pkg_files
@@ -37,6 +39,31 @@
 from vbench import VBench
 
 
+def _default_torch_load_to_full_pickles() -> None:
+    """Make torch.load default to weights_only=False in this subprocess.
+
+    VBench's reference checkpoints (e.g. motion_smoothness AMT, RAFT) are full
+    pickled objects, written before torch 2.6 flipped torch.load's default to
+    weights_only=True; loading them with the new default fails with
+    UnpicklingError (e.g. "Unsupported global: typing.OrderedDict"). They are
+    the VBench-sanctioned reference weights, so loading them unrestricted here
+    matches upstream VBench behavior on the torch versions it was written for.
+    Callers that pass weights_only explicitly (e.g. torch.hub) are unaffected;
+    scoped to this subprocess only, the parent benchmark keeps stock semantics.
+    """
+    orig_load = torch.load
+
+    @functools.wraps(orig_load)
+    def _load(*args, **kwargs):
+        kwargs.setdefault("weights_only", False)
+        return orig_load(*args, **kwargs)
+
+    torch.load = _load
+
+
+_default_torch_load_to_full_pickles()
+
+
 def _emit_error(exc: BaseException) -> None:
     """Print a structured JSON error line on stderr for the parent to surface."""
     payload = {
@@ -85,6 +112,31 @@ def main() -> int:
     )
     args = parser.parse_args()
 
+    # vbench.utils.init_submodules downloads per-dimension weights via literal
+    # `wget`/`unzip` subprocesses on a cold cache. Fail fast with a clear
+    # message instead of a FileNotFoundError mid-evaluation (after the videos
+    # have already been generated and staged).
+    missing_tools = [t for t in ("wget", "unzip") if shutil.which(t) is None]
+    if missing_tools:
+        print(
+            json.dumps(
+                {
+                    "status": "error",
+                    "type": "MissingSystemDependency",
+                    "message": (
+                        f"Required system tool(s) not found: {', '.join(missing_tools)}. "
+                        "VBench downloads its per-dimension model weights via "
+                        "wget/unzip on first use. Install them in the client "
+                        "environment (e.g. `apt-get install wget unzip`) or "
+                        "pre-populate the VBench cache directory."
+                    ),
+                }
+            ),
+            file=sys.stderr,
+            flush=True,
+        )
+        return 2
+
     if not torch.cuda.is_available() and not args.allow_cpu:
         print(
             json.dumps(

src/inference_endpoint/commands/benchmark/execute.py

@@ -452,6 +452,7 @@ def _build_phases(
 ) -> list[PhaseConfig]:
     """Build the phase list from BenchmarkContext."""
     phases: list[PhaseConfig] = []
+    drain_cfg = ctx.config.settings.drain
 
     # Warmup phase (optional, before performance)
     warmup_cfg = ctx.config.settings.warmup
@@ -479,6 +480,7 @@ def _build_phases(
                 warmup_dataset,
                 PhaseType.WARMUP,
                 drain_after=warmup_cfg.drain,
+                drain_timeout_s=drain_cfg.warmup_timeout_s,
             )
         )
 
@@ -489,6 +491,7 @@ def _build_phases(
             ctx.rt_settings,
             ctx.dataloader,
             PhaseType.PERFORMANCE,
+            drain_timeout_s=drain_cfg.performance_timeout_s,
             strategy=perf_strategy,
         )
     )
@@ -521,7 +524,13 @@ def _build_phases(
             load_pattern=acc_load_pattern,
         )
         phases.append(
-            PhaseConfig(eval_cfg.dataset_name, acc_settings, acc_ds, PhaseType.ACCURACY)
+            PhaseConfig(
+                eval_cfg.dataset_name,
+                acc_settings,
+                acc_ds,
+                PhaseType.ACCURACY,
+                drain_timeout_s=drain_cfg.accuracy_timeout_s,
+            )
         )
 
     return phases

src/inference_endpoint/config/schema.py

@@ -483,6 +483,48 @@ class WarmupConfig(BaseModel):
     ] = Field(42, description="RNG seed for warmup scheduling and sample ordering")
 
 
+class DrainConfig(BaseModel):
+    """Per-phase in-flight drain timeouts (seconds). ``None`` waits indefinitely.
+
+    After a phase finishes issuing, the session waits up to the phase's drain
+    timeout for in-flight responses before moving on. For slow generative
+    workloads (e.g. text-to-video at minutes per sample), the historical fixed
+    240 s bound silently abandoned every in-flight sample: a run could end
+    "successfully" with 0 samples completed, and accuracy phases (which issue
+    all samples up front) could never finish at all.
+    """
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    warmup_timeout_s: float | None = Field(
+        240.0,
+        gt=0,
+        description=(
+            "Drain timeout after the warmup phase. Raise it when warmup "
+            "requests take minutes each, so leftovers don't leak into the "
+            "measured performance phase. None = wait for all."
+        ),
+    )
+    performance_timeout_s: float | None = Field(
+        240.0,
+        gt=0,
+        description=(
+            "Drain timeout after the performance phase. None = wait for all "
+            "in-flight responses (recommended when per-sample latency can "
+            "exceed 240 s, e.g. video generation)."
+        ),
+    )
+    accuracy_timeout_s: float | None = Field(
+        None,
+        gt=0,
+        description=(
+            "Drain timeout after accuracy phases. Default None: wait for all "
+            "responses, because abandoning in-flight accuracy samples "
+            "silently invalidates the score."
+        ),
+    )
+
+
 @cyclopts.Parameter(name="*")
 class Settings(BaseModel):
     """Test settings."""
@@ -493,6 +535,7 @@ class Settings(BaseModel):
     load_pattern: LoadPattern = Field(default_factory=LoadPattern)
     client: HTTPClientConfig = Field(default_factory=HTTPClientConfig)
     warmup: WarmupConfig = Field(default_factory=WarmupConfig)
+    drain: DrainConfig = Field(default_factory=DrainConfig)
 
 
 class OfflineSettings(Settings):

src/inference_endpoint/evaluation/scoring.py

@@ -1009,7 +1009,12 @@ def _stage_videos(
             # strict=True surfaces missing/unmounted sources here, not as an
             # opaque decord read failure inside VBench 30 minutes later.
             resolved_src = src.resolve(strict=True)
-            dst = staged_dir / f"{safe_prompt}-{idx}{src.suffix or '.mp4'}"
+            # Always stage as .mp4: VBench's load_video dispatches purely on
+            # the file extension and raises NotImplementedError for anything
+            # else (e.g. the MJPEG .avi that trtllm-serve emits when ffmpeg
+            # is unavailable server-side). decord detects the real container
+            # by content, so a non-mp4 file under an .mp4 name decodes fine.
+            dst = staged_dir / f"{safe_prompt}-{idx}.mp4"
             dst.symlink_to(resolved_src)
 
     def _run_vbench_subprocess(

src/inference_endpoint/load_generator/session.py

@@ -90,6 +90,12 @@ class PhaseConfig:
     dataset: Dataset
     phase_type: PhaseType = PhaseType.PERFORMANCE
     drain_after: bool = True
+    drain_timeout_s: float | None = 240.0
+    """Max seconds to wait for in-flight responses after issuing ends.
+
+    ``None`` waits indefinitely. Configured per phase type via
+    ``settings.drain`` (see ``config.schema.DrainConfig``).
+    """
     strategy: LoadStrategy | None = field(default=None, compare=False)
 
 
@@ -417,7 +423,7 @@ async def _run_phase(self, phase: PhaseConfig) -> PhaseResult | None:
             self._strategy_task = None
 
         if phase.drain_after:
-            await self._drain_inflight(phase_issuer)
+            await self._drain_inflight(phase_issuer, phase.drain_timeout_s)
 
         if phase.phase_type == PhaseType.PERFORMANCE:
             self._publish_session_event(SessionEventType.STOP_PERFORMANCE_TRACKING)
@@ -442,21 +448,29 @@ async def _run_phase(self, phase: PhaseConfig) -> PhaseResult | None:
             end_time_ns=phase_end,
         )
 
-    async def _drain_inflight(self, phase_issuer: PhaseIssuer) -> None:
+    async def _drain_inflight(
+        self, phase_issuer: PhaseIssuer, timeout_s: float | None = 240.0
+    ) -> None:
         """Wait for all in-flight responses from this phase to complete.
 
-        Hard-bounded at 240 s; logs an error and returns if exceeded so the
-        next phase starts regardless of stuck requests."""
+        Bounded by ``timeout_s`` (``None`` = wait indefinitely); on timeout,
+        logs an error and returns so the next phase starts regardless of
+        stuck requests."""
         if phase_issuer.inflight <= 0 or self._stop_requested:
             return
-        logger.info("Draining %d in-flight responses...", phase_issuer.inflight)
+        logger.info(
+            "Draining %d in-flight responses (timeout: %s)...",
+            phase_issuer.inflight,
+            f"{timeout_s:.0f} s" if timeout_s is not None else "none",
+        )
         self._drain_event.clear()
         try:
-            await asyncio.wait_for(self._drain_event.wait(), timeout=240.0)
+            await asyncio.wait_for(self._drain_event.wait(), timeout=timeout_s)
         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,
             )
 

tests/unit/evaluation/test_scoring.py

@@ -466,6 +466,31 @@ def test_stage_clears_stale_files_from_prior_run(
         assert names == ["a cat-0.mp4", "a dog-0.mp4", "a tree-0.mp4"]
         assert not zombie.exists()
 
+    def test_stage_videos_renames_non_mp4_sources_to_mp4(
+        self, dataset, staged, vbench_project, tmp_path
+    ):
+        """Non-mp4 sources (e.g. trtllm-serve's MJPEG .avi) stage as .mp4.
+
+        VBench's load_video dispatches purely on the file extension and raises
+        NotImplementedError for anything but .mp4; decord detects the real
+        container by content, so an .avi under an .mp4 name decodes fine.
+        """
+        report_dir, _ = staged
+        avi = tmp_path / "video_x.avi"
+        avi.write_bytes(b"")
+        scorer = VBenchScorer(
+            dataset_name="vid_acc",
+            dataset=dataset,
+            report_dir=report_dir,
+            ground_truth_column="prompt",
+            vbench_project_path=vbench_project,
+        )
+        staged_dir = report_dir / "vbench_videos"
+        scorer._stage_videos(staged_dir, [str(avi)], ["a cat"])
+        (staged_file,) = staged_dir.iterdir()
+        assert staged_file.name == "a cat-0.mp4"
+        assert staged_file.resolve() == avi.resolve()
+
     def test_subprocess_failure_includes_stderr_tail(
         self, dataset, staged, vbench_project, monkeypatch, tmp_path
     ):

tests/unit/load_generator/test_async_session.py

@@ -444,6 +444,35 @@ async def test_recv_none_triggers_stop(self):
         result = await asyncio.wait_for(session.run(phases), timeout=10.0)
         assert result is not None
 
+    @pytest.mark.asyncio
+    async def test_drain_respects_per_phase_timeout(self, caplog):
+        """A phase's drain_timeout_s bounds the in-flight drain wait.
+
+        With a never-responding issuer and a sub-second timeout, the drain
+        must expire (logging an error) instead of hanging for the historical
+        hard-coded 240 s.
+        """
+        loop = asyncio.get_running_loop()
+        publisher = FakePublisher()
+
+        issuer = FakeIssuer()
+        issuer._loop = loop
+        issuer._auto_respond = False  # samples issue but never complete
+
+        session = BenchmarkSession(issuer, publisher, loop)
+        phases = [
+            PhaseConfig(
+                "perf",
+                _make_settings(n_samples=2),
+                FakeDataset(2),
+                drain_timeout_s=0.1,
+            ),
+        ]
+
+        result = await asyncio.wait_for(session.run(phases), timeout=10.0)
+        assert result is not None
+        assert "Drain timed out after 0 s" in caplog.text or "Drain timed out" in caplog.text
+
     @pytest.mark.asyncio
     async def test_streaming_query_completes_via_queryresult(self):
         """Streaming: StreamChunks publish timing events, QueryResult handles completion.