From 56d20647ca24fce2372777e208433854a32bf0b7 Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Fri, 17 Jul 2026 10:52:59 +0800 Subject: [PATCH] test(e2e): add dynamic-axes e2e tests for export, config, build, perf Cover the shared --dynamic-axes option across the four commands that support it, using microsoft/resnet-50: export marks the named input axis symbolic (dim_param) and propagates it to logits, config records the mapping under export.dynamic_axes, the symbolic batch dim survives a full build, and perf re-exports so the benchmarked model exposes a dynamic batch axis. --- tests/e2e/test_build_e2e.py | 65 +++++++++++++++++++++++++++ tests/e2e/test_config_e2e.py | 35 +++++++++++++++ tests/e2e/test_export_e2e.py | 85 +++++++++++++++++++++++++++++++++--- tests/e2e/test_perf_e2e.py | 62 ++++++++++++++++++++++++++ 4 files changed, 241 insertions(+), 6 deletions(-) diff --git a/tests/e2e/test_build_e2e.py b/tests/e2e/test_build_e2e.py index fc05c5275..a0b145e5b 100644 --- a/tests/e2e/test_build_e2e.py +++ b/tests/e2e/test_build_e2e.py @@ -30,6 +30,7 @@ from typing import TYPE_CHECKING from unittest.mock import MagicMock, patch +import onnx import pytest from click.testing import CliRunner @@ -293,3 +294,67 @@ def test_onnx_passthrough_no_optimize(self, tmp_path: Path, onnx_model_path: Pat ) assert result.exit_code == 0, f"build failed (exit {result.exit_code}):\n{result.output}" assert output_dir.exists() + + +# =========================================================================== +# Dynamic axes: --dynamic-axes survives the full build pipeline. +# =========================================================================== + + +@pytest.mark.slow +@pytest.mark.network +class TestBuildDynamicAxes: + """``--dynamic-axes`` is applied during export and survives optimization.""" + + def test_dynamic_batch_survives_build(self, tmp_path: Path): + """A dynamic batch axis stays symbolic in the built ONNX artifact.""" + config_path = _generate_config_file( + tmp_path, + "microsoft/resnet-50", + task="image-classification", + ) + axes = tmp_path / "axes.json" + axes.write_text(json.dumps({"pixel_values": {"0": "batch"}})) + output_dir = tmp_path / "output" + + result = CliRunner().invoke( + build, + [ + "-c", + config_path, + "-m", + "microsoft/resnet-50", + "-o", + str(output_dir), + "--no-quant", + "--no-compile", + "--no-analyze", + "--dynamic-axes", + str(axes), + ], + obj={"debug": True}, + catch_exceptions=False, + ) + assert result.exit_code == 0, f"build failed (exit {result.exit_code}):\n{result.output}" + + onnx_files = list(output_dir.rglob("*.onnx")) + assert onnx_files, ( + f"No ONNX files found in {output_dir}. Contents: " + f"{[str(p) for p in output_dir.rglob('*')]}" + ) + + # Locate the graph carrying the model input (skip any auxiliary files). + model = None + for path in onnx_files: + candidate = onnx.load(str(path)) + if any(i.name == "pixel_values" for i in candidate.graph.input): + model = candidate + break + assert model is not None, ( + f"no ONNX with a 'pixel_values' input among {[str(p) for p in onnx_files]}" + ) + + pixel_values = next(i for i in model.graph.input if i.name == "pixel_values") + dims = pixel_values.type.tensor_type.shape.dim + assert dims[0].dim_param == "batch", f"expected a symbolic batch dim, got {list(dims)}" + assert [d.dim_value for d in dims[1:]] == [3, 224, 224] diff --git a/tests/e2e/test_config_e2e.py b/tests/e2e/test_config_e2e.py index a8a1cb9c0..10892dcd2 100644 --- a/tests/e2e/test_config_e2e.py +++ b/tests/e2e/test_config_e2e.py @@ -539,3 +539,38 @@ def test_explicit_fp16_still_triggers_quant(self) -> None: quant = data.get("quant") assert quant is not None, "Explicit --precision fp16 should produce a quant config" assert quant.get("mode") == "fp16" + + +# =========================================================================== +# Dynamic axes: --dynamic-axes +# =========================================================================== + + +class TestConfigDynamicAxes: + """``--dynamic-axes`` is recorded in the generated ``export`` config section. + + JSON serialization keeps axis keys as strings, so the round-tripped mapping + is ``{"pixel_values": {"0": "batch"}}``. + """ + + MODEL = "microsoft/resnet-50" + + def test_dynamic_axes_recorded(self, tmp_path: Path) -> None: + axes = tmp_path / "axes.json" + axes.write_text(json.dumps({"pixel_values": {"0": "batch"}})) + data = _run_config("-m", self.MODEL, "--dynamic-axes", str(axes)) + _assert_hf_config_structure(data) + assert data["export"]["dynamic_axes"] == {"pixel_values": {"0": "batch"}} + + def test_dynamic_axes_absent_by_default(self) -> None: + data = _run_config("-m", self.MODEL) + _assert_hf_config_structure(data) + assert data["export"].get("dynamic_axes") is None + + def test_multiple_axes_recorded(self, tmp_path: Path) -> None: + mapping = {"pixel_values": {"0": "batch", "2": "height", "3": "width"}} + axes = tmp_path / "axes.json" + axes.write_text(json.dumps(mapping)) + data = _run_config("-m", self.MODEL, "--dynamic-axes", str(axes)) + _assert_hf_config_structure(data) + assert data["export"]["dynamic_axes"] == mapping diff --git a/tests/e2e/test_export_e2e.py b/tests/e2e/test_export_e2e.py index 42cec6b40..bc9fa9fb9 100644 --- a/tests/e2e/test_export_e2e.py +++ b/tests/e2e/test_export_e2e.py @@ -186,6 +186,30 @@ def _write_json(path: Path, payload: dict) -> Path: return path +def _symbolic_dims(tensors, name: str) -> dict[int, str]: + """Return ``{axis: dim_param}`` for the symbolic dims of the named tensor.""" + for tensor in tensors: + if tensor.name == name: + return { + i: d.dim_param + for i, d in enumerate(tensor.type.tensor_type.shape.dim) + if d.dim_param + } + raise AssertionError(f"tensor {name!r} not found in {[t.name for t in tensors]}") + + +def _static_dims(tensors, name: str) -> dict[int, int]: + """Return ``{axis: dim_value}`` for the static (non-symbolic) dims of the tensor.""" + for tensor in tensors: + if tensor.name == name: + return { + i: d.dim_value + for i, d in enumerate(tensor.type.tensor_type.shape.dim) + if not d.dim_param + } + raise AssertionError(f"tensor {name!r} not found in {[t.name for t in tensors]}") + + # =========================================================================== # --help # =========================================================================== @@ -220,17 +244,13 @@ def test_minimal_resnet50(self, tmp_path: Path): class TestExportDinoV2: - MODEL = "facebook/dinov2-base" def test_image_feature_extraction(self, tmp_path: Path): """``-t image-feature-extraction`` must produce a valid ONNX export.""" onnx_path = tmp_path / "model.onnx" - result = _invoke(["-m", self.MODEL, "-o", str(onnx_path), - "-t", "image-feature-extraction"]) - assert result.exit_code == 0, ( - f"export failed (exit {result.exit_code}):\n{result.output}" - ) + result = _invoke(["-m", self.MODEL, "-o", str(onnx_path), "-t", "image-feature-extraction"]) + assert result.exit_code == 0, f"export failed (exit {result.exit_code}):\n{result.output}" assert onnx_path.exists(), f"ONNX model not found at {onnx_path}" model = onnx.load(str(onnx_path)) @@ -430,3 +450,56 @@ def test_build_config_with_dynamo(self, tmp_path: Path): f"expected opset 18 with -c + --dynamo, got {_opset_version(model)}" ) _assert_some_node_has(model, "pkg.onnxscript.rewriter.rule_name") + + +# =========================================================================== +# Dynamic axes: --dynamic-axes +# =========================================================================== + + +class TestExportDynamicAxes: + """``--dynamic-axes`` marks the named tensor axes symbolic in the ONNX graph. + + ResNet-50 has a single input ``pixel_values`` of shape [1, 3, 224, 224] and + a ``logits`` output. Making axis 0 dynamic turns the static batch dim into a + symbolic ``dim_param`` that also propagates to the output. + """ + + def test_batch_axis_symbolic(self, tmp_path: Path): + onnx_path = tmp_path / "model.onnx" + axes = _write_json(tmp_path / "axes.json", {"pixel_values": {"0": "batch"}}) + model = _assert_succeeds(_happy_args(onnx_path, "--dynamic-axes", str(axes)), onnx_path) + # Axis 0 of pixel_values becomes symbolic; channel/spatial dims stay static. + assert _symbolic_dims(model.graph.input, "pixel_values") == {0: "batch"} + assert _static_dims(model.graph.input, "pixel_values") == {1: 3, 2: 224, 3: 224} + # The symbolic batch dim propagates to the logits output. + assert _symbolic_dims(model.graph.output, "logits") == {0: "batch"} + + def test_static_batch_without_flag(self, tmp_path: Path): + # Baseline contrast: absent the flag, every input dim is a fixed integer. + onnx_path = tmp_path / "model.onnx" + model = _assert_succeeds(_happy_args(onnx_path), onnx_path) + assert _symbolic_dims(model.graph.input, "pixel_values") == {} + assert _input_shape_dims(model) == [1, 3, 224, 224] + + def test_multiple_axes_symbolic(self, tmp_path: Path): + onnx_path = tmp_path / "model.onnx" + axes = _write_json( + tmp_path / "axes.json", + {"pixel_values": {"0": "batch", "2": "height", "3": "width"}}, + ) + model = _assert_succeeds(_happy_args(onnx_path, "--dynamic-axes", str(axes)), onnx_path) + assert _symbolic_dims(model.graph.input, "pixel_values") == { + 0: "batch", + 2: "height", + 3: "width", + } + # Only the channel dim remains static. + assert _static_dims(model.graph.input, "pixel_values") == {1: 3} + + def test_invalid_dynamic_axes_fails(self, tmp_path: Path): + # An empty symbolic dim name is rejected by WinMLExportConfig validation, + # so the command must fail cleanly without writing an ONNX file. + onnx_path = tmp_path / "model.onnx" + bad = _write_json(tmp_path / "axes.json", {"pixel_values": {"0": ""}}) + _assert_fails(_happy_args(onnx_path, "--dynamic-axes", str(bad)), onnx_path) diff --git a/tests/e2e/test_perf_e2e.py b/tests/e2e/test_perf_e2e.py index 3b0722f68..335ce2cb0 100644 --- a/tests/e2e/test_perf_e2e.py +++ b/tests/e2e/test_perf_e2e.py @@ -35,6 +35,7 @@ import sys from pathlib import Path +import onnx import pytest from click.testing import CliRunner @@ -741,6 +742,67 @@ def test_module_invalid_lists_available(self, tmp_path: Path): assert not output_file.exists(), "Output file should not be written on failure" +# =========================================================================== +# Dynamic axes: --dynamic-axes re-exports with a symbolic batch dim. +# =========================================================================== + + +class TestPerfDynamicAxes: + """``--dynamic-axes`` feeds the HF export so the benchmarked model is dynamic. + + ``--ignore-cache`` forces a fresh build in a throwaway folder so the cached + static export isn't reused, and ``--no-skip-build`` guarantees the export + actually runs. The benchmarked ResNet-50 then exposes a dynamic (``None``) + batch axis instead of the static ``1``. + """ + + def test_dynamic_axes_cpu(self, tmp_path: Path): + axes = tmp_path / "axes.json" + axes.write_text(json.dumps({"pixel_values": {"0": "batch"}})) + output_file = tmp_path / "perf_dynamic_axes.json" + + result = CliRunner().invoke( + perf, + [ + "-m", + "microsoft/resnet-50", + "--iterations", + "3", + "--warmup", + "1", + "-o", + str(output_file), + "--device", + "cpu", + "--dynamic-axes", + str(axes), + "--ignore-cache", + "--no-skip-build", + "--no-optimize", + "--no-quant", + "--no-memory", + ], + obj={}, + catch_exceptions=False, + ) + assert result.exit_code == 0, f"perf failed (exit {result.exit_code}):\n{result.output}" + assert "Dynamic axes:" in result.output + assert output_file.exists() + data = json.loads(output_file.read_text()) + + # The benchmarked model exposes a dynamic (None) batch axis; the + # remaining channel/spatial dims stay static. + input_shapes = data["model_info"]["input_shapes"] + assert input_shapes[0][0] is None, f"expected dynamic batch, got {input_shapes}" + assert input_shapes[0][1:] == [3, 224, 224] + assert data["latency_ms"]["mean"] > 0 + + # The ONNX graph ORT actually loaded carries the symbolic batch dim. + running = onnx.load(data["benchmark_info"]["running_model_path"]) + pixel_values = next(i for i in running.graph.input if i.name == "pixel_values") + assert pixel_values.type.tensor_type.shape.dim[0].dim_param == "batch" + + # =========================================================================== # GenAI runtime (winml-genai): --device / --ep override # ===========================================================================