diff --git a/docs/commands/config.md b/docs/commands/config.md index 63c5bdf0e..278fe3220 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 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. | @@ -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 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 7fc8ea922..1f2bcade0 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 @@ -276,6 +312,27 @@ def config( "--module is not supported with ONNX file input. " "Module discovery requires a HuggingFace model." ) + # 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." + ) + 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] if hf_model and _hf_is_onnx: @@ -306,6 +363,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. 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 " + "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, @@ -327,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. @@ -345,7 +426,10 @@ def config( ep=ep, ) if isinstance(result, list): - configs = 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] @@ -353,7 +437,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 +491,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"{input_specs.name} [green]\u2713[/green]" + ) + if export_config: + console.print( + f" \U0001f4c1 [bold]Export config:[/bold] " + f"{export_config.name} [green]\u2713[/green]" + ) + if dynamic_axes: + console.print( + f" \U0001f4c1 [bold]Dynamic axes:[/bold] " + f"{dynamic_axes.name} [green]\u2713[/green]" + ) console.print() 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.""" diff --git a/tests/unit/commands/test_config_cli.py b/tests/unit/commands/test_config_cli.py index 3a19fa6b7..47c51bdc0 100644 --- a/tests/unit/commands/test_config_cli.py +++ b/tests/unit/commands/test_config_cli.py @@ -693,3 +693,285 @@ 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.""" + + @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 + + 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_input_specs_patch_by_name_and_derive_dynamic_axes( + self, runner: CliRunner, tmp_path: Path + ) -> None: + """--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({"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({"input_ids": {"0": "batch"}})) + 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", + "--export-config", + str(export_config), + "--dynamic-axes", + str(dynamic_axes), + "-o", + str(out), + ], + ) + + assert result.exit_code == 0, result.output + 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 the generated export section is emitted as-is.""" + from winml.modelkit.commands.config import config + + 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 + 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_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"]}})) + + with patch( + "winml.modelkit.config.generate_hf_build_config", + ) as mock_generate: + result = runner.invoke( + config, + [ + "-m", + "bert-base-uncased", + "--module", + "BertLayer", + "--input-specs", + str(input_specs), + ], + ) + + 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 + ) -> 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_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() + + 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 + ) -> 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 + + 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