Show optimization capability probe progress

src/winml/modelkit/commands/optimize.py

@@ -28,6 +28,7 @@
 import click
 from rich.console import Console
 from rich.markup import escape
+from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn
 
 from ..onnx import load_onnx, save_onnx
 from ..utils import cli as cli_utils
@@ -256,8 +257,20 @@ def _run_check_optim(model: Path, all_caps: dict[str, Any], verbose: bool) -> No
         f"[bold]Probing {probe_count} optimization capabilities...[/bold] "
         "[dim](this can take a while on large models)[/dim]"
     )
-    with console.status("[bold]Analyzing...[/bold]", spinner="dots"):
-        findings = analyze_model(onnx_model, all_caps)
+    with Progress(
+        TextColumn("[bold]Checking[/bold] [cyan]{task.description}[/cyan]"),
+        BarColumn(),
+        MofNCompleteColumn(),
+        TimeElapsedColumn(),
+        console=console,
+    ) as progress:
+        task_id = progress.add_task("preparing capabilities", total=probe_count)
+        findings = analyze_model(
+            onnx_model,
+            all_caps,
+            on_probe_start=lambda name: progress.update(task_id, description=name),
+            on_probe_complete=lambda _: progress.advance(task_id),
+        )
 
     _render_check_optim(console, findings, verbose)
 

src/winml/modelkit/optim/analysis.py

@@ -32,7 +32,7 @@
 
 
 if TYPE_CHECKING:
-    from collections.abc import Iterator
+    from collections.abc import Callable, Iterator
 
     from .registry import CapabilityDef
 
