From dde59769cd4182f208cd7a21f93ec7693465f05a Mon Sep 17 00:00:00 2001 From: Zac <1221537+tezheng@users.noreply.github.com> Date: Mon, 13 Jul 2026 08:15:06 +0800 Subject: [PATCH 1/4] feat(session): unified-source EP refactor + follow-ups (src) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Consolidated src/ changes spanning cc7a6650..fa9986fb + fix: - cc7a6650 feat(session): unified-source EP refactor + P0 fixes + simplification pass - 5f101e01 fix: close out PR #1019 regressions and stale tests - 057e00ba fix(sys/ep): truthful isolated-register errors, hardened rendering, stale-test alignment - (local) refactor(session): rename _ep_short_or_none -> ep_short_or_none; shrink session facade __all__ (drop SessionState, InferenceError — session-lifecycle types now sourced from session.session where tested). Kept in facade: EPDeviceSpec, UnknownListingPick, WinMLEPMonitorMismatch, default_device_for_ep, known_ep_short_names, lookup_device_spec — the test_ep_device_import_rule fitness function requires them there. Test updates and non-src changes (docs, pyproject.toml, CI workflow) land in follow-up commits. --- src/winml/modelkit/__init__.py | 9 +- src/winml/modelkit/_warnings.py | 28 +- src/winml/modelkit/analyze/analyzer.py | 20 +- .../analyze/core/doc_constraint_checker.py | 2 +- .../model_validator_manager.py | 3 +- .../analyze/core/runtime_checker_query.py | 50 +- src/winml/modelkit/analyze/models/ihv_type.py | 4 +- .../modelkit/analyze/models/information.py | 4 +- .../modelkit/analyze/models/onnx_model.py | 4 +- .../modelkit/analyze/models/runtime_checks.py | 6 +- .../modelkit/analyze/models/support_level.py | 4 +- .../analyze/pattern/check_patterns.py | 85 +- .../analyze/runtime_checker/check_ops.py | 22 +- .../analyze/runtime_checker/ep_checker.py | 29 +- src/winml/modelkit/analyze/utils/__init__.py | 2 - src/winml/modelkit/analyze/utils/ep_utils.py | 12 +- .../analyze/utils/pattern_matching.py | 118 -- src/winml/modelkit/build/common.py | 204 ++- src/winml/modelkit/build/hf.py | 153 +- src/winml/modelkit/build/onnx.py | 150 +- src/winml/modelkit/commands/_ep_arg.py | 134 ++ src/winml/modelkit/commands/_live_chart.py | 100 +- src/winml/modelkit/commands/_pre_bench.py | 160 ++ src/winml/modelkit/commands/analyze.py | 6 +- src/winml/modelkit/commands/build.py | 54 +- src/winml/modelkit/commands/compile.py | 89 +- src/winml/modelkit/commands/config.py | 29 +- src/winml/modelkit/commands/eval.py | 8 +- src/winml/modelkit/commands/perf.py | 745 +++++--- src/winml/modelkit/commands/sys.py | 773 +++++++-- src/winml/modelkit/compiler/__init__.py | 2 +- src/winml/modelkit/compiler/cli.py | 397 ----- src/winml/modelkit/compiler/compiler.py | 12 +- src/winml/modelkit/compiler/configs.py | 208 +-- src/winml/modelkit/compiler/stages/compile.py | 45 +- src/winml/modelkit/config/build.py | 17 +- src/winml/modelkit/config/precision.py | 72 +- src/winml/modelkit/core/time_utils.py | 4 +- src/winml/modelkit/ep_path.py | 1526 +++++++++++++++++ src/winml/modelkit/eval/evaluate.py | 11 +- .../modelkit/export/htp/config_generator.py | 345 ---- .../modelkit/export/htp/metadata_writer.py | 4 +- src/winml/modelkit/export/htp/monitor.py | 1 - src/winml/modelkit/inference/engine.py | 32 +- src/winml/modelkit/models/auto.py | 47 +- src/winml/modelkit/models/winml/base.py | 20 +- .../modelkit/models/winml/composite_model.py | 24 +- src/winml/modelkit/onnx/__init__.py | 18 +- src/winml/modelkit/onnx/domains.py | 4 +- src/winml/modelkit/optracing/__init__.py | 34 - src/winml/modelkit/optracing/base.py | 35 - .../modelkit/optracing/qnn/csv_parser.py | 227 --- src/winml/modelkit/optracing/qnn/profiler.py | 351 ---- .../modelkit/optracing/qnn/qhas_parser.py | 113 -- src/winml/modelkit/optracing/registry.py | 64 - src/winml/modelkit/optracing/result.py | 99 -- src/winml/modelkit/pattern/config.py | 2 +- src/winml/modelkit/pattern/models.py | 4 +- src/winml/modelkit/serve/app.py | 4 +- src/winml/modelkit/serve/cli_api.py | 4 +- src/winml/modelkit/serve/manager.py | 2 +- src/winml/modelkit/serve/schema.py | 2 +- src/winml/modelkit/session/__init__.py | 73 +- src/winml/modelkit/session/ep_device.py | 864 ++++++++++ src/winml/modelkit/session/ep_registry.py | 718 ++++++-- .../qnn => session/monitor}/__init__.py | 2 +- src/winml/modelkit/session/monitor/_pdh.py | 128 ++ .../modelkit/session/monitor/_xrt_smi.py | 24 +- .../modelkit/session/monitor/ep_monitor.py | 125 +- .../modelkit/session/monitor/hw_monitor.py | 29 +- .../modelkit/session/monitor/live_display.py | 207 --- .../modelkit/session/monitor/op_metrics.py | 168 ++ .../session/monitor/openvino_monitor.py | 112 +- .../modelkit/session/monitor/qnn/__init__.py | 27 + .../modelkit/session/monitor/qnn/_internal.py | 425 +++++ .../monitor}/qnn/viewer.py | 93 +- .../modelkit/session/monitor/qnn_monitor.py | 656 ++++++- .../{optracing => session/monitor}/report.py | 149 +- .../session/monitor/vitisai_monitor.py | 8 +- .../modelkit/session/qairt/qairt_session.py | 18 +- src/winml/modelkit/session/session.py | 608 ++++--- src/winml/modelkit/sysinfo/__init__.py | 6 +- src/winml/modelkit/sysinfo/device.py | 191 --- src/winml/modelkit/sysinfo/hardware.py | 29 + src/winml/modelkit/sysinfo/sysinfo.py | 38 - .../modelkit/telemetry/deviceid/deviceid.py | 4 +- .../modelkit/telemetry/library/exporter.py | 4 +- src/winml/modelkit/transformers_compat.py | 238 +++ src/winml/modelkit/utils/__init__.py | 16 - src/winml/modelkit/utils/cli.py | 83 +- src/winml/modelkit/utils/constants.py | 120 -- src/winml/modelkit/utils/hub_utils.py | 344 ---- src/winml/modelkit/utils/optimum_loader.py | 160 -- src/winml/modelkit/winml.py | 248 ++- 94 files changed, 7979 insertions(+), 4673 deletions(-) delete mode 100644 src/winml/modelkit/analyze/utils/pattern_matching.py create mode 100644 src/winml/modelkit/commands/_ep_arg.py create mode 100644 src/winml/modelkit/commands/_pre_bench.py delete mode 100644 src/winml/modelkit/compiler/cli.py create mode 100644 src/winml/modelkit/ep_path.py delete mode 100644 src/winml/modelkit/export/htp/config_generator.py delete mode 100644 src/winml/modelkit/optracing/__init__.py delete mode 100644 src/winml/modelkit/optracing/base.py delete mode 100644 src/winml/modelkit/optracing/qnn/csv_parser.py delete mode 100644 src/winml/modelkit/optracing/qnn/profiler.py delete mode 100644 src/winml/modelkit/optracing/qnn/qhas_parser.py delete mode 100644 src/winml/modelkit/optracing/registry.py delete mode 100644 src/winml/modelkit/optracing/result.py create mode 100644 src/winml/modelkit/session/ep_device.py rename src/winml/modelkit/{optracing/qnn => session/monitor}/__init__.py (82%) delete mode 100644 src/winml/modelkit/session/monitor/live_display.py create mode 100644 src/winml/modelkit/session/monitor/op_metrics.py create mode 100644 src/winml/modelkit/session/monitor/qnn/__init__.py create mode 100644 src/winml/modelkit/session/monitor/qnn/_internal.py rename src/winml/modelkit/{optracing => session/monitor}/qnn/viewer.py (63%) rename src/winml/modelkit/{optracing => session/monitor}/report.py (52%) delete mode 100644 src/winml/modelkit/sysinfo/device.py create mode 100644 src/winml/modelkit/transformers_compat.py delete mode 100644 src/winml/modelkit/utils/constants.py delete mode 100644 src/winml/modelkit/utils/hub_utils.py delete mode 100644 src/winml/modelkit/utils/optimum_loader.py diff --git a/src/winml/modelkit/__init__.py b/src/winml/modelkit/__init__.py index 59a45c313..bfb0855c6 100644 --- a/src/winml/modelkit/__init__.py +++ b/src/winml/modelkit/__init__.py @@ -34,7 +34,14 @@ logging.getLogger(__name__).addHandler(logging.NullHandler()) -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 # noqa: I001 + +transformers_compat.arm() try: diff --git a/src/winml/modelkit/_warnings.py b/src/winml/modelkit/_warnings.py index fbe08288d..f40c7b5ec 100644 --- a/src/winml/modelkit/_warnings.py +++ b/src/winml/modelkit/_warnings.py @@ -19,6 +19,7 @@ import logging import os +import sys import warnings @@ -79,13 +80,14 @@ def filter(self, record: logging.LogRecord) -> bool: warnings.filterwarnings("ignore", category=_cat, module=r"torch\..*") # TracerWarning (from torch.jit, inherits Warning not UserWarning) - # fires during ONNX export tracing — safe to suppress in both torch and transformers - try: - from torch.jit import TracerWarning - - warnings.filterwarnings("ignore", category=TracerWarning) - except ImportError: - pass # torch not installed + # 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( @@ -93,5 +95,17 @@ def filter(self, record: logging.LogRecord) -> bool: ) +def install_torch_tracer_filter() -> None: + """Register the ``TracerWarning`` filter — call after ``import torch``. + + Idempotent — ``warnings.filterwarnings`` de-duplicates identical entries. + """ + try: + from torch.jit import TracerWarning + 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 60e5126f0..0348d8cf5 100644 --- a/src/winml/modelkit/analyze/analyzer.py +++ b/src/winml/modelkit/analyze/analyzer.py @@ -17,7 +17,7 @@ from typing import TYPE_CHECKING from ..optim.config import WinMLOptimizationConfig -from ..utils.constants import normalize_ep_name +from ..utils.cli import normalize_ep_name from .models.information import Information from .models.support_level import SupportLevel @@ -667,22 +667,20 @@ def analyze_from_proto( # Determine which EPs to analyze eps_to_analyze: list[str] = [] if ep_normalized is None: - # Analyze all supported EPs - eps_to_analyze = [ - "QNNExecutionProvider", - "OpenVINOExecutionProvider", - "VitisAIExecutionProvider", - ] - logger.info("No EP specified, analyzing all supported EPs: %s", eps_to_analyze) + # Derive the NPU EP list from the catalog so future EP additions + # are automatically included. sorted() gives deterministic order. + from ..session import eps_for_device + + eps_to_analyze = sorted(eps_for_device("npu")) + logger.info("No EP specified, analyzing all NPU-capable EPs: %s", eps_to_analyze) else: eps_to_analyze = [ep_normalized] # 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") - 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" diff --git a/src/winml/modelkit/analyze/core/doc_constraint_checker.py b/src/winml/modelkit/analyze/core/doc_constraint_checker.py index 107fe516d..3fab5a53a 100644 --- a/src/winml/modelkit/analyze/core/doc_constraint_checker.py +++ b/src/winml/modelkit/analyze/core/doc_constraint_checker.py @@ -124,7 +124,7 @@ def _load_constraints(self) -> dict[str, pd.DataFrame]: df = pd.DataFrame(op_data) op_dfs[op_type] = df logger.debug(f"Loaded constraints for operator: {op_type} ({len(df)} records)") - except Exception as e: # noqa: PERF203 + except Exception as e: logger.error(f"Failed to load constraints for {op_type}: {e}") return op_dfs 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 167b8eb7d..2882b2356 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 @@ -89,7 +89,6 @@ def __init__( self.model_proto = model.get_model() self.op_runtime_results = op_runtime_results or [] self.device = device or "NPU" - self.device = device self.enabled_validators = enabled_validators or list(self.VALIDATORS.keys()) # Instantiate enabled validators @@ -142,7 +141,7 @@ def run_all_validators(self) -> list[Information]: if info: logger.info(f"{validator.validator_name} found issue: {info.pattern_id}") information_list.append(info) - except Exception as e: # noqa: PERF203 + except Exception as e: logger.exception( f"Validator {validator.validator_name} failed with exception: " f"{type(e).__name__}", diff --git a/src/winml/modelkit/analyze/core/runtime_checker_query.py b/src/winml/modelkit/analyze/core/runtime_checker_query.py index 8ac93183b..4a23f5286 100644 --- a/src/winml/modelkit/analyze/core/runtime_checker_query.py +++ b/src/winml/modelkit/analyze/core/runtime_checker_query.py @@ -15,7 +15,6 @@ import numpy as np import onnx -import onnxruntime as ort import pandas as pd from onnx import numpy_helper @@ -1286,31 +1285,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 - - winml.register_execution_providers(ort=True) - - 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 = ort.get_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 @@ -1322,11 +1334,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 48d77433a..f5f096b14 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 afd6d0d5a..7793e72b0 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 d66d487b2..f40c34ef3 100644 --- a/src/winml/modelkit/analyze/models/onnx_model.py +++ b/src/winml/modelkit/analyze/models/onnx_model.py @@ -6,7 +6,7 @@ from __future__ import annotations -from enum import Enum +from enum import StrEnum import onnx from pydantic import BaseModel, Field, field_validator @@ -14,7 +14,7 @@ from ...onnx import ONNXDomain -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 5356534f3..32e31f07f 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 @@ -16,14 +16,14 @@ from .support_level import SupportLevel -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 05c98d131..90f90cf19 100644 --- a/src/winml/modelkit/analyze/pattern/check_patterns.py +++ b/src/winml/modelkit/analyze/pattern/check_patterns.py @@ -18,24 +18,23 @@ from pathlib import Path from typing import Any -import onnxruntime as ort - -from ... import winml from ...onnx import ONNXDomain from ...pattern.base import ( PatternInputGenerator, get_pattern_input_generator, get_registered_pattern_input_generators, ) +from ...session import DEVICE_TYPE_TO_DEVICE from ...sysinfo import SysInfo -from ...utils import constants +from ..runtime_checker.check_ops import ( + OpenVINONPUChecker, + QNNNPUChecker, + get_ep_checker, +) from ..runtime_checker.ep_checker import EPChecker from ..utils import CheckResultWriter -winml.register_execution_providers(ort=True) - - def check_patterns( ep_checker: EPChecker, patterns: list[str], @@ -126,7 +125,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 @@ -198,50 +197,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) - - -def get_ep_checker(ep_name: str, 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, Any] = { - "QNNExecutionProvider": QNNNPUChecker, - "OpenVINOExecutionProvider": OpenVINONPUChecker, - # 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" - ) - return ep_name_to_checker[ep_name](device_type=device_type) +# NPU 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(): @@ -274,6 +237,10 @@ def build_parser(): "--ep", type=str, required=True, + # CARVE-OUT: This subprocess tool intentionally supports only a curated subset of + # NPU EPs. VitisAI and future NPU EPs are excluded because this pattern-checking + # tool has not been validated against them. Do NOT derive from eps_for_device("npu") + # or EP_DEVICE_SPECS — this is an explicit opt-in allowlist, not catalog drift. choices=["QNNExecutionProvider", "OpenVINOExecutionProvider"], help=( "Execution Provider names to test. " @@ -293,10 +260,7 @@ def build_parser(): "--opset_mapping", type=str, nargs="+", - help=( - "Domain:version pairs for ONNX opset versions, " - "e.g., ai.onnx:17 com.microsoft:1" - ), + help=("Domain:version pairs for ONNX opset versions, e.g., ai.onnx:17 com.microsoft:1"), ) opset_group.add_argument( "--opset_version", @@ -311,8 +275,7 @@ def build_parser(): type=str, default=ONNXDomain.AI_ONNX.value, help=( - "ONNX opset domain to use with --opset_version " - f"(default: {ONNXDomain.AI_ONNX.value})" + f"ONNX opset domain to use with --opset_version (default: {ONNXDomain.AI_ONNX.value})" ), ) parser.add_argument( @@ -392,8 +355,7 @@ def _parse_opset_mapping(args: Any) -> dict[str, int]: for pair in args.opset_mapping: if ":" not in pair: raise ValueError( - "Invalid --opset_mapping value " - f"'{pair}'. Expected format: domain:version" + f"Invalid --opset_mapping value '{pair}'. Expected format: domain:version" ) domain, version_text = pair.split(":", 1) if not domain: @@ -402,8 +364,7 @@ def _parse_opset_mapping(args: Any) -> dict[str, int]: opset_mapping[domain] = int(version_text) except ValueError as exc: raise ValueError( - "Invalid --opset_mapping value " - f"'{pair}'. Version must be an integer" + f"Invalid --opset_mapping value '{pair}'. Version must be an integer" ) from exc return opset_mapping diff --git a/src/winml/modelkit/analyze/runtime_checker/check_ops.py b/src/winml/modelkit/analyze/runtime_checker/check_ops.py index 1a0b9f094..8d1a0d53a 100644 --- a/src/winml/modelkit/analyze/runtime_checker/check_ops.py +++ b/src/winml/modelkit/analyze/runtime_checker/check_ops.py @@ -21,7 +21,6 @@ import onnxruntime as ort from onnx.defs import SchemaError -from ... import winml from ...onnx import ONNXDomain from ...pattern.op_input_gen import ( OpInputGenerator, @@ -29,18 +28,13 @@ get_runtime_checker_op, ) from ...pattern.op_input_gen.qdq_gen import QDQGenerator +from ...session import DEVICE_TO_DEVICE_TYPE, DEVICE_TYPE_TO_DEVICE from ...sysinfo import SysInfo -from ...utils import constants from ..utils import CheckResultWriter 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], @@ -131,7 +125,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}" @@ -261,10 +255,10 @@ class RTXChecker(EPChecker): def __init__(self, device_type: ort.OrtHardwareDeviceType) -> None: if device_type != ort.OrtHardwareDeviceType.GPU: - raise ValueError("NvTensorRTRTXExecutionProvider only supports GPU device type") + raise ValueError("NvTensorRtRtxExecutionProvider only supports GPU device type") """Initialize RTX checker.""" super().__init__( - ep_name="NvTensorRTRTXExecutionProvider", device_type=ort.OrtHardwareDeviceType.GPU + ep_name="NvTensorRtRtxExecutionProvider", device_type=ort.OrtHardwareDeviceType.GPU ) @@ -280,13 +274,13 @@ def get_ep_checker(ep_name: str, 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, "VitisAIExecutionProvider": VitisAIChecker, "MIGraphXExecutionProvider": MIGraphXChecker, - "NvTensorRTRTXExecutionProvider": RTXChecker, + "NvTensorRtRtxExecutionProvider": RTXChecker, # Add other EPChecker subclasses here as needed } if ep_name not in ep_name_to_checker: @@ -332,7 +326,7 @@ def build_parser(): "OpenVINOExecutionProvider", "VitisAIExecutionProvider", "MIGraphXExecutionProvider", - "NvTensorRTRTXExecutionProvider", + "NvTensorRtRtxExecutionProvider", ], help=( "Execution Provider names to test. " @@ -340,7 +334,7 @@ def build_parser(): "OpenVINOExecutionProvider, " "VitisAIExecutionProvider, " "MIGraphXExecutionProvider, " - "NvTensorRTRTXExecutionProvider" + "NvTensorRtRtxExecutionProvider" ), ) parser.add_argument( diff --git a/src/winml/modelkit/analyze/runtime_checker/ep_checker.py b/src/winml/modelkit/analyze/runtime_checker/ep_checker.py index 94a073965..c42ddbbd0 100644 --- a/src/winml/modelkit/analyze/runtime_checker/ep_checker.py +++ b/src/winml/modelkit/analyze/runtime_checker/ep_checker.py @@ -11,8 +11,6 @@ import onnx import onnxruntime as ort -from ... import winml - # TODO: allow test case iter to take dtypes as inputs # TODO: define dataclass for result @@ -41,9 +39,34 @@ def __init__( self._provider_options = provider_options def _get_sess_options(self) -> ort.SessionOptions: + 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 15b015306..4826a39bc 100644 --- a/src/winml/modelkit/analyze/utils/__init__.py +++ b/src/winml/modelkit/analyze/utils/__init__.py @@ -7,7 +7,6 @@ from .ep_utils import get_devices_with_rule_data, has_rule_data_for_ep, infer_ihv_from_ep_name from .json_utils import validate_json_schema from .op_utils import CheckResultWriter -from .pattern_matching import match_pattern_with_wildcards from .rule_loader import RuleLoader, get_runtime_rules_search_dirs, resolve_rule_zip_path @@ -18,7 +17,6 @@ "get_runtime_rules_search_dirs", "has_rule_data_for_ep", "infer_ihv_from_ep_name", - "match_pattern_with_wildcards", "resolve_rule_zip_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 10c4b377a..86d5168c4 100644 --- a/src/winml/modelkit/analyze/utils/ep_utils.py +++ b/src/winml/modelkit/analyze/utils/ep_utils.py @@ -70,7 +70,7 @@ def get_devices_with_rule_data(ep_name: str) -> list[str]: First probes rule zip search directories for files matching ``{ep_name}_{device}_*.zip``. 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"``). @@ -79,10 +79,10 @@ def get_devices_with_rule_data(ep_name: str) -> 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 @@ -91,9 +91,9 @@ def get_devices_with_rule_data(ep_name: str) -> 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_rule_data_for_ep(ep_name: str, device: str) -> bool: 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 d01527d84..a77909d51 100644 --- a/src/winml/modelkit/build/common.py +++ b/src/winml/modelkit/build/common.py @@ -4,21 +4,26 @@ # -------------------------------------------------------------------------- """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 ..onnx import copy_onnx_model +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: @@ -27,6 +32,199 @@ 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, + 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, + **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, + **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 run_optimize_analyze_loop( model_path: Path, optimized_path: Path, diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 73d4f5d1d..4c4eac9ca 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -27,11 +27,8 @@ 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, is_quantized_onnx -from ..quant import quantize_onnx -from .common import run_optimize_analyze_loop +from .common import run_build_stages if TYPE_CHECKING: @@ -228,130 +225,32 @@ def _name(base: str) -> str: 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. + # [3]-[6] OPTIMIZE -> QUANTIZE -> COMPILE -> FINALIZE + # Shared with build_onnx_model via ``common.run_build_stages``. # ========================================================================= skip_optimize: bool = kwargs.pop("skip_optimize", False) - # Defensive fallback: when called through the unified pipeline, - # generate_*_build_config() already detects QDQ models and sets - # config.quant=None. This is_quantized_onnx() check is redundant in that - # path but kept for backward compatibility when build_hf_model() - # is called directly with a hand-built config. - is_pre_quantized = is_quantized_onnx(current_path) or skip_optimize - - if is_pre_quantized: - logger.info( - "Pre-quantized model detected (QDQ nodes present). " - "Skipping optimize + quantize, running analyze-only." - ) - stages_skipped.append("optimize") - # Optimize+analyze only, no autoconf re-optimization - 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, - **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, - **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) - # ========================================================================= - 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: - # Defensive fallback: catches the edge case where a direct caller - # provides config.quant != None but the model already has QDQ nodes - # (e.g., hand-built config without running generate_*_build_config). - if is_quantized_onnx(current_path): - logger.warning( - "Model already contains QDQ nodes, skipping quantization. " - "Set config.quant=None to silence this warning." - ) - stages_skipped.append("quantize") - else: - 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) + 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, + onnx_kwargs=onnx_kwargs, + ) + stages_completed.extend(stages.stages_completed) + stages_skipped.extend(stages.stages_skipped) + stage_timings.update(stages.stage_timings) + current_path = stages.current_path + 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) @@ -365,7 +264,7 @@ def _name(base: str) -> str: "task": task, "cache_key": cache_key, "config_hash": cache_key.rsplit("_", 1)[-1] if cache_key and "_" in cache_key else None, - "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "timestamp": datetime.datetime.now(datetime.UTC).isoformat(), "elapsed_seconds": round(elapsed, 3), "stages": [], "final_artifact": final_path.name, diff --git a/src/winml/modelkit/build/onnx.py b/src/winml/modelkit/build/onnx.py index 9bccd8761..ce5c77eeb 100644 --- a/src/winml/modelkit/build/onnx.py +++ b/src/winml/modelkit/build/onnx.py @@ -20,10 +20,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from ..compiler import compile_onnx -from ..onnx import copy_onnx_model, is_quantized_onnx -from ..quant import quantize_onnx -from .common import run_optimize_analyze_loop +from ..onnx import copy_onnx_model +from .common import run_build_stages from .hf import BuildResult @@ -133,126 +131,32 @@ def build_onnx_model( copy_onnx_model(onnx_path, current_path) # ========================================================================= - # [1] OPTIMIZE + ANALYZE (or ANALYZE-ONLY 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``. # ========================================================================= skip_optimize: bool = kwargs.pop("skip_optimize", False) - # Defensive fallback: when called through the unified pipeline, - # generate_onnx_build_config() already detects QDQ models and sets - # config.quant=None. This is_quantized_onnx() check is redundant in that - # path but kept for backward compatibility when build_onnx_model() - # is called directly with a hand-built config. - is_pre_quantized = is_quantized_onnx(current_path) or skip_optimize - - if is_pre_quantized: - logger.info( - "Pre-quantized model detected (QDQ nodes present). " - "Skipping optimize + quantize, running analyze-only." - ) - stages_skipped.append("optimize") - # Optimize+analyze only, no autoconf re-optimization - 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, - **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, - **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) - # ========================================================================= - 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: - # Defensive fallback: catches the edge case where a direct caller - # provides config.quant != None but the model already has QDQ nodes - # (e.g., hand-built config without running generate_*_build_config). - if is_quantized_onnx(current_path): - logger.warning( - "Model already contains QDQ nodes, skipping quantization. " - "Set config.quant=None to silence this warning." - ) - stages_skipped.append("quantize") - else: - 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) + 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, + onnx_kwargs=onnx_kwargs, + ) + stages_completed.extend(stages.stages_completed) + stages_skipped.extend(stages.stages_skipped) + stage_timings.update(stages.stage_timings) + current_path = stages.current_path + 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) @@ -264,7 +168,7 @@ def build_onnx_model( "schema_version": 1, "source": "onnx", "input_onnx": str(onnx_path), - "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(), + "timestamp": datetime.datetime.now(datetime.UTC).isoformat(), "elapsed_seconds": round(elapsed, 3), "stages": [], "final_artifact": final_path.name, diff --git a/src/winml/modelkit/commands/_ep_arg.py b/src/winml/modelkit/commands/_ep_arg.py new file mode 100644 index 000000000..7730e4b5c --- /dev/null +++ b/src/winml/modelkit/commands/_ep_arg.py @@ -0,0 +1,134 @@ +# ------------------------------------------------------------------------- +# 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 + +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( # type: ignore[override] + self, + value, + param, + ctx, + ) -> 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: + 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 ef462a7a0..9ea28e86d 100644 --- a/src/winml/modelkit/commands/_live_chart.py +++ b/src/winml/modelkit/commands/_live_chart.py @@ -17,12 +17,28 @@ # 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 +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. @@ -35,7 +51,7 @@ 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, ) -> None: @@ -75,13 +91,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, @@ -90,6 +108,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 @@ -106,12 +127,15 @@ 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 NPU (green) and CPU (cyan) with distinct colors. + Plots NPU (green), CPU (cyan), and GPU (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. """ @@ -155,6 +179,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 %") @@ -175,12 +215,12 @@ def _render_chart( from rich.text import Text # Rich-colored title line with legend swatches + legend_parts = ["[green]\u2588\u2588[/green] NPU %"] if has_cpu: - title = Text.from_markup( - " Utilization ([green]\u2588\u2588[/green] NPU % [cyan]\u2588\u2588[/cyan] CPU %)" - ) - else: - title = Text.from_markup(" Utilization ([green]\u2588\u2588[/green] NPU %)") + legend_parts.append("[cyan]\u2588\u2588[/cyan] CPU %") + if has_gpu: + legend_parts.append("[bright_yellow]\u2588\u2588[/bright_yellow] GPU %") + 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()] @@ -195,15 +235,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 (NPU / 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)}]" @@ -215,23 +265,31 @@ def _render_status( throughput = 1000.0 / latency_ms if latency_ms > 0 else 0.0 + npu_avg, npu_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) - npu_cell = f"NPU: {mean_util:.1f}% avg ({current_util:.1f}% now)" - cpu_cell = f"CPU: {cpu_pct:.1f}%" + # Row 2: Compute (unified now/avg format across all three) + npu_cell = f"NPU: {npu_now:.1f}%/{npu_avg:.1f}%" + cpu_cell = f"CPU: {cpu_now:.1f}%/{cpu_avg:.1f}%" + gpu_cell = f"GPU: {gpu_now:.1f}%/{gpu_avg:.1f}%" + row2 = f" {npu_cell:<20}| {cpu_cell:<20}| {gpu_cell:<20}" + + # 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)" - row2 = f" {npu_cell:<30}| {cpu_cell:<12}| {ram_cell} | {mem_cell}" + 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/_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 ``