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/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..2ad3d45f0 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,17 @@ 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: + 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 890f98540..83c4cdc89 100644 --- a/src/winml/modelkit/analyze/analyzer.py +++ b/src/winml/modelkit/analyze/analyzer.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any 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 @@ -186,7 +186,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: @@ -227,7 +227,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: @@ -260,7 +260,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: @@ -293,7 +293,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: @@ -374,7 +374,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: @@ -409,7 +409,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: @@ -452,7 +452,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 +594,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, @@ -728,7 +728,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, @@ -793,25 +793,27 @@ 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" logger.info("Using device: %s", device_to_use) # Determine which EPs to analyze - eps_to_analyze: list[EPName] = [] + eps_to_analyze: list[str] = [] 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_to_analyze = 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] @@ -834,8 +836,8 @@ def analyze_from_proto( extraction_ms = int((time.perf_counter() - extraction_start) * 1000) # Step 2: Check runtime support for each EP - check_op_results: dict[EPName, list[PatternRuntime]] = {} - information_list: dict[EPName, list[Information]] = {} + check_op_results: dict[str, list[PatternRuntime]] = {} + information_list: dict[str, list[Information]] = {} runtime_debug_details_summary: dict[ str, dict[str, list[str] | dict[str, RuntimeDebugSummaryEntry]] ] = {} @@ -964,7 +966,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 71b8babea..6b8c21024 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 e5594f6fc..23d0aac42 100644 --- a/src/winml/modelkit/analyze/pattern/check_patterns.py +++ b/src/winml/modelkit/analyze/pattern/check_patterns.py @@ -20,31 +20,28 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from ... import winml from ...onnx import ONNXDomain + + +if TYPE_CHECKING: + import argparse + + 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 ...sysinfo import SysInfo - - -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 ..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], @@ -144,7 +141,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 @@ -216,50 +213,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: 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, - # 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() -> argparse.ArgumentParser: @@ -292,6 +253,10 @@ def build_parser() -> argparse.ArgumentParser: "--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. " diff --git a/src/winml/modelkit/analyze/runtime_checker/check_ops.py b/src/winml/modelkit/analyze/runtime_checker/check_ops.py index b4fbcc022..71cf6fe9b 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,21 +33,14 @@ if TYPE_CHECKING: import argparse - - from ...utils.constants import EPName 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, 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], @@ -148,7 +140,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}" @@ -285,7 +277,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: @@ -297,7 +289,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 9edd0ed11..679490876 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 @@ -35,20 +31,20 @@ 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, ) -> None: @@ -57,10 +53,34 @@ def __init__( self._provider_options = provider_options 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..ef80a2760 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, @@ -68,12 +67,12 @@ def infer_ihv_from_ep_name(ep_name: EPNameOrAlias) -> IHVType: return ep_name_to_ihv.get(canonical, IHVType.UNKNOWN) # type: ignore[arg-type] -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..045858b63 100644 --- a/src/winml/modelkit/build/common.py +++ b/src/winml/modelkit/build/common.py @@ -4,30 +4,227 @@ # -------------------------------------------------------------------------- """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, + 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 ensure_pre_quantized_stamped( config: WinMLBuildConfig, onnx_path: Path, *, force: bool = False ) -> None: @@ -72,7 +269,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 +381,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..2ab2323e8 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,7 +85,7 @@ 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, **kwargs: Any, @@ -140,7 +135,8 @@ def build_hf_model( """ # TODO: Move hack_max_optim_iterations to global env config hack_max_optim_iterations: int = kwargs.pop("hack_max_optim_iterations", 3) - allow_unsupported_nodes: bool = kwargs.pop("allow_unsupported_nodes", False) + # Consumed so it doesn't leak into **kwargs; not yet threaded through run_build_stages. + kwargs.pop("allow_unsupported_nodes", False) # ONNX-level kwargs forwarded to export, optimize, quantize stages onnx_kwargs = { @@ -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: @@ -239,130 +234,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``. # ========================================================================= - # 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) - # ========================================================================= - # 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=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) diff --git a/src/winml/modelkit/build/onnx.py b/src/winml/modelkit/build/onnx.py index 8e0095e5e..358920115 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, @@ -80,7 +76,8 @@ def build_onnx_model( RuntimeError: If a pipeline stage fails. """ hack_max_optim_iterations: int = kwargs.pop("hack_max_optim_iterations", 3) - allow_unsupported_nodes: bool = kwargs.pop("allow_unsupported_nodes", False) + # Consumed so it doesn't leak into **kwargs; not yet threaded through run_build_stages. + kwargs.pop("allow_unsupported_nodes", False) onnx_kwargs = { "use_external_data": kwargs.get("use_external_data", True), } @@ -118,7 +115,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 +148,32 @@ 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, + 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) 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 e75628195..e16c876ac 100644 --- a/src/winml/modelkit/commands/_live_chart.py +++ b/src/winml/modelkit/commands/_live_chart.py @@ -19,12 +19,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. @@ -37,7 +53,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, device_kind: str | None = None, @@ -90,13 +106,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 +123,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,12 +142,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 adapter (NPU/GPU, 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. """ @@ -180,6 +204,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 %") @@ -201,16 +241,13 @@ 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 = ["[green]\u2588\u2588[/green] NPU %"] + if has_cpu: + 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()] @@ -225,15 +262,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)}]" @@ -245,28 +292,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). - # 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" - 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}" + # 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)" + 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 b40838055..98a6caf22 100644 --- a/src/winml/modelkit/commands/_perf_genai.py +++ b/src/winml/modelkit/commands/_perf_genai.py @@ -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, @@ -94,14 +94,14 @@ 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 - return EP_NAME_TO_ALIAS[eps[0]] + return short_ep_name(eps[0]) def genai_output_path(bundle_dir: str | Path) -> Path: 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 ``