From cf18420a122f638cafd39e0b9481008331f98fea Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Tue, 14 Jul 2026 13:15:03 +0800 Subject: [PATCH 1/5] feat(config): add dynamic export controls to config command Mirror build/perf by adding --input-specs, --export-config, and --dynamic-axes to `winml config`. Overrides are applied via the shared merge_export_overrides helper so --input-specs patches auto-resolved input tensors by name (preserving unlisted inputs and their dtype/value_range) and symbolic dims re-derive dynamic axes. Wired into the single, --module list, and composite-pipeline paths; rejected for pre-exported ONNX inputs and configs with export=null. Docs and config CLI tests updated. --- docs/commands/config.md | 12 ++++ src/winml/modelkit/commands/config.py | 73 ++++++++++++++++++++- tests/unit/commands/test_config_cli.py | 89 ++++++++++++++++++++++++++ 3 files changed, 172 insertions(+), 2 deletions(-) diff --git a/docs/commands/config.md b/docs/commands/config.md index 63c5bdf0e..301e3930e 100644 --- a/docs/commands/config.md +++ b/docs/commands/config.md @@ -23,6 +23,9 @@ $ winml config [options] | `--module` | | `TEXT` | *(none)* | Generate configs for every submodule whose class name matches the given string (e.g., `ResNetConvLayer`). The output is a JSON array instead of a single object. | | `--config` | `-c` | `PATH` | *(none)* | JSON override file in `WinMLBuildConfig` format. Fields present in this file take precedence over auto-detected values. | | `--shape-config` | | `PATH` | *(none)* | JSON file with input shape overrides for dummy input generation. Valid keys by modality — text: `sequence_length`; vision: `height`, `width`, `num_channels`; audio: `feature_size`, `nb_max_frames`, `audio_sequence_length`. | +| `--input-specs` | | `PATH` | *(none)* | JSON file with input specifications for the HuggingFace export. Fields are patched onto the auto-resolved input tensors *by name* (unlisted inputs and their `dtype`/`value_range` are preserved); unknown names are appended. Only valid when generating a HuggingFace export config. | +| `--export-config` | | `PATH` | *(none)* | ONNX export configuration JSON (`opset_version`, `do_constant_folding`, etc.) merged into the generated `export` section. Only valid when generating a HuggingFace export config. | +| `--dynamic-axes` | | `PATH` | *(none)* | JSON dynamic axes mapping for the ONNX export (e.g., `{"input_ids": {"0": "batch", "1": "sequence"}}`). Symbolic string dimensions in `--input-specs` shapes also infer dynamic axes. Only valid when generating a HuggingFace export config. | | `--device` | `-d` | `auto\|npu\|gpu\|cpu` | `auto` | Target device. Affects the generated quantization and compilation sub-configs. `auto` leaves those sections unchanged from the kit defaults. | | `--ep` | | `TEXT` | *(none)* | Force a specific execution provider (`qnn`, `dml`, `migraphx`, `tensorrt`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. When used without `--device`, the device is inferred from the EP. | | `--precision` | `-p` | `TEXT` | `auto` | Target precision: `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed format such as `w8a16`. `auto` selects the precision based on the chosen device. | @@ -81,6 +84,14 @@ Generate a config from an already-exported ONNX file, skipping quantization (com $ winml config -m facebook/convnext-tiny-224.onnx --no-quant -o convnext_optim_only.json ``` +Generate a config with dynamic axes so the exported model accepts a variable batch dimension: + +```bash +$ winml config -m microsoft/resnet-50 --dynamic-axes dynamic_axes.json -o resnet_dyn.json +``` + +The `dynamic_axes.json` maps each input to its dynamic dimensions, e.g. `{"pixel_values": {"0": "batch"}}`. Symbolic string dimensions in `--input-specs` shapes (e.g. `{"input_ids": {"shape": ["batch", "sequence"]}}`) infer dynamic axes automatically without a separate `--dynamic-axes` file. + ## Common pitfalls - **At least one of `-m`, `--model-type`, or `--model-class` is required** — calling `winml config` with none of these three flags raises a usage error immediately. @@ -88,6 +99,7 @@ $ winml config -m facebook/convnext-tiny-224.onnx --no-quant -o convnext_optim_o - **`--module` changes the output shape** — with `--module` the JSON output is an array of configs, not a single object. Scripts that expect a single object will fail to parse this output. - **`--trust-remote-code` has security implications** — only use this flag with model repositories you own or explicitly trust; it allows arbitrary Python execution from the remote model card. - **Shape overrides in `--shape-config` are modality-specific** — passing a `sequence_length` key for a vision model has no effect. Check the `--help` description for valid keys per modality. +- **Export controls require a HuggingFace export** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. ## See also diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 7fc8ea922..34441faa9 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -59,6 +59,27 @@ def _apply_stage_overrides(cfg: Any, *, no_quant: bool, no_compile: bool) -> Non cfg.compile = None +def _merge_export_overrides(cfg: Any, export_overrides: dict[str, Any]) -> Any: + """Apply --export-config/--dynamic-axes/--input-specs onto a generated config. + + Returns ``cfg`` unchanged when no export overrides were supplied. Mirrors the + ``build`` command: ``--input-specs`` patches the auto-resolved input tensors + by name (preserving unlisted inputs and their dtype/value_range) and any + symbolic dims re-derive dynamic axes via ``WinMLExportConfig.__post_init__``. + """ + if not export_overrides: + return cfg + if cfg.export is None: + raise click.UsageError( + "--input-specs, --export-config, and --dynamic-axes require a " + "HuggingFace export config; they are not supported when the " + "generated build config has export=null." + ) + from ..config import merge_export_overrides + + return merge_export_overrides(cfg, export_overrides) + + @click.command("config") @cli_utils.model_option(required=False, optional_message="Optional when --model-type is provided.") @click.option( @@ -99,6 +120,14 @@ def _apply_stage_overrides(cfg: Any, *, no_quant: bool, no_compile: bool) -> Non "vision: height, width, num_channels; " "audio: feature_size, nb_max_frames, audio_sequence_length.", ) +@cli_utils.input_specs_option() +@cli_utils.export_config_option() +@cli_utils.dynamic_axes_option( + help_text=( + "JSON dynamic axes mapping for HuggingFace ONNX export " + '(e.g., {"input_ids": {"0": "batch", "1": "sequence"}}).' + ) +) @cli_utils.device_option( required=False, optional_message="Affects quant/compile config.", @@ -140,6 +169,9 @@ def config( module: str | None, config_file: str | None, shape_config_file: str | None, + input_specs: Path | None, + export_config: Path | None, + dynamic_axes: Path | None, device: str, ep: EPNameOrAlias | None, precision: str, @@ -190,6 +222,10 @@ def config( # Vision model with shape overrides ({"height": 224, "width": 224}) winml config --model-type resnet -t image-classification --shape-config shapes.json + # Dynamic export controls (mirrors ``winml build`` / ``winml export``) + winml config -m bert-base-uncased --dynamic-axes dynamic_axes.json + winml config -m bert-base-uncased --input-specs inputs.json --export-config export.json + # Save to file winml config -m bert-base-uncased -o config.json @@ -266,6 +302,15 @@ def config( ) from e _shape_config_file = shape_config_path.name + # Load export CLI overrides (--input-specs/--export-config/--dynamic-axes). + # Returned sparse so unspecified fields don't clobber auto-detected values; + # applied per generated config via merge_export_overrides below. + export_overrides = cli_utils.load_export_overrides( + export_config=export_config, + input_specs=input_specs, + dynamic_axes=dynamic_axes, + ) + # ONNX file detection: generate simpler config without loader/export _model_input = classify_model_input(hf_model) if hf_model else None if _model_input is not None and _model_input.kind is ModelInputKind.INVALID: @@ -276,6 +321,12 @@ def config( "--module is not supported with ONNX file input. " "Module discovery requires a HuggingFace model." ) + if hf_model and _hf_is_onnx and export_overrides: + raise click.UsageError( + "--input-specs, --export-config, and --dynamic-axes are only " + "supported when generating a HuggingFace export config, not " + "pre-exported ONNX files." + ) config_obj: WinMLBuildConfig | None = None output_data: dict[str, Any] | list[Any] if hf_model and _hf_is_onnx: @@ -314,6 +365,7 @@ def config( model_type=model_type, override=override, shape_config=shape_config, + export_overrides=export_overrides, library_name=library_name, device=device, precision=precision, @@ -345,7 +397,7 @@ def config( ep=ep, ) if isinstance(result, list): - configs = result + configs = [_merge_export_overrides(cfg, export_overrides) for cfg in result] for cfg in configs: _apply_stage_overrides(cfg, no_quant=not quant, no_compile=no_compile) output_data = [cfg.to_dict() for cfg in configs] @@ -353,7 +405,7 @@ def config( # Use first config for display metadata config_obj = configs[0] if configs else None else: - config_obj = result + config_obj = _merge_export_overrides(result, export_overrides) configs = [] _apply_stage_overrides(config_obj, no_quant=not quant, no_compile=no_compile) output_data = config_obj.to_dict() @@ -407,6 +459,21 @@ def config( f" \U0001f4c1 [bold]Shape config:[/bold] " f"{_shape_config_file} [green]\u2713[/green]" ) + if input_specs: + console.print( + f" \U0001f4c1 [bold]Input specs:[/bold] " + f"{Path(input_specs).name} [green]\u2713[/green]" + ) + if export_config: + console.print( + f" \U0001f4c1 [bold]Export config:[/bold] " + f"{Path(export_config).name} [green]\u2713[/green]" + ) + if dynamic_axes: + console.print( + f" \U0001f4c1 [bold]Dynamic axes:[/bold] " + f"{Path(dynamic_axes).name} [green]\u2713[/green]" + ) console.print() @@ -509,6 +576,7 @@ def _generate_pipeline_configs( model_type: str | None, override: Any, shape_config: dict | None, + export_overrides: dict[str, Any], library_name: str, device: str, precision: str, @@ -542,6 +610,7 @@ def _generate_pipeline_configs( trust_remote_code=trust_remote_code, ep=ep, ) + cfg = _merge_export_overrides(cfg, export_overrides) _apply_stage_overrides(cfg, no_quant=no_quant, no_compile=no_compile) config_json = json.dumps(cfg.to_dict(), indent=2) diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 3a19fa6b7..60d5e14a8 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -693,3 +693,92 @@ def test_non_qdq_onnx_has_default_quant(self, runner: CliRunner, tmp_path: Path) assert data.get("quant") is not None, ( f"Non-QDQ model should have default quant settings, got: {data.get('quant')}" ) + + +class TestConfigExportControls: + """Export CLI overrides (--input-specs/--export-config/--dynamic-axes) on config.""" + + def test_help_shows_export_control_options(self, runner: CliRunner) -> None: + from winml.modelkit.commands.config import config + + result = runner.invoke(config, ["--help"]) + assert result.exit_code == 0 + for opt in ("--input-specs", "--export-config", "--dynamic-axes"): + assert opt in result.output, f"Expected '{opt}' in help output" + + def test_export_overrides_forwarded_to_merge( + self, + runner: CliRunner, + tmp_path: Path, + mock_generate_config: MagicMock, + ) -> None: + """--input-specs/--export-config/--dynamic-axes are parsed and merged onto the config.""" + from winml.modelkit.commands.config import config + + input_specs = tmp_path / "inputs.json" + input_specs.write_text( + json.dumps({"pixel_values": {"dtype": "float32", "shape": ["batch", 3, 224, 224]}}) + ) + export_config = tmp_path / "export.json" + export_config.write_text(json.dumps({"opset_version": 18})) + dynamic_axes = tmp_path / "dynamic_axes.json" + dynamic_axes.write_text(json.dumps({"pixel_values": {"0": "batch"}})) + + with patch( + "winml.modelkit.config.merge_export_overrides", + side_effect=lambda cfg, overrides: cfg, + ) as mock_merge: + result = runner.invoke( + config, + [ + "-m", + "microsoft/resnet-50", + "--input-specs", + str(input_specs), + "--export-config", + str(export_config), + "--dynamic-axes", + str(dynamic_axes), + ], + ) + + assert result.exit_code == 0, result.output + assert mock_merge.call_count == 1 + overrides = mock_merge.call_args.args[1] + assert overrides["opset_version"] == 18 + assert overrides["dynamic_axes"] == {"pixel_values": {"0": "batch"}} + assert overrides["input_tensors"][0].name == "pixel_values" + assert overrides["input_tensors"][0].shape == ("batch", 3, 224, 224) + + def test_no_export_overrides_skips_merge( + self, + runner: CliRunner, + mock_generate_config: MagicMock, + ) -> None: + """Without export flags, merge_export_overrides is not invoked.""" + from winml.modelkit.commands.config import config + + with patch( + "winml.modelkit.config.merge_export_overrides", + side_effect=lambda cfg, overrides: cfg, + ) as mock_merge: + result = runner.invoke(config, ["-m", "microsoft/resnet-50"]) + + assert result.exit_code == 0, result.output + assert mock_merge.call_count == 0 + + def test_export_overrides_rejected_for_onnx_input( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """--dynamic-axes on a pre-exported ONNX input is a usage error.""" + from winml.modelkit.commands.config import config + + onnx_file = tmp_path / "model.onnx" + onnx_file.write_bytes(b"fake-onnx") + dynamic_axes = tmp_path / "dynamic_axes.json" + dynamic_axes.write_text(json.dumps({"input": {"0": "batch"}})) + + result = runner.invoke(config, ["-m", str(onnx_file), "--dynamic-axes", str(dynamic_axes)]) + + assert result.exit_code != 0 + assert "pre-exported ONNX" in result.output From c902f0e802012bb6234335c0bf0e8876575df945 Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Wed, 15 Jul 2026 16:03:02 +0800 Subject: [PATCH 2/5] fix(config): address PR review on export controls - Reject --input-specs/--export-config/--dynamic-axes for composite (multi-component) models instead of fanning the same overrides across sub-components with distinct inputs, which could append bogus inputs. - Move the pre-exported-ONNX rejection ahead of loading/validating the override JSON so the error names the ONNX input (not a downstream "Invalid export configuration") and skips needless file I/O. - Replace mock-based tests with integration coverage that lets merge_export_overrides run and asserts on to_dict(): by-name patching, unlisted-input preservation, symbolic-dim -> dynamic_axes derivation, export=null rejection, the --module list path, composite rejection, and ONNX rejection ordering. - Document the composite rejection in config.md. --- docs/commands/config.md | 2 +- src/winml/modelkit/commands/config.py | 42 +++-- tests/unit/commands/test_config_cli.py | 245 +++++++++++++++++++++---- 3 files changed, 240 insertions(+), 49 deletions(-) diff --git a/docs/commands/config.md b/docs/commands/config.md index 301e3930e..b53c7ce85 100644 --- a/docs/commands/config.md +++ b/docs/commands/config.md @@ -99,7 +99,7 @@ The `dynamic_axes.json` maps each input to its dynamic dimensions, e.g. `{"pixel - **`--module` changes the output shape** — with `--module` the JSON output is an array of configs, not a single object. Scripts that expect a single object will fail to parse this output. - **`--trust-remote-code` has security implications** — only use this flag with model repositories you own or explicitly trust; it allows arbitrary Python execution from the remote model card. - **Shape overrides in `--shape-config` are modality-specific** — passing a `sequence_length` key for a vision model has no effect. Check the `--help` description for valid keys per modality. -- **Export controls require a HuggingFace export** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. +- **Export controls require a HuggingFace export** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. They are also rejected for composite (multi-component) models, whose sub-components each have their own export inputs — generate the per-component configs first, then edit their export sections individually. ## See also diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index 34441faa9..cbd9df53f 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -302,15 +302,6 @@ def config( ) from e _shape_config_file = shape_config_path.name - # Load export CLI overrides (--input-specs/--export-config/--dynamic-axes). - # Returned sparse so unspecified fields don't clobber auto-detected values; - # applied per generated config via merge_export_overrides below. - export_overrides = cli_utils.load_export_overrides( - export_config=export_config, - input_specs=input_specs, - dynamic_axes=dynamic_axes, - ) - # ONNX file detection: generate simpler config without loader/export _model_input = classify_model_input(hf_model) if hf_model else None if _model_input is not None and _model_input.kind is ModelInputKind.INVALID: @@ -321,12 +312,26 @@ def config( "--module is not supported with ONNX file input. " "Module discovery requires a HuggingFace model." ) - if hf_model and _hf_is_onnx and export_overrides: + # Reject export controls for pre-exported ONNX up front — before loading + # and validating the JSON — so the error names the real problem (ONNX + # input, no export stage) instead of a downstream validation failure, and + # so we skip the needless file I/O. + if hf_model and _hf_is_onnx and (input_specs or export_config or dynamic_axes): raise click.UsageError( "--input-specs, --export-config, and --dynamic-axes are only " "supported when generating a HuggingFace export config, not " "pre-exported ONNX files." ) + + # Load export CLI overrides (--input-specs/--export-config/--dynamic-axes). + # Returned sparse so unspecified fields don't clobber auto-detected values; + # applied per generated config via merge_export_overrides below. + export_overrides = cli_utils.load_export_overrides( + export_config=export_config, + input_specs=input_specs, + dynamic_axes=dynamic_axes, + ) + config_obj: WinMLBuildConfig | None = None output_data: dict[str, Any] | list[Any] if hf_model and _hf_is_onnx: @@ -357,6 +362,20 @@ def config( hf_model, model_type, task, trust_remote_code=trust_remote_code ) if pipeline_components: + # Export controls target a single export graph; a composite model + # has one export per sub-component with distinct inputs. Applying a + # shared --input-specs/--dynamic-axes across them would append bogus + # inputs on components that lack the named input. Reject instead of + # emitting invalid per-component configs (unlike build, config never + # fans these overrides out). + if export_overrides: + raise click.UsageError( + "--input-specs, --export-config, and --dynamic-axes are not " + "supported for composite (multi-component) models, whose " + "sub-components each have their own export inputs. Generate " + "the per-component configs first, then edit their export " + "sections individually." + ) # composite model: generate one config per sub-component _generate_pipeline_configs( pipeline_components, @@ -365,7 +384,6 @@ def config( model_type=model_type, override=override, shape_config=shape_config, - export_overrides=export_overrides, library_name=library_name, device=device, precision=precision, @@ -576,7 +594,6 @@ def _generate_pipeline_configs( model_type: str | None, override: Any, shape_config: dict | None, - export_overrides: dict[str, Any], library_name: str, device: str, precision: str, @@ -610,7 +627,6 @@ def _generate_pipeline_configs( trust_remote_code=trust_remote_code, ep=ep, ) - cfg = _merge_export_overrides(cfg, export_overrides) _apply_stage_overrides(cfg, no_quant=no_quant, no_compile=no_compile) config_json = json.dumps(cfg.to_dict(), indent=2) diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 60d5e14a8..05b265744 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -698,6 +698,27 @@ def test_non_qdq_onnx_has_default_quant(self, runner: CliRunner, tmp_path: Path) class TestConfigExportControls: """Export CLI overrides (--input-specs/--export-config/--dynamic-axes) on config.""" + @staticmethod + def _real_config(): + """Build a real WinMLBuildConfig with an export section for integration tests.""" + from winml.modelkit.config import WinMLBuildConfig + from winml.modelkit.export import InputTensorSpec, WinMLExportConfig + from winml.modelkit.loader import WinMLLoaderConfig + + return WinMLBuildConfig( + loader=WinMLLoaderConfig( + task="fill-mask", model_class="BertForMaskedLM", model_type="bert" + ), + export=WinMLExportConfig( + input_tensors=[ + InputTensorSpec( + name="input_ids", dtype="int64", shape=(1, 16), value_range=(0, 30522) + ), + InputTensorSpec(name="attention_mask", dtype="int64", shape=(1, 16)), + ], + ), + ) + def test_help_shows_export_control_options(self, runner: CliRunner) -> None: from winml.modelkit.commands.config import config @@ -706,66 +727,199 @@ def test_help_shows_export_control_options(self, runner: CliRunner) -> None: for opt in ("--input-specs", "--export-config", "--dynamic-axes"): assert opt in result.output, f"Expected '{opt}' in help output" - def test_export_overrides_forwarded_to_merge( - self, - runner: CliRunner, - tmp_path: Path, - mock_generate_config: MagicMock, + def test_input_specs_patch_by_name_and_derive_dynamic_axes( + self, runner: CliRunner, tmp_path: Path ) -> None: - """--input-specs/--export-config/--dynamic-axes are parsed and merged onto the config.""" + """--input-specs patches by name; symbolic dims re-derive dynamic_axes. + + Lets merge_export_overrides actually run and asserts on the emitted config: + the unlisted attention_mask is preserved (with its int64 dtype), and the + symbolic ["batch", "seq"] shape on input_ids produces dynamic_axes without + a separate --dynamic-axes file. + """ from winml.modelkit.commands.config import config input_specs = tmp_path / "inputs.json" - input_specs.write_text( - json.dumps({"pixel_values": {"dtype": "float32", "shape": ["batch", 3, 224, 224]}}) - ) + input_specs.write_text(json.dumps({"input_ids": {"shape": ["batch", "seq"]}})) + out = tmp_path / "out.json" + + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value=None, + ), + patch( + "winml.modelkit.config.generate_hf_build_config", + return_value=self._real_config(), + ), + ): + result = runner.invoke( + config, + ["-m", "bert-base-uncased", "--input-specs", str(input_specs), "-o", str(out)], + ) + + assert result.exit_code == 0, result.output + export = json.loads(out.read_text())["export"] + names = [t["name"] for t in export["input_tensors"]] + assert names == ["input_ids", "attention_mask"] # unlisted input preserved + ids = next(t for t in export["input_tensors"] if t["name"] == "input_ids") + assert ids["dtype"] == "int64" # preserved, not forced to float32 + axes = export["dynamic_axes"]["input_ids"] # symbolic dims derived dynamic axes + assert set(axes.values()) == {"batch", "seq"} + + def test_export_config_and_dynamic_axes_applied( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """--export-config and --dynamic-axes are merged onto the generated config.""" + from winml.modelkit.commands.config import config + export_config = tmp_path / "export.json" export_config.write_text(json.dumps({"opset_version": 18})) dynamic_axes = tmp_path / "dynamic_axes.json" - dynamic_axes.write_text(json.dumps({"pixel_values": {"0": "batch"}})) + dynamic_axes.write_text(json.dumps({"input_ids": {"0": "batch"}})) + out = tmp_path / "out.json" - with patch( - "winml.modelkit.config.merge_export_overrides", - side_effect=lambda cfg, overrides: cfg, - ) as mock_merge: + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value=None, + ), + patch( + "winml.modelkit.config.generate_hf_build_config", + return_value=self._real_config(), + ), + ): result = runner.invoke( config, [ "-m", - "microsoft/resnet-50", - "--input-specs", - str(input_specs), + "bert-base-uncased", "--export-config", str(export_config), "--dynamic-axes", str(dynamic_axes), + "-o", + str(out), ], ) assert result.exit_code == 0, result.output - assert mock_merge.call_count == 1 - overrides = mock_merge.call_args.args[1] - assert overrides["opset_version"] == 18 - assert overrides["dynamic_axes"] == {"pixel_values": {"0": "batch"}} - assert overrides["input_tensors"][0].name == "pixel_values" - assert overrides["input_tensors"][0].shape == ("batch", 3, 224, 224) - - def test_no_export_overrides_skips_merge( - self, - runner: CliRunner, - mock_generate_config: MagicMock, + export = json.loads(out.read_text())["export"] + assert export["opset_version"] == 18 + assert export["dynamic_axes"]["input_ids"] == {"0": "batch"} + + def test_no_export_overrides_leaves_export_unchanged( + self, runner: CliRunner, tmp_path: Path ) -> None: - """Without export flags, merge_export_overrides is not invoked.""" + """Without export flags the generated export section is emitted as-is.""" from winml.modelkit.commands.config import config - with patch( - "winml.modelkit.config.merge_export_overrides", - side_effect=lambda cfg, overrides: cfg, - ) as mock_merge: - result = runner.invoke(config, ["-m", "microsoft/resnet-50"]) + out = tmp_path / "out.json" + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value=None, + ), + patch( + "winml.modelkit.config.generate_hf_build_config", + return_value=self._real_config(), + ), + ): + result = runner.invoke(config, ["-m", "bert-base-uncased", "-o", str(out)]) assert result.exit_code == 0, result.output - assert mock_merge.call_count == 0 + export = json.loads(out.read_text())["export"] + # Concrete auto-resolved shapes -> no dynamic axes are invented. + assert not export.get("dynamic_axes") + + def test_module_path_applies_overrides_to_each_config( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """--module list path patches --input-specs onto every generated submodule config.""" + from winml.modelkit.commands.config import config + + input_specs = tmp_path / "inputs.json" + input_specs.write_text(json.dumps({"input_ids": {"shape": ["batch", "seq"]}})) + out = tmp_path / "out.json" + + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value=None, + ), + patch( + "winml.modelkit.config.generate_hf_build_config", + return_value=[self._real_config(), self._real_config()], + ), + ): + result = runner.invoke( + config, + [ + "-m", + "bert-base-uncased", + "--module", + "BertLayer", + "--input-specs", + str(input_specs), + "-o", + str(out), + ], + ) + + assert result.exit_code == 0, result.output + configs = json.loads(out.read_text()) + assert isinstance(configs, list) and len(configs) == 2 + for cfg in configs: + axes = cfg["export"]["dynamic_axes"]["input_ids"] + assert set(axes.values()) == {"batch", "seq"} + + def test_composite_model_rejects_export_overrides( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """Composite models reject export overrides instead of fanning them out.""" + from winml.modelkit.commands.config import config + + dynamic_axes = tmp_path / "dynamic_axes.json" + dynamic_axes.write_text(json.dumps({"input_ids": {"0": "batch"}})) + + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value={"encoder": "feature-extraction", "decoder": "text-generation"}, + ), + patch( + "winml.modelkit.commands.config._generate_pipeline_configs", + ) as mock_pipeline, + ): + result = runner.invoke( + config, ["-m", "some/seq2seq", "--dynamic-axes", str(dynamic_axes)] + ) + + assert result.exit_code != 0 + assert "composite" in result.output + mock_pipeline.assert_not_called() + + def test_merge_export_overrides_rejects_export_null(self) -> None: + """_merge_export_overrides raises when the generated config has export=null.""" + import click + + from winml.modelkit.commands.config import _merge_export_overrides + from winml.modelkit.config import WinMLBuildConfig + from winml.modelkit.loader import WinMLLoaderConfig + + cfg = WinMLBuildConfig( + loader=WinMLLoaderConfig(task="fill-mask", model_class="X", model_type="bert"), + export=None, + ) + with pytest.raises(click.UsageError, match="export=null"): + _merge_export_overrides(cfg, {"dynamic_axes": {"input_ids": {"0": "batch"}}}) + + def test_merge_export_overrides_noop_when_empty(self) -> None: + """_merge_export_overrides returns the config untouched when no overrides given.""" + from winml.modelkit.commands.config import _merge_export_overrides + + cfg = self._real_config() + assert _merge_export_overrides(cfg, {}) is cfg def test_export_overrides_rejected_for_onnx_input( self, runner: CliRunner, tmp_path: Path @@ -782,3 +936,24 @@ def test_export_overrides_rejected_for_onnx_input( assert result.exit_code != 0 assert "pre-exported ONNX" in result.output + + def test_onnx_rejection_precedes_json_validation( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """ONNX inputs are rejected before the override JSON is loaded/validated. + + A malformed --dynamic-axes on an ONNX input should surface the relevant + ONNX error, not a downstream "Invalid export configuration" from parsing. + """ + from winml.modelkit.commands.config import config + + onnx_file = tmp_path / "model.onnx" + onnx_file.write_bytes(b"fake-onnx") + dynamic_axes = tmp_path / "dynamic_axes.json" + dynamic_axes.write_text(json.dumps({"input_ids": {"not-an-int": "batch"}})) + + result = runner.invoke(config, ["-m", str(onnx_file), "--dynamic-axes", str(dynamic_axes)]) + + assert result.exit_code != 0 + assert "pre-exported ONNX" in result.output + assert "Invalid export configuration" not in result.output From e0b0580c1f438dcc494eb4b83d18bd5e079d35bb Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 16 Jul 2026 13:48:36 +0800 Subject: [PATCH 3/5] fix: reject export controls for --module and align composite/module rejection Address PR #1106 review comments: - --module now rejects --input-specs/--export-config/--dynamic-axes up front instead of fanning them onto every submodule config (would append bogus inputs on submodules with different input names). - Reject the ONNX / --module / composite paths on raw flag presence before loading the override JSON, so the path-specific error wins over a downstream "Invalid export configuration" (mirrors the ONNX guard). - Move load_export_overrides below composite resolution; the single HF config path is now the only consumer. - Simplify Path(input_specs).name to input_specs.name (already Path|None). - Docs + tests updated: --module rejection, composite-precedes-JSON test. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- docs/commands/config.md | 8 ++-- src/winml/modelkit/commands/config.py | 62 +++++++++++++++---------- tests/unit/commands/test_config_cli.py | 64 ++++++++++++++++---------- 3 files changed, 82 insertions(+), 52 deletions(-) diff --git a/docs/commands/config.md b/docs/commands/config.md index b53c7ce85..278fe3220 100644 --- a/docs/commands/config.md +++ b/docs/commands/config.md @@ -23,9 +23,9 @@ $ winml config [options] | `--module` | | `TEXT` | *(none)* | Generate configs for every submodule whose class name matches the given string (e.g., `ResNetConvLayer`). The output is a JSON array instead of a single object. | | `--config` | `-c` | `PATH` | *(none)* | JSON override file in `WinMLBuildConfig` format. Fields present in this file take precedence over auto-detected values. | | `--shape-config` | | `PATH` | *(none)* | JSON file with input shape overrides for dummy input generation. Valid keys by modality — text: `sequence_length`; vision: `height`, `width`, `num_channels`; audio: `feature_size`, `nb_max_frames`, `audio_sequence_length`. | -| `--input-specs` | | `PATH` | *(none)* | JSON file with input specifications for the HuggingFace export. Fields are patched onto the auto-resolved input tensors *by name* (unlisted inputs and their `dtype`/`value_range` are preserved); unknown names are appended. Only valid when generating a HuggingFace export config. | -| `--export-config` | | `PATH` | *(none)* | ONNX export configuration JSON (`opset_version`, `do_constant_folding`, etc.) merged into the generated `export` section. Only valid when generating a HuggingFace export config. | -| `--dynamic-axes` | | `PATH` | *(none)* | JSON dynamic axes mapping for the ONNX export (e.g., `{"input_ids": {"0": "batch", "1": "sequence"}}`). Symbolic string dimensions in `--input-specs` shapes also infer dynamic axes. Only valid when generating a HuggingFace export config. | +| `--input-specs` | | `PATH` | *(none)* | JSON file with input specifications for the HuggingFace export. Fields are patched onto the auto-resolved input tensors *by name* (unlisted inputs and their `dtype`/`value_range` are preserved); unknown names are appended. Only valid for a single HuggingFace export config (not `--module` or composite models). | +| `--export-config` | | `PATH` | *(none)* | ONNX export configuration JSON (`opset_version`, `do_constant_folding`, etc.) merged into the generated `export` section. Only valid for a single HuggingFace export config (not `--module` or composite models). | +| `--dynamic-axes` | | `PATH` | *(none)* | JSON dynamic axes mapping for the ONNX export (e.g., `{"input_ids": {"0": "batch", "1": "sequence"}}`). Symbolic string dimensions in `--input-specs` shapes also infer dynamic axes. Only valid for a single HuggingFace export config (not `--module` or composite models). | | `--device` | `-d` | `auto\|npu\|gpu\|cpu` | `auto` | Target device. Affects the generated quantization and compilation sub-configs. `auto` leaves those sections unchanged from the kit defaults. | | `--ep` | | `TEXT` | *(none)* | Force a specific execution provider (`qnn`, `dml`, `migraphx`, `tensorrt`, `vitisai`, `openvino`, `cpu`). Overrides the device-to-provider mapping. When used without `--device`, the device is inferred from the EP. | | `--precision` | `-p` | `TEXT` | `auto` | Target precision: `auto`, `fp32`, `fp16`, `int8`, `int16`, or a mixed format such as `w8a16`. `auto` selects the precision based on the chosen device. | @@ -99,7 +99,7 @@ The `dynamic_axes.json` maps each input to its dynamic dimensions, e.g. `{"pixel - **`--module` changes the output shape** — with `--module` the JSON output is an array of configs, not a single object. Scripts that expect a single object will fail to parse this output. - **`--trust-remote-code` has security implications** — only use this flag with model repositories you own or explicitly trust; it allows arbitrary Python execution from the remote model card. - **Shape overrides in `--shape-config` are modality-specific** — passing a `sequence_length` key for a vision model has no effect. Check the `--help` description for valid keys per modality. -- **Export controls require a HuggingFace export** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. They are also rejected for composite (multi-component) models, whose sub-components each have their own export inputs — generate the per-component configs first, then edit their export sections individually. +- **Export controls require a single HuggingFace export graph** — `--input-specs`, `--export-config`, and `--dynamic-axes` are rejected for pre-exported `.onnx` inputs (which set `export` to `null`); use them only when config generates the export section. They are also rejected for composite (multi-component) models and for `--module` (which fan out one config per sub-component / submodule, each with their own export inputs) — generate those per-config outputs first, then edit their export sections individually. ## See also diff --git a/src/winml/modelkit/commands/config.py b/src/winml/modelkit/commands/config.py index cbd9df53f..1f2bcade0 100644 --- a/src/winml/modelkit/commands/config.py +++ b/src/winml/modelkit/commands/config.py @@ -312,25 +312,26 @@ def config( "--module is not supported with ONNX file input. " "Module discovery requires a HuggingFace model." ) - # Reject export controls for pre-exported ONNX up front — before loading - # and validating the JSON — so the error names the real problem (ONNX - # input, no export stage) instead of a downstream validation failure, and - # so we skip the needless file I/O. - if hf_model and _hf_is_onnx and (input_specs or export_config or dynamic_axes): + # Export controls (--input-specs/--export-config/--dynamic-axes) target a + # single export graph, so they only apply to the plain single-config + # HuggingFace path. Reject the multi-graph / no-export paths up front — on + # raw flag presence, before loading/validating the JSON — so the error + # names the real problem (ONNX input, --module fan-out, composite) instead + # of a downstream "Invalid export configuration", and we skip needless I/O. + _export_flags_given = bool(input_specs or export_config or dynamic_axes) + if hf_model and _hf_is_onnx and _export_flags_given: raise click.UsageError( "--input-specs, --export-config, and --dynamic-axes are only " "supported when generating a HuggingFace export config, not " "pre-exported ONNX files." ) - - # Load export CLI overrides (--input-specs/--export-config/--dynamic-axes). - # Returned sparse so unspecified fields don't clobber auto-detected values; - # applied per generated config via merge_export_overrides below. - export_overrides = cli_utils.load_export_overrides( - export_config=export_config, - input_specs=input_specs, - dynamic_axes=dynamic_axes, - ) + if module and _export_flags_given: + raise click.UsageError( + "--input-specs, --export-config, and --dynamic-axes are not " + "supported with --module, which generates one config per matched " + "submodule. Generate the per-module configs first, then edit their " + "export sections individually." + ) config_obj: WinMLBuildConfig | None = None output_data: dict[str, Any] | list[Any] @@ -363,12 +364,12 @@ def config( ) if pipeline_components: # Export controls target a single export graph; a composite model - # has one export per sub-component with distinct inputs. Applying a - # shared --input-specs/--dynamic-axes across them would append bogus - # inputs on components that lack the named input. Reject instead of - # emitting invalid per-component configs (unlike build, config never - # fans these overrides out). - if export_overrides: + # has one export per sub-component with distinct inputs. Reject on + # raw flag presence — before loading/validating the JSON — so the + # composite-specific error wins (mirroring the ONNX path) instead + # of a downstream "Invalid export configuration". config never fans + # these overrides out across heterogeneous components. + if _export_flags_given: raise click.UsageError( "--input-specs, --export-config, and --dynamic-axes are not " "supported for composite (multi-component) models, whose " @@ -397,6 +398,16 @@ def config( ) return + # Load export CLI overrides now that the multi-graph paths (ONNX, + # --module, composite) have all been rejected — the single HF config + # path below is the only one that applies them. Returned sparse so + # unspecified fields don't clobber auto-detected values. + export_overrides = cli_utils.load_export_overrides( + export_config=export_config, + input_specs=input_specs, + dynamic_axes=dynamic_axes, + ) + # Generate config(s). The ``module: str | None`` overload of # generate_hf_build_config returns WinMLBuildConfig | list[...], # which isinstance(result, list) narrows for the branches below. @@ -415,7 +426,10 @@ def config( ep=ep, ) if isinstance(result, list): - configs = [_merge_export_overrides(cfg, export_overrides) for cfg in result] + # --module + export overrides is rejected up front, so + # export_overrides is empty here; emit the submodule configs as + # generated without fanning any overrides across them. + configs = list(result) for cfg in configs: _apply_stage_overrides(cfg, no_quant=not quant, no_compile=no_compile) output_data = [cfg.to_dict() for cfg in configs] @@ -480,17 +494,17 @@ def config( if input_specs: console.print( f" \U0001f4c1 [bold]Input specs:[/bold] " - f"{Path(input_specs).name} [green]\u2713[/green]" + f"{input_specs.name} [green]\u2713[/green]" ) if export_config: console.print( f" \U0001f4c1 [bold]Export config:[/bold] " - f"{Path(export_config).name} [green]\u2713[/green]" + f"{export_config.name} [green]\u2713[/green]" ) if dynamic_axes: console.print( f" \U0001f4c1 [bold]Dynamic axes:[/bold] " - f"{Path(dynamic_axes).name} [green]\u2713[/green]" + f"{dynamic_axes.name} [green]\u2713[/green]" ) console.print() diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 05b265744..ecee0589b 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -832,26 +832,16 @@ def test_no_export_overrides_leaves_export_unchanged( # Concrete auto-resolved shapes -> no dynamic axes are invented. assert not export.get("dynamic_axes") - def test_module_path_applies_overrides_to_each_config( - self, runner: CliRunner, tmp_path: Path - ) -> None: - """--module list path patches --input-specs onto every generated submodule config.""" + def test_module_path_rejects_export_overrides(self, runner: CliRunner, tmp_path: Path) -> None: + """--module rejects export overrides instead of fanning them onto each config.""" from winml.modelkit.commands.config import config input_specs = tmp_path / "inputs.json" input_specs.write_text(json.dumps({"input_ids": {"shape": ["batch", "seq"]}})) - out = tmp_path / "out.json" - with ( - patch( - "winml.modelkit.commands.config._resolve_composite_model_components", - return_value=None, - ), - patch( - "winml.modelkit.config.generate_hf_build_config", - return_value=[self._real_config(), self._real_config()], - ), - ): + with patch( + "winml.modelkit.config.generate_hf_build_config", + ) as mock_generate: result = runner.invoke( config, [ @@ -861,17 +851,13 @@ def test_module_path_applies_overrides_to_each_config( "BertLayer", "--input-specs", str(input_specs), - "-o", - str(out), ], ) - assert result.exit_code == 0, result.output - configs = json.loads(out.read_text()) - assert isinstance(configs, list) and len(configs) == 2 - for cfg in configs: - axes = cfg["export"]["dynamic_axes"]["input_ids"] - assert set(axes.values()) == {"batch", "seq"} + assert result.exit_code != 0 + assert "--module" in result.output + # Rejected up front, before any config generation. + mock_generate.assert_not_called() def test_composite_model_rejects_export_overrides( self, runner: CliRunner, tmp_path: Path @@ -899,7 +885,37 @@ def test_composite_model_rejects_export_overrides( assert "composite" in result.output mock_pipeline.assert_not_called() - def test_merge_export_overrides_rejects_export_null(self) -> None: + def test_composite_rejection_precedes_json_validation( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """Composite models are rejected before the override JSON is loaded/validated. + + A malformed --dynamic-axes on a composite model should surface the + composite-specific error, not a downstream "Invalid export configuration" + from parsing — mirroring the ONNX path. + """ + from winml.modelkit.commands.config import config + + dynamic_axes = tmp_path / "dynamic_axes.json" + dynamic_axes.write_text(json.dumps({"input_ids": {"not-an-int": "batch"}})) + + with ( + patch( + "winml.modelkit.commands.config._resolve_composite_model_components", + return_value={"encoder": "feature-extraction", "decoder": "text-generation"}, + ), + patch( + "winml.modelkit.commands.config._generate_pipeline_configs", + ) as mock_pipeline, + ): + result = runner.invoke( + config, ["-m", "some/seq2seq", "--dynamic-axes", str(dynamic_axes)] + ) + + assert result.exit_code != 0 + assert "composite" in result.output + assert "Invalid export configuration" not in result.output + mock_pipeline.assert_not_called() """_merge_export_overrides raises when the generated config has export=null.""" import click From ae412b3cded7ac0d31b33ba616a367a2c9da207f Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 16 Jul 2026 16:00:21 +0800 Subject: [PATCH 4/5] test: restore test_merge_export_overrides_rejects_export_null The previous refactor fused two tests: the export=null assertion body was left orphaned at the tail of test_composite_rejection_precedes_json_validation, so test_merge_export_overrides_rejects_export_null no longer existed as a discrete test. Split it back into its own method. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/unit/commands/test_config_cli.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index ecee0589b..47c51bdc0 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -916,6 +916,8 @@ def test_composite_rejection_precedes_json_validation( assert "composite" in result.output assert "Invalid export configuration" not in result.output mock_pipeline.assert_not_called() + + def test_merge_export_overrides_rejects_export_null(self) -> None: """_merge_export_overrides raises when the generated config has export=null.""" import click From d7aeaa8a1a72baadcb122e6bde43868b5329439a Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Thu, 16 Jul 2026 16:43:14 +0800 Subject: [PATCH 5/5] test: stub composite resolution in build unit tests to avoid network hang build calls resolve_composite_components before dispatching, which resolves the model config from HuggingFace Hub (AutoConfig.from_pretrained) over the network. The build CLI plumbing tests only mocked _run_single_build, so this detection call stayed live: repeated metadata requests get throttled and huggingface_hub's retry backoff sleeps, making a later unrelated test appear to hang mid-run. Add an autouse fixture that defaults resolve_composite_components to None (non-composite), matching the existing mock_task_model_compatibility_validator pattern that keeps these tests off HF resolution paths. Composite-specific tests already override it with their own patch, which nests inside and wins. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- tests/unit/commands/test_build.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index ac0b376fa..03a2ce127 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -88,6 +88,25 @@ def mock_task_model_compatibility_validator(): yield +@pytest.fixture(autouse=True) +def mock_composite_resolution(): + """Default composite detection to "not composite" without any network call. + + ``build`` calls ``resolve_composite_components`` before dispatching, which + resolves the model config from HuggingFace Hub (``AutoConfig.from_pretrained``) + over the network. Left live in these CLI plumbing tests, repeated metadata + requests get throttled and ``huggingface_hub``'s retry backoff sleeps — making + a later, unrelated test appear to hang. Returning ``None`` keeps the plain + single-build path. Composite-specific tests override this with their own + ``patch(...)``, which nests inside and wins for the duration of that test. + """ + with patch( + "winml.modelkit.loader.resolution.resolve_composite_components", + return_value=None, + ): + yield + + @pytest.fixture def runner() -> CliRunner: """Create a CLI test runner."""