diff --git a/docs/commands/export.md b/docs/commands/export.md index ff2cf42ec..f3e86cac3 100644 --- a/docs/commands/export.md +++ b/docs/commands/export.md @@ -22,7 +22,7 @@ $ winml export [options] | `--output` | `-o` | path | *(required)* | Output ONNX file path (e.g., `model.onnx`). | | `--with-report/--no-with-report` | | flag | `false` | Generate full export reports: Markdown, JSON, and a console tree. | | `--hierarchy/--no-hierarchy` | | flag | `true` | Preserve `hierarchy_tag` metadata in ONNX nodes (use `--no-hierarchy` for a clean ONNX file). | -| `--dynamo/--no-dynamo` | | flag | `false` | Enable PyTorch 2.9+ dynamo export for richer node metadata. (Experimental — currently logs a warning.) | +| `--dynamo/--no-dynamo` | | flag | `true` | Use PyTorch's TorchDynamo ONNX exporter (default) for richer per-node module metadata. Pass `--no-dynamo` for the legacy TorchScript exporter, whose opset-17 op decomposition is the validated path for QNN/NPU compilation today. | | `--torch-module` | | string | `None` | Comma-separated list of `torch.nn` module types to include in hierarchy (e.g., `LayerNorm,Embedding`). (Experimental — currently logs a warning.) | | `--input-specs` | | path | `None` | JSON file with explicit input tensor specifications. Auto-generated when omitted. | | `--task` | `-t` | string | `None` | Override auto-detected Hugging Face task (e.g., `image-feature-extraction`). | @@ -37,9 +37,12 @@ $ winml export [options] `winml export` loads the model via Hugging Face `transformers`, then runs the eight-step Hierarchy-preserving Tags Protocol (HTP): model preparation, input -generation, module-hierarchy tracing, TorchScript ONNX export, node-tagger +generation, module-hierarchy recovery, ONNX export, node-tagger creation, per-node tagging, tag injection into ONNX `metadata_props`, and -optional report generation. The hierarchy metadata allows downstream tools to +optional report generation. By default the hierarchy is reconstructed from the +TorchDynamo exporter's per-node module metadata; with `--no-dynamo` it is +captured by tracing the model's forward pass under the legacy TorchScript +exporter instead. The hierarchy metadata allows downstream tools to reason about operators grouped by their originating module rather than flat graph position. When `--no-hierarchy` is specified, hierarchy steps are bypassed and a bare ONNX file is written, useful for third-party tools that do not @@ -79,6 +82,15 @@ winml export -m bert-base-uncased -o bert.onnx \ winml export -m bert-base-uncased -o bert.onnx --input-specs inputs.json ``` +```bash +# Export a fully static-shaped model (the default) for NPU/QNN compilation. +# The default dummy inputs already produce static shapes; use --shape-config to +# pin any symbolic dimensions to concrete sizes, or --input-specs to fully +# specify every input tensor. +winml export -m bert-base-uncased -o bert.onnx --shape-config shape_config.json +# shape_config.json: {"sequence_length": 128} +``` + ```bash # Export with dynamic batch and sequence dimensions winml export -m bert-base-uncased -o bert.onnx --dynamic-axes dynamic_axes.json @@ -114,9 +126,26 @@ winml export -m microsoft/resnet-50 -o resnet50_clean.onnx --no-hierarchy - **Dynamic dimensions can reduce QNN optimization coverage.** Static batch and static shapes remain the default because some QNN fusions require them. Use `--dynamic-axes` only when downstream runtime scenarios need variable sizes. -- **`--dynamo` and `--torch-module` are experimental.** Both flags emit a - warning and have no effect in the current release. Do not rely on them in - automated pipelines yet. +- **Dynamo is the default exporter.** `winml export` uses PyTorch's TorchDynamo + ONNX exporter, which records rich per-node module metadata that drives the + hierarchy tags. Pass `--no-dynamo` to select the legacy TorchScript exporter. + Shape staticness is independent of the exporter: both default to static shapes + and only emit dynamic axes when you ask for them. The QNN-relevant + difference is opset and op decomposition: torch's dynamo op library is built + against a minimum opset of 18. When you request a lower opset (17 by default), + dynamo attempts to down-convert the graph to it; this can succeed (the graph is + saved at the requested opset) or fail (the graph stays at opset 18). `winml + export` reports the opset actually produced and warns when it differs from the + one you requested, so pass `--no-dynamo` if you need a graph natively at opset + 17. The TorchScript path always exports at the configured opset (17 by + default). Dynamo also + lowers some ops differently -- for example ResNet's classification head becomes + `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under + TorchScript -- which can change or reduce QNN fusion coverage. Prefer + `--no-dynamo` for hand exports targeting QNN/NPU compilation until the dynamo + graph is validated for your model. +- **`--torch-module` is experimental.** The flag emits a warning and has no + effect in the current release. Do not rely on it in automated pipelines yet. - **Output directory must be writable.** The command creates parent directories automatically, but will fail with a permission error on read-only paths. - **Model weights are downloaded to the Hugging Face cache.** Set `HF_HOME` or diff --git a/docs/concepts/load-and-export.md b/docs/concepts/load-and-export.md index b195adcd6..8fbc378d7 100644 --- a/docs/concepts/load-and-export.md +++ b/docs/concepts/load-and-export.md @@ -2,13 +2,13 @@ The first stage of the winml-cli pipeline is the most deterministic: bring a model into memory and convert it to ONNX. Everything that follows — optimization, quantization, compilation — operates on that ONNX artifact. A well-exported graph with accurate metadata travels cleanly through the rest of the pipeline without requiring patching or re-export. -Loading is an internal operation: the loader module resolves model provenance, selects the right HuggingFace model class, and prepares the weights for tracing. The `winml export` command is the surface users interact with directly. +Loading is an internal operation: the loader module resolves model provenance, selects the right HuggingFace model class, and prepares the weights for export. The `winml export` command is the surface users interact with directly. ## Loading a model When you point winml-cli at a model identifier, the internal loader resolves it in one of two ways. If the identifier looks like a HuggingFace Hub path (e.g., `prajjwal1/bert-tiny`), the loader downloads the model weights and configuration to the standard HuggingFace cache at `~/.cache/huggingface`. Subsequent runs are served from that cache without re-downloading. If the identifier is a path to a local PyTorch checkpoint directory, the loader reads it directly without network access. -In both cases the loader auto-detects the task — image classification, text feature extraction, and so on — and selects a corresponding HuggingFace model class. The result is a PyTorch model object ready for tracing. +In both cases the loader auto-detects the task — image classification, text feature extraction, and so on — and selects a corresponding HuggingFace model class. The result is a PyTorch model object ready for export. Before committing to a full export you can verify that the loader resolved everything correctly with `winml inspect`. It prints the detected task, the HuggingFace model class, the export configuration, and the WinML inference class — all without downloading weights. Add `--hierarchy` to reconstruct the PyTorch module tree from random-weight tracing. @@ -16,9 +16,9 @@ Some community models host custom Python code in their repositories. The loader ## Exporting to ONNX -`winml export` converts the loaded model to ONNX. The conversion uses TorchScript tracing by default, which follows actual execution paths and tends to produce compact, inference-oriented graphs. A `--dynamo` flag exists for the PyTorch 2.x dynamo exporter; however, **Note:** the `--dynamo` flag is reserved for the PyTorch 2.x dynamo exporter but is **not yet functional** in the current release — passing it logs a warning and the flag is ignored. +`winml export` converts the loaded model to ONNX. By default it uses PyTorch's TorchDynamo ONNX exporter (`torch.onnx.export(dynamo=True)`), which records rich per-node module metadata that is used to derive the `winml.hierarchy.*` node tags. Pass `--no-dynamo` to fall back to the legacy TorchScript exporter, which follows actual execution paths and tends to produce compact inference-oriented graphs. Both exporters default to static shapes; the QNN-relevant difference is opset and op decomposition: torch's dynamo op library is built against a minimum opset of 18. When a lower opset is requested (17 by default), dynamo attempts to down-convert to it, which may succeed (the graph is saved at the requested opset) or fail (it stays at opset 18); `winml export` reports the opset actually produced and warns when it differs from the requested value. The TorchScript path always exports at the configured opset (17 by default) and lowers some ops differently (e.g. ResNet's head becomes `ReduceMean` + `Reshape` under dynamo but `GlobalAveragePool` + `Flatten` under TorchScript). Because the opset-17 TorchScript graph is what the QNN/NPU toolchain has been validated against, `--no-dynamo` remains the validated choice for QNN/NPU hand exports today. -By default the exporter runs an eight-step process that includes hierarchy tracing and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. +By default the exporter runs an eight-step process that includes hierarchy recovery and tag injection. The result is an ONNX file enriched with structural metadata that powers downstream features such as per-module benchmarking, inspector views, and optimizer scoping. ### Hierarchy tagging in detail @@ -31,7 +31,9 @@ During export the HTP (Hierarchy-preserving Tags Protocol) exporter attaches two #### How tags are built -The exporter registers PyTorch forward hooks on each module. When a module executes, a pre-hook pushes its class name onto a tag stack; the post-hook pops it. This produces hierarchical paths that mirror the PyTorch module tree: +The module path attached to each node is obtained differently depending on the exporter. Under the default TorchDynamo exporter, dynamo records originating-module information as ONNX metadata, and the exporter reconstructs the hierarchy directly from that metadata after export — no forward hooks are involved. Under `--no-dynamo`, the TorchScript path instead registers PyTorch forward hooks on each module: when a module executes, a pre-hook pushes its class name onto a tag stack and the post-hook pops it. Either way the result is hierarchical paths that mirror the PyTorch module tree. + +The hook flow below applies only to the `--no-dynamo` TorchScript path: ```mermaid flowchart LR @@ -42,11 +44,11 @@ flowchart LR E --> F[Tag stack → path] ``` -Only modules that are actually executed during tracing receive tags — unused modules are excluded. For example, `prajjwal1/bert-tiny` has 48 registered modules but only 18 are reached during a forward pass. +Under `--no-dynamo`, only modules that are actually executed during tracing receive tags — unused modules are excluded. For example, `prajjwal1/bert-tiny` has 48 registered modules but only 18 are reached during a forward pass. (Under the default dynamo path the hierarchy instead reflects the modules present in the exported graph's metadata.) #### Concrete example: BERT-tiny -Running `winml export -m prajjwal1/bert-tiny -o model.onnx -v` produces the following hierarchy tree (18 traced modules, 132 ONNX nodes, 100 % coverage): +Running `winml export -m prajjwal1/bert-tiny -o model.onnx -v --no-dynamo` produces the following hierarchy tree (18 traced modules, 132 ONNX nodes, 100 % coverage) — the numbers below are from the TorchScript trace path: ``` BertModel (132 nodes) @@ -75,7 +77,7 @@ Each ONNX node gets its tag from the module it belongs to. Here are a few exampl #### Node-to-module mapping -After the ONNX graph is produced by `torch.onnx.export`, a 4-priority system assigns each ONNX node to the closest matching module: +Under the default dynamo path, each node is assigned directly from its `pkg.torch.onnx.name_scopes` and `pkg.torch.onnx.class_hierarchy` metadata; nodes without usable metadata receive the model-root fallback. Under `--no-dynamo`, after the ONNX graph is produced by `torch.onnx.export`, the legacy tagger instead uses a 4-priority system to assign each ONNX node to the closest traced module: 1. **Direct match** (61 %) — the node's scope name maps exactly to a traced module. 2. **Parent match** (24 %) — walk up the scope hierarchy until a traced module is found. @@ -100,7 +102,7 @@ These I/O specs enable tools like `winml perf` to generate correct dummy inputs Alongside the `.onnx` file, the exporter writes a `*_htp_metadata.json` sidecar containing: - **`nodes`** — complete mapping of every ONNX node name → hierarchy tag -- **`modules`** — traced module information (class name, tag, execution order) +- **`modules`** — hierarchy module information (class name, tag, discovery order) - **`statistics`** — export time, node counts, coverage percentage - **`outputs`** — I/O tensor specifications diff --git a/docs/concepts/perf-and-monitoring.md b/docs/concepts/perf-and-monitoring.md index 235018e0f..3a30ae92a 100644 --- a/docs/concepts/perf-and-monitoring.md +++ b/docs/concepts/perf-and-monitoring.md @@ -139,12 +139,13 @@ In JSON output (`-f json`), these metrics appear under the `hw_monitor` key: "adapter_luid": null, "cpu": { "mean_pct": 15.8, "peak_pct": 16.71, "sample_count": 2 }, "ram": { "used_mb": 640.21, "peak_mb": 640.21 }, + "gpu": { "mean_pct": 0.0, "peak_pct": 0.0, "sample_count": 0, "luids": [] }, "device_memory": { "local_peak_mb": 0.0, "shared_peak_mb": 0.0 }, "running_time_ns": 0 } ``` -When a hardware accelerator is active, `device_kind` will be `"npu"` or `"gpu"`, and an additional key (e.g. `"npu"`) appears with device utilisation: +The `gpu` block is always present because GPU telemetry is collected independently of the selected inference adapter. When an accelerator is active, `device_kind` will be `"npu"` or `"gpu"` and the selected-adapter metrics appear under the stable `adapter` block. For compatibility, an additional dynamic block is only added when the selected kind does not already have a reserved aggregate block — today that means NPU selections still include `npu`, while GPU selections keep `gpu` reserved for aggregate telemetry. ```json "hw_monitor": { @@ -153,12 +154,16 @@ When a hardware accelerator is active, `device_kind` will be `"npu"` or `"gpu"`, "adapter_luid": "0x0000abcd12340000", "cpu": { "mean_pct": 12.1, "peak_pct": 34.5, "sample_count": 50 }, "ram": { "used_mb": 1842.0, "peak_mb": 1910.0 }, + "gpu": { "mean_pct": 3.2, "peak_pct": 9.7, "sample_count": 50, "luids": ["0x00001234abcd0000"] }, + "adapter": { "mean_pct": 87.3, "peak_pct": 100.0, "sample_count": 50 }, "device_memory": { "local_peak_mb": 245.0, "shared_peak_mb": 0.0 }, "npu": { "mean_pct": 87.3, "peak_pct": 100.0, "sample_count": 50 }, "running_time_ns": 4820000000 } ``` +When the selected device is a GPU, `adapter` contains the selected GPU's live metrics while `gpu` continues to report the aggregate GPU view (including `luids`). + This makes it straightforward to track memory consumption across model revisions or compare devices programmatically. ## Per-module benchmarking diff --git a/docs/naming-convention.md b/docs/naming-convention.md index f1cd3a9a5..d7f33d52b 100644 --- a/docs/naming-convention.md +++ b/docs/naming-convention.md @@ -11,7 +11,7 @@ Domain acronyms in PascalCase class names **retain their uppercase form**, excep | Acronym | Meaning | Class Casing | Example | |---------|---------|--------------|---------| | ONNX | Open Neural Network Exchange | `ONNX` | `ONNXStaticAnalyzer`, `ONNXLoader` | -| EP | Execution Provider | `EP` | `EPChecker`, `EPConfig`, `EPMonitor` | +| EP | Execution Provider | `EP` | `EPChecker`, `EPConfig`, `WinMLEPMonitor` | | QDQ | Quantize-Dequantize | `QDQ` | `QDQParameterConfig`, `QDQGenerator` | | QNN | Qualcomm Neural Network | `QNN` | `QNNMonitor` | | Op | Operator (2-letter prefix) | `Op` | `OpUnsupportedError` | diff --git a/docs/reference/index.md b/docs/reference/index.md index 086f12cf0..568ad9ede 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -54,7 +54,7 @@ stages based on the target device and precision. | `export_params` | `bool` | `true` | Include model parameters in ONNX. | | `do_constant_folding` | `bool` | `true` | Fold constants during export. | | `verbose` | `bool` | `false` | Verbose export logging. | -| `dynamo` | `bool` | `false` | Use PyTorch 2.x Dynamo exporter. | +| `dynamo` | `bool` | `true` | Use PyTorch's TorchDynamo ONNX exporter (default); set `false` for the legacy TorchScript exporter. | | `enable_hierarchy_tags` | `bool` | `true` | Add module hierarchy tags to ONNX nodes. | | `clean_onnx` | `bool` | `false` | Strip hierarchy tags after export. | | `hierarchy_tag_format` | `"full" \| "module_only"` | `"full"` | Tag detail level. | diff --git a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json new file mode 100644 index 000000000..6ac10fb42 --- /dev/null +++ b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp16_config.json @@ -0,0 +1,60 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_values", + "dtype": "float32", + "shape": [ + 1, + 1024, + 128 + ], + "value_range": [ + -1, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "audio-classification", + "model_class": "AutoModelForAudioClassification", + "model_type": "audio-spectrogram-transformer" + } +} diff --git a/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json new file mode 100644 index 000000000..439f9edee --- /dev/null +++ b/examples/recipes/MIT_ast-finetuned-audioset-10-10-0.4593/cpu/cpu/audio-classification_fp32_config.json @@ -0,0 +1,41 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_values", + "dtype": "float32", + "shape": [ + 1, + 1024, + 128 + ], + "value_range": [ + -1, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "audio-classification", + "model_class": "AutoModelForAudioClassification", + "model_type": "audio-spectrogram-transformer" + } +} diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json new file mode 100644 index 000000000..0fe7d1a21 --- /dev/null +++ b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp16_config.json @@ -0,0 +1,81 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 50265 + ] + }, + { + "name": "bbox", + "dtype": "int32", + "shape": [ + 1, + 512, + 4 + ], + "value_range": [ + 0, + 1000 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "start_logits" + }, + { + "name": "end_logits" + } + ] + }, + "optim": { + "clamp_constant_values": true + }, + "quant": null, + "loader": { + "task": "question-answering", + "model_class": "LayoutLMForQuestionAnswering", + "model_type": "layoutlm" + } +} diff --git a/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json new file mode 100644 index 000000000..0fe7d1a21 --- /dev/null +++ b/examples/recipes/impira_layoutlm-document-qa/cpu/cpu/question-answering_fp32_config.json @@ -0,0 +1,81 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 50265 + ] + }, + { + "name": "bbox", + "dtype": "int32", + "shape": [ + 1, + 512, + 4 + ], + "value_range": [ + 0, + 1000 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 2 + ] + }, + { + "name": "token_type_ids", + "dtype": "int32", + "shape": [ + 1, + 512 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "start_logits" + }, + { + "name": "end_logits" + } + ] + }, + "optim": { + "clamp_constant_values": true + }, + "quant": null, + "loader": { + "task": "question-answering", + "model_class": "LayoutLMForQuestionAnswering", + "model_type": "layoutlm" + } +} diff --git a/pyproject.toml b/pyproject.toml index 4ffda3c89..785ee1d2b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,8 +45,12 @@ dependencies = [ "numpy>=2.2,<2.3", "onnx==1.18", "onnx-ir>=0.1", - "onnxruntime-genai-winml==0.14.1", - "onnxruntime-windowsml==1.24.5.202604171637", + "onnxruntime-genai-winml==0.14.1 ; sys_platform == 'win32'", + # Range, not an exact pin: the exact build `1.24.5.202604171637` is the PyPI + # form of what WinML's wasdk metadata pins as `1.24.5.post202604171637` (PEP 440 + # normalized) — an exact `==` can't reconcile the two. A bounded range keeps the + # 1.24.x series while resolving to the published wheel. + "onnxruntime-windowsml>=1.24.5,<1.25 ; sys_platform == 'win32'", "onnxscript>=0.2", "opentelemetry-sdk>=1.39.1", "optimum>=2", @@ -75,9 +79,9 @@ dependencies = [ "torchmetrics[detection]>=1.0", "torchvision>=0.24", "tqdm>=4.67", - "transformers>=4.57", + "transformers>=5", "uvicorn[standard]>=0.42.0", - "windowsml==2.0.300", + "windowsml==2.0.300 ; sys_platform == 'win32'", ] optional-dependencies.dev = [ @@ -133,6 +137,19 @@ environments = [ "sys_platform == 'linux' and platform_machine == 'x86_64'", "sys_platform == 'win32' and platform_machine == 'AMD64'", ] +# Overrides: +# 1. Force transformers 5.x over optimum-onnx 0.1.0's `transformers<4.58.0` ceiling +# (it predates the transformers-5 API; the transformers_compat shim bridges the +# few 4.x internals optimum-onnx reaches for). +# 2. Neutralize onnxruntime-qnn's transitive standard `onnxruntime` requirement. +# onnxruntime-windowsml ships the same `onnxruntime/` module; installing plain +# onnxruntime alongside it overwrites windowsml's files on disk and kills the +# DirectML EP. The `sys_platform == 'unobtainium'` marker is never true, so uv +# never installs plain onnxruntime while still keeping onnxruntime-qnn. +override-dependencies = [ + "transformers>=5,<6", + "onnxruntime ; sys_platform == 'unobtainium'", +] [dependency-groups] dev = [ @@ -536,6 +553,13 @@ module = [ # Not a dependency of the main/CI environment, so it has no stubs here. "qairt", "qairt.*", + # WinUI3 / Windows App SDK bootstrap, WinRT projections, and pywin32 — Windows + # runtime packages imported in ep_path.py / ep_device.py. Installed and usable + # at runtime but ship no type stubs, so treat as untyped here. + "winui3", + "winui3.*", + "winrt.*", + "win32api", ] ignore_missing_imports = true diff --git a/src/winml/modelkit/__init__.py b/src/winml/modelkit/__init__.py index 463a60b7f..0419050c1 100644 --- a/src/winml/modelkit/__init__.py +++ b/src/winml/modelkit/__init__.py @@ -43,7 +43,6 @@ logging.getLogger(__name__).addHandler(logging.NullHandler()) - def _preload_bundled_onnxruntime_dll() -> None: # Windows ships C:\Windows\System32\onnxruntime.dll (older API version) # as part of the system WindowsML component. When WinML EP plugin DLLs @@ -74,7 +73,14 @@ def _preload_bundled_onnxruntime_dll() -> None: _preload_bundled_onnxruntime_dll() -from . import _warnings # Configure warning filters before importing subpackages +# _warnings configures filters before any subpackage imports. +# transformers_compat arms a sys.meta_path hook — the shim fires lazily +# the first time anything imports optimum.*; lightweight commands +# (``winml sys``, ``winml --help``) never pay the transformers cost. +from . import _warnings # noqa: I001 +from . import transformers_compat + +transformers_compat.arm() try: diff --git a/src/winml/modelkit/_warnings.py b/src/winml/modelkit/_warnings.py index cd19e5d08..aad0a6da9 100644 --- a/src/winml/modelkit/_warnings.py +++ b/src/winml/modelkit/_warnings.py @@ -20,6 +20,7 @@ import logging import os +import sys import warnings from ._env import env_flag_enabled @@ -91,12 +92,15 @@ def filter(self, record: logging.LogRecord) -> bool: for _cat in (FutureWarning, DeprecationWarning, UserWarning): warnings.filterwarnings("ignore", category=_cat, module=r"torch\..*") - # NOTE: TracerWarning (from torch.jit) is intentionally NOT filtered here. - # Importing torch.jit at startup would pull all of torch (~1.7s) into - # `winml --help` and violate the CLI import budget (tests/cli/test_import_time.py). - # During ONNX export, export_pytorch() wraps torch.onnx.export in - # `warnings.catch_warnings()` + `filterwarnings("ignore")`, which is strictly - # broader than a TracerWarning-only filter. + # TracerWarning (from torch.jit, inherits Warning not UserWarning) + # fires during ONNX export tracing — safe to suppress in both torch and + # transformers. Only register the filter if torch has ALREADY been + # imported; otherwise loading torch here would add ~2s to every + # lightweight command (``winml sys`` etc.) that never touches ONNX + # export. The export path re-triggers this by calling + # :func:`install_torch_tracer_filter` after loading torch. + if "torch" in sys.modules: + install_torch_tracer_filter() # Diffusers warnings.filterwarnings( @@ -140,5 +144,18 @@ def filter(self, record: logging.LogRecord) -> bool: logging.getLogger("transformers.modeling_utils").addFilter(_TransformersWeightsFilter()) +def install_torch_tracer_filter() -> None: + """Register the ``TracerWarning`` filter — call after ``import torch``. + + Idempotent — ``warnings.filterwarnings`` de-duplicates identical entries. + """ + try: + # torch.jit exposes TracerWarning at runtime but its stubs don't export it. + from torch.jit import TracerWarning # type: ignore[attr-defined] + except ImportError: + return # torch not installed + warnings.filterwarnings("ignore", category=TracerWarning) + + # Auto-configure on import _configure() diff --git a/src/winml/modelkit/analyze/analyzer.py b/src/winml/modelkit/analyze/analyzer.py index 890f98540..de1f50392 100644 --- a/src/winml/modelkit/analyze/analyzer.py +++ b/src/winml/modelkit/analyze/analyzer.py @@ -15,10 +15,10 @@ import time from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast from ..optim.config import WinMLOptimizationConfig -from ..utils.constants import EP_SUPPORTED_DEVICES, EPName, EPNameOrAlias, normalize_ep_name +from ..utils.constants import normalize_ep_name from .models.information import Information from .models.output import RuntimeDebugSummaryEntry from .models.support_level import SupportLevel @@ -30,6 +30,7 @@ import onnx + from ..utils.constants import EPName, EPNameOrAlias from .models.information import Action from .models.output import AnalysisOutput from .models.runtime_checks import PatternRuntime, RuntimeTestResult @@ -186,7 +187,7 @@ def __repr__(self) -> str: pattern_count = sum(self.output.metadata.detected_pattern_count.values()) return f"AnalysisResult(patterns={pattern_count})" - def is_fully_supported(self, ep: EPNameOrAlias | None = None) -> bool: + def is_fully_supported(self, ep: str | None = None) -> bool: """Check if model is fully supported on the target EP and device. Args: @@ -214,7 +215,7 @@ def is_fully_supported(self, ep: EPNameOrAlias | None = None) -> bool: return False # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None # Track if we found the target EP when filtering found_target = ep_normalized is None # True if not filtering @@ -227,7 +228,7 @@ def is_fully_supported(self, ep: EPNameOrAlias | None = None) -> bool: return False return found_target - def has_errors(self, ep: EPNameOrAlias | None = None) -> bool: + def has_errors(self, ep: str | None = None) -> bool: """Check if there are any unsupported patterns (blocking errors). Args: @@ -251,7 +252,7 @@ def has_errors(self, ep: EPNameOrAlias | None = None) -> bool: return False # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None for ep_support in self.output.results: if ep_normalized and ep_support.ep_type != ep_normalized: @@ -260,7 +261,7 @@ def has_errors(self, ep: EPNameOrAlias | None = None) -> bool: return True return False - def has_warnings(self, ep: EPNameOrAlias | None = None) -> bool: + def has_warnings(self, ep: str | None = None) -> bool: """Check if there are any partial patterns (warnings/optimization opportunities). Args: @@ -284,7 +285,7 @@ def has_warnings(self, ep: EPNameOrAlias | None = None) -> bool: return False # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None for ep_support in self.output.results: if ep_normalized and ep_support.ep_type != ep_normalized: @@ -293,7 +294,7 @@ def has_warnings(self, ep: EPNameOrAlias | None = None) -> bool: return True return False - def get_lint_result(self, ep: EPNameOrAlias | None = None) -> LintResult: + def get_lint_result(self, ep: str | None = None) -> LintResult: """Get lint-style result with error/warning/info counts. Args: @@ -332,7 +333,7 @@ def get_lint_result(self, ep: EPNameOrAlias | None = None) -> LintResult: ) # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None # Aggregate counts and lists error_patterns: list[str] = [] @@ -374,7 +375,7 @@ def get_lint_result(self, ep: EPNameOrAlias | None = None) -> LintResult: optimization_config=optimization_config, ) - def get_unsupported_operators(self, ep: EPNameOrAlias | None = None) -> list[str]: + def get_unsupported_operators(self, ep: str | None = None) -> list[str]: """Get list of unsupported operators for the target EP and device. Args: @@ -395,7 +396,7 @@ def get_unsupported_operators(self, ep: EPNameOrAlias | None = None) -> list[str ... print(f"Unsupported: {op_name}") """ # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None unsupported = [] for ep_support in self.output.results: @@ -409,7 +410,7 @@ def get_unsupported_operators(self, ep: EPNameOrAlias | None = None) -> list[str return unsupported - def get_optimization_opportunities(self, ep: EPNameOrAlias | None = None) -> list[Action]: + def get_optimization_opportunities(self, ep: str | None = None) -> list[Action]: """Get actions for patterns that could be optimized (UNSUPPORTED or PARTIAL status). Args: @@ -432,7 +433,7 @@ def get_optimization_opportunities(self, ep: EPNameOrAlias | None = None) -> lis ... print(f"Optimize: {action.pattern_from_id} -> {action.action}") """ # Normalize EP if specified - ep_normalized = normalize_ep_name(ep) if ep else None + ep_normalized = normalize_ep_name(cast("EPNameOrAlias", ep)) if ep else None actions: list[Action] = [] seen_patterns: set[tuple[str, str]] = set() @@ -452,7 +453,7 @@ def get_optimization_opportunities(self, ep: EPNameOrAlias | None = None) -> lis seen_patterns.add(pattern_key) return actions - def get_optimization_config(self, ep: EPNameOrAlias | None = None) -> WinMLOptimizationConfig: + def get_optimization_config(self, ep: str | None = None) -> WinMLOptimizationConfig: """Generate WinML optimization configuration based on action items. This method extracts optimization settings from action_items in Actions, @@ -594,7 +595,7 @@ def __init__(self, config: AnalyzerConfig | None = None) -> None: def analyze( self, model_path: str, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, enable_information: bool = True, htp_metadata_path: str | None = None, @@ -675,7 +676,7 @@ def analyze( total_start = time.perf_counter() # Normalize EP name (convert aliases to full names) - ep_normalized = normalize_ep_name(ep) + ep_normalized = normalize_ep_name(cast("EPNameOrAlias | None", ep)) if ep != ep_normalized: logger.debug("EP alias '%s' normalized to '%s'", ep, ep_normalized) @@ -728,7 +729,7 @@ def analyze( def analyze_from_proto( self, model_proto: onnx.ModelProto, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, enable_information: bool = True, model_path: str | None = None, @@ -785,7 +786,7 @@ def analyze_from_proto( # Normalize EP name (convert aliases to full names) total_start = time.perf_counter() - ep_normalized = normalize_ep_name(ep) + ep_normalized = normalize_ep_name(cast("EPNameOrAlias | None", ep)) if ep != ep_normalized: logger.debug("EP alias '%s' normalized to '%s'", ep, ep_normalized) @@ -793,10 +794,9 @@ def analyze_from_proto( # Resolve device — rule files are device-specific (CPU/GPU/NPU). if device is not None and device.lower() == "auto": - from ..sysinfo import resolve_device + from ..session import auto_detect_device - resolved, _ = resolve_device("auto", ep=ep_normalized) - device_to_use = resolved.upper() + device_to_use = auto_detect_device().upper() logger.info("Device 'auto' resolved to: %s", device_to_use) else: device_to_use = device if device is not None else "NPU" @@ -805,13 +805,18 @@ def analyze_from_proto( # Determine which EPs to analyze eps_to_analyze: list[EPName] = [] if ep_normalized is None: - # Analyze all EPs that support the target device - eps_to_analyze = [ - ep_name - for ep_name, supported_devices in EP_SUPPORTED_DEVICES.items() - if device_to_use.lower() in supported_devices - ] - logger.info("No EP specified, analyzing all supported EPs: %s", eps_to_analyze) + # Derive the EP list from the catalog so future EP additions + # are automatically included. sorted() gives deterministic order. + from ..session import eps_for_device + + # eps_for_device returns EP full names as ``str``; they are members of + # the ``EPName`` Literal by construction (catalog parity is test-enforced). + eps_to_analyze = cast("list[EPName]", sorted(eps_for_device(device_to_use.lower()))) + logger.info( + "No EP specified, analyzing all %s-capable EPs: %s", + device_to_use, + eps_to_analyze, + ) else: eps_to_analyze = [ep_normalized] @@ -964,7 +969,7 @@ def has_errors(self) -> bool: def analyze_onnx( model: str | Path, *, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, autoconf: bool = True, run_unknown_op: bool = False, diff --git a/src/winml/modelkit/analyze/core/model_validators/model_validator_manager.py b/src/winml/modelkit/analyze/core/model_validators/model_validator_manager.py index bcf8af563..f56ec80f1 100644 --- a/src/winml/modelkit/analyze/core/model_validators/model_validator_manager.py +++ b/src/winml/modelkit/analyze/core/model_validators/model_validator_manager.py @@ -24,7 +24,6 @@ if TYPE_CHECKING: - from ....utils.constants import EPName from ...models.information import Information from ...models.onnx_model import ONNXModel from ...models.runtime_checks import PatternRuntime @@ -78,8 +77,8 @@ def __init__( enabled_validators: list[str] | None = None, op_runtime_results: list[PatternRuntime] | None = None, *, - device: str, - ep: EPName, + device: str | None = None, + ep: str | None = None, ) -> None: """Initialize validator manager. @@ -100,7 +99,7 @@ def __init__( self.model = model self.model_proto = model.get_model() self.op_runtime_results = op_runtime_results or [] - self.device = device + self.device = device or "NPU" self.ep = ep self.enabled_validators = enabled_validators or list(self.VALIDATORS.keys()) diff --git a/src/winml/modelkit/analyze/core/runtime_checker_query.py b/src/winml/modelkit/analyze/core/runtime_checker_query.py index 18382f57d..454a67422 100644 --- a/src/winml/modelkit/analyze/core/runtime_checker_query.py +++ b/src/winml/modelkit/analyze/core/runtime_checker_query.py @@ -81,7 +81,6 @@ if TYPE_CHECKING: from winml.modelkit.pattern.match import PatternMatchResult - from ...utils.constants import EPName from .node_checkers.base import NodeChecker @@ -983,7 +982,7 @@ class RuntimeCheckerQuery: def __init__( self, model_proto: onnx.ModelProto, - ep_name: EPName, + ep_name: str, device_type: str, model_path: str | Path | None = None, dynamic_axis_strict_mode: bool = False, @@ -1146,28 +1145,44 @@ def _collect_qdq_types(self) -> None: def _is_ep_available_locally(self) -> bool: """Check if the target EP is available on the local machine. + Targeted probe through :meth:`WinMLEPRegistry.auto_device`: the + registry handles plugin discovery + DLL load lazily, so we no + longer need a separate module-level pre-register pass. Negative + outcomes (EP not discovered, registration failed, no matching + device) are caught and reported as "not available locally" + rather than propagated. + Returns: True if the EP+device combination is available locally. """ if self._ep_available_locally is not None: return self._ep_available_locally - from ... import winml - from ...utils.constants import DEVICE_TO_DEVICE_TYPE - - device_type_enum = DEVICE_TO_DEVICE_TYPE.get(self.device_type) - if device_type_enum is None: - self._ep_available_locally = False - return False + from ...session import ( + DeviceNotFound, + EPDeviceTarget, + WinMLEPNotDiscovered, + WinMLEPRegistrationFailed, + WinMLEPRegistry, + resolve_device, + short_ep_name, + ) try: - ep_devices = winml.get_registered_ep_devices() - self._ep_available_locally = any( - ep_dev.ep_name == self.ep_name and ep_dev.device.type == device_type_enum - for ep_dev in ep_devices + target = EPDeviceTarget( + ep=short_ep_name(self.ep_name), + device=self.device_type.lower(), ) - except Exception as e: - logger.debug("Failed to query EP devices: %s", e) + resolved = resolve_device(target) + WinMLEPRegistry.instance().auto_device(resolved) + self._ep_available_locally = True + except ( + WinMLEPNotDiscovered, + WinMLEPRegistrationFailed, + DeviceNotFound, + ValueError, + ) as e: + logger.debug("EP %s on %s not available locally: %s", self.ep_name, self.device_type, e) self._ep_available_locally = False return self._ep_available_locally @@ -1179,11 +1194,11 @@ def _get_ep_checker(self) -> EPChecker: EPChecker instance configured for the target EP+device. """ if self._ep_checker is None: - from ...utils.constants import DEVICE_TO_DEVICE_TYPE + from ...session import DEVICE_TO_DEVICE_TYPE self._ep_checker = EPChecker( ep_name=self.ep_name, - device_type=DEVICE_TO_DEVICE_TYPE[self.device_type], + device_type=DEVICE_TO_DEVICE_TYPE[self.device_type.lower()], ) return self._ep_checker diff --git a/src/winml/modelkit/analyze/models/ihv_type.py b/src/winml/modelkit/analyze/models/ihv_type.py index f517c948f..340ea6c2b 100644 --- a/src/winml/modelkit/analyze/models/ihv_type.py +++ b/src/winml/modelkit/analyze/models/ihv_type.py @@ -4,10 +4,10 @@ # -------------------------------------------------------------------------- """IHV type enum.""" -from enum import Enum +from enum import StrEnum -class IHVType(str, Enum): +class IHVType(StrEnum): """IHV (Independent Hardware Vendor) type.""" QC = "QC" diff --git a/src/winml/modelkit/analyze/models/information.py b/src/winml/modelkit/analyze/models/information.py index 481b9419d..309126829 100644 --- a/src/winml/modelkit/analyze/models/information.py +++ b/src/winml/modelkit/analyze/models/information.py @@ -7,7 +7,7 @@ Represents actionable information for fixing unsupported patterns. """ -from enum import Enum +from enum import StrEnum from uuid import uuid4 from pydantic import BaseModel, Field @@ -16,7 +16,7 @@ from .support_level import SupportLevel -class ActionLevel(str, Enum): +class ActionLevel(StrEnum): """Action priority level.""" REQUIRED = "required" diff --git a/src/winml/modelkit/analyze/models/onnx_model.py b/src/winml/modelkit/analyze/models/onnx_model.py index b5482cbbf..b418537be 100644 --- a/src/winml/modelkit/analyze/models/onnx_model.py +++ b/src/winml/modelkit/analyze/models/onnx_model.py @@ -7,7 +7,7 @@ from __future__ import annotations import logging -from enum import Enum +from enum import StrEnum import onnx from pydantic import BaseModel, Field, field_validator @@ -23,7 +23,7 @@ logger = logging.getLogger(__name__) -class ModelTag(str, Enum): +class ModelTag(StrEnum): """Tags for marking model-level issues and validation states. These tags are stored in ONNXModel.model_tags to record various diff --git a/src/winml/modelkit/analyze/models/runtime_checks.py b/src/winml/modelkit/analyze/models/runtime_checks.py index d396aedda..3329acdad 100644 --- a/src/winml/modelkit/analyze/models/runtime_checks.py +++ b/src/winml/modelkit/analyze/models/runtime_checks.py @@ -6,7 +6,7 @@ from __future__ import annotations -from enum import Enum +from enum import StrEnum from typing import Any from uuid import uuid4 @@ -43,14 +43,14 @@ class RuntimeDebugDetails(TypedDict): error_message: NotRequired[str] -class NodeTag(str, Enum): +class NodeTag(StrEnum): """Node tag enum for classifying nodes based on their properties.""" ALL_INPUTS_CONSTANT = "all_inputs_constant" MISSING_SHAPE_INFERENCE = "missing_shape_inference" -class AlternativeType(str, Enum): +class AlternativeType(StrEnum): """Alternative pattern relationship type enum.""" EQUIVALENT = "equivalent" diff --git a/src/winml/modelkit/analyze/models/support_level.py b/src/winml/modelkit/analyze/models/support_level.py index 2ba30af17..f4104787f 100644 --- a/src/winml/modelkit/analyze/models/support_level.py +++ b/src/winml/modelkit/analyze/models/support_level.py @@ -4,10 +4,10 @@ # -------------------------------------------------------------------------- """Support level classification enum.""" -from enum import Enum +from enum import StrEnum -class SupportLevel(str, Enum): +class SupportLevel(StrEnum): """Support level classification.""" SUPPORTED = "supported" diff --git a/src/winml/modelkit/analyze/pattern/check_patterns.py b/src/winml/modelkit/analyze/pattern/check_patterns.py index ebcb64bf9..020fcc34d 100644 --- a/src/winml/modelkit/analyze/pattern/check_patterns.py +++ b/src/winml/modelkit/analyze/pattern/check_patterns.py @@ -18,33 +18,31 @@ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast -from ... import winml from ...onnx import ONNXDomain -from ...pattern.base import ( - PatternInputGenerator, - get_pattern_input_generator, - get_registered_pattern_input_generators, -) if TYPE_CHECKING: import argparse - from collections.abc import Callable - - import onnxruntime as ort from ...utils.constants import EPName -from ...utils import constants + from ..runtime_checker.ep_checker import EPChecker +from ...pattern.base import ( + PatternInputGenerator, + get_pattern_input_generator, + get_registered_pattern_input_generators, +) +from ...session import DEVICE_TYPE_TO_DEVICE from ..core.rules_prefilter import RuntimeCheckerRulesPrefilter -from ..runtime_checker.ep_checker import EPChecker +from ..runtime_checker.check_ops import ( + OpenVINONPUChecker, + QNNNPUChecker, + get_ep_checker, +) from ..utils import CheckResultWriter, load_case_indices_from_conflict_file -winml.register_execution_providers(ort=True) - - def check_patterns( ep_checker: EPChecker, patterns: list[str], @@ -112,7 +110,7 @@ def check_patterns( ep_checker.set_rules_prefilter( RuntimeCheckerRulesPrefilter( - ep_name=ep_checker.ep_name, + ep_name=cast("EPName", ep_checker.ep_name), device_type=ep_checker.device_type.name, ) ) @@ -149,7 +147,7 @@ def check_patterns( opset_suffix = f"_{first_domain.value}_opset{first_version}" # Prepare output file - device = constants.DEVICE_TYPE_TO_DEVICE[ep_checker.device_type] + device = DEVICE_TYPE_TO_DEVICE[ep_checker.device_type].upper() output_filename = f"{pattern_name}_{ep_checker.ep_name}_{device}{opset_suffix}.json" output_path = output_dir / output_filename @@ -223,65 +221,14 @@ def check_patterns( return all_results -# don't use EPChecker directly as there is a bug with pytest in subprocess -class OpenVINONPUChecker(EPChecker): - """OpenVINO NPU execution provider checker wrapper for pytest compatibility.""" - - def __init__(self, device_type: ort.OrtHardwareDeviceType) -> None: - """Initialize OpenVINO NPU checker.""" - super().__init__(ep_name="OpenVINOExecutionProvider", device_type=device_type) - - -# don't use EPChecker directly as there is a bug with pytest in subprocess -class QNNNPUChecker(EPChecker): - """QNN NPU execution provider checker wrapper for pytest compatibility.""" - - def __init__(self, device_type: ort.OrtHardwareDeviceType) -> None: - """Initialize QNN NPU checker.""" - super().__init__(ep_name="QNNExecutionProvider", device_type=device_type) - - -class RTXChecker(EPChecker): - """NVIDIA TensorRT RTX execution provider checker wrapper for pytest compatibility.""" - - def __init__(self, device_type: ort.OrtHardwareDeviceType) -> None: - """Initialize RTX checker.""" - if device_type != constants.DEVICE_TO_DEVICE_TYPE["GPU"]: - raise ValueError("NvTensorRTRTXExecutionProvider only supports GPU device type") - super().__init__( - ep_name="NvTensorRTRTXExecutionProvider", - device_type=constants.DEVICE_TO_DEVICE_TYPE["GPU"], - ) - - -def get_ep_checker(ep_name: EPName, device: str) -> EPChecker: - """Get EPChecker for given execution provider name. - - Args: - ep_name: Execution provider name (e.g., "QNNExecutionProvider") - device: Target device type (CPU, GPU, NPU) - - Returns: - EPChecker corresponding to the execution provider. - - Raises: - ValueError: If the execution provider name is not supported. - """ - device_type = constants.DEVICE_TO_DEVICE_TYPE[device] - ep_name_to_checker: dict[str, Callable[..., EPChecker]] = { - "QNNExecutionProvider": QNNNPUChecker, - "OpenVINOExecutionProvider": OpenVINONPUChecker, - "NvTensorRTRTXExecutionProvider": RTXChecker, - # Add other EPChecker subclasses here as needed - } - if ep_name not in ep_name_to_checker: - raise ValueError( - f"Unsupported execution provider: {ep_name}. " - f"Available: QNNExecutionProvider, " - f"OpenVINOExecutionProvider, " - f"NvTensorRTRTXExecutionProvider" - ) - return ep_name_to_checker[ep_name](device_type=device_type) +# EPCheckers and get_ep_checker are re-exported from +# ..runtime_checker.check_ops to keep a single source of truth. +__all__ = [ + "OpenVINONPUChecker", + "QNNNPUChecker", + "check_patterns", + "get_ep_checker", +] def build_parser() -> argparse.ArgumentParser: @@ -314,6 +261,11 @@ def build_parser() -> argparse.ArgumentParser: "--ep", type=str, required=True, + # CARVE-OUT: This subprocess tool intentionally supports only a curated allowlist + # of EPs. VitisAI and other unvalidated NPU EPs are excluded because this + # pattern-checking tool has not been validated against them. Do NOT derive from + # eps_for_device(...) or EP_DEVICE_SPECS — this is an explicit opt-in allowlist, + # not catalog drift. choices=[ "QNNExecutionProvider", "OpenVINOExecutionProvider", diff --git a/src/winml/modelkit/analyze/runtime_checker/check_ops.py b/src/winml/modelkit/analyze/runtime_checker/check_ops.py index ea66a1d7e..f0384cd0f 100644 --- a/src/winml/modelkit/analyze/runtime_checker/check_ops.py +++ b/src/winml/modelkit/analyze/runtime_checker/check_ops.py @@ -23,7 +23,6 @@ import onnxruntime as ort from onnx.defs import SchemaError -from ... import winml from ...onnx import ONNXDomain from ...pattern.op_input_gen import ( OpInputGenerator, @@ -34,20 +33,13 @@ if TYPE_CHECKING: import argparse - - from ...utils.constants import EPName from ...pattern.op_input_gen.qdq_gen import QDQGenerator -from ...utils import constants +from ...session import DEVICE_TO_DEVICE_TYPE, DEVICE_TYPE_TO_DEVICE from ..utils import CheckResultWriter, load_case_indices_from_conflict_file from ..utils.model_utils import get_op_since_version from .ep_checker import EPChecker -# Register WinML EPs at module level before any ORT session is created. -# This must stay at the top of the file so EPs are available for all downstream usage. -winml.register_execution_providers(ort=True) - - def check_ops( ep_checker: EPChecker, ops: list[str], @@ -146,7 +138,7 @@ def check_ops( # Prepare output file since_version = get_op_since_version(op_name, current_opset_version, opset_domain) - device = constants.DEVICE_TYPE_TO_DEVICE[ep_checker.device_type] + device = DEVICE_TYPE_TO_DEVICE[ep_checker.device_type].upper() qdq_suffix = "_qdq" if use_qdq else "" output_filename = ( f"{op_name}_{ep_checker.ep_name}_{device}" @@ -279,7 +271,7 @@ def __init__(self, device_type: ort.OrtHardwareDeviceType) -> None: ) -def get_ep_checker(ep_name: EPName, device: str) -> EPChecker: +def get_ep_checker(ep_name: str, device: str) -> EPChecker: """Get EPChecker for given execution provider name. Args: @@ -291,7 +283,7 @@ def get_ep_checker(ep_name: EPName, device: str) -> EPChecker: Raises: ValueError: If the execution provider name is not supported. """ - device_type = constants.DEVICE_TO_DEVICE_TYPE[device] + device_type = DEVICE_TO_DEVICE_TYPE[device.lower()] ep_name_to_checker = { "QNNExecutionProvider": QNNNPUChecker, "OpenVINOExecutionProvider": OpenVINONPUChecker, diff --git a/src/winml/modelkit/analyze/runtime_checker/ep_checker.py b/src/winml/modelkit/analyze/runtime_checker/ep_checker.py index 14ed48a4c..94ab82d42 100644 --- a/src/winml/modelkit/analyze/runtime_checker/ep_checker.py +++ b/src/winml/modelkit/analyze/runtime_checker/ep_checker.py @@ -11,15 +11,11 @@ import onnx import onnxruntime as ort -from ... import winml - if TYPE_CHECKING: from collections.abc import Sequence from os import PathLike - from ...utils.constants import EPName - # TODO: allow test case iter to take dtypes as inputs # TODO: define dataclass for result @@ -45,34 +41,58 @@ class EPChecker: # EPs that require a file path (not in-memory bytes) for compilation. # VitisAI EP fails with "ep.context_file_path and model_path are both empty" # when given in-memory model bytes. - EPS_REQUIRING_FILE_PATH: ClassVar[set[EPName]] = {"VitisAIExecutionProvider"} + EPS_REQUIRING_FILE_PATH: ClassVar[set[str]] = {"VitisAIExecutionProvider"} # EP/device combinations that are known to leak resources/state across many # sequential checks inside a single worker process. Running each case in an # isolated process avoids "first case passes, later cases fail" behavior. EPS_REQUIRING_CASE_ISOLATION_BY_DEVICE: ClassVar[ - dict[EPName, set[ort.OrtHardwareDeviceType]] + dict[str, set[ort.OrtHardwareDeviceType]] ] = { "OpenVINOExecutionProvider": {ort.OrtHardwareDeviceType.NPU}, } def __init__( self, - ep_name: EPName, + ep_name: str, device_type: ort.OrtHardwareDeviceType, provider_options: Sequence[dict[Any, Any]] | None = None, rules_prefilter: _RulesPrefilterProtocol | None = None, ) -> None: self.device_type = device_type - self.ep_name: EPName = ep_name + self.ep_name: str = ep_name self._provider_options = provider_options self._rules_prefilter = rules_prefilter def _get_sess_options(self) -> ort.SessionOptions: - winml.register_execution_providers(ort=True) + from ...session import ( + EPDeviceTarget, + WinMLEPRegistry, + resolve_device, + short_ep_name, + ) + sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL - winml.add_ep_for_device(sess_options, self.ep_name, self.device_type) + + # self.device_type is ort.OrtHardwareDeviceType (CPU/GPU/NPU enum). + # self.ep_name is the full EP name (e.g. "QNNExecutionProvider"). + target = EPDeviceTarget( + ep=short_ep_name(self.ep_name), + device=self.device_type.name.lower(), + ) + resolved = resolve_device(target) + ep_device = WinMLEPRegistry.instance().auto_device(resolved) + + options: dict[str, str] = {} + if self._provider_options: + # _provider_options is Sequence[dict[Any, Any]] | None; take the first. + options = dict(self._provider_options[0]) + + sess_options.add_provider_for_devices( + [ep_device.device.ort_handle], + options, + ) return sess_options def _needs_file_path(self) -> bool: diff --git a/src/winml/modelkit/analyze/utils/__init__.py b/src/winml/modelkit/analyze/utils/__init__.py index 05b780f95..3d34ff729 100644 --- a/src/winml/modelkit/analyze/utils/__init__.py +++ b/src/winml/modelkit/analyze/utils/__init__.py @@ -13,7 +13,6 @@ from .json_utils import validate_json_schema from .model_utils import encode_rule_condition_value_for_parquet from .op_utils import CheckResultWriter, load_case_indices_from_conflict_file -from .pattern_matching import match_pattern_with_wildcards from .rule_loader import ( RuleLoader, get_runtime_rules_search_dirs, @@ -31,7 +30,6 @@ "has_rule_data_for_ep", "infer_ihv_from_ep_name", "load_case_indices_from_conflict_file", - "match_pattern_with_wildcards", "resolve_rule_parquet_path", "validate_json_schema", ] diff --git a/src/winml/modelkit/analyze/utils/ep_utils.py b/src/winml/modelkit/analyze/utils/ep_utils.py index e4ab29864..0a553ebed 100644 --- a/src/winml/modelkit/analyze/utils/ep_utils.py +++ b/src/winml/modelkit/analyze/utils/ep_utils.py @@ -13,24 +13,23 @@ if TYPE_CHECKING: from pathlib import Path - from ...utils.constants import EPName, EPNameOrAlias from ..models.ihv_type import IHVType logger = logging.getLogger(__name__) -def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: +def infer_ihv_from_ep_name(ep_name: str) -> IHVType: """Infer IHVType from an Execution Provider name or alias. - Accepts either a canonical ``EPName`` or a shorthand ``EPAlias`` (e.g. + Accepts either a canonical EP name or a shorthand alias (e.g. ``"openvino"``); aliases are normalized to their canonical name before the exact lookup, which covers every member of the canonical set. Names that are neither a known EP nor a known alias resolve to ``IHVType.UNKNOWN`` rather than raising, so callers can treat inference as total. Args: - ep_name: Execution Provider name or alias (see ``utils.constants``). + ep_name: Execution Provider name or alias. Returns: IHVType: Inferred IHV type (QC, INTEL, AMD, NVIDIA, MICROSOFT, or @@ -53,7 +52,7 @@ def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: from ...utils.constants import normalize_ep_name from ..models.ihv_type import IHVType - ep_name_to_ihv: dict[EPName, IHVType] = { + ep_name_to_ihv: dict[str, IHVType] = { "QNNExecutionProvider": IHVType.QC, "OpenVINOExecutionProvider": IHVType.INTEL, "VitisAIExecutionProvider": IHVType.AMD, @@ -65,15 +64,15 @@ def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: } canonical = normalize_ep_name(ep_name) - return ep_name_to_ihv.get(canonical, IHVType.UNKNOWN) # type: ignore[arg-type] + return ep_name_to_ihv.get(canonical, IHVType.UNKNOWN) -def get_devices_with_rule_data(ep_name: EPName) -> list[str]: +def get_devices_with_rule_data(ep_name: str) -> list[str]: """Return all devices supported by an EP. First probes runtime-rule directories for parquet artifacts for each ``EP + device`` pair. If no rule data is found, falls - back to the EP→device mapping from :func:`sysinfo.get_ep_device_map`. + back to the EP→device mapping in :data:`session.EP_DEVICE_SPECS`. Args: ep_name: Full execution provider name (e.g., ``"QNNExecutionProvider"``). @@ -82,10 +81,10 @@ def get_devices_with_rule_data(ep_name: EPName) -> list[str]: List of device strings (e.g., ``["NPU", "GPU"]``), empty if the EP is completely unknown. """ - from ...sysinfo.device import get_ep_device_map + from ...session import EP_DEVICE_SPECS # Priority order: NPU > GPU > CPU (first match used as default device) - known_devices = {d.upper() for v in get_ep_device_map().values() for d in v.split("/") if d} + known_devices = {spec.device.upper() for spec in EP_DEVICE_SPECS} priority = ["NPU", "GPU", "CPU"] probe_order = [d for d in priority if d in known_devices] # Append any devices not in the priority list @@ -94,9 +93,9 @@ def get_devices_with_rule_data(ep_name: EPName) -> list[str]: devices = [d for d in probe_order if has_rule_data_for_ep(ep_name, d)] if devices: return devices - # Fallback: derive from the authoritative EP→device mapping - device_str = get_ep_device_map().get(ep_name, "") - return [d.upper() for d in device_str.split("/") if d] + # Fallback: derive from the authoritative EP→device catalog. Preserve + # catalog order (NPU rows precede GPU/CPU rows for each vendor EP). + return [spec.device.upper() for spec in EP_DEVICE_SPECS if spec.ep == ep_name] def has_any_rule_data() -> bool: @@ -119,7 +118,7 @@ def has_any_rule_data() -> bool: return False -def has_rule_data_for_ep(ep_name: EPName, device: str) -> bool: +def has_rule_data_for_ep(ep_name: str, device: str) -> bool: """Check whether runtime check rule data exists for a given EP and device. Probes runtime-rule search directories for provider subdirectory layout only: @@ -137,7 +136,7 @@ def has_rule_data_for_ep(ep_name: EPName, device: str) -> bool: """ from .rule_loader import get_runtime_rules_search_dirs - def _has_parquet_in_search_dir(search_dir: Path, ep: EPName, device_upper: str) -> bool: + def _has_parquet_in_search_dir(search_dir: Path, ep: str, device_upper: str) -> bool: provider_dir = search_dir / f"{ep}_{device_upper}" return provider_dir.is_dir() and any(provider_dir.glob("*.parquet")) diff --git a/src/winml/modelkit/analyze/utils/pattern_matching.py b/src/winml/modelkit/analyze/utils/pattern_matching.py deleted file mode 100644 index b06b1fb57..000000000 --- a/src/winml/modelkit/analyze/utils/pattern_matching.py +++ /dev/null @@ -1,118 +0,0 @@ -# ------------------------------------------------------------------------- -# Copyright (c) Microsoft Corporation. All rights reserved. -# Licensed under the MIT License. -# -------------------------------------------------------------------------- -"""Pattern matching with wildcard support.""" - -from typing import Any - - -def match_pattern_with_wildcards(pattern: dict[str, Any], attributes: dict[str, Any]) -> bool: - """Match detected pattern attributes against rule with wildcard support. - - Implements universal wildcard matching where "*" in pattern matches any value. - - Args: - pattern: Pattern/rule attributes (may contain "*" wildcards) - attributes: Actual attributes from detected pattern - - Returns: - True if pattern matches attributes, False otherwise - - Examples: - >>> match_pattern_with_wildcards( - ... {"kernel_shape": [3, 3]}, {"kernel_shape": [3, 3]} - ... ) - True - >>> match_pattern_with_wildcards( - ... {"kernel_shape": "*"}, {"kernel_shape": [3, 3]} - ... ) - True - >>> match_pattern_with_wildcards( - ... {"kernel_shape": [3, 3]}, {"kernel_shape": [5, 5]} - ... ) - False - """ - # Iterate through each attribute in pattern - for attr_name, expected_value in pattern.items(): - # Wildcard matches any value - if expected_value == "*": - continue - - # Get actual value from detected pattern - actual_value = attributes.get(attr_name) - - # If attribute missing or doesn't match, fail - if actual_value != expected_value: - return False - - return True - - -def match_type_vars_with_wildcards(pattern: dict[str, str], types: dict[str, str]) -> bool: - """Match detected type variables against pattern with wildcard support. - - Supports: - - Exact match: "float32" matches only "float32" - - Wildcard: "*" matches any type - - Alternatives: "float32|float16" matches either "float32" or "float16" - - Args: - pattern: Pattern type constraints (may contain "*" wildcards or "|" - alternatives) - types: Actual data types from detected pattern (e.g., {"T": "float32"}) - - Returns: - True if types match pattern, False otherwise - - Examples: - >>> match_type_vars_with_wildcards({"T": "float32"}, {"T": "float32"}) - True - >>> match_type_vars_with_wildcards({"T": "*"}, {"T": "float32"}) - True - >>> match_type_vars_with_wildcards({"T": "float32|float16"}, {"T": "float16"}) - True - >>> match_type_vars_with_wildcards({"T": "float32"}, {"T": "int8"}) - False - """ - for type_var, expected_type in pattern.items(): - # Wildcard matches any type - if expected_type == "*": - continue - - actual_type = types.get(type_var) - - # Check if expected_type contains alternatives (pipe-separated) - if "|" in expected_type: - allowed_types = [t.strip() for t in expected_type.split("|")] - if actual_type not in allowed_types: - return False - else: - # Exact match - if actual_type != expected_type: - return False - - return True - - -def match_version_with_wildcard(actual_version: str, rule_version: str) -> bool: - """Match version string with wildcard support. - - Args: - actual_version: Actual version string - rule_version: Rule version ("*" matches any) - - Returns: - True if versions match, False otherwise - - Examples: - >>> match_version_with_wildcard("2.3.1", "2.3.1") - True - >>> match_version_with_wildcard("2.3.1", "*") - True - >>> match_version_with_wildcard("2.3.1", "2.3.0") - False - """ - if rule_version == "*": - return True - return actual_version == rule_version diff --git a/src/winml/modelkit/build/common.py b/src/winml/modelkit/build/common.py index 7e221d2ee..99835a4ee 100644 --- a/src/winml/modelkit/build/common.py +++ b/src/winml/modelkit/build/common.py @@ -4,30 +4,235 @@ # -------------------------------------------------------------------------- """Shared build pipeline utilities. -Provides the optimize-analyze loop reused by both build_hf_model() and -build_onnx_model(). +Provides the optimize-analyze loop and the stage runner +(optimize -> quantize -> compile -> finalize) reused by both +:func:`build_hf_model` and :func:`build_onnx_model`. """ from __future__ import annotations +import json import logging import tempfile import time +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any from ..analyze import analyze_onnx +from ..compiler import compile_onnx from ..onnx import copy_onnx_model, is_quantized_onnx from ..optim import optimize_onnx +from ..quant import quantize_onnx if TYPE_CHECKING: from ..config import WinMLBuildConfig - from ..utils.constants import EPNameOrAlias logger = logging.getLogger(__name__) +@dataclass +class StagesResult: + """Outcome of :func:`run_build_stages`. + + Bundles the fields both build entry points need to persist into their + build manifest and return via :class:`BuildResult`. + """ + + current_path: Path + is_pre_quantized: bool + stages_completed: list[str] = field(default_factory=list) + stages_skipped: list[str] = field(default_factory=list) + stage_timings: dict[str, float] = field(default_factory=dict) + analyze_iterations: int = 0 + analyze_unsupported_nodes: int = 0 + analyze_details: dict = field(default_factory=dict) + quant_result: Any = None + + +def run_build_stages( + *, + current_path: Path, + optimized_path: Path, + quantized_path: Path, + compiled_path: Path, + final_path: Path, + config: WinMLBuildConfig, + config_path: Path, + ep: str | None = None, + device: str | None = None, + hack_max_optim_iterations: int = 3, + skip_optimize: bool = False, + allow_unsupported_nodes: bool = False, + analyze_result_path: Path | None = None, + onnx_kwargs: dict[str, Any] | None = None, +) -> StagesResult: + """Run the shared build stages: optimize -> quantize -> compile -> finalize. + + Extracted from :func:`build_hf_model` and :func:`build_onnx_model` per + the FIXMEs in ``build/hf.py`` and ``build/onnx.py``. Behaviour is + identical to the previous inline blocks. + + Args: + current_path: Path to the ONNX model at the start of the stages + (post-export for HF, post-copy for ONNX). Mutated in place as + the pipeline advances. + optimized_path: Destination path for the optimize stage. + quantized_path: Destination path for the quantize stage. + compiled_path: Destination path for the compile stage. + final_path: Destination for the finalize stage (``model.onnx``). + config: Full build config; ``config.optim`` / ``config.quant`` / + ``config.compile`` gate the individual stages. + config_path: Where the resolved config is persisted after + optimize (autoconf may have expanded ``config.optim``). + ep, device: Passed through to :func:`run_optimize_analyze_loop`. + hack_max_optim_iterations: Max analyzer autoconf rounds. + skip_optimize: When True, bypasses optimize and only runs analyze + (used for pre-quantized inputs). + onnx_kwargs: ONNX-level kwargs forwarded to optimize/quantize. + """ + onnx_kwargs = onnx_kwargs or {} + result = StagesResult( + current_path=current_path, + is_pre_quantized=is_quantized_onnx(current_path) or skip_optimize, + ) + + # ========================================================================= + # OPTIMIZE + ANALYZE (or ANALYZE-ONLY for pre-quantized) + # ========================================================================= + if result.is_pre_quantized: + logger.info( + "Pre-quantized model detected (QDQ nodes present). " + "Skipping optimize + quantize, running analyze-only." + ) + result.stages_skipped.append("optimize") + ( + result.current_path, + _, + result.analyze_iterations, + result.analyze_unsupported_nodes, + result.analyze_details, + ) = run_optimize_analyze_loop( + model_path=result.current_path, + optimized_path=optimized_path, + config=config, + ep=ep, + device=device, + max_optim_iterations=hack_max_optim_iterations, + allow_unsupported_nodes=allow_unsupported_nodes, + skip_optimize=True, + analyze_output_path=analyze_result_path, + **onnx_kwargs, + ) + else: + logger.info("Optimizing ONNX model...") + ( + result.current_path, + opt_elapsed, + result.analyze_iterations, + result.analyze_unsupported_nodes, + result.analyze_details, + ) = run_optimize_analyze_loop( + model_path=result.current_path, + optimized_path=optimized_path, + config=config, + ep=ep, + device=device, + max_optim_iterations=hack_max_optim_iterations, + allow_unsupported_nodes=allow_unsupported_nodes, + analyze_output_path=analyze_result_path, + **onnx_kwargs, + ) + result.stage_timings["optimize"] = opt_elapsed + result.stages_completed.append("optimize") + logger.info("Optimize done (%.1fs) -> %s", opt_elapsed, optimized_path) + + # Persist config AFTER autoconf — includes discovered optimization flags + config_path.write_text(json.dumps(config.to_dict(), indent=2)) + logger.debug("Config persisted: %s", config_path) + + # ========================================================================= + # QUANTIZE (optional — config.quant=None means skip) + # ========================================================================= + if result.is_pre_quantized: + if "quantize" not in result.stages_skipped: + result.stages_skipped.append("quantize") + logger.info("Quantize skipped (pre-quantized model)") + elif config.quant is not None: + # Defensive fallback: catches the edge case where a direct caller + # provides config.quant != None but the model already has QDQ nodes. + if is_quantized_onnx(result.current_path): + logger.warning( + "Model already contains QDQ nodes, skipping quantization. " + "Set config.quant=None to silence this warning." + ) + result.stages_skipped.append("quantize") + else: + logger.info("Quantizing model...") + t0 = time.monotonic() + result.quant_result = quantize_onnx( + model_path=result.current_path, + output_path=quantized_path, + config=config.quant, + **onnx_kwargs, + ) + if not result.quant_result.success: + errors = ( + ", ".join(result.quant_result.errors) + if result.quant_result.errors + else "Unknown" + ) + raise RuntimeError(f"Quantization failed: {errors}") + result.current_path = quantized_path + result.stage_timings["quantize"] = time.monotonic() - t0 + result.stages_completed.append("quantize") + logger.info( + "Quantize done (%.1fs) -> %s", + result.stage_timings["quantize"], + quantized_path, + ) + else: + result.stages_skipped.append("quantize") + logger.info("Quantize skipped (config.quant is None)") + + # ========================================================================= + # COMPILE (optional — config.compile=None means skip) + # ========================================================================= + if config.compile is not None: + logger.info("Compiling model...") + t0 = time.monotonic() + compile_result = compile_onnx( + model_path=result.current_path, + output_path=compiled_path, + config=config.compile, + ) + if hasattr(compile_result, "success") and not compile_result.success: + errors = ", ".join(compile_result.errors) if compile_result.errors else "Unknown" + raise RuntimeError(f"Compilation failed: {errors}") + if compile_result.output_path and Path(compile_result.output_path) != compiled_path: + copy_onnx_model(compile_result.output_path, compiled_path) + if compiled_path.exists(): + result.current_path = compiled_path + result.stage_timings["compile"] = time.monotonic() - t0 + result.stages_completed.append("compile") + logger.info( + "Compile done (%.1fs) -> %s", result.stage_timings["compile"], result.current_path + ) + else: + result.stages_skipped.append("compile") + logger.info("Compile skipped (config.compile is None)") + + # ========================================================================= + # FINALIZE — Copy last stage output as model.onnx + # ========================================================================= + if result.current_path != final_path: + copy_onnx_model(result.current_path, final_path) + result.current_path = final_path + + return result + + def ensure_pre_quantized_stamped( config: WinMLBuildConfig, onnx_path: Path, *, force: bool = False ) -> None: @@ -72,7 +277,7 @@ def run_optimize_analyze_loop( optimized_path: Path, config: WinMLBuildConfig, *, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, max_optim_iterations: int = 0, allow_unsupported_nodes: bool = False, @@ -184,7 +389,7 @@ def run_optimize_analyze_loop( def _run_analyze_loop( *, optimized_path: Path, - ep: EPNameOrAlias | None, + ep: str | None, device: str | None, max_optim_iterations: int, config: WinMLBuildConfig, diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 888c65df8..ed253ef3d 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -20,26 +20,21 @@ from __future__ import annotations import datetime -import json import logging import time from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any -from ..compiler import compile_onnx from ..export import export_onnx -from ..onnx import copy_onnx_model -from ..quant import quantize_onnx from ..utils import MANIFEST_FILENAME, ManifestStage, WinMLManifest -from .common import ensure_pre_quantized_stamped, run_optimize_analyze_loop +from .common import run_build_stages if TYPE_CHECKING: import torch.nn as nn from ..config import WinMLBuildConfig - from ..utils.constants import EPNameOrAlias logger = logging.getLogger(__name__) @@ -90,9 +85,10 @@ def build_hf_model( trust_remote_code: bool = False, random_init: bool = False, cache_key: str | None = None, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, model_type: str | None = None, + hf_config: Any | None = None, **kwargs: Any, ) -> BuildResult: """Build an ONNX model from a HuggingFace model architecture. @@ -176,7 +172,6 @@ def _name(base: str) -> str: final_path = output_dir / _name("model.onnx") config_path = output_dir / _name("winml_build_config.json") manifest_path = output_dir / _name(MANIFEST_FILENAME) - analyze_result_path = output_dir / _name("analyze_result.json") # Check for existing artifact (skip build if present and not rebuilding) if final_path.exists() and not rebuild: @@ -213,6 +208,7 @@ def _name(base: str) -> str: model_id, trust_remote_code, random_init=random_init, + hf_config=hf_config, model_type=model_type, ) @@ -233,136 +229,38 @@ def _name(base: str) -> str: verbose=False, **onnx_kwargs, ) - current_path = export_path stage_timings["export"] = time.monotonic() - t0 stages_completed.append("export") logger.info("Export done (%.1fs) -> %s", stage_timings["export"], export_path) # ========================================================================= - # [3] OPTIMIZE — ONNX graph optimization + autoconf loop - # FIXME: Stages [3]-[6] (optimize, quantize, compile, finalize) are - # duplicated between build_hf_model() and build_onnx_model(). Extract - # into a shared run_build_stages() function in common.py. - # ========================================================================= - # Single defensive detection on the freshly exported ONNX. No-op when - # the caller already stamped ``config.skip_optimize``. HF export rarely - # produces a pre-quantized ONNX (Optimum exports float weights), but a - # direct caller could plausibly hand a pre-quantized - # ``pytorch_model`` and reach this branch. - skip_optimize_kwarg: bool = kwargs.pop("skip_optimize", False) - ensure_pre_quantized_stamped(config, current_path, force=skip_optimize_kwarg) - is_pre_quantized = config.skip_optimize - - if is_pre_quantized: - logger.info("Skipping optimize + quantize stages (config.skip_optimize=True)") - stages_skipped.append("optimize") - # Skip the ORT-based graph optimization (no kernel for QOperator - # ops like ConvInteger on the host EP). The autoconf re-optim/ - # analyze loop is disabled too -- ``run_optimize_analyze_loop`` - # forces ``max_optim_iterations=0`` when ``skip_optimize=True``, - # so ``_run_analyze_loop`` is not invoked. The model still flows - # through later stages (quantize-skip + compile) for validation. - current_path, _, analyze_iterations, analyze_unsupported_nodes, analyze_details = ( - run_optimize_analyze_loop( - model_path=current_path, - optimized_path=optimized_path, - config=config, - ep=ep, - device=device, - skip_optimize=True, - **onnx_kwargs, - ) - ) - else: - logger.info("Optimizing ONNX model...") - ( - current_path, - opt_elapsed, - analyze_iterations, - analyze_unsupported_nodes, - analyze_details, - ) = run_optimize_analyze_loop( - model_path=current_path, - optimized_path=optimized_path, - config=config, - ep=ep, - device=device, - max_optim_iterations=hack_max_optim_iterations, - allow_unsupported_nodes=allow_unsupported_nodes, - analyze_output_path=analyze_result_path, - **onnx_kwargs, - ) - stage_timings["optimize"] = opt_elapsed - stages_completed.append("optimize") - logger.info("Optimize done (%.1fs) -> %s", opt_elapsed, optimized_path) - logger.info("Portable ONNX ready: %s", current_path) - - # Persist config AFTER autoconf — includes discovered optimization flags - config_path.write_text(json.dumps(config.to_dict(), indent=2)) - logger.debug("Config persisted: %s", config_path) - - # ========================================================================= - # [4] QUANTIZE (optional — config.quant=None means skip) + # [3]-[6] OPTIMIZE -> QUANTIZE -> COMPILE -> FINALIZE + # Shared with build_onnx_model via ``common.run_build_stages``. # ========================================================================= - # No defensive ``is_quantized_onnx`` re-check here: when the model is - # pre-quantized, ``ensure_pre_quantized_stamped`` has already set - # ``config.quant = None`` at stage [3], so this branch naturally - # falls through to the ``quant is None`` skip path. - quant_result = None - if is_pre_quantized: - if "quantize" not in stages_skipped: - stages_skipped.append("quantize") - logger.info("Quantize skipped (pre-quantized model)") - elif config.quant is not None: - logger.info("Quantizing model...") - t0 = time.monotonic() - quant_result = quantize_onnx( - model_path=current_path, - output_path=quantized_path, - config=config.quant, - **onnx_kwargs, - ) - if not quant_result.success: - errors = ", ".join(quant_result.errors) if quant_result.errors else "Unknown" - raise RuntimeError(f"Quantization failed: {errors}") - current_path = quantized_path - stage_timings["quantize"] = time.monotonic() - t0 - stages_completed.append("quantize") - logger.info("Quantize done (%.1fs) -> %s", stage_timings["quantize"], quantized_path) - else: - stages_skipped.append("quantize") - logger.info("Quantize skipped (config.quant is None)") - - # ========================================================================= - # [5] COMPILE (optional — config.compile=None means skip) - # ========================================================================= - if config.compile is not None: - logger.info("Compiling model...") - t0 = time.monotonic() - compile_result = compile_onnx( - model_path=current_path, - output_path=compiled_path, - config=config.compile, - ) - if hasattr(compile_result, "success") and not compile_result.success: - errors = ", ".join(compile_result.errors) if compile_result.errors else "Unknown" - raise RuntimeError(f"Compilation failed: {errors}") - if compile_result.output_path and Path(compile_result.output_path) != compiled_path: - copy_onnx_model(compile_result.output_path, compiled_path) - if compiled_path.exists(): - current_path = compiled_path - stage_timings["compile"] = time.monotonic() - t0 - stages_completed.append("compile") - logger.info("Compile done (%.1fs) -> %s", stage_timings["compile"], current_path) - else: - stages_skipped.append("compile") - logger.info("Compile skipped (config.compile is None)") - - # ========================================================================= - # [6] FINALIZE — Copy last stage output as model.onnx - # ========================================================================= - if current_path != final_path: - copy_onnx_model(current_path, final_path) + skip_optimize: bool = kwargs.pop("skip_optimize", False) + stages = run_build_stages( + current_path=export_path, + optimized_path=optimized_path, + quantized_path=quantized_path, + compiled_path=compiled_path, + final_path=final_path, + config=config, + config_path=config_path, + ep=ep, + device=device, + hack_max_optim_iterations=hack_max_optim_iterations, + skip_optimize=skip_optimize or config.skip_optimize, + allow_unsupported_nodes=allow_unsupported_nodes, + analyze_result_path=output_dir / _name("analyze_result.json"), + onnx_kwargs=onnx_kwargs, + ) + stages_completed.extend(stages.stages_completed) + stages_skipped.extend(stages.stages_skipped) + stage_timings.update(stages.stage_timings) + analyze_iterations = stages.analyze_iterations + analyze_unsupported_nodes = stages.analyze_unsupported_nodes + analyze_details = stages.analyze_details + quant_result = stages.quant_result elapsed = time.monotonic() - start_time logger.info("Build complete in %.1fs -> %s", elapsed, final_path) diff --git a/src/winml/modelkit/build/onnx.py b/src/winml/modelkit/build/onnx.py index 8e0095e5e..f21a7b53c 100644 --- a/src/winml/modelkit/build/onnx.py +++ b/src/winml/modelkit/build/onnx.py @@ -14,23 +14,19 @@ from __future__ import annotations import datetime -import json import logging import time from pathlib import Path from typing import TYPE_CHECKING, Any -from ..compiler import compile_onnx from ..onnx import copy_onnx_model -from ..quant import quantize_onnx from ..utils import MANIFEST_FILENAME, ManifestStage, WinMLManifest -from .common import ensure_pre_quantized_stamped, run_optimize_analyze_loop +from .common import run_build_stages from .hf import BuildResult if TYPE_CHECKING: from ..config import WinMLBuildConfig - from ..utils.constants import EPNameOrAlias logger = logging.getLogger(__name__) @@ -41,7 +37,7 @@ def build_onnx_model( config: WinMLBuildConfig, output_dir: Path | str, rebuild: bool = False, - ep: EPNameOrAlias | None = None, + ep: str | None = None, device: str | None = None, cache_key: str | None = None, **kwargs: Any, @@ -118,7 +114,6 @@ def _name(base: str) -> str: final_path = output_dir / _name("model.onnx") config_path = output_dir / _name("winml_build_config.json") manifest_path = output_dir / _name(MANIFEST_FILENAME) - analyze_result_path = output_dir / _name("analyze_result.json") # Check for existing artifact (skip build if present and not rebuilding) if final_path.exists() and not rebuild: @@ -152,124 +147,33 @@ def _name(base: str) -> str: copy_onnx_model(onnx_path, current_path) # ========================================================================= - # [1] OPTIMIZE + ANALYZE (or SKIP-BOTH for pre-quantized) - # FIXME: Stages [1]-[4] (optimize, quantize, compile, finalize) are - # duplicated between build_onnx_model() and build_hf_model(). Extract - # into a shared run_build_stages() function in common.py. + # [1]-[4] OPTIMIZE -> QUANTIZE -> COMPILE -> FINALIZE + # Shared with build_hf_model via ``common.run_build_stages``. # ========================================================================= - # Single defensive detection. No-op when the CLI path (via - # ``generate_onnx_build_config``) already stamped ``config.skip_optimize``. - # Direct callers who hand-built a config trigger the one detection here. - skip_optimize_kwarg: bool = kwargs.pop("skip_optimize", False) - ensure_pre_quantized_stamped(config, current_path, force=skip_optimize_kwarg) - is_pre_quantized = config.skip_optimize - - if is_pre_quantized: - logger.info("Skipping optimize + quantize stages (config.skip_optimize=True)") - stages_skipped.append("optimize") - # Skip the ORT-based graph optimization (no kernel for QOperator - # ops like ConvInteger on the host EP). The autoconf re-optim/ - # analyze loop is disabled too -- ``run_optimize_analyze_loop`` - # forces ``max_optim_iterations=0`` when ``skip_optimize=True``, - # so ``_run_analyze_loop`` is not invoked. The model still flows - # through later stages (quantize-skip + compile) for validation. - current_path, _, analyze_iters, analyze_unsupported, analyze_details = ( - run_optimize_analyze_loop( - model_path=current_path, - optimized_path=optimized_path, - config=config, - ep=ep, - device=device, - skip_optimize=True, - **onnx_kwargs, - ) - ) - else: - logger.info("Optimizing ONNX model...") - current_path, opt_elapsed, analyze_iters, analyze_unsupported, analyze_details = ( - run_optimize_analyze_loop( - model_path=current_path, - optimized_path=optimized_path, - config=config, - ep=ep, - device=device, - max_optim_iterations=hack_max_optim_iterations, - allow_unsupported_nodes=allow_unsupported_nodes, - analyze_output_path=analyze_result_path, - **onnx_kwargs, - ) - ) - stage_timings["optimize"] = opt_elapsed - stages_completed.append("optimize") - logger.info("Optimize done (%.1fs) -> %s", opt_elapsed, optimized_path) - - # Persist config AFTER autoconf — includes discovered optimization flags - config_path.write_text(json.dumps(config.to_dict(), indent=2)) - logger.debug("Config persisted: %s", config_path) - - # ========================================================================= - # [2] QUANTIZE (optional — config.quant=None means skip) - # ========================================================================= - # No defensive ``is_quantized_onnx`` re-check here: when the model is - # pre-quantized, ``ensure_pre_quantized_stamped`` has already set - # ``config.quant = None`` at stage [1], so this branch naturally - # falls through to the ``quant is None`` skip path. - quant_result = None - if is_pre_quantized: - # Already handled above -- skip quantize for pre-quantized models - if "quantize" not in stages_skipped: - stages_skipped.append("quantize") - logger.info("Quantize skipped (pre-quantized model)") - elif config.quant is not None: - logger.info("Quantizing model...") - t0 = time.monotonic() - quant_result = quantize_onnx( - model_path=current_path, - output_path=quantized_path, - config=config.quant, - **onnx_kwargs, - ) - if not quant_result.success: - errors = ", ".join(quant_result.errors) if quant_result.errors else "Unknown" - raise RuntimeError(f"Quantization failed: {errors}") - current_path = quantized_path - stage_timings["quantize"] = time.monotonic() - t0 - stages_completed.append("quantize") - logger.info("Quantize done (%.1fs) -> %s", stage_timings["quantize"], quantized_path) - else: - stages_skipped.append("quantize") - logger.info("Quantize skipped (config.quant is None)") - - # ========================================================================= - # [3] COMPILE (optional — config.compile=None means skip) - # ========================================================================= - if config.compile is not None: - logger.info("Compiling model...") - t0 = time.monotonic() - compile_result = compile_onnx( - model_path=current_path, - output_path=compiled_path, - config=config.compile, - ) - if hasattr(compile_result, "success") and not compile_result.success: - errors = ", ".join(compile_result.errors) if compile_result.errors else "Unknown" - raise RuntimeError(f"Compilation failed: {errors}") - if compile_result.output_path and Path(compile_result.output_path) != compiled_path: - copy_onnx_model(compile_result.output_path, compiled_path) - if compiled_path.exists(): - current_path = compiled_path - stage_timings["compile"] = time.monotonic() - t0 - stages_completed.append("compile") - logger.info("Compile done (%.1fs) -> %s", stage_timings["compile"], current_path) - else: - stages_skipped.append("compile") - logger.info("Compile skipped (config.compile is None)") - - # ========================================================================= - # [4] FINALIZE — Copy last stage output as model.onnx - # ========================================================================= - if current_path != final_path: - copy_onnx_model(current_path, final_path) + skip_optimize: bool = kwargs.pop("skip_optimize", False) + stages = run_build_stages( + current_path=current_path, + optimized_path=optimized_path, + quantized_path=quantized_path, + compiled_path=compiled_path, + final_path=final_path, + config=config, + config_path=config_path, + ep=ep, + device=device, + hack_max_optim_iterations=hack_max_optim_iterations, + skip_optimize=skip_optimize or config.skip_optimize, + allow_unsupported_nodes=allow_unsupported_nodes, + analyze_result_path=output_dir / _name("analyze_result.json"), + onnx_kwargs=onnx_kwargs, + ) + stages_completed.extend(stages.stages_completed) + stages_skipped.extend(stages.stages_skipped) + stage_timings.update(stages.stage_timings) + analyze_iters = stages.analyze_iterations + analyze_unsupported = stages.analyze_unsupported_nodes + analyze_details = stages.analyze_details + quant_result = stages.quant_result elapsed = time.monotonic() - start_time logger.info("Build complete in %.1fs -> %s", elapsed, final_path) diff --git a/src/winml/modelkit/commands/_ep_arg.py b/src/winml/modelkit/commands/_ep_arg.py new file mode 100644 index 000000000..d5309c393 --- /dev/null +++ b/src/winml/modelkit/commands/_ep_arg.py @@ -0,0 +1,136 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Module-level helper for the ``--ep [@]`` CLI syntax. + +Splits the raw click argument into ``(ep, source)`` and validates the +source-tag at parse time so malformed input surfaces at the CLI layer +instead of further down the chain inside ``EPDeviceTarget.__post_init__``. + +Per design doc ``docs/design/session/2_coreloop.md`` §6.2 Scenarios A.5/A.6. +""" + +from __future__ import annotations + +from typing import Any + +import click + +from ..session import VALID_SOURCE_TAGS + + +def split_ep_at_source(value: str) -> tuple[str, str | None]: + """Split ``"openvino@pypi"`` into ``("openvino", "pypi")``. + + Without ``@`` returns ``(value, None)`` (unqualified ``--ep openvino`` + form). The source tag is normalized to lowercase before being matched + against :data:`VALID_SOURCE_TAGS`. The EP name is returned verbatim + (case preserved) so full names like ``"OpenVINOExecutionProvider"`` + survive intact for :func:`expand_ep_name` (which lowercases + short-name lookups itself) and :class:`EPDeviceTarget`'s + case-sensitive full-name match. + + Raises: + ValueError: On whitespace, multiple ``@``, empty ep or source + substring, or unknown source tag. + """ + if any(c.isspace() for c in value): + raise ValueError(f"Invalid --ep value {value!r}: whitespace is not allowed") + + if value.count("@") > 1: + raise ValueError(f"Invalid --ep value {value!r}: expected at most one '@'") + + if "@" not in value: + return value, None + + ep, source = value.split("@", 1) + if not ep or not source: + raise ValueError( + f"Invalid --ep value {value!r}: expected '@' " + f"with non-empty ep and source" + ) + + source = source.lower() + if source not in VALID_SOURCE_TAGS: + raise ValueError( + f"Unknown source tag {source!r}; expected one of {sorted(VALID_SOURCE_TAGS)}" + ) + + return ep, source + + +class EpAtSourceParamType(click.ParamType): + """Click ParamType wrapping :func:`split_ep_at_source` for ``--ep``. + + Converts ``"openvino@pypi"`` into ``("openvino", "pypi")`` at click + parse time. Empty / unset values pass through as ``None`` (click's + standard "option not provided" shape). + + Failures from :func:`split_ep_at_source` surface as + :class:`click.UsageError` via :meth:`click.ParamType.fail`, so each + CLI command stops re-wrapping ``ValueError`` in its own try/except + block. + + Commands wire this up as + ``@click.option("--ep", type=EpAtSourceParamType())`` and receive + the option value pre-split as a ``tuple[str, str | None] | None``. + """ + + name = "ep_at_source" + + def convert( + self, + value: Any, + param: click.Parameter | None, + ctx: click.Context | None, + ) -> tuple[str, str | None] | None: + if value is None or value == "": + return None + # Idempotency: click may invoke convert() twice (e.g. when the + # value came from a callback that already returned the parsed + # tuple). Return pre-split tuples unchanged. + if isinstance(value, tuple): + return value + try: + return split_ep_at_source(value) + except ValueError as e: + return self.fail(str(e), param, ctx) + + +def _reject_ep_source( + ep: tuple[str, str | None] | None, + command_name: str, +) -> str | None: + """Reject the ``--ep @`` form at the CLI boundary. + + Used by commands that route ``--ep`` through + :class:`EpAtSourceParamType` but whose downstream pipeline does not + yet honor the source-tag (build, config). Collapses the verbatim + try/except block that those commands previously duplicated. + + Args: + ep: The pre-split value from :class:`EpAtSourceParamType` — + ``None`` when ``--ep`` was not given, otherwise + ``(ep_short_name, source_tag_or_None)``. + command_name: User-visible command string for the error message + (e.g. ``"winml build"``, ``"winml config"``). + + Returns: + The bare ``ep`` short name (``str``) when given, or ``None`` + when ``--ep`` was not supplied. + + Raises: + click.UsageError: when ``ep`` carries a non-``None`` source tag + (e.g. ``--ep openvino@pypi``). + """ + if ep is None: + return None + ep_part, ep_source = ep + if ep_source is not None: + raise click.UsageError( + f"`{command_name}` does not yet support source pinning " + f"(got --ep {ep_part}@{ep_source!r}); " + f"use --ep {ep_part!r} without '@'." + ) + return ep_part diff --git a/src/winml/modelkit/commands/_live_chart.py b/src/winml/modelkit/commands/_live_chart.py index e75628195..b88f8faa4 100644 --- a/src/winml/modelkit/commands/_live_chart.py +++ b/src/winml/modelkit/commands/_live_chart.py @@ -4,27 +4,51 @@ # -------------------------------------------------------------------------- """Live hardware monitor display for performance benchmarking. -Renders a live NPU/CPU utilization chart during benchmarking using +Renders a live adapter/CPU utilization chart during benchmarking using plotext for chart rendering and Rich Live for terminal refresh. """ from __future__ import annotations -from typing import Any +from typing import Any, Final from rich.console import Console from rich.panel import Panel from ..session.monitor.hw_monitor import adapter_label +from ..utils.constants import ACCELERATOR_DEVICE_TYPES # Moving window size for the x-axis (seconds) -_CHART_WINDOW_SECONDS = 10.0 +_CHART_WINDOW_SECONDS = 15.0 # Display refresh rate (frames per second) _REFRESH_FPS = 5 +class _OmittedDeviceKind: + """Sentinel for callers that omitted the resolved adapter kind.""" + + +_DEVICE_KIND_OMITTED: Final = _OmittedDeviceKind() + + +def _avg_now( + samples: list[float] | None, + fallback_now: float = 0.0, +) -> tuple[float, float]: + """Return ``(avg, now)`` for a samples list. + + ``fallback_now`` is used for the ``now`` value when ``samples`` is empty + or ``None`` (e.g. when a caller has a scalar current reading but no + time-series to compute an average from — in that case ``avg`` mirrors + the scalar so the display stays honest rather than reading 0.0). + """ + if not samples: + return (fallback_now, fallback_now) + return (sum(samples) / len(samples), samples[-1]) + + class LiveMonitorDisplay: """Renders a live hardware utilization chart during benchmarking. @@ -37,10 +61,10 @@ def __init__( warmup: int, model_id: str, device: str, - chart_width: int = 80, + chart_width: int = 120, chart_height: int = 15, poll_interval_ms: int = 100, - device_kind: str | None = None, + device_kind: str | None | _OmittedDeviceKind = _DEVICE_KIND_OMITTED, ) -> None: self._total = total_iterations self._warmup = warmup @@ -50,9 +74,9 @@ def __init__( # in when you want the legend to reflect what's actually polled (e.g. # "auto" that resolved to GPU). Falls back to the requested string # when the caller doesn't know the resolved kind yet. - if device_kind is None: + if isinstance(device_kind, _OmittedDeviceKind): requested = (device or "").lower() - device_kind = requested if requested in ("npu", "gpu") else None + device_kind = requested if requested in ACCELERATOR_DEVICE_TYPES else None # When no adapter is polled (CPU-only / auto resolved to nothing), # hide the adapter line + status cell entirely instead of drawing # a flat zero series labelled "Adapter". @@ -65,6 +89,28 @@ def __init__( # Track the last rendered panel for transient=False final display self._last_panel: Any = None + def _resolved_device_label(self) -> str: + """Return the display label for the requested or resolved device.""" + return self._adapter_label if self._show_adapter else self._device + + def _uses_gpu_role_labels(self) -> bool: + """Whether selected-adapter and aggregate GPU labels would collide.""" + return self._show_adapter and self._adapter_label == adapter_label("gpu") + + def _selected_adapter_label(self) -> str: + """Return the selected-adapter label for the chart/status display.""" + label = self._adapter_label + if self._uses_gpu_role_labels(): + return f"{label} (selected)" + return label + + def _aggregate_gpu_label(self) -> str: + """Return the aggregate GPU label for the chart/status display.""" + label = adapter_label("gpu") + if self._uses_gpu_role_labels(): + return f"{label} (aggregate)" + return label + def __enter__(self) -> LiveMonitorDisplay: from rich.live import Live @@ -90,13 +136,15 @@ def update( cpu_pct: float = 0.0, ram_mb: float = 0.0, cpu_samples: list[float] | None = None, + gpu_samples: list[float] | None = None, + gpu_pct: float = 0.0, ) -> None: """Update the live display with current metrics.""" if self._live is None: return try: - chart_renderable = self._render_chart(util_samples, cpu_samples) + chart_renderable = self._render_chart(util_samples, cpu_samples, gpu_samples) status_line = self._render_status( iteration, latency_ms, @@ -105,6 +153,9 @@ def update( memory_shared_mb, cpu_pct, ram_mb, + gpu_pct=gpu_pct, + cpu_samples=cpu_samples, + gpu_samples=gpu_samples, ) from rich.console import Group @@ -121,16 +172,21 @@ def update( pass # Don't let display errors interrupt the benchmark def _render_chart( - self, util_samples: list[float], cpu_samples: list[float] | None = None + self, + util_samples: list[float], + cpu_samples: list[float] | None = None, + gpu_samples: list[float] | None = None, ) -> Any: """Render utilization chart as a Rich renderable. Uses plotext with AnsiDecoder for flicker-free Rich Live integration. - Plots adapter (NPU/GPU, green) and CPU (cyan) with distinct colors. + Plots the selected adapter (green), CPU (cyan), and aggregate GPU + telemetry (yellow) with distinct colors. X-axis is a moving window of the last N seconds. Y-axis has fixed ticks: 0, 20, 40, 60, 80, 100. """ - adapter = self._adapter_label + adapter = self._selected_adapter_label() + gpu_label = self._aggregate_gpu_label() show_adapter = self._show_adapter try: import plotext as plt @@ -157,16 +213,16 @@ def _render_chart( # Compute moving window: keep last N seconds of samples window_samples = int(_CHART_WINDOW_SECONDS / self._poll_interval_s) - total_npu = len(util_samples) if util_samples else 0 + total_adapter = len(util_samples) if util_samples else 0 # Plot the adapter line only when an adapter is actually being polled. if show_adapter: - npu_window = util_samples[-window_samples:] if util_samples else [0] - window_start_idx = max(0, total_npu - len(npu_window)) - npu_times = [ - (window_start_idx + i) * self._poll_interval_s for i in range(len(npu_window)) + adapter_window = util_samples[-window_samples:] if util_samples else [0] + window_start_idx = max(0, total_adapter - len(adapter_window)) + adapter_times = [ + (window_start_idx + i) * self._poll_interval_s for i in range(len(adapter_window)) ] - plt.plot(npu_times, npu_window, marker="braille", color="green") + plt.plot(adapter_times, adapter_window, marker="braille", color="green") # Plot CPU in cyan (distinct from adapter) has_cpu = False @@ -180,6 +236,22 @@ def _render_chart( ] plt.plot(cpu_times, cpu_window, marker="braille", color="cyan") + # Plot GPU in yellow (distinct from NPU green and CPU cyan) + has_gpu = False + if gpu_samples: + has_gpu = True + total_gpu = len(gpu_samples) + gpu_window = gpu_samples[-window_samples:] + gpu_start_idx = max(0, total_gpu - len(gpu_window)) + gpu_times = [ + (gpu_start_idx + i) * self._poll_interval_s for i in range(len(gpu_window)) + ] + # plotext's palette exposes 'orange+' (ANSI bright yellow, code 11) + # but has no 'yellow' key — `color="yellow"` silently falls through + # to default (white). `orange+` matches Rich's `[bright_yellow]` + # legend swatch below. + plt.plot(gpu_times, gpu_window, marker="braille", color="orange+") + # No plotext title -- we render our own Rich-colored title with legend plt.ylabel("Usage %") @@ -189,7 +261,7 @@ def _render_chart( # X-axis: absolute elapsed time, sliding window. Use whichever series # we have to anchor the timeline so a CPU-only chart still scrolls. - sample_count = total_npu if show_adapter else total_cpu + sample_count = total_adapter if show_adapter else total_cpu elapsed = sample_count * self._poll_interval_s x_min = max(0.0, elapsed - _CHART_WINDOW_SECONDS) x_max = max(elapsed, _CHART_WINDOW_SECONDS) @@ -201,16 +273,15 @@ def _render_chart( from rich.console import Group from rich.text import Text - # Rich-colored title line with legend swatches. - if show_adapter and has_cpu: - title = Text.from_markup( - f" Utilization ([green]\u2588\u2588[/green] {adapter} % " - f"[cyan]\u2588\u2588[/cyan] CPU %)" - ) - elif show_adapter: - title = Text.from_markup(f" Utilization ([green]\u2588\u2588[/green] {adapter} %)") - else: - title = Text.from_markup(" Utilization ([cyan]\u2588\u2588[/cyan] CPU %)") + # Rich-colored title line with legend swatches + legend_parts = [] + if show_adapter: + legend_parts.append(f"[green]\u2588\u2588[/green] {adapter} %") + if has_cpu: + legend_parts.append("[cyan]\u2588\u2588[/cyan] CPU %") + if has_gpu: + legend_parts.append(f"[bright_yellow]\u2588\u2588[/bright_yellow] {gpu_label} %") + title = Text.from_markup(f" Utilization ({' '.join(legend_parts)})") ansi_output = plt.build() chart_lines = [Text.from_ansi(line) for line in ansi_output.splitlines()] @@ -225,15 +296,25 @@ def _render_status( memory_shared_mb: float = 0.0, cpu_pct: float = 0.0, ram_mb: float = 0.0, + gpu_pct: float = 0.0, + cpu_samples: list[float] | None = None, + gpu_samples: list[float] | None = None, ) -> str: - """Render 3-row status below the chart.""" + """Render 4-row status below the chart. + + Row 1: progress bar + phase counter + device label. + Row 2: compute utilization (adapter / CPU / GPU) — unified ``now%/avg%``. + Row 3: memory (Sys Mem + Device Mem local/shared). + Row 4: inference latency + throughput. + + CPU and GPU accept a samples list to compute ``avg`` — the ``cpu_pct`` + and ``gpu_pct`` scalars remain as the ``now`` value (and as fallbacks + for ``avg`` when no samples were supplied). + """ phase = "warmup" if iteration <= self._warmup else "benchmark" effective_iter = iteration - self._warmup if phase == "benchmark" else iteration total_bench = self._total - self._warmup - current_util = util_samples[-1] if util_samples else 0.0 - mean_util = sum(util_samples) / len(util_samples) if util_samples else 0.0 - pct = iteration / self._total if self._total > 0 else 0 bar_len = int(pct * 20) bar = f"[{'=' * bar_len}{' ' * (20 - bar_len)}]" @@ -245,28 +326,36 @@ def _render_status( throughput = 1000.0 / latency_ms if latency_ms > 0 else 0.0 + adapter_avg, adapter_now = _avg_now(util_samples) + cpu_avg, cpu_now = _avg_now(cpu_samples, fallback_now=cpu_pct) + gpu_avg, gpu_now = _avg_now(gpu_samples, fallback_now=gpu_pct) + # Row 1: Progress pct_cell = f"{bar} {pct:.0%}" - row1 = f" {pct_cell:<30}| {progress} | Device: {self._device}" - - # Row 2: Hardware (pad each cell to fixed width, spaces before divider). - # CPU-only mode drops the adapter cell + device-memory cell since we - # have no live values to populate them with. - cpu_cell = f"CPU: {cpu_pct:.1f}%" - ram_cell = f"RAM: {ram_mb:.0f} MB" + row1 = f" {pct_cell:<30}| {progress} | Device: {self._resolved_device_label()}" + + # Row 2: Compute (unified now/avg format across all three) + adapter_label_text = self._selected_adapter_label() + gpu_label_text = self._aggregate_gpu_label() + adapter_cell = f"{adapter_label_text}: {adapter_now:.1f}%/{adapter_avg:.1f}%" + cpu_cell = f"CPU: {cpu_now:.1f}%/{cpu_avg:.1f}%" + gpu_cell = f"{gpu_label_text}: {gpu_now:.1f}%/{gpu_avg:.1f}%" + row2_cells = [cpu_cell, gpu_cell] if self._show_adapter: - adapter_cell = f"{self._adapter_label}: {mean_util:.1f}% avg ({current_util:.1f}% now)" - mem_cell = f"VRAM: {memory_local_mb:.0f}/{memory_shared_mb:.0f} MB (local/shared)" - row2 = f" {adapter_cell:<30}| {cpu_cell:<12}| {ram_cell} | {mem_cell}" - else: - row2 = f" {cpu_cell:<12}| {ram_cell}" + row2_cells.insert(0, adapter_cell) + row2 = " " + "| ".join(f"{cell:<28}" for cell in row2_cells) + + # Row 3: Memory + ram_cell = f"Sys Mem: {ram_mb:.0f} MB" + mem_cell = f"Device Mem: {memory_local_mb:.0f}/{memory_shared_mb:.0f} MB (local/shared)" + row3 = f" {ram_cell:<24}| {mem_cell}" - # Row 3: Inference (pad each cell to fixed width, spaces before divider) + # Row 4: Inference lat_cell = f"Latency: {latency_ms:.2f} ms" thr_cell = f"Throughput: ~{throughput:.0f} smp/s" - row3 = f" {lat_cell:<24}| {thr_cell}" + row4 = f" {lat_cell:<24}| {thr_cell}" - return f"{row1}\n{row2}\n{row3}" + return f"{row1}\n{row2}\n{row3}\n{row4}" def print_final_snapshot( self, diff --git a/src/winml/modelkit/commands/_perf_genai.py b/src/winml/modelkit/commands/_perf_genai.py index 6ae0ede18..bf9e102d2 100644 --- a/src/winml/modelkit/commands/_perf_genai.py +++ b/src/winml/modelkit/commands/_perf_genai.py @@ -30,7 +30,7 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, cast import click @@ -40,9 +40,9 @@ GenaiSession, GenaiSessionError, GenerationConfig, + short_ep_name, ) from ..utils.constants import ( - EP_NAME_TO_ALIAS, EP_SUPPORTED_DEVICES, EPNameOrAlias, normalize_ep_name, @@ -52,6 +52,8 @@ if TYPE_CHECKING: from rich.console import Console + from ..utils.constants import EPName + logger = logging.getLogger(__name__) @@ -94,11 +96,11 @@ def resolve_genai_ep(device: str) -> EPNameOrAlias | None: return None # Function-local import mirrors the ONNX path (perf.py) and avoids a - # module-level cycle: ``sysinfo`` pulls in heavier device-probing deps. - from ..sysinfo import resolve_device, resolve_eps + # module-level cycle. + from ..session import EPDeviceTarget, available_eps_for_device, resolve_device - resolved_device, _ = resolve_device(device=device, ep=None) - eps = resolve_eps(resolved_device) + resolved_device = resolve_device(EPDeviceTarget(ep="auto", device=device)).device + eps = available_eps_for_device(resolved_device) if not eps: return None @@ -107,9 +109,11 @@ def resolve_genai_ep(device: str) -> EPNameOrAlias | None: # cpu/gpu support but their primary target is npu — when the user says # ``--device cpu`` or ``--device gpu`` they expect the native EP for that # device, not a cross-device accelerator that also happens to support it. - native = [ep for ep in eps if EP_SUPPORTED_DEVICES[ep][0] == resolved_device] + native = [ep for ep in eps if EP_SUPPORTED_DEVICES[cast("EPName", ep)][0] == resolved_device] best = native[0] if native else eps[0] - return EP_NAME_TO_ALIAS[best] + # short_ep_name returns a plain ``str``; the value is a canonical EP short + # alias (a member of EPAlias) that GenaiSession accepts as an override. + return cast("EPNameOrAlias", short_ep_name(best)) def genai_output_path(bundle_dir: str | Path) -> Path: @@ -336,8 +340,6 @@ def _session_device(self) -> str | None: if device and device not in ("config", "auto"): return device canonical = normalize_ep_name(ep) - if canonical is None: - return None devices = EP_SUPPORTED_DEVICES.get(canonical) return devices[0] if devices else None diff --git a/src/winml/modelkit/commands/_pre_bench.py b/src/winml/modelkit/commands/_pre_bench.py new file mode 100644 index 000000000..726450545 --- /dev/null +++ b/src/winml/modelkit/commands/_pre_bench.py @@ -0,0 +1,160 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Pre-benchmark identity block. + +Renders two bordered panels before the benchmark loop: a **Model** panel +(identity + surface: task/opset/I/O) and a **Device** panel (resolved +device + EP provenance + DLL path). Option B semantics — no arrows, no +requested-vs-resolved leak — inside the classic two-Panel shell. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from rich.console import Group +from rich.panel import Panel +from rich.text import Text + + +if TYPE_CHECKING: + from collections.abc import Sequence + + from rich.console import Console + + +# Every label column is padded to LABEL_WIDTH so the values on the right +# line up in a single column. Matches the width used in the mockup at +# ``docs/design/perf/console_mockup.py::build_pre_bench_block``. +_LABEL_WIDTH = 10 + + +def print_pre_bench_block( + console: Console, + *, + model_id: str | None, + task: str | None, + opset: int | None, + inputs: Sequence[tuple[str, str, tuple[int | str, ...]]] | None, + outputs: Sequence[tuple[str, str, tuple[int | str, ...]]] | None, + cached_onnx_path: str | None, + onnx_file: str | None, + device: str, + hardware_name: str, + ep: str, + ep_source: str, + ep_version: str | None, + ep_dll_path: str, +) -> None: + """Print the pre-benchmark identity block. + + Layout (Option B — see the mockup design doc for the target shape): + + - Identity: ``Model:`` (bold cyan; ``(HF)`` / ``(local)`` suffix), plus + an ``ONNX:`` line when a cached artifact path is supplied. + - Surface: ``Task:``, ``Opset:``, ``Inputs:``, ``Outputs:`` — each + omitted when the source field is empty / ``None``. + - Device: ``Device:`` (resolved short name + hardware name in dim + parens), ``EP:`` (``@`` + optional ``v``), + ``EP DLL:`` (full plugin path; ``(bundled with ORT)`` when the EP + is built into ORT and has no plugin DLL). + + Args: + console: Rich console sink (usually ``Console(stderr=True)``). + model_id: HF model identifier when the user passed one; ``None`` + selects the raw-``.onnx`` branch. + task: Resolved task string (e.g. ``"image-classification"``). + opset: ONNX opset for the surface block. + inputs / outputs: I/O spec triples. Empty / ``None`` skips the row. + cached_onnx_path: Path of the compiled ONNX cached on disk (HF + path only). Rendered on a dedicated ``ONNX:`` line. + onnx_file: Raw ``.onnx`` file path when the user bypassed HF. + device: Resolved device short name (``"npu"`` / ``"gpu"`` / + ``"cpu"``). Never the literal ``"auto"`` — callers are + responsible for passing the resolved value. + hardware_name: Human-readable hardware label (e.g. ``"Intel(R) AI + Boost"``); rendered in dim parens after the device. + ep: Short EP alias (e.g. ``"qnn"`` / ``"openvino"``). Never + ``"auto"``. + ep_source: Canonical origin tag (``"bundled"`` / ``"directory"`` / + ``"pypi"`` / etc.). + ep_version: Per-source EP version string, or ``None`` (omits the + ``v`` chunk when absent). + ep_dll_path: Full path to the plugin DLL. Empty string signals a + built-in EP and renders as ``(bundled with ORT)``. + """ + # --- Model panel: identity + surface --------------------------------- + model_lines: list[Text] = [] + if model_id: + model_lines.append( + _labeled_line( + "Model:", f"[bold cyan]{model_id}[/bold cyan] [dim](HF)[/dim]" + ) + ) + if cached_onnx_path: + model_lines.append(_labeled_line("ONNX:", f"[dim]{cached_onnx_path}[/dim]")) + elif onnx_file: + model_lines.append( + _labeled_line( + "Model:", f"[bold cyan]{onnx_file}[/bold cyan] [dim](local)[/dim]" + ) + ) + + if task: + model_lines.append(_labeled_line("Task:", f"[cyan]{task}[/cyan]")) + if opset is not None: + model_lines.append(_labeled_line("Opset:", f"[green]{opset}[/green]")) + if inputs: + model_lines.extend(_io_lines("Inputs:", inputs)) + if outputs: + model_lines.extend(_io_lines("Outputs:", outputs)) + + if model_lines: + console.print(Panel(Group(*model_lines), title="Model", expand=True)) + + # --- Device panel: resolved device + EP + DLL ------------------------- + hw_suffix = f" [dim]({hardware_name})[/dim]" if hardware_name else "" + ep_line = f"[cyan]{ep}[/cyan]@[cyan]{ep_source}[/cyan]" + if ep_version: + ep_line += f" [green]v{ep_version}[/green]" + dll_display = ep_dll_path if ep_dll_path else "(bundled with ORT)" + + device_lines: list[Text] = [ + _labeled_line("Device:", f"[cyan]{device}[/cyan]{hw_suffix}"), + _labeled_line("EP:", ep_line), + _labeled_line("EP DLL:", f"[dim]{dll_display}[/dim]"), + ] + console.print(Panel(Group(*device_lines), title="Device", expand=True)) + + +def _labeled_line(label: str, value_markup: str) -> Text: + """Render one ``