diff --git a/src/winml/modelkit/commands/optimize.py b/src/winml/modelkit/commands/optimize.py index 0b5a7917b..f86092d0e 100644 --- a/src/winml/modelkit/commands/optimize.py +++ b/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) diff --git a/src/winml/modelkit/optim/analysis.py b/src/winml/modelkit/optim/analysis.py index 649a5d2d3..cbea1b3a1 100644 --- a/src/winml/modelkit/optim/analysis.py +++ b/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( diff --git a/tests/unit/commands/test_optimize_cli.py b/tests/unit/commands/test_optimize_cli.py index d5b6507aa..98dda97ca 100644 --- a/tests/unit/commands/test_optimize_cli.py +++ b/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.""" diff --git a/tests/unit/optim/test_analysis.py b/tests/unit/optim/test_analysis.py index a69856d17..6c778ee4f 100644 --- a/tests/unit/optim/test_analysis.py +++ b/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())