@@ -283,6 +283,8 @@ def _run_pipe(pipe: Any, model: ModelProto, config: Any) -> ModelProto:
 def _iter_findings(
     model: ModelProto,
     capabilities: dict[str, CapabilityDef],
+    on_probe_start: Callable[[str], None] | None = None,
+    on_probe_complete: Callable[[str], None] | None = None,
 ) -> Iterator[tuple[CapabilityFinding, ModelProto]]:
     """Yield ``(finding, produced_model)`` for every applicable optimization.
 
@@ -303,6 +305,18 @@ def _iter_findings(
     from .pipes import PIPES
     from .registry import BoolCapability, auto_enable_dependencies
 
+    remaining_probes = Counter(
+        name
+        for pipe_class in PIPES
+        for name, cap in pipe_class.capabilities.items()
+        if isinstance(cap, BoolCapability) and not cap.default
+    )
+
+    def complete_probe(cap_name: str) -> None:
+        remaining_probes[cap_name] -= 1
+        if remaining_probes[cap_name] == 0 and on_probe_complete is not None:
+            on_probe_complete(cap_name)
+
     # Baseline kwargs = every capability at its default value.
     default_kwargs = {cap.python_name: cap.default for cap in capabilities.values()}
     kebab_defaults = {name: cap.default for name, cap in capabilities.items()}
@@ -315,6 +329,11 @@ def _iter_findings(
 
     for pipe_class in PIPES:
         pipe = pipe_class()
+        probe_caps = [
+            (name, cap)
+            for name, cap in pipe.capabilities.items()
+            if isinstance(cap, BoolCapability) and not cap.default
+        ]
 
         base_config = pipe.build_config(**default_kwargs)
         try:
@@ -328,76 +347,76 @@ def _iter_findings(
                 getattr(pipe, "name", pipe_class.__name__),
                 exc,
             )
+            for cap_name, _ in probe_caps:
+                complete_probe(cap_name)
             continue
         base_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {}
         _collect_nodes(base_out.graph, (), base_nodes)
         base_inits = _collect_initializers(base_out)
 
-        probe_caps = [
-            (name, cap)
-            for name, cap in pipe.capabilities.items()
-            if isinstance(cap, BoolCapability) and not cap.default
-        ]
-
         for cap_name, cap in probe_caps:
-            # Enable only this capability (plus its dependencies) on top of the
-            # all-defaults configuration.
-            kebab = dict(kebab_defaults)
-            kebab[cap_name] = True
-            kebab = auto_enable_dependencies(kebab, capabilities)
-            probe_kwargs = {
-                capabilities[name].python_name: value
-                for name, value in kebab.items()
-                if name in capabilities
-            }
-
-            probe_config = pipe.build_config(**probe_kwargs)
-            should_process = getattr(pipe, "should_process", None)
-            if callable(should_process) and not should_process(probe_config):
-                # Pipe would not run for this capability — nothing to apply.
-                continue
-
-            # NB: this deliberately calls pipe.process directly rather than
-            # reusing _run_pipe. The probe must distinguish "pipe opts out"
-            # (handled above by continuing without emitting a finding) from
-            # "pipe ran but changed nothing", and it must isolate a failing
-            # capability in try/except so one bad probe cannot abort the scan —
-            # neither of which _run_pipe's return-unchanged contract expresses.
+            if on_probe_start is not None:
+                on_probe_start(cap_name)
             try:
-                probe_out = pipe.process(_clone(current), probe_config)
-            except Exception as exc:
-                logger.warning(
-                    "Could not evaluate capability '%s' on pipe '%s': %s",
-                    cap_name,
-                    pipe.name,
-                    exc,
+                # Enable only this capability (plus its dependencies) on top of
+                # the all-defaults configuration.
+                kebab = dict(kebab_defaults)
+                kebab[cap_name] = True
+                kebab = auto_enable_dependencies(kebab, capabilities)
+                probe_kwargs = {
+                    capabilities[name].python_name: value
+                    for name, value in kebab.items()
+                    if name in capabilities
+                }
+
+                probe_config = pipe.build_config(**probe_kwargs)
+                should_process = getattr(pipe, "should_process", None)
+                if callable(should_process) and not should_process(probe_config):
+                    # Pipe would not run for this capability — nothing to apply.
+                    continue
+
+                # NB: this deliberately calls pipe.process directly rather than
+                # reusing _run_pipe. The probe must distinguish "pipe opts out"
+                # (handled above by continuing without emitting a finding) from
+                # "pipe ran but changed nothing", and it must isolate a failing
+                # capability in try/except so one bad probe cannot abort the scan.
+                try:
+                    probe_out = pipe.process(_clone(current), probe_config)
+                except Exception as exc:
+                    logger.warning(
+                        "Could not evaluate capability '%s' on pipe '%s': %s",
+                        cap_name,
+                        pipe.name,
+                        exc,
+                    )
+                    continue
+
+                probe_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {}
+                _collect_nodes(probe_out.graph, (), probe_nodes)
+                removed, added, modified = _diff_nodes(base_nodes, probe_nodes)
+
+                probe_inits = _collect_initializers(probe_out)
+                rem_init, add_init, mod_init = _diff_initializers(base_inits, probe_inits)
+
+                finding = CapabilityFinding(
+                    name=cap.name,
+                    python_name=cap.python_name,
+                    enable_flag=f"--enable-{cap.name}",
+                    category=cap.category.value,
+                    description=cap.description,
+                    pipe_name=pipe.name,
+                    removed_nodes=removed,
+                    added_nodes=added,
+                    modified_nodes=modified,
+                    removed_initializers=rem_init,
+                    added_initializers=add_init,
+                    modified_initializers=mod_init,
                 )
-                continue
-
-            probe_nodes: dict[tuple[Any, ...], tuple[bytes, NodeRef]] = {}
-            _collect_nodes(probe_out.graph, (), probe_nodes)
-            removed, added, modified = _diff_nodes(base_nodes, probe_nodes)
-
-            probe_inits = _collect_initializers(probe_out)
-            rem_init, add_init, mod_init = _diff_initializers(base_inits, probe_inits)
-
-            finding = CapabilityFinding(
-                name=cap.name,
-                python_name=cap.python_name,
-                enable_flag=f"--enable-{cap.name}",
-                category=cap.category.value,
-                description=cap.description,
-                pipe_name=pipe.name,
-                removed_nodes=removed,
-                added_nodes=added,
-                modified_nodes=modified,
-                removed_initializers=rem_init,
-                added_initializers=add_init,
-                modified_initializers=mod_init,
-            )
 
-            if finding.applicable:
-                yield finding, probe_out
+                if finding.applicable:
+                    yield finding, probe_out
+            finally:
+                complete_probe(cap_name)
 
         # Advance the pipeline exactly as the real optimizer would.
         current = base_out
@@ -406,6 +425,9 @@ def _iter_findings(
 def analyze_model(
     model: ModelProto,
     capabilities: dict[str, CapabilityDef],
+    *,
+    on_probe_start: Callable[[str], None] | None = None,
+    on_probe_complete: Callable[[str], None] | None = None,
 ) -> list[CapabilityFinding]:
     """Probe every applicable optimization capability against ``model``.
 
