fix(perf): unify HF and ONNX paths through PerfBenchmark

src/winml/modelkit/commands/perf.py

@@ -318,13 +318,25 @@ def run(self) -> BenchmarkResult:
         return self._collect_results(stats)
 
     def _load_model(self) -> None:
-        """Load model via WinMLAutoModel (handles both HF and ONNX)."""
+        """Load model via WinMLAutoModel.
+
+        Both HF model IDs and pre-exported .onnx files flow through this
+        single path so latency numbers stay comparable: HF runs export →
+        optimize → [quantize] → [compile], and ONNX runs the same pipeline
+        minus export.
+        """
         from ..config import WinMLBuildConfig
         from ..models import WinMLAutoModel
 
         model_id = self.config.model_id
         model_path = Path(model_id)
-        is_onnx = model_path.suffix.lower() == ".onnx" and model_path.exists()
+        is_onnx = model_path.suffix.lower() == ".onnx"
+        if is_onnx and not model_path.exists():
+            # Surface a clear error for programmatic callers. The CLI guards
+            # this earlier, but without this check from_pretrained would fall
+            # through to HF loading and produce a confusing "not a valid JSON
+            # file" error from AutoConfig.
+            raise FileNotFoundError(f"ONNX file not found: {model_path}")
 
         # Only override config when user explicitly passes --no-quantize
         override = None
@@ -943,10 +955,7 @@ def _run_monitored_loop(
     model_id: str,
     device: str,
 ) -> None:
