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( 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/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/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/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 ): 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.