From f655f2ddf342738db99a6b99cc3a6276c49581c7 Mon Sep 17 00:00:00 2001 From: Tin-Yin Lai Date: Wed, 8 Jul 2026 14:05:00 -0700 Subject: [PATCH 1/3] Make per-phase in-flight drain timeouts configurable The benchmark session hard-coded a 240 s drain timeout after every phase. For slow generative workloads (e.g. wan2.2 text-to-video at ~2 min per sample on a single GPU), this silently abandoned every in-flight sample: a performance run could exit green with 0 samples completed (observed on GB300: 144 issued, 0 completed), and accuracy phases, which issue all samples up front, could never finish at all. Add settings.drain (DrainConfig) with per-phase timeouts: - warmup_timeout_s: 240 (raise when warmup requests take minutes each, so leftovers don't leak into the measured performance phase) - performance_timeout_s: 240 (unchanged default) - accuracy_timeout_s: None = wait for all responses, because abandoning in-flight accuracy samples silently invalidates the score PhaseConfig gains drain_timeout_s and the phase builder wires the config through. None waits indefinitely. --- .../commands/benchmark/execute.py | 11 ++++- src/inference_endpoint/config/schema.py | 43 +++++++++++++++++++ .../load_generator/session.py | 28 +++++++++--- .../unit/load_generator/test_async_session.py | 29 +++++++++++++ 4 files changed, 103 insertions(+), 8 deletions(-) diff --git a/src/inference_endpoint/commands/benchmark/execute.py b/src/inference_endpoint/commands/benchmark/execute.py index 88e337f89..42bddfc45 100644 --- a/src/inference_endpoint/commands/benchmark/execute.py +++ b/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 diff --git a/src/inference_endpoint/config/schema.py b/src/inference_endpoint/config/schema.py index 462c5a146..77893b214 100644 --- a/src/inference_endpoint/config/schema.py +++ b/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): diff --git a/src/inference_endpoint/load_generator/session.py b/src/inference_endpoint/load_generator/session.py index cdfa28415..ab3898f68 100644 --- a/src/inference_endpoint/load_generator/session.py +++ b/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, ) diff --git a/tests/unit/load_generator/test_async_session.py b/tests/unit/load_generator/test_async_session.py index 331d930c2..a6dcf673a 100644 --- a/tests/unit/load_generator/test_async_session.py +++ b/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. From 408df7194a4e0aa2b8d011e31626e4e9b1bc1045 Mon Sep 17 00:00:00 2001 From: Tin-Yin Lai Date: Wed, 8 Jul 2026 14:05:00 -0700 Subject: [PATCH 2/3] Stage VBench videos as .mp4 regardless of source suffix VBench's load_video dispatches purely on the file extension and raises NotImplementedError for anything but .mp4/.gif/frame dirs. trtllm-serve emits MJPEG .avi when ffmpeg is unavailable server-side, so a full accuracy run generated 248 videos and then failed at scoring. decord (libav) detects the real container by content, so an .avi symlinked under an .mp4 name decodes fine (verified on GB300: 81-frame 720x1280 MJPEG-AVI reads correctly). Always name the staged symlinks .mp4. --- src/inference_endpoint/evaluation/scoring.py | 7 +++++- tests/unit/evaluation/test_scoring.py | 25 ++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/inference_endpoint/evaluation/scoring.py b/src/inference_endpoint/evaluation/scoring.py index 911ec6293..a46faa3cc 100644 --- a/src/inference_endpoint/evaluation/scoring.py +++ b/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( diff --git a/tests/unit/evaluation/test_scoring.py b/tests/unit/evaluation/test_scoring.py index 3b7a5a90a..8dfa9d9cd 100644 --- a/tests/unit/evaluation/test_scoring.py +++ b/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 ): From e28dc9aa206ca7f0a5506175398feb0388494c59 Mon Sep 17 00:00:00 2001 From: Tin-Yin Lai Date: Wed, 8 Jul 2026 14:05:00 -0700 Subject: [PATCH 3/3] vbench_runner: survive torch>=2.6 checkpoints and missing wget Two cold-start failures observed on aarch64 GB300 runs: - The accuracy venv resolves torch 2.12, where torch.load defaults to weights_only=True (since 2.6). VBench's reference checkpoints (motion_smoothness AMT, RAFT) are full pickles and fail with UnpicklingError. Default torch.load to weights_only=False in this subprocess only; explicit weights_only callers are unaffected. - vbench.utils.init_submodules downloads per-dimension weights via literal wget/unzip subprocesses; without them the run dies with an opaque FileNotFoundError after videos were already generated. Fail fast preflight with a structured error naming the missing tools. --- .../accuracy/vbench_runner.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/examples/09_Wan22_VideoGen_Example/accuracy/vbench_runner.py b/examples/09_Wan22_VideoGen_Example/accuracy/vbench_runner.py index db3c41229..3818cbe12 100644 --- a/examples/09_Wan22_VideoGen_Example/accuracy/vbench_runner.py +++ b/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(