-    """Run the benchmark iteration loop with live hardware monitoring.
-
-    Shared by both HF-path (PerfBenchmark) and ONNX-path (_run_onnx_benchmark).
-    """
+    """Run the benchmark iteration loop with live hardware monitoring."""
     display = LiveMonitorDisplay(
         total_iterations=total_iterations,
         warmup=warmup,
@@ -976,120 +985,14 @@ def _run_simple_loop(
     inputs: dict[str, Any],
     total_iterations: int,
 ) -> None:
-    """Run the benchmark iteration loop with periodic debug logging.
-
-    Shared by both HF-path (PerfBenchmark) and ONNX-path (_run_onnx_benchmark).
-    """
+    """Run the benchmark iteration loop with periodic debug logging."""
     for i in range(total_iterations):
         session.run(inputs)
 
         if (i + 1) % max(1, total_iterations // 10) == 0:
             logger.debug("Progress: %d/%d", i + 1, total_iterations)
 
 
-# =============================================================================
-# ONNX Direct Benchmark
-# =============================================================================
-
-
-def _run_onnx_benchmark(
-    onnx_path: Path,
-    *,
-    device: str,
-    iterations: int,
-    warmup: int,
-    batch_size: int,
-    config: BenchmarkConfig,
-) -> BenchmarkResult:
-    """Benchmark an ONNX file directly via WinMLSession (no HF build).
-
-    Creates a WinMLSession, reads io_config for input shapes,
-    generates random inputs, and runs the standard benchmark loop.
-    """
-    from ..session import WinMLSession
-
-    session = WinMLSession(onnx_path=onnx_path, device=device, ep=config.ep)
-
-    # Generate random inputs from session's I/O config
-    io_cfg = session.io_config
-    inputs = generate_random_inputs(io_config=io_cfg, batch_size=batch_size)
-
-    # Compile session early so session.device is resolved for display
-    session.compile()
-
-    # Print model info before benchmark starts
-    _print_model_info(io_cfg, req_device=device, act_device=session.device, ep_name=session.ep_name)
-
-    # Run benchmark
-    total_iterations = warmup + iterations
-    hw_metrics = None
-    hw_ctx = None
-
-    # Determine if hardware monitoring is available
-    if config.monitor:
-        from ..session.monitor.hw_monitor import HWMonitor
-
-        if HWMonitor.is_available():
-            hw_ctx = HWMonitor(
-                poll_interval_ms=_HW_POLL_INTERVAL_MS,
-                device=session.device or device,
-                ep_name=session.ep_name,
-            )
-        else:
-            Console(stderr=True).print(
-                "[yellow]Warning:[/yellow] HWMonitor unavailable. "
-                "Running ONNX benchmark without monitoring."
-            )
-
-    if hw_ctx:
-        with session.perf(warmup=warmup) as stats, hw_ctx as hw:
-            _run_monitored_loop(
-                session,
-                inputs,
-                stats,
-                hw,
-                total_iterations=total_iterations,
-                warmup=warmup,
-                model_id=str(onnx_path.name),
-                device=session.device or device,
-            )
-            hw_metrics = hw.to_dict()
-    else:
-        with session.perf(warmup=warmup) as stats:
-            _run_simple_loop(session, inputs, total_iterations)
-
-    # Collect results
-    mean_latency_sec = stats.mean_ms / 1000.0
-    samples_per_sec = batch_size / mean_latency_sec if mean_latency_sec > 0 else 0
-    batches_per_sec = 1.0 / mean_latency_sec if mean_latency_sec > 0 else 0
-    samples = stats.samples_ms
-    std_ms = float(np.std(samples)) if samples else 0.0
-
-    return BenchmarkResult(
-        config=config,
-        input_names=io_cfg["input_names"],
-        input_shapes=[list(s) if s else [] for s in io_cfg["input_shapes"]],
-        input_types=[str(t) for t in io_cfg["input_types"]],
-        output_names=io_cfg["output_names"],
-        output_shapes=[list(s) if s else [] for s in io_cfg["output_shapes"]],
-        mean_ms=stats.mean_ms,
-        min_ms=stats.min_ms,
-        max_ms=stats.max_ms,
-        p50_ms=stats.p50_ms,
-        p90_ms=stats.p90_ms,
-        p95_ms=stats.p95_ms,
-        p99_ms=stats.p99_ms,
-        std_ms=std_ms,
-        raw_samples_ms=stats.samples_ms,
-        samples_per_sec=samples_per_sec,
-        batches_per_sec=batches_per_sec,
-        actual_device=session.device,
-        actual_task="n/a (direct ONNX)",
-        actual_ep=session.ep_name,
-        hw_monitor=hw_metrics,
-    )
-
-
 # =============================================================================
 # CLI Command
 # =============================================================================
@@ -1224,8 +1127,10 @@ def perf(
     Measures latency and throughput using random input data generated
     from the model's I/O configuration.
 
-    Accepts both HuggingFace model IDs and local .onnx files.
-    HF models go through PerfBenchmark; .onnx files use _run_onnx_benchmark.
+    Accepts both HuggingFace model IDs and local .onnx files. Both flow
+    through the same PerfBenchmark pipeline (optimize → [quantize] → [compile]
+    minus export for ONNX inputs), so latency numbers are directly comparable
+    between the two inputs.
 
     \b
     Examples:
@@ -1354,33 +1259,25 @@ def perf(
         is_onnx = model_path.suffix.lower() == ".onnx"
 
         if is_onnx:
-            # ONNX direct path -- skip HF build, benchmark via WinMLSession
+            # Validate file existence up front; otherwise WinMLAutoModel would
+            # fall through to HF loading and surface a confusing
+            # "not a valid JSON file" error from AutoConfig.
+            if not model_path.exists():
+                raise FileNotFoundError(f"ONNX file not found: {model_path}")
             if shape_config:
                 console.print(
                     "[yellow]Warning:[/yellow] --shape-config is ignored for "
                     "pre-exported ONNX files (shapes are baked into the model)."
                 )
                 config.shape_config = None
-            if not model_path.exists():
-                raise FileNotFoundError(f"ONNX file not found: {model_path}")
             console.print(f"[dim]Benchmarking ONNX:[/dim] {model_path}")
-
-            result = _run_onnx_benchmark(
-                model_path,
-                device=config.device,
-                iterations=iterations,
-                warmup=warmup,
-                batch_size=batch_size,
-                config=config,
-            )
         else:
-            # HF model path -- full build + benchmark via PerfBenchmark
             if precision != "auto":
                 console.print(f"[dim]Precision: {precision} (applied during model build)[/dim]")
             console.print(f"[dim]Loading model:[/dim] {hf_model}")
 
-            benchmark = PerfBenchmark(config)
-            result = benchmark.run()
+        benchmark = PerfBenchmark(config)
+        result = benchmark.run()
 
         # Display console report
         display_console_report(result, console)

tests/unit/commands/test_perf_cli.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import json
 import re
 from pathlib import Path
 from unittest.mock import MagicMock, patch
@@ -279,14 +280,22 @@ def test_no_quantize_false_passes_no_override(self) -> None:
         override = mock_from_pretrained.call_args.kwargs["config"]
         assert override is None
 
-    def test_cli_onnx_goes_through_onnx_benchmark(self, runner: CliRunner, tmp_path: Path) -> None:
-        """CLI with .onnx file should route through _run_onnx_benchmark."""
+    def test_cli_onnx_routes_through_perf_benchmark(
+        self, runner: CliRunner, tmp_path: Path
+    ) -> None:
+        """CLI with .onnx file should route through the same PerfBenchmark as HF.
+
+        Both paths must share the build+benchmark pipeline so latency numbers
+        from `winml perf -m hf/id` and `winml perf -m <built.onnx>` are
+        comparable (issue #596).
+        """
         onnx_file = tmp_path / "model.onnx"
         onnx_file.write_bytes(b"fake onnx")
 
         with (
-            patch(
-                "winml.modelkit.commands.perf._run_onnx_benchmark",
+            patch.object(
+                PerfBenchmark,
+                "run",
                 return_value=MagicMock(),
             ) as mock_run,
             patch(
@@ -305,6 +314,55 @@ def test_cli_onnx_goes_through_onnx_benchmark(self, runner: CliRunner, tmp_path:
         assert result.exit_code == 0, result.output
         mock_run.assert_called_once()
 
+    def test_cli_onnx_clears_shape_config_with_warning(
+        self, runner: CliRunner, tmp_path: Path
+    ) -> None:
+        """ONNX input with --shape-config: warn + clear shape_config before PerfBenchmark.
+
+        Shapes are baked into a pre-exported ONNX, so --shape-config is silently
+        ignored; we want to be sure the CLI both surfaces the warning to the
+        user and actually drops the override before constructing PerfBenchmark.
+        """
+        onnx_file = tmp_path / "model.onnx"
+        onnx_file.write_bytes(b"fake onnx")
+
+        shape_cfg_file = tmp_path / "shapes.json"
+        shape_cfg_file.write_text(json.dumps({"input_ids": [1, 128]}))
+
+        captured: dict[str, BenchmarkConfig] = {}
+
+        def capture_config(config: BenchmarkConfig) -> MagicMock:
+            captured["config"] = config
+            mock = MagicMock()
+            mock.run.return_value = MagicMock()
+            return mock
+
+        with (
+            patch(
+                "winml.modelkit.commands.perf.PerfBenchmark",
+                side_effect=capture_config,
+            ),
+            patch("winml.modelkit.commands.perf.display_console_report"),
+            patch("winml.modelkit.commands.perf.write_json_report"),
+        ):
+            result = runner.invoke(
+                perf,
+                [
+                    "-m",
+                    str(onnx_file),
+                    "--shape-config",
+                    str(shape_cfg_file),
+                    "-o",
+                    str(tmp_path / "out.json"),
+                ],
+                obj={},
+            )
+
+        assert result.exit_code == 0, result.output
+        assert "shape-config is ignored" in result.output
+        assert "Benchmarking ONNX" in result.output
+        assert captured["config"].shape_config is None
+
     def test_cli_onnx_not_found_error(self, runner: CliRunner, tmp_path: Path) -> None:
         """CLI with non-existent .onnx file should raise FileNotFoundError."""
         missing = tmp_path / "missing.onnx"