@@ -418,12 +440,24 @@ def analyze_model(
         model: The input ONNX model (never modified).
         capabilities: The full capability registry (kebab-case keyed), e.g.
             from ``optim.pipes.get_all_capabilities()``.
+        on_probe_start: Optional callback invoked with the capability name
+            before each probe begins.
+        on_probe_complete: Optional callback invoked with the capability name
+            after all of its probes finish or are skipped.
 
     Returns:
         Applicable findings in pipeline order, each naming the affected nodes
         and constants.
     """
-    return [finding for finding, _ in _iter_findings(model, capabilities)]
+    return [
+        finding
+        for finding, _ in _iter_findings(
+            model,
+            capabilities,
+            on_probe_start=on_probe_start,
+            on_probe_complete=on_probe_complete,
+        )
+    ]
 
 
 def iter_optimization_outputs(

tests/unit/commands/test_optimize_cli.py

@@ -270,6 +270,45 @@ def test_check_optim_no_findings_message(self, runner: CliRunner, tmp_path: Path
         assert result.exit_code == 0, result.output
         assert "No registered optimizations" in result.output
 
+    def test_check_optim_shows_probe_name_and_completed_progress(
+        self, runner: CliRunner, tmp_path: Path
+    ) -> None:
+        from winml.modelkit.optim import BoolCapability, get_all_capabilities
+
+        model_file = tmp_path / "model.onnx"
+        model_file.touch()
+        probe_names = [
+            name
+            for name, cap in get_all_capabilities().items()
+            if isinstance(cap, BoolCapability) and not cap.default
+        ]
+
+        def analyze_with_progress(
+            model: MagicMock,
+            capabilities: dict[str, object],
+            *,
+            on_probe_start: object,
+            on_probe_complete: object,
+        ) -> list[object]:
+            del model
+            assert callable(on_probe_start)
+            assert callable(on_probe_complete)
+            for name, cap in capabilities.items():
+                if isinstance(cap, BoolCapability) and not cap.default:
+                    on_probe_start(name)
+                    on_probe_complete(name)
+            return []
+
+        with (
+            patch(_LOAD_ONNX, return_value=_make_mock_model()),
+            patch(_ANALYZE_MODEL, side_effect=analyze_with_progress),
+        ):
+            result = runner.invoke(optimize, ["-m", str(model_file), "--check-optim"])
+
+        assert result.exit_code == 0, result.output
+        assert f"Checking {probe_names[-1]}" in result.output
+        assert f"{len(probe_names)}/{len(probe_names)}" in result.output
+
 
 class TestConfigFile:
     """Config file loading and precedence."""

tests/unit/optim/test_analysis.py

@@ -16,9 +16,11 @@
 from onnx import GraphProto, ModelProto, TensorProto, helper, numpy_helper
 
 from winml.modelkit.optim import (
+    BoolCapability,
     CapabilityFinding,
     NodeRef,
     analyze_model,
+    get_all_capabilities,
     iter_optimization_outputs,
 )
 from winml.modelkit.optim.analysis import (
@@ -28,7 +30,6 @@
     _diff_initializers,
     _diff_nodes,
 )
-from winml.modelkit.optim.pipes import get_all_capabilities
 
 
 # =============================================================================
@@ -373,9 +374,23 @@ def test_clamp_reported_for_extreme_constant(self) -> None:
         assert "BIG" in clamp.modified_initializers
 
     def test_clamp_not_reported_for_benign_constant(self) -> None:
-        findings = analyze_model(_benign_model(), get_all_capabilities())
+        capabilities = get_all_capabilities()
+        started: list[str] = []
+        completed: list[str] = []
+        findings = analyze_model(
+            _benign_model(),
+            capabilities,
+            on_probe_start=started.append,
+            on_probe_complete=completed.append,
+        )
         names = {f.name for f in findings}
         assert "clamp-constant-values" not in names
+        assert set(started) == set(completed)
+        assert sorted(completed) == sorted(
+            name
+            for name, capability in capabilities.items()
+            if isinstance(capability, BoolCapability) and not capability.default
+        )
 
     def test_matmul_add_fusion_detected(self) -> None:
         findings = analyze_model(_matmul_add_model(), get_all_capabilities())