Add --duration flag to perf for time-budgeted benchmarking

src/winml/modelkit/commands/_live_chart.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import time
 from typing import Any
 
 from rich.console import Console
@@ -41,11 +42,17 @@ def __init__(
         chart_height: int = 15,
         poll_interval_ms: int = 100,
         device_kind: str | None = None,
+        duration_sec: float | None = None,
     ) -> None:
         self._total = total_iterations
         self._warmup = warmup
         self._model_id = model_id
         self._device = device
+        # When set, the benchmark phase runs on a wall-clock budget instead of a
+        # fixed iteration count, so progress is reported as elapsed/total time.
+        # ``_bench_start`` is stamped on the first benchmark-phase update().
+        self._duration_sec = duration_sec
+        self._bench_start: float | None = None
         # `device_kind` is the value HWMonitor resolved at start() — pass it
         # in when you want the legend to reflect what's actually polled (e.g.
         # "auto" that resolved to GPU). Falls back to the requested string
@@ -95,6 +102,15 @@ def update(
         if self._live is None:
             return
 
+        # Stamp the start of the timed benchmark phase so duration-based progress
+        # is measured from the first post-warmup iteration.
+        if (
+            self._duration_sec is not None
+            and self._bench_start is None
+            and iteration > self._warmup
+        ):
+            self._bench_start = time.perf_counter()
+
         try:
             chart_renderable = self._render_chart(util_samples, cpu_samples)
             status_line = self._render_status(
@@ -234,15 +250,23 @@ def _render_status(
         current_util = util_samples[-1] if util_samples else 0.0
         mean_util = sum(util_samples) / len(util_samples) if util_samples else 0.0
 
-        pct = iteration / self._total if self._total > 0 else 0
+        if self._duration_sec is not None and phase == "benchmark":
+            # Duration mode: base progress on elapsed wall-clock time, since the
+            # benchmark iteration count is not known ahead of time.
+            elapsed = time.perf_counter() - self._bench_start if self._bench_start else 0.0
+            pct = min(elapsed / self._duration_sec, 1.0) if self._duration_sec > 0 else 0.0
+            shown = min(elapsed, self._duration_sec)
+            progress = f"[green]Time: {shown:.1f}/{self._duration_sec:.0f}s[/green]"
+        else:
+            pct = iteration / self._total if self._total > 0 else 0
+            if phase == "warmup":
+                progress = f"[yellow]Warmup: {iteration}/{self._warmup}[/yellow]"
+            else:
+                progress = f"[green]Iter: {effective_iter}/{total_bench}[/green]"
+
         bar_len = int(pct * 20)
         bar = f"[{'=' * bar_len}{' ' * (20 - bar_len)}]"
 
-        if phase == "warmup":
-            progress = f"[yellow]Warmup: {iteration}/{self._warmup}[/yellow]"
-        else:
-            progress = f"[green]Iter: {effective_iter}/{total_bench}[/green]"
-
         throughput = 1000.0 / latency_ms if latency_ms > 0 else 0.0
 
         # Row 1: Progress

src/winml/modelkit/commands/perf.py

@@ -17,6 +17,7 @@
 
 import json
 import logging
+import time
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from pathlib import Path
@@ -37,6 +38,7 @@
 
 if TYPE_CHECKING:
     import contextlib
+    from collections.abc import Iterator
 
     from ..models.winml.base import WinMLPreTrainedModel
     from ..models.winml.composite_model import WinMLCompositeModel
@@ -88,6 +90,10 @@ class BenchmarkConfig:
     precision: str = "auto"
     iterations: int = 100
     warmup: int = 10
+    # When set, the benchmark phase runs for this many wall-clock seconds (after
+    # warmup) instead of a fixed ``iterations`` count. Ideal with --monitor,
+    # whose PDH counters need time to emit real utilization data.
+    duration: float | None = None
     batch_size: int = 1
     output_path: Path | None = None
     no_quantize: bool = False
@@ -181,8 +187,16 @@ def to_dict(self) -> dict[str, Any]:
                 "ep": self.actual_ep,
                 "ep_options": self.config.ep_options,
                 "precision": self.config.precision,
-                "iterations": self.config.iterations,
+                # In duration mode the run isn't bounded by a fixed count, so
+                # report the actual number of timed (post-warmup) samples rather
+                # than the unused ``--iterations`` default.
+                "iterations": (
+                    len(self.raw_samples_ms)
+                    if self.config.duration is not None
+                    else self.config.iterations
+                ),
                 "warmup": self.config.warmup,
+                "duration_sec": self.config.duration,
                 "batch_size": self.config.batch_size,
                 "effective_batch_size": self.effective_batch_size,
                 "timestamp": self.timestamp,
@@ -927,7 +941,13 @@ def _run_benchmark_simple(self) -> PerfStats:
 
         session = self._single._session
         with session.perf(warmup=self.config.warmup) as stats:
-            _run_simple_loop(session, self._inputs, total_iterations)
+            _run_simple_loop(
+                session,
+                self._inputs,
+                total_iterations,
+                warmup=self.config.warmup,
+                duration_sec=self.config.duration,
+            )
 
         return stats
 
@@ -990,6 +1010,7 @@ def _run_benchmark_monitored(self) -> PerfStats:
                 warmup=self.config.warmup,
                 model_id=self.config.model_id,
                 device=monitor_device,
+                duration_sec=self.config.duration,
             )
 
             # Store hardware metrics
@@ -1067,6 +1088,7 @@ def _perf_modules(
     task: str | None,
     iterations: int,
     warmup: int,
+    duration: float | None = None,
     batch_size: int,
     no_quantize: bool,
     no_optimize: bool,
@@ -1097,6 +1119,8 @@ def _perf_modules(
         task: Explicit task override, or None for auto-detection.
         iterations: Number of benchmark iterations.
         warmup: Number of warmup iterations.
+        duration: When set, run the benchmark phase for this many wall-clock
+            seconds (after warmup) instead of a fixed ``iterations`` count.
         batch_size: Batch size for input generation.
         no_quantize: If True, skip quantization during the per-module build.
         no_optimize: If True, skip graph optimization during the per-module build.
@@ -1312,15 +1336,21 @@ def _perf_modules(
                             warmup=warmup,
                             model_id=label,
                             device=resolved_device,
+                            duration_sec=duration,
                         )
                         # Collect inside the `with` block: hw_ctx.__exit__
                         # stops the monitor, so to_dict() must read while it's
                         # still live (mirrors the single-model path).
                         hw_metrics = hw.to_dict()
                 else:
                     with session.perf(warmup=warmup) as stats:
-                        for _ in range(total_iters):
-                            session.run(inputs)
+                        _run_simple_loop(
+                            session,
+                            inputs,
+                            total_iters,
+                            warmup=warmup,
+                            duration_sec=duration,
+                        )
 
                 mod_stats = stats
                 result_entry: dict[str, Any] = {
@@ -1340,6 +1370,10 @@ def _perf_modules(
                     "throughput_sps": (
                         round(1000.0 / mod_stats.mean_ms, 2) if mod_stats.mean_ms > 0 else 0.0
                     ),
+                    # Actual timed sample count. Under a --duration budget each
+                    # module runs a different number of iterations, so this is
+                    # recorded per instance rather than as one top-level value.
+                    "iterations": len(mod_stats.samples_ms),
                 }
                 if hw_metrics:
                     result_entry["hw_monitor"] = hw_metrics
@@ -1389,8 +1423,11 @@ def _perf_modules(
         "model_id": hf_model,
         "module_class": module_class,
         "instance_count": len(all_results),
-        "iterations": iterations,
+        # Duration mode has no single iteration count (each instance runs its
+        # own — see per-instance "iterations"), so the top-level value is null.
+        "iterations": None if duration is not None else iterations,
         "warmup": warmup,
+        "duration_sec": duration,
         "instances": all_results,
     }
     output.parent.mkdir(parents=True, exist_ok=True)
@@ -1707,6 +1744,38 @@ def _print_model_info(
     console.print()
 
 
+def _benchmark_indices(
+    total_iterations: int,
+    warmup: int,
+    duration_sec: float | None,
+) -> Iterator[int]:
+    """Yield 0-based iteration indices for a benchmark run.
+
+    Always yields ``warmup`` warmup indices first (these fall inside the
+    PerfStats warmup window and are excluded from statistics). Then:
+
+    * ``duration_sec is None`` -> yields indices up to ``total_iterations``.
+    * ``duration_sec`` set -> keeps yielding until that many wall-clock seconds
+      elapse, measured from the end of warmup, always running at least one
+      benchmark iteration so stats are never empty.
+    """
+    i = 0
+    while i < warmup:
+        yield i
+        i += 1
+    if duration_sec is None:
+        while i < total_iterations:
+            yield i
+            i += 1
+        return
+    start = time.perf_counter()
+    while True:
+        yield i
+        i += 1
+        if time.perf_counter() - start >= duration_sec:
+            return
+
+
 def _run_monitored_loop(
     session: Any,
     inputs: dict[str, Any],
@@ -1717,6 +1786,7 @@ def _run_monitored_loop(
     warmup: int,
     model_id: str,
     device: str,
+    duration_sec: float | None = None,
 ) -> None:
     """Run the benchmark iteration loop with live hardware monitoring."""
     display = LiveMonitorDisplay(
@@ -1725,9 +1795,10 @@ def _run_monitored_loop(
         model_id=model_id,
         device=device,
         device_kind=getattr(hw, "device_kind", None),
+        duration_sec=duration_sec,
     )
     with display:
-        for i in range(total_iterations):
+        for i in _benchmark_indices(total_iterations, warmup, duration_sec):
             session.run(inputs)
 
             latest_latency = stats.all_samples_ms[-1] if stats.all_samples_ms else 0
@@ -1747,9 +1818,16 @@ def _run_simple_loop(
     session: Any,
     inputs: dict[str, Any],
     total_iterations: int,
+    *,
+    warmup: int = 0,
+    duration_sec: float | None = None,
 ) -> None:
-    """Run the benchmark iteration loop with periodic debug logging."""
-    for i in range(total_iterations):
+    """Run the benchmark iteration loop with periodic debug logging.
+
+    When ``duration_sec`` is set, the benchmark phase (after ``warmup``) runs
+    until the wall-clock duration elapses instead of a fixed iteration count.
+    """
+    for i in _benchmark_indices(total_iterations, warmup, duration_sec):
         session.run(inputs)
 
         if (i + 1) % max(1, total_iterations // 10) == 0:
@@ -1785,6 +1863,7 @@ def _run_simple_loop(
     "memory": "--memory",
     "op_tracing": "--op-tracing",
     "batch_size": "--batch-size",
+    "duration": "--duration",
 }
 
 # Subsets of the above that the model-id auto-build path honors, so they are
@@ -2129,6 +2208,14 @@ def _resolve_composite_components_for_perf(model: str, task: str | None) -> dict
     show_default=True,
     help="Number of warmup iterations (excluded from statistics; must be >= 0)",
 )
+@click.option(
+    "--duration",
+    type=click.FloatRange(min=0, min_open=True),
+    default=None,
+    help="Run the benchmark for this many seconds (after warmup) instead of a "
+    "fixed --iterations count. Ideal with --monitor, whose PDH counters need "
+    "time to emit real utilization data. Not valid with --op-tracing.",
+)
 @cli_utils.device_option(
     required=False,
     default="auto",
@@ -2246,6 +2333,7 @@ def perf(
     submodel: str | None,
     iterations: int,
     warmup: int,
+    duration: float | None,
     device: str,
     precision: str,
     ep: EPNameOrAlias | None,
@@ -2346,6 +2434,15 @@ def perf(
     # the inference session for both HF model IDs and ONNX file inputs.
     ep_provider_options = cli_utils.parse_ep_options(ep_options)
 
+    # --duration replaces the fixed iteration count with a wall-clock budget.
+    # Op-tracing runs its own fixed, small iteration count, so the two are
+    # mutually exclusive.
+    if duration is not None and op_tracing:
+        raise click.UsageError(
+            "--duration is not valid with --op-tracing "
+            "(op-tracing runs a fixed, small iteration count)."
+        )
+
     json_mode = output_format == "json"
     console = Console(stderr=True) if json_mode else Console()
 
@@ -2470,6 +2567,7 @@ def perf(
             task=task,
             iterations=iterations,
             warmup=warmup,
+            duration=duration,
             batch_size=batch_size,
             no_quantize=not quant,
             no_optimize=not optimize,
@@ -2577,6 +2675,7 @@ def perf(
         precision=precision.lower(),
         iterations=iterations,
         warmup=warmup,
+        duration=duration,
         batch_size=batch_size,
         output_path=output,
         no_quantize=not quant,