From 4da82cde6c93894199f00e108f80be13f5dfe3bd Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 8 Sep 2026 13:43:41 +0800 Subject: [PATCH 1/2] fix(build): honor CLI optimization bypass --- src/winml/modelkit/commands/build.py | 5 +- tests/unit/commands/test_build.py | 252 +++++++++++++++++++++++++++ 2 files changed, 256 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index a56d9f52e..c17946189 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -2068,6 +2068,7 @@ def _build_hf_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + skip_optimize: bool = extra_kwargs.pop("skip_optimize", False) model_label = model_id or "random-init" # ── Validate + setup ───────────────────────────────────────── @@ -2153,6 +2154,7 @@ def _name(base: str) -> str: show_io_first=False, analyze_output_path=analyze_result_path, allow_unsupported_nodes=allow_unsupported_nodes, + skip_optimize=skip_optimize or config.skip_optimize, ) # Persist config after autoconf @@ -2209,6 +2211,7 @@ def _build_onnx_pipeline( max_iters: int = extra_kwargs.pop("hack_max_optim_iterations", 3) allow_unsupported_nodes: bool = extra_kwargs.pop("allow_unsupported_nodes", False) + skip_optimize: bool = extra_kwargs.pop("skip_optimize", False) # ── Validate + setup ───────────────────────────────────────── if not onnx_path.exists(): @@ -2264,7 +2267,7 @@ def _build_onnx_pipeline( show_io_first=True, analyze_output_path=analyze_result_path, allow_unsupported_nodes=allow_unsupported_nodes, - skip_optimize=config.skip_optimize, + skip_optimize=skip_optimize or config.skip_optimize, ) config_path.write_text(json.dumps(config.to_dict(), indent=2)) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 37f6bde79..50ae7ea76 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -1791,6 +1791,258 @@ def test_no_optimize_default_not_present( assert "skip_optimize" not in extra +class TestBuildPipelineStageControls: + """Exercise real CLI sinks and stage control, mocking only expensive payloads.""" + + @staticmethod + def _write_model(path: Path, graph_kind: str) -> None: + import onnx + from onnx import TensorProto, helper + + input_type = output_type = TensorProto.FLOAT + initializers = [] + if graph_kind == "raw": + nodes = [helper.make_node("Identity", ["x"], ["y"])] + elif graph_kind == "qdq": + initializers = [ + helper.make_tensor("scale", TensorProto.FLOAT, [], [0.125]), + helper.make_tensor("zero", TensorProto.UINT8, [], [0]), + ] + nodes = [ + helper.make_node("QuantizeLinear", ["x", "scale", "zero"], ["q"]), + helper.make_node("DequantizeLinear", ["q", "scale", "zero"], ["y"]), + ] + else: + assert graph_kind == "qoperator" + input_type, output_type = TensorProto.UINT8, TensorProto.INT32 + initializers = [helper.make_tensor("weight", TensorProto.UINT8, [1, 1], [1])] + # No Q/DQ nodes: detection must recognize the integer operator itself. + nodes = [helper.make_node("MatMulInteger", ["x", "weight"], ["y"])] + graph = helper.make_graph( + nodes, + "stage-control-fixture", + [helper.make_tensor_value_info("x", input_type, [1, 1])], + [helper.make_tensor_value_info("y", output_type, [1, 1])], + initializer=initializers, + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + onnx.checker.check_model(model) + onnx.save(model, path) + + @staticmethod + def _config(input_kind: str, quant_mode: str, config_skip: bool = False) -> dict: + return { + "loader": {"task": "feature-extraction"}, + "export": {"opset_version": 17} if input_kind == "hf" else None, + "optim": {}, + "quant": { + "mode": quant_mode, + "samples": 1, + "task": "feature-extraction", + "model_id": "test-model", + }, + "compile": {"execution_provider": "cpu", "enable_ep_context": False}, + "auto": True, + "skip_optimize": config_skip, + } + + @pytest.fixture + def payloads(self): + from winml.modelkit.analyze import AnalyzeResult, LintResult + from winml.modelkit.onnx import copy_onnx_model + from winml.modelkit.optim import WinMLOptimizationConfig + + def export_payload(*, model, output_path, **kwargs): + copy_onnx_model(model, output_path) + + def optimize_payload(*, model, output, **kwargs): + copy_onnx_model(model, output) + + def convert_payload(*, model_path, output_path, **kwargs): + copy_onnx_model(model_path, output_path) + return MagicMock(success=True, output_path=output_path, errors=[]) + + optim = WinMLOptimizationConfig() + analysis = AnalyzeResult( + lint=LintResult( + errors=0, + warnings=0, + info=0, + passed=True, + error_patterns=[], + warning_patterns=[], + information=[], + optimization_config=optim, + ), + optimization_config=optim, + ) + with ( + patch("winml.modelkit.build.hf._load_model") as load, + patch("winml.modelkit.export.export_onnx", side_effect=export_payload) as export, + patch( + "winml.modelkit.build.common.optimize_onnx", side_effect=optimize_payload + ) as optimize, + patch("winml.modelkit.build.common.analyze_onnx", return_value=analysis) as analyze, + patch("winml.modelkit.quant.quantize_onnx", side_effect=convert_payload) as quantize, + patch("winml.modelkit.compiler.compile_onnx", side_effect=convert_payload) as compile_, + patch("winml.modelkit.utils.console.StageLive"), + patch( + "winml.modelkit.session.resolve_device", + return_value=EPDeviceTarget(ep="CPUExecutionProvider", device="cpu"), + ), + ): + yield { + "load": load, + "export": export, + "optimize": optimize, + "analyze": analyze, + "quantize": quantize, + "compile": compile_, + } + + def _invoke_pipeline( + self, + tmp_path: Path, + payloads: dict, + input_kind: str, + quant_mode: str, + flags: tuple[str, ...], + *, + config_skip: bool = False, + graph_kind: str = "raw", + ) -> Path: + from winml.modelkit.onnx import is_quantized_onnx + + model_path = tmp_path / "input.onnx" + self._write_model(model_path, graph_kind) + assert is_quantized_onnx(model_path) is (graph_kind != "raw") + payloads["load"].return_value = model_path + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(self._config(input_kind, quant_mode, config_skip))) + output_dir = tmp_path / "out" + result = _invoke( + [ + "-c", + str(config_path), + "-m", + "test-model" if input_kind == "hf" else str(model_path), + "-o", + str(output_dir), + "--ep", + "cpu", + "--export-type", + "generic", + *flags, + ] + ) + assert result.exit_code == 0, result.output + return output_dir + + @pytest.mark.parametrize("input_kind", ["hf", "onnx"]) + @pytest.mark.parametrize("quant_mode", ["fp16", "static"]) + @pytest.mark.parametrize( + ("flags", "config_skip"), + [ + pytest.param((), False, id="default"), + pytest.param(("--optimize",), False, id="explicit-optimize"), + pytest.param(("--no-optimize",), False, id="cli-skip"), + pytest.param((), True, id="config-skip"), + pytest.param(("--optimize",), True, id="config-skip-with-cli-optimize"), + pytest.param( + ("--no-optimize", "--max-optim-iterations", "5"), + False, + id="cli-skip-with-autoconf", + ), + ], + ) + def test_raw_pipeline_preserves_optimization_and_quantization_controls( + self, tmp_path: Path, payloads, input_kind, quant_mode, flags, config_skip + ) -> None: + output_dir = self._invoke_pipeline( + tmp_path, payloads, input_kind, quant_mode, flags, config_skip=config_skip + ) + + skip = config_skip or "--no-optimize" in flags + assert payloads["optimize"].call_count == (0 if skip else 1) + assert payloads["analyze"].call_count == (0 if skip else 2) + assert payloads["export"].call_count == (1 if input_kind == "hf" else 0) + payloads["quantize"].assert_called_once() + assert payloads["quantize"].call_args.kwargs["config"].mode == quant_mode + payloads["compile"].assert_called_once() + assert ( + payloads["compile"].call_args.kwargs["model_path"] + == payloads["quantize"].call_args.kwargs["output_path"] + ) + saved = json.loads((output_dir / "winml_build_config.json").read_text()) + assert saved["quant"]["mode"] == quant_mode + # Explicit bypass is not evidence of prequantization: do not force-stamp it. + assert saved.get("skip_optimize", False) is config_skip + optimized_name = "optimized.onnx" if input_kind == "hf" else "input_optimized.onnx" + assert (output_dir / optimized_name).is_file() + assert (output_dir / "model.onnx").is_file() + + @pytest.mark.parametrize("graph_kind", ["qdq", "qoperator"]) + @pytest.mark.parametrize("quant_mode", ["fp16", "static"]) + @pytest.mark.parametrize( + ("flags", "config_skip"), + [((), False), (("--no-optimize",), False), ((), True)], + ids=["default", "cli-skip", "config-skip"], + ) + def test_graph_proven_prequantization_suppresses_requantization( + self, tmp_path: Path, payloads, graph_kind, quant_mode, flags, config_skip + ) -> None: + output_dir = self._invoke_pipeline( + tmp_path, + payloads, + "onnx", + quant_mode, + flags, + config_skip=config_skip, + graph_kind=graph_kind, + ) + + for stage in ("export", "optimize", "analyze", "quantize"): + payloads[stage].assert_not_called() + payloads["compile"].assert_called_once() + saved = json.loads((output_dir / "winml_build_config.json").read_text()) + assert saved["skip_optimize"] is True + assert saved["quant"] is None + assert (output_dir / "model.onnx").is_file() + + @pytest.mark.parametrize("quant_mode", ["fp16", "static"]) + @pytest.mark.parametrize("flags", [(), ("--no-optimize",)], ids=["default", "cli-skip"]) + def test_composite_controls_reach_every_actual_hf_sink( + self, tmp_path: Path, payloads, quant_mode, flags + ) -> None: + from winml.modelkit.config import WinMLBuildConfig + + components = dict.fromkeys(("part_a", "part_b"), "feature-extraction") + + def component_config(_model_id, *, task, **kwargs): + data = self._config("hf", quant_mode) + data["loader"]["task"] = task + return WinMLBuildConfig.from_dict(data) + + with ( + patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=components, + ), + patch("winml.modelkit.config.generate_build_config", side_effect=component_config), + ): + output_dir = self._invoke_pipeline(tmp_path, payloads, "hf", quant_mode, flags) + + assert payloads["optimize"].call_count == (0 if flags else 2) + assert payloads["analyze"].call_count == (0 if flags else 4) + for stage in ("load", "export", "quantize", "compile"): + assert payloads[stage].call_count == len(components) + quant_configs = [call.kwargs["config"] for call in payloads["quantize"].call_args_list] + assert quant_configs[0] is not quant_configs[1] + assert all(config.mode == quant_mode for config in quant_configs) + for name in components: + assert (output_dir / f"{name}_model.onnx").is_file() + + # ============================================================================= # _run_compile_stage UNIT TESTS # ============================================================================= From c29257b4834d99c0c0c91203471a29ed930a2a92 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 15 Sep 2026 20:32:25 +0800 Subject: [PATCH 2/2] test(build): use consistent ONNX imports --- tests/unit/commands/test_build.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 50ae7ea76..e9706299d 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -1796,8 +1796,7 @@ class TestBuildPipelineStageControls: @staticmethod def _write_model(path: Path, graph_kind: str) -> None: - import onnx - from onnx import TensorProto, helper + from onnx import TensorProto, checker, helper, save input_type = output_type = TensorProto.FLOAT initializers = [] @@ -1826,8 +1825,8 @@ def _write_model(path: Path, graph_kind: str) -> None: initializer=initializers, ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) - onnx.checker.check_model(model) - onnx.save(model, path) + checker.check_model(model) + save(model, path) @staticmethod def _config(input_kind: str, quant_mode: str, config_skip: bool = False) -> dict: