From af7c24eb35e67aefe626ccde5898f6e2ba5b150e Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 6 Aug 2026 09:46:43 +0800 Subject: [PATCH] Serialize check-optim findings in analyze JSON --- src/winml/modelkit/analyze/optim_output.py | 62 ++++- src/winml/modelkit/commands/analyze.py | 86 ++++-- tests/unit/analyze/test_optim_output.py | 44 ++++ .../unit/analyze/test_static_analyzer_cli.py | 244 ++++++++++++++++++ 4 files changed, 407 insertions(+), 29 deletions(-) diff --git a/src/winml/modelkit/analyze/optim_output.py b/src/winml/modelkit/analyze/optim_output.py index ea0a9e25e..3cd868d45 100644 --- a/src/winml/modelkit/analyze/optim_output.py +++ b/src/winml/modelkit/analyze/optim_output.py @@ -24,6 +24,7 @@ from dataclasses import dataclass, field from typing import TYPE_CHECKING, cast +from ..optim.analysis import NodeRef from .models.onnx_model import ONNXModel from .models.support_level import SupportLevel @@ -33,7 +34,7 @@ from onnx import ModelProto, NodeProto - from ..optim import CapabilityFinding, NodeRef + from ..optim import CapabilityFinding from ..utils.constants import EPName @@ -68,6 +69,18 @@ class ProducedOperatorSupport: support: SupportLevel reason: str | None = None + def to_dict(self) -> dict[str, str]: + """Return the stable JSON representation of this produced operator.""" + data = { + "op_type": self.op_type, + "label": self.label, + "change": self.change, + "support": self.support.value, + } + if self.reason is not None: + data["reason"] = self.reason + return data + @dataclass class OptimizationOutputSupport: @@ -88,6 +101,12 @@ class OptimizationOutputSupport: category: str description: str pipe_name: str + removed_nodes: list[NodeRef] = field(default_factory=list) + added_nodes: list[NodeRef] = field(default_factory=list) + modified_nodes: list[NodeRef] = field(default_factory=list) + removed_initializers: list[str] = field(default_factory=list) + added_initializers: list[str] = field(default_factory=list) + modified_initializers: list[str] = field(default_factory=list) operators: list[ProducedOperatorSupport] = field(default_factory=list) error: str | None = None @@ -105,6 +124,41 @@ def support_counts(self) -> dict[SupportLevel, int]: """Return a count of produced operators per support level.""" return dict(Counter(op.support for op in self.operators)) + @staticmethod + def _node_ref_dict(ref: NodeRef) -> dict[str, object]: + """Return the stable JSON representation of one graph-delta node.""" + return { + "op_type": ref.op_type, + "name": ref.name, + "outputs": list(ref.outputs), + } + + def to_dict(self) -> dict[str, object]: + """Return actionable graph-delta and target-support evidence.""" + data: dict[str, object] = { + "name": self.name, + "enable_flag": self.enable_flag, + "category": self.category, + "description": self.description, + "pipe_name": self.pipe_name, + "worst_support": self.worst_support.value, + "support_counts": { + level.value: count for level, count in self.support_counts().items() + }, + "graph_delta": { + "removed_nodes": [self._node_ref_dict(ref) for ref in self.removed_nodes], + "added_nodes": [self._node_ref_dict(ref) for ref in self.added_nodes], + "modified_nodes": [self._node_ref_dict(ref) for ref in self.modified_nodes], + "removed_initializers": self.removed_initializers, + "added_initializers": self.added_initializers, + "modified_initializers": self.modified_initializers, + }, + "operators": [operator.to_dict() for operator in self.operators], + } + if self.error is not None: + data["error"] = self.error + return data + def _produced_node_refs( finding: CapabilityFinding, @@ -137,6 +191,12 @@ def _check_one( category=finding.category, description=finding.description, pipe_name=finding.pipe_name, + removed_nodes=list(finding.removed_nodes), + added_nodes=list(finding.added_nodes), + modified_nodes=list(finding.modified_nodes), + removed_initializers=list(finding.removed_initializers), + added_initializers=list(finding.added_initializers), + modified_initializers=list(finding.modified_initializers), ) produced = _produced_node_refs(finding) diff --git a/src/winml/modelkit/commands/analyze.py b/src/winml/modelkit/commands/analyze.py index dfa18bb6f..882485c2e 100644 --- a/src/winml/modelkit/commands/analyze.py +++ b/src/winml/modelkit/commands/analyze.py @@ -1196,18 +1196,20 @@ def analyze( # Optionally probe which optimizations would change the model and what # operators they would introduce. This probe is target-independent, so # materialize it once here and only re-run the (cheap) support lookup per - # EP/device below. Console-only feature — skipped in quiet mode. + # EP/device below. Quiet mode disables rendering, not data collection. optim_outputs: list[tuple[Any, Any]] = [] - if check_optim and not quiet: + optim_probe_error: str | None = None + if check_optim: try: import onnx from ..optim import get_all_capabilities, iter_optimization_outputs - console.print( - "[dim]Probing optimization outputs " - "(this can take a while on large models)…[/dim]" - ) + if not quiet: + console.print( + "[dim]Probing optimization outputs " + "(this can take a while on large models)…[/dim]" + ) optim_proto = onnx.load(str(model)) optim_outputs = list(iter_optimization_outputs(optim_proto, get_all_capabilities())) # Each entry retains a full produced-model clone so the (cheap) @@ -1221,6 +1223,7 @@ def analyze( ) except Exception as exc: logger.warning("Could not probe optimization outputs: %s", exc) + optim_probe_error = str(exc) optim_outputs = [] # Model info header @@ -1273,9 +1276,38 @@ def analyze( ep_counter = 0 _no_data_eps: set[tuple[str, str]] = set() # EP/device pairs with no op rule data analysis_results: list = [] + optimization_support_payloads: list[dict[str, object] | None] = [] current_run_unknown_op = False current_op_check_skipped = False + def _collect_optimization_support( + target_ep: EPName, target_device: str + ) -> tuple[list[Any], dict[str, object]]: + """Check one target and return renderable results plus JSON data.""" + from ..analyze.optim_output import check_optimization_output_support + + support_error: str | None = None + optim_support = [] + if optim_probe_error is None: + try: + optim_support = check_optimization_output_support( + optim_outputs, + ep=target_ep, + device=target_device, + model_path=str(model), + ) + except Exception as exc: + support_error = str(exc) + logger.warning("Could not check optimization output support: %s", exc) + payload: dict[str, object] = { + "ep_type": target_ep, + "device_type": target_device, + "probe_error": optim_probe_error, + "support_error": support_error, + "optimizations": [item.to_dict() for item in optim_support], + } + return optim_support, payload + def _current_ep_device_pair_display_name() -> str: """Return current EP/device display label, or empty when unset.""" if current_ep_device_pair is None: @@ -1523,14 +1555,10 @@ def on_node_result(pattern_runtime: PatternRuntime) -> None: # Optimization output support section (per-EP), when opted in. if check_optim: - from ..analyze.optim_output import check_optimization_output_support - - optim_support = check_optimization_output_support( - optim_outputs, - ep=target_ep, - device=target_device, - model_path=str(model), + optim_support, payload = _collect_optimization_support( + target_ep, target_device ) + optimization_support_payloads.append(payload) _render_optim_output_support( console, optim_support, @@ -1561,21 +1589,29 @@ def on_node_result(pattern_runtime: PatternRuntime) -> None: ) analysis_results.append(result) + if check_optim: + _, payload = _collect_optimization_support(target_ep, target_device) + optimization_support_payloads.append(payload) + result = analysis_results[-1] + serialized_results: list[dict[str, object]] = [] + json_mode = output_format == "json" + if output or json_mode: + for index, run_result in enumerate(analysis_results): + payload = json.loads(run_result.to_json()) + if check_optim: + payload["optimization_output_support"] = optimization_support_payloads[index] + serialized_results.append(payload) + # Save JSON if requested if output: try: output.parent.mkdir(parents=True, exist_ok=True) if len(analysis_results) == 1: - output.write_text(result.to_json(), encoding="utf-8") + output.write_text(json.dumps(serialized_results[0], indent=2), encoding="utf-8") else: - output.write_text( - json.dumps( - [json.loads(run_result.to_json()) for run_result in analysis_results] - ), - encoding="utf-8", - ) + output.write_text(json.dumps(serialized_results), encoding="utf-8") logger.info("JSON results saved to: %s", output) except OSError as e: logger.error("Failed to write JSON output to %s: %s", output, e) @@ -1654,17 +1690,11 @@ def on_node_result(pattern_runtime: PatternRuntime) -> None: logger.debug("Config generation traceback:", exc_info=True) # Emit JSON to stdout if requested - json_mode = output_format == "json" if json_mode: if len(analysis_results) == 1: - click.echo(result.to_json()) + click.echo(json.dumps(serialized_results[0], indent=2)) else: - click.echo( - json.dumps( - [json.loads(run_result.to_json()) for run_result in analysis_results], - indent=2, - ) - ) + click.echo(json.dumps(serialized_results, indent=2)) # Exit code: 0 = fully supported, 1 = partial support overall_supported = all(run_result.is_fully_supported() for run_result in analysis_results) diff --git a/tests/unit/analyze/test_optim_output.py b/tests/unit/analyze/test_optim_output.py index b15ac4aa2..04cc9fd52 100644 --- a/tests/unit/analyze/test_optim_output.py +++ b/tests/unit/analyze/test_optim_output.py @@ -164,6 +164,50 @@ def test_support_counts(self) -> None: assert counts[SupportLevel.SUPPORTED] == 2 assert counts[SupportLevel.PARTIAL] == 1 + def test_to_dict_includes_graph_delta_and_target_support(self) -> None: + """Structured JSON retains actionable graph and support evidence.""" + from winml.modelkit.optim import NodeRef + + opt = OptimizationOutputSupport( + name="static-split-to-slice", + enable_flag="--enable-static-split-to-slice", + category="rewrite", + description="Replace static Split with Slice.", + pipe_name="algebraic", + removed_nodes=[NodeRef("Split", "split", ("a", "b"))], + added_nodes=[NodeRef("Slice", "slice_0", ("a",))], + modified_initializers=["starts"], + operators=[ + ProducedOperatorSupport("Slice", "Slice 'slice_0'", "added", SupportLevel.SUPPORTED) + ], + ) + + assert opt.to_dict() == { + "name": "static-split-to-slice", + "enable_flag": "--enable-static-split-to-slice", + "category": "rewrite", + "description": "Replace static Split with Slice.", + "pipe_name": "algebraic", + "worst_support": "supported", + "support_counts": {"supported": 1}, + "graph_delta": { + "removed_nodes": [{"op_type": "Split", "name": "split", "outputs": ["a", "b"]}], + "added_nodes": [{"op_type": "Slice", "name": "slice_0", "outputs": ["a"]}], + "modified_nodes": [], + "removed_initializers": [], + "added_initializers": [], + "modified_initializers": ["starts"], + }, + "operators": [ + { + "op_type": "Slice", + "label": "Slice 'slice_0'", + "change": "added", + "support": "supported", + } + ], + } + # ============================================================================= # END-TO-END SUPPORT CHECK diff --git a/tests/unit/analyze/test_static_analyzer_cli.py b/tests/unit/analyze/test_static_analyzer_cli.py index b7a078dd3..129d19438 100644 --- a/tests/unit/analyze/test_static_analyzer_cli.py +++ b/tests/unit/analyze/test_static_analyzer_cli.py @@ -2266,3 +2266,247 @@ def _should_not_run(*_a: object, **_k: object) -> list: assert result.exit_code == 0 assert "OPTIMIZATION OUTPUT SUPPORT" not in result.output + + @patch("winml.modelkit.analyze.ONNXStaticAnalyzer") + def test_quiet_json_includes_structured_optimization_support( + self, + mock_analyzer_class: MagicMock, + runner: CliRunner, + tmp_path: Path, + mock_analyzer_result: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """--check-optim produces JSON even when Rich rendering is quiet.""" + from winml.modelkit.analyze.models.support_level import SupportLevel + from winml.modelkit.analyze.optim_output import ( + OptimizationOutputSupport, + ProducedOperatorSupport, + ) + + model_file = tmp_path / "model.onnx" + output_file = tmp_path / "analysis.json" + self._write_model(model_file) + + mock_instance = Mock() + mock_instance.analyze.return_value = mock_analyzer_result + mock_analyzer_class.return_value = mock_instance + + monkeypatch.setattr( + "winml.modelkit.optim.iter_optimization_outputs", + lambda *_args, **_kwargs: [(object(), object())], + ) + monkeypatch.setattr("winml.modelkit.optim.get_all_capabilities", dict) + monkeypatch.setattr( + "winml.modelkit.analyze.optim_output.check_optimization_output_support", + lambda *_args, **_kwargs: [ + OptimizationOutputSupport( + name="static-split-to-slice", + enable_flag="--enable-static-split-to-slice", + category="rewrite", + description="Replace static Split with Slice.", + pipe_name="algebraic", + operators=[ + ProducedOperatorSupport( + "Slice", + "Slice 'slice_0'", + "added", + SupportLevel.SUPPORTED, + ) + ], + ) + ], + ) + + result = runner.invoke( + analyze, + [ + "--model", + str(model_file), + "--ep", + "qnn", + "--device", + "NPU", + "--check-optim", + "--format", + "json", + "--output", + str(output_file), + "--quiet", + ], + ) + + assert result.exit_code == 0, result.output + stdout_data = json.loads(result.output) + file_data = json.loads(output_file.read_text(encoding="utf-8")) + assert file_data == stdout_data + support = stdout_data["optimization_output_support"] + assert support["ep_type"] == "QNNExecutionProvider" + assert support["device_type"] == "NPU" + assert support["probe_error"] is None + assert support["support_error"] is None + assert support["optimizations"][0]["enable_flag"] == ("--enable-static-split-to-slice") + assert support["optimizations"][0]["worst_support"] == "supported" + + @patch("winml.modelkit.analyze.ONNXStaticAnalyzer") + def test_quiet_json_preserves_optimization_probe_error( + self, + mock_analyzer_class: MagicMock, + runner: CliRunner, + tmp_path: Path, + mock_analyzer_result: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Optimization probe failures remain visible to JSON consumers.""" + model_file = tmp_path / "model.onnx" + self._write_model(model_file) + + mock_instance = Mock() + mock_instance.analyze.return_value = mock_analyzer_result + mock_analyzer_class.return_value = mock_instance + + def _fail_probe(*_args: object, **_kwargs: object) -> list: + raise RuntimeError("probe failed") + + monkeypatch.setattr("winml.modelkit.optim.iter_optimization_outputs", _fail_probe) + + result = runner.invoke( + analyze, + [ + "--model", + str(model_file), + "--ep", + "qnn", + "--device", + "NPU", + "--check-optim", + "--format", + "json", + "--quiet", + ], + ) + + assert result.exit_code == 0, result.output + support = json.loads(result.output)["optimization_output_support"] + assert support["probe_error"] == "probe failed" + assert support["support_error"] is None + assert support["optimizations"] == [] + + @patch("winml.modelkit.analyze.ONNXStaticAnalyzer") + def test_quiet_json_preserves_target_support_error( + self, + mock_analyzer_class: MagicMock, + runner: CliRunner, + tmp_path: Path, + mock_analyzer_result: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Target support failures are distinct from graph probe failures.""" + model_file = tmp_path / "model.onnx" + self._write_model(model_file) + + mock_instance = Mock() + mock_instance.analyze.return_value = mock_analyzer_result + mock_analyzer_class.return_value = mock_instance + + monkeypatch.setattr( + "winml.modelkit.optim.iter_optimization_outputs", + lambda *_args, **_kwargs: [(object(), object())], + ) + monkeypatch.setattr("winml.modelkit.optim.get_all_capabilities", dict) + + def _fail_support(*_args: object, **_kwargs: object) -> list: + raise RuntimeError("support check failed") + + monkeypatch.setattr( + "winml.modelkit.analyze.optim_output.check_optimization_output_support", + _fail_support, + ) + + result = runner.invoke( + analyze, + [ + "--model", + str(model_file), + "--ep", + "qnn", + "--device", + "NPU", + "--check-optim", + "--format", + "json", + "--quiet", + ], + ) + + assert result.exit_code == 0, result.output + support = json.loads(result.output)["optimization_output_support"] + assert support["probe_error"] is None + assert support["support_error"] == "support check failed" + assert support["optimizations"] == [] + + @patch("winml.modelkit.analyze.ONNXStaticAnalyzer") + def test_multi_device_json_aligns_optimization_support_with_each_target( + self, + mock_analyzer_class: MagicMock, + runner: CliRunner, + tmp_path: Path, + mock_analyzer_result: Mock, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Each result in a fan-out carries support for its own EP/device.""" + from winml.modelkit.analyze.optim_output import OptimizationOutputSupport + + model_file = tmp_path / "model.onnx" + self._write_model(model_file) + + mock_instance = Mock() + mock_instance.analyze.return_value = mock_analyzer_result + mock_analyzer_class.return_value = mock_instance + + monkeypatch.setattr( + "winml.modelkit.optim.iter_optimization_outputs", + lambda *_args, **_kwargs: [(object(), object())], + ) + monkeypatch.setattr("winml.modelkit.optim.get_all_capabilities", dict) + + def _support_for_target(*_args: object, **kwargs: object) -> list: + device = str(kwargs["device"]) + return [ + OptimizationOutputSupport( + name=f"optimization-for-{device.lower()}", + enable_flag=f"--enable-for-{device.lower()}", + category="rewrite", + description="Target-specific test result.", + pipe_name="algebraic", + ) + ] + + monkeypatch.setattr( + "winml.modelkit.analyze.optim_output.check_optimization_output_support", + _support_for_target, + ) + + result = runner.invoke( + analyze, + [ + "--model", + str(model_file), + "--ep", + "qnn", + "--device", + "all", + "--check-optim", + "--format", + "json", + "--quiet", + ], + ) + + assert result.exit_code == 0, result.output + payloads = json.loads(result.output) + assert len(payloads) == 2 + for payload in payloads: + support = payload["optimization_output_support"] + device = support["device_type"] + assert support["ep_type"] == "QNNExecutionProvider" + assert support["optimizations"][0]["name"] == (f"optimization-for-{device.lower()}")