From e176488a7b0c6e856cdfd4d2eadd33d92621f495 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 29 Jul 2026 14:26:58 +0800 Subject: [PATCH 1/5] fix: build Qwen3 genai bundles without runtime sessions Preserve quantization when raw graphs skip ORT optimization, and make genai bundle assembly use private artifact builds so bundle export does not create runtime sessions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/winml/modelkit/build/common.py | 49 ++----- src/winml/modelkit/build/hf.py | 3 + src/winml/modelkit/commands/build.py | 18 +-- src/winml/modelkit/config/build.py | 10 +- src/winml/modelkit/models/auto.py | 136 +++++++++++++----- .../models/hf/qwen3/qwen_transformer_only.py | 1 + .../modelkit/models/winml/genai_bundle.py | 62 ++++---- tests/unit/build/test_hf.py | 50 ++++++- tests/unit/build/test_onnx.py | 6 +- tests/unit/commands/test_build.py | 67 +++++++-- .../models/auto/test_from_pretrained_ep.py | 41 ++++++ .../unit/models/qwen3/test_qwen3_modeling.py | 7 + .../winml/test_genai_bundle_orchestrator.py | 99 ++++++++----- 13 files changed, 372 insertions(+), 177 deletions(-) diff --git a/src/winml/modelkit/build/common.py b/src/winml/modelkit/build/common.py index 99835a4ee..863a74a21 100644 --- a/src/winml/modelkit/build/common.py +++ b/src/winml/modelkit/build/common.py @@ -88,24 +88,23 @@ def run_build_stages( 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). + skip_optimize: When True, bypasses optimize. 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, + is_pre_quantized=is_quantized_onnx(current_path), ) # ========================================================================= - # OPTIMIZE + ANALYZE (or ANALYZE-ONLY for pre-quantized) + # OPTIMIZE + ANALYZE # ========================================================================= - if result.is_pre_quantized: - logger.info( - "Pre-quantized model detected (QDQ nodes present). " - "Skipping optimize + quantize, running analyze-only." - ) + if result.is_pre_quantized or skip_optimize: + if result.is_pre_quantized: + logger.info("Pre-quantized model detected (QDQ nodes present).") + else: + logger.info("Optimize skipped (skip_optimize=True).") result.stages_skipped.append("optimize") ( result.current_path, @@ -236,17 +235,7 @@ def run_build_stages( def ensure_pre_quantized_stamped( config: WinMLBuildConfig, onnx_path: Path, *, force: bool = False ) -> None: - """Stamp ``config.skip_optimize`` (and clear ``config.quant``) once. - - Sets ``config.skip_optimize = True`` and clears ``config.quant`` if the - input ONNX is already quantized. - - This is the **single defensive detection point** for the library entry - points (``build_onnx_model``, ``build_hf_model``). When - ``config.skip_optimize`` is already True (i.e. the unified CLI path - via :func:`generate_onnx_build_config` already stamped the config), it - still enforces ``config.quant = None`` without re-running - ``is_quantized_onnx()``. + """Stamp ``config.skip_optimize`` and clear ``config.quant`` for quantized ONNX. Args: config: Build config to stamp in place. @@ -255,9 +244,6 @@ def ensure_pre_quantized_stamped( ``is_quantized_onnx`` (used to honor the legacy ``skip_optimize=True`` kwarg from direct callers). """ - if config.skip_optimize: - config.quant = None - return if force: config.skip_optimize = True config.quant = None @@ -318,11 +304,7 @@ def run_optimize_analyze_loop( analyze_output_path: Optional path to write the full analysis result as JSON. Written after every analyze pass; each pass overwrites the previous one so the file always reflects the most recent analysis. - skip_optimize: When True, skip the initial ``optimize_onnx`` call and - just copy the input model to ``optimized_path``. Used for - pre-quantized models (QDQ or QOperator format) where ORT-based - graph optimization would fail because the runtime lacks kernels - for ops like ``ConvInteger`` on the host EP. + skip_optimize: When True, skip the initial ``optimize_onnx`` call. **onnx_kwargs: Additional ONNX-level kwargs. Returns: @@ -335,21 +317,14 @@ def run_optimize_analyze_loop( if not config.auto: max_optim_iterations = 0 - # Enforce the skip_optimize invariant: autoconf re-optimize would - # crash on pre-quantized models for the same reason the initial - # optimize was skipped (ORT lacks kernels for the integer ops on the - # host EP). Drop iterations to 0 so callers can pass any value safely. + # If optimize is bypassed, autoconf must not re-run it. if skip_optimize: max_optim_iterations = 0 t0 = time.monotonic() - # 1. Optimize (or skip for pre-quantized models) + # 1. Optimize if skip_optimize: - # Pre-quantized models (QOperator format with ConvInteger / - # MatMulInteger) cannot pass through ORT graph optimization on - # hosts that lack kernels for those integer ops. Simply forward - # the input as the "optimized" artifact. if model_path.resolve() != optimized_path.resolve(): copy_onnx_model(model_path, optimized_path) else: diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 9ce3989ad..6e64e3d34 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -20,6 +20,7 @@ from __future__ import annotations import datetime +import gc import logging import time from dataclasses import dataclass, field @@ -232,6 +233,8 @@ def _name(base: str) -> str: stage_timings["export"] = time.monotonic() - t0 stages_completed.append("export") logger.info("Export done (%.1fs) -> %s", stage_timings["export"], export_path) + pytorch_model = None + gc.collect() # ========================================================================= # [3]-[6] OPTIMIZE -> QUANTIZE -> COMPILE -> FINALIZE diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index 51b5abb49..a90b86428 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -1047,7 +1047,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: resolved_quant, _ = resolve_quant_compile_config( device=device, precision=precision, ep=ep_value ) - if cfg.skip_optimize or not quant or resolved_quant is None: + if not quant or resolved_quant is None or (cfg.skip_optimize and cfg.quant is None): cfg.quant = None elif cfg.quant is None: # Populate calibration identifiers from the loader/model @@ -1641,8 +1641,6 @@ def _run_optimize_stage( show_io_first: If True, show I/O tensors at the start of the stage (used in ONNX mode where there is no export stage). skip_optimize: When True, skip the ORT graph-optimization pass. - Used for pre-quantized models (QDQ or QOperator format) whose - integer ops have no kernel on the host EP. Returns: Tuple of (current_path, opt_elapsed). @@ -1797,10 +1795,6 @@ def _run_quantize_stage( from ..quant import quantize_onnx from ..utils.console import StageLive - if config.skip_optimize: - config.quant = None - return current_path - if config.quant is None: # ``generate_onnx_build_config`` and ``ensure_pre_quantized_stamped`` # (in build/common.py) both clear ``config.quant`` for pre-quantized @@ -2177,14 +2171,6 @@ def _build_onnx_pipeline( # run on integer ops and the quantize stage may try to re-quantize. ensure_pre_quantized_stamped(config, current_path) - # Pre-quantized models (QDQ or QOperator format) cannot pass through - # ORT-based graph optimization on hosts that lack kernels for ops like - # ``ConvInteger``. The unified pipeline stamps ``config.skip_optimize`` - # exactly once in ``generate_onnx_build_config`` -- downstream stages - # (here and inside ``build_onnx_model``) read the flag instead of - # re-running ``is_quantized_onnx`` on the same file. - is_pre_quantized = config.skip_optimize - # ── Optimize stage (first stage for ONNX — show I/O here) ──── current_path, _ = _run_optimize_stage( config=config, @@ -2197,7 +2183,7 @@ def _build_onnx_pipeline( show_io_first=True, analyze_output_path=analyze_result_path, allow_unsupported_nodes=allow_unsupported_nodes, - skip_optimize=is_pre_quantized, + skip_optimize=config.skip_optimize, ) config_path.write_text(json.dumps(config.to_dict(), indent=2)) diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index 96b702cdf..f458ac27f 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -138,13 +138,7 @@ class WinMLBuildConfig: compile: WinMLCompileConfig | None = field(default_factory=WinMLCompileConfig) eval: WinMLEvaluationConfig | None = None auto: bool = True - # Stamped True by generate_*_build_config (or by the build_*_model - # entry-point defensive fallback) when the input ONNX is already - # quantized (QDQ or QOperator format). When True, the optimize stage - # is bypassed for downstream pipelines (no ORT graph optimization, - # no autoconf analyze loop). This is the SINGLE source of truth for - # "is this model pre-quantized?" — downstream stages must read this - # flag instead of calling ``is_quantized_onnx`` again. + # Skip ORT optimization. Pre-quantized inputs also clear ``quant``. skip_optimize: bool = False def __post_init__(self) -> None: @@ -597,8 +591,6 @@ def generate_onnx_build_config( if is_quantized_onnx(onnx_path_resolved): # Skip optimize+quantize, compile with resolved policy. - # ``skip_optimize`` is the single source of truth — downstream - # pipelines must read this flag and not re-detect. config.quant = None config.skip_optimize = True config.compile = resolved_compile diff --git a/src/winml/modelkit/models/auto.py b/src/winml/modelkit/models/auto.py index f24c8229b..922aa1254 100644 --- a/src/winml/modelkit/models/auto.py +++ b/src/winml/modelkit/models/auto.py @@ -25,6 +25,7 @@ import logging from collections.abc import Mapping +from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -42,6 +43,7 @@ from transformers import PretrainedConfig + from ..build import BuildResult from ..session import WinMLEPDevice from .winml.base import WinMLPreTrainedModel from .winml.composite_model import WinMLCompositeModel @@ -71,6 +73,15 @@ def _resolved_ep_short_name(ep_device: WinMLEPDevice) -> str: return short_ep_name(ep_device.device.ep_name) +@dataclass(frozen=True) +class _PretrainedArtifact: + result: "BuildResult" + build_config: WinMLBuildConfig + hf_config: "PretrainedConfig" + task: str + model_type: str + + # ============================================================================= # WinMLAutoModel Factory # ============================================================================= @@ -456,12 +467,83 @@ def from_pretrained( ) # ===================================================================== - # [1] CONFIG PHASE - Generate complete config with I/O specs (Lightweight, ~2s) + # [1]-[3] BUILD PHASE - Create or reuse artifact files. # ===================================================================== + artifact = cls._build_pretrained_artifact( + model_id, + task=task, + config=config, + ep_device=ep_device, + ep=ep, + precision=precision, + cache_dir=cache_dir, + use_cache=use_cache, + force_rebuild=force_rebuild, + trust_remote_code=trust_remote_code, + shape_config=shape_config, + model_type=model_type, + allow_unsupported_nodes=allow_unsupported_nodes, + no_compile=no_compile, + skip_optimize=skip_optimize, + hack_max_optim_iterations=hack_max_optim_iterations, + **kwargs, + ) + onnx_path = artifact.result.final_onnx_path + + # ===================================================================== + # [4] RUNTIME PHASE - Return inference wrapper + # ===================================================================== + winml_class = get_winml_class(artifact.model_type, artifact.task) + logger.info("Creating inference wrapper: %s", winml_class.__name__) + + model = winml_class( + onnx_path=onnx_path, + config=artifact.hf_config, + ep_device=ep_device, + provider_options=provider_options, + session_options=session_options, + ) + model._build_config = artifact.build_config + return model + + @classmethod + def _build_pretrained_artifact( + cls, + model_id_or_path: str | Path, + ep_device: WinMLEPDevice | None = None, + *, + device: str | None = None, + ep: str | None = None, + task: str | None = None, + config: WinMLBuildConfig | dict[str, Any] | None = None, + precision: str = "auto", + cache_dir: str | Path | None = None, + use_cache: bool = True, + force_rebuild: bool = False, + trust_remote_code: bool = False, + shape_config: dict | None = None, + model_type: str | None = None, + allow_unsupported_nodes: bool = False, + no_compile: bool = False, + skip_optimize: bool = False, + hack_max_optim_iterations: int = 3, + **_kwargs: Any, + ) -> _PretrainedArtifact: + from ..utils.model_input import resolve_model_input + + model_input = resolve_model_input(str(model_id_or_path)) + model_id = model_input.local_path or model_input.raw + + if ep_device is None: + from ..session import EPDeviceTarget, WinMLEPRegistry, resolve_device + + target = resolve_device( + EPDeviceTarget(ep=ep or "auto", device=(device or "auto").lower()) + ) + ep_device = WinMLEPRegistry.instance().auto_device(target) + from ..config import generate_hf_build_config - # Config fields merge on top of defaults, while the already resolved - # runtime target remains authoritative for quant/compile policy. build_config = generate_hf_build_config( model_id, task=task, @@ -476,13 +558,9 @@ def from_pretrained( no_compile=no_compile, ) - resolved_task = build_config.loader.task + resolved_task = cast("str", build_config.loader.task) logger.debug("Generated config with task: %s", resolved_task) - # ===================================================================== - # [2] CONFIG PHASE - Load only the HF config for wrapper compatibility. - # The build API performs its own cache check before it loads weights. - # ===================================================================== effective_trust = trust_remote_code or ( build_config.loader.trust_remote_code if build_config.loader else False ) @@ -495,19 +573,12 @@ def from_pretrained( model_id, trust_remote_code=effective_trust, ) - resolved_model_type = model_type or getattr(hf_config, "model_type", "unknown") + resolved_model_type = model_type or getattr(hf_config, "model_type", None) or "unknown" logger.debug("Model type: %s, task: %s", resolved_model_type, resolved_task) - config = build_config - task = resolved_task - - # ===================================================================== - # [3] CACHE + BUILD PHASE -- delegate to build_hf_model() - # ===================================================================== if use_cache: cache_dir_path = get_cache_dir(override=cache_dir) else: - # No cache -- use temp directory, always rebuild import tempfile cache_dir_path = Path(tempfile.mkdtemp(prefix="winml_")) @@ -515,8 +586,8 @@ def from_pretrained( logger.info("Cache disabled -- using temp directory: %s", cache_dir_path) cache_key = get_cache_key( - get_task_abbrev(cast("str", task)), - config.generate_cache_key(), + get_task_abbrev(resolved_task), + build_config.generate_cache_key(), _get_cache_build_controls( skip_optimize=skip_optimize, hack_max_optim_iterations=hack_max_optim_iterations, @@ -526,14 +597,13 @@ def from_pretrained( from ..build import build_hf_model - # An explicit EP takes precedence over the compile-derived provider. resolved_ep = ep - if resolved_ep is None and config.compile is not None: - resolved_ep = config.compile.ep_config.provider + if resolved_ep is None and build_config.compile is not None: + resolved_ep = build_config.compile.ep_config.provider if resolved_ep is None: resolved_ep = _resolved_ep_short_name(ep_device) result = build_hf_model( - config=config, + config=build_config, output_dir=output_dir, model_id=model_id, rebuild=force_rebuild, @@ -547,23 +617,13 @@ def from_pretrained( skip_optimize=skip_optimize, hack_max_optim_iterations=hack_max_optim_iterations, ) - onnx_path = result.final_onnx_path - - # ===================================================================== - # [4] RUNTIME PHASE - Return inference wrapper - # ===================================================================== - winml_class = get_winml_class(resolved_model_type, task) - logger.info("Creating inference wrapper: %s", winml_class.__name__) - - model = winml_class( - onnx_path=onnx_path, - config=hf_config, # HF PretrainedConfig for pipeline compatibility - ep_device=ep_device, - provider_options=provider_options, - session_options=session_options, + return _PretrainedArtifact( + result=result, + build_config=build_config, + hf_config=hf_config, + task=resolved_task, + model_type=resolved_model_type, ) - model._build_config = config # resolved build config (task, quant, compile) - return model @classmethod def supported_tasks(cls) -> list[str]: diff --git a/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py b/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py index 282017fe3..0fe0cf054 100644 --- a/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py +++ b/src/winml/modelkit/models/hf/qwen3/qwen_transformer_only.py @@ -333,6 +333,7 @@ def outputs(self) -> dict[str, dict[int, str]]: # Pure graph (no post-export RMSNorm fusion / matmul-add fusion): the default # WinMLOptimizationConfig() leaves every fusion flag off. optim=WinMLOptimizationConfig(), + skip_optimize=True, ) diff --git a/src/winml/modelkit/models/winml/genai_bundle.py b/src/winml/modelkit/models/winml/genai_bundle.py index 3031bdaa2..3c0976af1 100644 --- a/src/winml/modelkit/models/winml/genai_bundle.py +++ b/src/winml/modelkit/models/winml/genai_bundle.py @@ -24,7 +24,7 @@ from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING import onnx @@ -35,9 +35,6 @@ if TYPE_CHECKING: from collections.abc import Mapping - from .base import WinMLPreTrainedModel - from .composite_model import WinMLCompositeModel - # ========================================================================= # Recipe specs (pure data) @@ -219,6 +216,21 @@ def _node_summary(path: str | Path, *, top: int = 6) -> str: return f"{total} nodes, {len(counts)} op types: {top_ops}" +def _transformer_sub_model_task(transformer: GenaiTransformerSpec, sub_model: str) -> str: + from .composite_model import COMPOSITE_MODEL_REGISTRY + + composite_cls = COMPOSITE_MODEL_REGISTRY.get((transformer.model_type, transformer.task)) + if composite_cls is None: + raise ValueError( + f"No composite model registered for ({transformer.model_type!r}, {transformer.task!r})" + ) + component_task = composite_cls._SUB_MODEL_CONFIG.get(sub_model) + if component_task is None: + valid = list(composite_cls._SUB_MODEL_CONFIG) + raise ValueError(f"Unknown transformer sub-model {sub_model!r}; valid sub-models: {valid}") + return component_task + + def build_genai_bundle( model_id: str, output_dir: str | Path, @@ -299,11 +311,11 @@ def build_genai_bundle( from ..auto import WinMLAutoModel - # --- Transformer (context + iterator) via the composite builder --- + # --- Transformer (context + iterator) via the composite build spec --- _emit(f"building transformer stages (device={device}, precision={transformer_precision})") - built = WinMLAutoModel.from_pretrained( + context_onnx = WinMLAutoModel._build_pretrained_artifact( model_id, - task=transformer.task, + task=_transformer_sub_model_task(transformer, transformer.context_sub_model), model_type=transformer.model_type, device=device, precision=transformer_precision, @@ -312,20 +324,21 @@ def build_genai_bundle( use_cache=True, force_rebuild=force_rebuild, cache_dir=cache_dir, - sub_model_kwargs={ - transformer.context_sub_model: { - "shape_config": {"max_cache_len": max_cache_len, "seq_len": prefill_seq_len} - }, - transformer.iterator_sub_model: { - "shape_config": {"max_cache_len": max_cache_len, "seq_len": 1} - }, - }, - ) - # The composite transformer build returns a WinMLCompositeModel, whose - # sub_models attribute is not declared on the WinMLPreTrainedModel base. - sub_models = cast("WinMLCompositeModel", built).sub_models - context_onnx = Path(sub_models[transformer.context_sub_model].onnx_path) - iterator_onnx = Path(sub_models[transformer.iterator_sub_model].onnx_path) + shape_config={"max_cache_len": max_cache_len, "seq_len": prefill_seq_len}, + ).result.final_onnx_path + iterator_onnx = WinMLAutoModel._build_pretrained_artifact( + model_id, + task=_transformer_sub_model_task(transformer, transformer.iterator_sub_model), + model_type=transformer.model_type, + device=device, + precision=transformer_precision, + ep=transformer_ep, + no_compile=True, + use_cache=True, + force_rebuild=force_rebuild, + cache_dir=cache_dir, + shape_config={"max_cache_len": max_cache_len, "seq_len": 1}, + ).result.final_onnx_path for label, model_path in (("ctx", context_onnx), ("iter", iterator_onnx)): _emit(f" [{label}] {model_path}") _emit(f" {_node_summary(model_path)}") @@ -340,7 +353,7 @@ def build_genai_bundle( _emit(f"using provided {spec.role}: {companion_srcs[spec.role]}") continue _emit(f"building {spec.role} (model_type={spec.model_type}, precision={spec.precision})") - companion = WinMLAutoModel.from_pretrained( + companion_path = WinMLAutoModel._build_pretrained_artifact( model_id, task=spec.task, model_type=spec.model_type, @@ -351,10 +364,7 @@ def build_genai_bundle( use_cache=True, force_rebuild=force_rebuild, cache_dir=cache_dir, - ) - # A companion build produces a single-model WinMLPreTrainedModel (not a - # composite), so onnx_path is present; narrow the from_pretrained union. - companion_path = Path(cast("WinMLPreTrainedModel", companion).onnx_path) + ).result.final_onnx_path _emit(f" [{spec.role}] {companion_path}") _emit(f" {_node_summary(companion_path)}") companion_srcs[spec.role] = companion_path diff --git a/tests/unit/build/test_hf.py b/tests/unit/build/test_hf.py index 4c529edc5..2e9ea6b46 100644 --- a/tests/unit/build/test_hf.py +++ b/tests/unit/build/test_hf.py @@ -9,8 +9,11 @@ from __future__ import annotations +import gc import json +import weakref from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -224,6 +227,47 @@ def test_produces_build_result(self, tmp_path: Path, sample_config, mock_pipelin assert result.reused is False assert result.elapsed >= 0 + def test_releases_loaded_model_before_build_stages( + self, tmp_path: Path, sample_config_no_quant_compile + ) -> None: + model_ref: dict[str, weakref.ReferenceType[object]] = {} + + class _Model: + pass + + def load_model(*_args: object, **_kwargs: object) -> object: + model = _Model() + model_ref["model"] = weakref.ref(model) + return model + + def export_model(*_args: object, **kwargs: object) -> None: + Path(kwargs["output_path"]).write_text("mock") + + def run_stages(**kwargs: object) -> SimpleNamespace: + gc.collect() + assert model_ref["model"]() is None + Path(kwargs["final_path"]).write_text("mock") + return SimpleNamespace( + stages_completed=[], + stages_skipped=["optimize", "quantize", "compile"], + stage_timings={}, + analyze_iterations=0, + analyze_unsupported_nodes=0, + analyze_details=None, + quant_result=None, + ) + + with ( + patch("winml.modelkit.build.hf._load_model", new=load_model), + patch("winml.modelkit.build.hf.export_onnx", new=export_model), + patch("winml.modelkit.build.hf.run_build_stages", new=run_stages), + ): + build_hf_model( + config=sample_config_no_quant_compile, + output_dir=tmp_path, + model_id="test", + ) + def test_persists_config(self, tmp_path: Path, sample_config, mock_pipeline) -> None: build_hf_model(config=sample_config, output_dir=tmp_path, model_id="test") config_path = tmp_path / "winml_build_config.json" @@ -880,7 +924,7 @@ def test_post_export_qdq_runs_analyze_only( mock_pipeline["optimize"].assert_not_called() def test_skip_optimize_kwarg(self, tmp_path: Path, sample_config, mock_pipeline) -> None: - """skip_optimize=True forces optimize+quantize skip.""" + """skip_optimize=True skips optimize but still quantizes raw exports.""" mock_pipeline["is_quantized_onnx"].return_value = False output_dir = tmp_path / "output" @@ -891,9 +935,9 @@ def test_skip_optimize_kwarg(self, tmp_path: Path, sample_config, mock_pipeline) skip_optimize=True, ) assert "optimize" in result.stages_skipped - assert "quantize" in result.stages_skipped + assert "quantize" in result.stages_completed mock_pipeline["optimize"].assert_not_called() - mock_pipeline["quantize"].assert_not_called() + mock_pipeline["quantize"].assert_called_once() # ============================================================================= diff --git a/tests/unit/build/test_onnx.py b/tests/unit/build/test_onnx.py index 36ca38b9c..74f4beb82 100644 --- a/tests/unit/build/test_onnx.py +++ b/tests/unit/build/test_onnx.py @@ -423,7 +423,7 @@ def test_pre_quantized_runs_analyze_only( def test_skip_optimize_kwarg( self, tmp_path: Path, fake_onnx: Path, sample_onnx_config, mock_onnx_pipeline ) -> None: - """skip_optimize=True forces optimize+quantize skip even without QDQ.""" + """skip_optimize=True skips optimize but still quantizes raw models.""" mock_onnx_pipeline["is_quantized_onnx"].return_value = False output_dir = tmp_path / "output" @@ -434,9 +434,9 @@ def test_skip_optimize_kwarg( skip_optimize=True, ) assert "optimize" in result.stages_skipped - assert "quantize" in result.stages_skipped + assert "quantize" in result.stages_completed mock_onnx_pipeline["optimize"].assert_not_called() - mock_onnx_pipeline["quantize"].assert_not_called() + mock_onnx_pipeline["quantize"].assert_called_once() # ============================================================================= diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index ac49e615a..46eb62f9f 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2051,6 +2051,49 @@ def test_auto_config_onnx_model_uses_single_generate_call( assert mock_gen.call_count == 1 assert mock_gen.call_args.kwargs["onnx_path"] == str(onnx_file) + def test_explicit_device_preserves_quant_for_raw_skip_optimize_config( + self, tmp_path: Path, mock_run_single_build: MagicMock + ) -> None: + from winml.modelkit.config import WinMLBuildConfig + + cfg = WinMLBuildConfig.from_dict( + { + "loader": {"task": "text-generation"}, + "export": {"opset_version": 18}, + "optim": {}, + "quant": { + "mode": "static", + "samples": 1, + "task": "text-generation", + "model_id": "Qwen/Qwen3-1.7B", + }, + "compile": None, + "skip_optimize": True, + } + ) + + with ( + patch("winml.modelkit.config.generate_build_config", return_value=cfg), + patch( + "winml.modelkit.commands.build._validate_loader_tasks_for_model", + return_value=None, + ), + ): + result = _invoke( + [ + "-m", + "Qwen/Qwen3-1.7B", + "-o", + str(tmp_path / "out"), + "--device", + "npu", + "--no-compile", + ] + ) + + assert result.exit_code == 0, result.output + assert mock_run_single_build.call_args.kwargs["config"].quant is not None + def test_shape_config_rejected_for_onnx_input( self, tmp_path: Path, mock_run_single_build: MagicMock ): @@ -2084,7 +2127,7 @@ class TestBuildOnnxPipelineRegressions: def test_pre_quantized_stamp_clears_quant_when_already_skip_optimize( self, tmp_path: Path ) -> None: - """skip_optimize is the pre-quantized invariant and must imply no quant.""" + """Pre-quantized detection clears quant even when skip_optimize is already set.""" from winml.modelkit.build.common import ensure_pre_quantized_stamped from winml.modelkit.config import WinMLBuildConfig from winml.modelkit.quant.config import WinMLQuantizationConfig @@ -2094,18 +2137,18 @@ def test_pre_quantized_stamp_clears_quant_when_already_skip_optimize( config = WinMLBuildConfig(quant=WinMLQuantizationConfig()) config.skip_optimize = True - ensure_pre_quantized_stamped(config, onnx_file) + with patch("winml.modelkit.build.common.is_quantized_onnx", return_value=True): + ensure_pre_quantized_stamped(config, onnx_file) assert config.skip_optimize is True assert config.quant is None @patch("winml.modelkit.quant.quantize_onnx") - def test_quantize_stage_skips_when_config_skip_optimize_is_set( + def test_quantize_stage_runs_when_raw_model_only_skips_optimize( self, mock_quantize: MagicMock, tmp_path: Path, ) -> None: - """The quantize stage must not double-quantize pre-quantized inputs.""" from winml.modelkit.commands.build import _run_quantize_stage from winml.modelkit.config import WinMLBuildConfig from winml.modelkit.quant.config import WinMLQuantizationConfig @@ -2115,18 +2158,24 @@ def test_quantize_stage_skips_when_config_skip_optimize_is_set( config = WinMLBuildConfig(quant=WinMLQuantizationConfig()) config.skip_optimize = True timings: list[tuple[str, float | None]] = [] + quantized_path = tmp_path / "quantized.onnx" + mock_result = MagicMock() + mock_result.success = True + mock_result.output_path = quantized_path + mock_result.errors = [] + mock_quantize.return_value = mock_result result = _run_quantize_stage( config=config, current_path=input_path, - quantized_path=tmp_path / "quantized.onnx", + quantized_path=quantized_path, stage_timings=timings, ) - assert result == input_path - assert config.quant is None - assert timings == [] - mock_quantize.assert_not_called() + assert result == quantized_path + assert config.quant is not None + assert timings and timings[0][0] == "Quantize" + mock_quantize.assert_called_once() def test_pre_quantized_stamp_runs_before_optimize(self, tmp_path: Path) -> None: """_build_onnx_pipeline must stamp config before optimize/quantize stages. diff --git a/tests/unit/models/auto/test_from_pretrained_ep.py b/tests/unit/models/auto/test_from_pretrained_ep.py index c4cfcea90..abc86415c 100644 --- a/tests/unit/models/auto/test_from_pretrained_ep.py +++ b/tests/unit/models/auto/test_from_pretrained_ep.py @@ -285,6 +285,47 @@ def reuse_build(**kwargs: Any) -> SimpleNamespace: assert received.get("pytorch_model") is None +def test_build_pretrained_artifact_does_not_create_runtime_wrapper( + monkeypatch: pytest.MonkeyPatch, tmp_path +) -> None: + from winml.modelkit.models import WinMLAutoModel + from winml.modelkit.models import auto as auto_module + + build_config = MagicMock() + build_config.loader.task = "image-classification" + build_config.loader.trust_remote_code = False + build_config.compile = None + build_config.generate_cache_key.return_value = "cache-key" + hf_config = SimpleNamespace(model_type="unit-type") + build_result = SimpleNamespace(final_onnx_path=tmp_path / "cached.onnx") + build_result.final_onnx_path.write_bytes(b"cached") + + monkeypatch.setattr( + "winml.modelkit.config.generate_hf_build_config", + lambda *_args, **_kwargs: build_config, + ) + monkeypatch.setattr( + "winml.modelkit.loader.load_hf_config", + lambda *_args, **_kwargs: hf_config, + ) + monkeypatch.setattr(auto_module, "get_cache_dir", lambda **_kwargs: tmp_path) + monkeypatch.setattr(auto_module, "get_model_dir", lambda *_args, **_kwargs: tmp_path) + monkeypatch.setattr("winml.modelkit.build.build_hf_model", lambda **_kwargs: build_result) + monkeypatch.setattr( + auto_module, + "get_winml_class", + lambda *_args: pytest.fail("artifact build must not create a runtime wrapper"), + ) + + result = WinMLAutoModel._build_pretrained_artifact( + "unit/model", + ep_device=MagicMock(device=MagicMock(device_type="CPU", ep_name="CPUExecutionProvider")), + task="image-classification", + ) + + assert result.result is build_result + + def test_cache_key_matches_shared_helper(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: """API cache keys must match the shared helper and ignore non-artifact flags.""" from winml.modelkit.cache import get_cache_key diff --git a/tests/unit/models/qwen3/test_qwen3_modeling.py b/tests/unit/models/qwen3/test_qwen3_modeling.py index 6826a95bf..6fe6acacf 100644 --- a/tests/unit/models/qwen3/test_qwen3_modeling.py +++ b/tests/unit/models/qwen3/test_qwen3_modeling.py @@ -18,6 +18,7 @@ import torch from winml.modelkit.models.hf.qwen3.qwen3_modeling import WinMLQwen3Attention +from winml.modelkit.models.hf.qwen3.qwen_transformer_only import QWEN_TRANSFORMER_ONLY_CONFIG # --------------------------------------------------------------------------- @@ -122,6 +123,12 @@ def test_fallback_to_max_position_embeddings(self): WinMLQwen3Attention.prepare_for_onnx_export(mod, matmul_to_conv=False) assert mod._max_rope_len == 512 + +class TestTransformerOnlyBuildConfig: + def test_skips_ort_optimize_without_disabling_quantization(self): + assert QWEN_TRANSFORMER_ONLY_CONFIG.skip_optimize is True + assert QWEN_TRANSFORMER_ONLY_CONFIG.quant is not None + def test_max_rope_len_is_plain_int(self): """_max_rope_len must be a plain int, not a tensor or other type.""" mod = _make_attention_module() diff --git a/tests/unit/models/winml/test_genai_bundle_orchestrator.py b/tests/unit/models/winml/test_genai_bundle_orchestrator.py index 194a62f79..20d3ddd2b 100644 --- a/tests/unit/models/winml/test_genai_bundle_orchestrator.py +++ b/tests/unit/models/winml/test_genai_bundle_orchestrator.py @@ -2,17 +2,13 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -"""Unit tests for the genai-bundle orchestrator (``build_genai_bundle``). - -Wiring-only: ``WinMLAutoModel.from_pretrained`` and the recipe assembler are -stubbed so no model is downloaded. The tests verify the orchestrator maps -recipe data + caller overrides onto the component builders and the assembler, -and that it stays architecture-agnostic (a synthetic, non-Qwen recipe is used). -""" +"""Unit tests for the genai-bundle orchestrator (``build_genai_bundle``).""" from __future__ import annotations from pathlib import Path +from types import SimpleNamespace +from typing import ClassVar import pytest from onnx import TensorProto, helper, save @@ -35,12 +31,6 @@ def _write_tiny_onnx(path: Path) -> None: save(helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]), str(path)) -class _StubModel: - def __init__(self, onnx_path: Path, sub_models: dict | None = None) -> None: - self.onnx_path = str(onnx_path) - self.sub_models = sub_models or {} - - def _dummy_pass(model): return model @@ -76,6 +66,14 @@ def _by_model_type(calls: list[dict], model_type: str) -> dict: return next(c for c in calls if c.get("model_type") == model_type) +def _by_model_type_and_task(calls: list[dict], model_type: str, task: str) -> dict: + return next(c for c in calls if c.get("model_type") == model_type and c.get("task") == task) + + +def _by_transformer_task(calls: list[dict], task: str) -> dict: + return _by_model_type_and_task(calls, "T-transformer", task) + + @pytest.fixture def harness(tmp_path, monkeypatch): onnx_file = tmp_path / "tiny.onnx" @@ -83,16 +81,32 @@ def harness(tmp_path, monkeypatch): calls: list[dict] = [] - def fake_from_pretrained(model_id, **kwargs): + def fake_build_artifact(model_id, **kwargs): calls.append({"model_id": model_id, **kwargs}) - if kwargs.get("model_type") == "T-transformer": - return _StubModel( - onnx_file, - sub_models={"ctx_sub": _StubModel(onnx_file), "iter_sub": _StubModel(onnx_file)}, - ) - return _StubModel(onnx_file) + return SimpleNamespace(result=SimpleNamespace(final_onnx_path=onnx_file)) - monkeypatch.setattr(WinMLAutoModel, "from_pretrained", staticmethod(fake_from_pretrained)) + monkeypatch.setattr( + WinMLAutoModel, "_build_pretrained_artifact", staticmethod(fake_build_artifact) + ) + monkeypatch.setattr( + WinMLAutoModel, + "from_pretrained", + staticmethod(lambda *_args, **_kwargs: pytest.fail("runtime wrapper was created")), + ) + + import winml.modelkit.models.winml.composite_model as cm_mod + + class _FakeTransformerComposite: + _SUB_MODEL_CONFIG: ClassVar[dict[str, str]] = { + "ctx_sub": "feature-extraction", + "iter_sub": "text2text-generation", + } + + monkeypatch.setattr( + cm_mod, + "COMPOSITE_MODEL_REGISTRY", + {("T-transformer", "text-generation"): _FakeTransformerComposite}, + ) assemble_kwargs: dict = {} @@ -126,21 +140,28 @@ def test_transformer_built_with_recipe_defaults(harness): out = harness["tmp_path"] / "bundle" build_genai_bundle("some/model", out, harness["recipe"], ep="qnn", device="npu") - t = _by_model_type(harness["calls"], "T-transformer") - assert t["task"] == "text-generation" - assert t["device"] == "npu" - assert t["precision"] == "w8a16" - assert t["ep"] == "QNNExecutionProvider" # normalized from short "qnn" - assert t["sub_model_kwargs"]["ctx_sub"]["shape_config"] == { + ctx = _by_transformer_task(harness["calls"], "feature-extraction") + it = _by_transformer_task(harness["calls"], "text2text-generation") + for transformer_call in (ctx, it): + assert transformer_call["device"] == "npu" + assert transformer_call["precision"] == "w8a16" + assert transformer_call["ep"] == "QNNExecutionProvider" + assert ctx["shape_config"] == { "max_cache_len": 2048, "seq_len": 64, } - assert t["sub_model_kwargs"]["iter_sub"]["shape_config"] == { + assert it["shape_config"] == { "max_cache_len": 2048, "seq_len": 1, } +def test_components_use_artifact_builds_without_runtime_wrappers(harness): + build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"]) + + assert all("build_only" not in call for call in harness["calls"]) + + def test_companions_built_on_cpu(harness): build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"]) emb = _by_model_type(harness["calls"], "T-emb") @@ -173,7 +194,8 @@ def test_precision_override_only_affects_transformer(harness): # ``T-transformer`` has no registered quant finalizer, so a precision # override is honored (forwarded to the transformer build). build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"], precision="w4a16") - assert _by_model_type(harness["calls"], "T-transformer")["precision"] == "w4a16" + assert _by_transformer_task(harness["calls"], "feature-extraction")["precision"] == "w4a16" + assert _by_transformer_task(harness["calls"], "text2text-generation")["precision"] == "w4a16" assert _by_model_type(harness["calls"], "T-emb")["precision"] == "fp32" assert _by_model_type(harness["calls"], "T-lmh")["precision"] == "w4a32" @@ -201,7 +223,8 @@ def test_matching_precision_override_allowed_when_finalizer_pinned(harness, monk monkeypatch.setattr(gb, "has_quant_finalizer", lambda model_type: True) build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"], precision="w8a16") - assert _by_model_type(harness["calls"], "T-transformer")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "feature-extraction")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "text2text-generation")["precision"] == "w8a16" def test_no_precision_override_uses_recipe_default_when_finalizer_pinned(harness, monkeypatch): @@ -210,7 +233,8 @@ def test_no_precision_override_uses_recipe_default_when_finalizer_pinned(harness monkeypatch.setattr(gb, "has_quant_finalizer", lambda model_type: True) build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"]) - assert _by_model_type(harness["calls"], "T-transformer")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "feature-extraction")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "text2text-generation")["precision"] == "w8a16" def test_auto_precision_override_allowed_when_finalizer_pinned(harness, monkeypatch): @@ -225,7 +249,8 @@ def test_auto_precision_override_allowed_when_finalizer_pinned(harness, monkeypa monkeypatch.setattr(gb, "has_quant_finalizer", lambda model_type: True) build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"], precision="auto") - assert _by_model_type(harness["calls"], "T-transformer")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "feature-extraction")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "text2text-generation")["precision"] == "w8a16" def test_case_variant_precision_override_allowed_when_finalizer_pinned(harness, monkeypatch): @@ -239,19 +264,21 @@ def test_case_variant_precision_override_allowed_when_finalizer_pinned(harness, monkeypatch.setattr(gb, "has_quant_finalizer", lambda model_type: True) build_genai_bundle("m", harness["tmp_path"] / "b", harness["recipe"], precision="W8A16") - assert _by_model_type(harness["calls"], "T-transformer")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "feature-extraction")["precision"] == "w8a16" + assert _by_transformer_task(harness["calls"], "text2text-generation")["precision"] == "w8a16" def test_length_overrides_flow_to_shapes_and_assembler(harness): build_genai_bundle( "m", harness["tmp_path"] / "b", harness["recipe"], max_cache_len=1024, prefill_seq_len=32 ) - t = _by_model_type(harness["calls"], "T-transformer") - assert t["sub_model_kwargs"]["ctx_sub"]["shape_config"] == { + ctx = _by_transformer_task(harness["calls"], "feature-extraction") + it = _by_transformer_task(harness["calls"], "text2text-generation") + assert ctx["shape_config"] == { "max_cache_len": 1024, "seq_len": 32, } - assert t["sub_model_kwargs"]["iter_sub"]["shape_config"] == { + assert it["shape_config"] == { "max_cache_len": 1024, "seq_len": 1, } From 0f72a3d23687c6d3839ce3be83b2788ef2642cfc Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Wed, 29 Jul 2026 20:53:41 +0800 Subject: [PATCH 2/5] fix: avoid WinML catalog shutdown cleanup --- src/winml/modelkit/ep_path.py | 16 --------- .../unit/ep_path/test_winml_catalog_source.py | 36 +++++-------------- 2 files changed, 9 insertions(+), 43 deletions(-) diff --git a/src/winml/modelkit/ep_path.py b/src/winml/modelkit/ep_path.py index 81691d582..66e92ec7d 100644 --- a/src/winml/modelkit/ep_path.py +++ b/src/winml/modelkit/ep_path.py @@ -40,7 +40,6 @@ from __future__ import annotations -import atexit import dataclasses import functools import logging @@ -713,24 +712,10 @@ def iter_eps(self) -> Iterable[str]: return self.dll_patterns.keys() -# --------------------------------------------------------------------------- -# windowsml EpCatalog singleton. -# --------------------------------------------------------------------------- # Lazy, process-wide initialization of the ``windowsml.EpCatalog``. -# We register an ``atexit`` cleanup so the catalog is closed on interpreter -# shutdown. ``__del__`` is intentionally NOT used — Python does not -# guarantee it is invoked on shutdown. _winml_catalog_warned_keys: set[str] = set() -def _release_winml_catalog(catalog: Any) -> None: - """Close the process-wide Windows ML catalog during interpreter shutdown.""" - try: - catalog.close() - except Exception as e: # pragma: no cover - shutdown best-effort - logger.debug("Windows ML catalog cleanup raised: %s", e) - - @functools.cache def _get_catalog() -> Any | None: """Return the cached ``windowsml.EpCatalog`` or ``None``. @@ -762,7 +747,6 @@ def _get_catalog() -> Any | None: logger.warning("WinMLCatalogSource: EpCatalog() failed: %s", e) return None - atexit.register(_release_winml_catalog, catalog) return catalog diff --git a/tests/unit/ep_path/test_winml_catalog_source.py b/tests/unit/ep_path/test_winml_catalog_source.py index 345e40e3c..13acfe9a7 100644 --- a/tests/unit/ep_path/test_winml_catalog_source.py +++ b/tests/unit/ep_path/test_winml_catalog_source.py @@ -15,12 +15,13 @@ Also covers: - The default EP source list includes the 5 ``WinMLCatalogSource`` rows with the canonical EP names from the design doc. - - ``atexit`` cleanup is registered exactly once across many - ``_get_catalog()`` calls. + - ``_get_catalog()`` does not register native catalog cleanup at process + exit. """ from __future__ import annotations +import atexit import logging import os import sys @@ -542,27 +543,26 @@ def test_is_ready(value: Any, expected: bool) -> None: # --------------------------------------------------------------------------- -# atexit cleanup. +# catalog lifecycle. # --------------------------------------------------------------------------- -class TestAtexitCleanup: - """The catalog is registered for cleanup exactly once.""" +class TestCatalogLifecycle: + """The catalog singleton avoids process-exit native cleanup hooks.""" - def test_atexit_registered_once_across_multiple_calls( + def test_get_catalog_does_not_register_atexit_cleanup( self, reset_catalog_singleton: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, ) -> None: - # Track atexit.register calls within ep_path. registered: list[Any] = [] def fake_register(func: Any, *args: Any, **kwargs: Any) -> Any: registered.append((func, args, kwargs)) return func - monkeypatch.setattr(_ep.atexit, "register", fake_register) + monkeypatch.setattr(atexit, "register", fake_register) # Install a working fake module. dll = tmp_path / "x.dll" @@ -586,22 +586,4 @@ def fake_register(func: Any, *args: Any, **kwargs: Any) -> Any: assert c1 is not None assert c2 is c1 assert c3 is c1 - # Exactly one atexit registration. - cleanup_callbacks = [r for r in registered if r[0] is _ep._release_winml_catalog] - assert len(cleanup_callbacks) == 1 - # The catalog object itself is what gets registered for cleanup. - assert cleanup_callbacks[0][1][0] is catalog - - def test_release_catalog_swallows_exceptions(self, caplog: pytest.LogCaptureFixture) -> None: - # Cleanup must not propagate exceptions during interpreter shutdown. - class _BoomCatalog: - def close(self) -> None: - raise RuntimeError("cleanup failure") - - # Should not raise. - _ep._release_winml_catalog(_BoomCatalog()) - - def test_release_catalog_calls_close(self) -> None: - catalog = _FakeCatalog([]) - _ep._release_winml_catalog(catalog) - assert catalog.closed is True + assert registered == [] From e921bead2ed60241c73e5af2ff1e8e8c44187ae1 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 30 Jul 2026 11:16:27 +0800 Subject: [PATCH 3/5] fix: preserve skip optimize quant policy --- src/winml/modelkit/commands/build.py | 10 ++++++- src/winml/modelkit/config/build.py | 1 + tests/unit/commands/test_build.py | 42 ++++++++++++++++++++++++++++ tests/unit/config/test_build.py | 41 +++++++++++++++++++++++++++ 4 files changed, 93 insertions(+), 1 deletion(-) diff --git a/src/winml/modelkit/commands/build.py b/src/winml/modelkit/commands/build.py index a90b86428..b679d6095 100644 --- a/src/winml/modelkit/commands/build.py +++ b/src/winml/modelkit/commands/build.py @@ -1040,6 +1040,14 @@ def build( # to honor the requested policy. fp16/fp32 clear quant; npu/int8 etc set it. if cli_utils.is_cli_provided(ctx, "device") or cli_utils.is_cli_provided(ctx, "precision"): from ..compiler.configs import WinMLCompileConfig + from ..onnx import is_quantized_onnx + + is_pre_quantized_onnx_input = ( + model_input is not None + and model_input.kind is ModelInputKind.ONNX_FILE + and model is not None + and is_quantized_onnx(Path(model)) + ) def _patch_device(cfg: WinMLBuildConfig) -> None: from ..config import resolve_quant_compile_config @@ -1047,7 +1055,7 @@ def _patch_device(cfg: WinMLBuildConfig) -> None: resolved_quant, _ = resolve_quant_compile_config( device=device, precision=precision, ep=ep_value ) - if not quant or resolved_quant is None or (cfg.skip_optimize and cfg.quant is None): + if not quant or resolved_quant is None or is_pre_quantized_onnx_input: cfg.quant = None elif cfg.quant is None: # Populate calibration identifiers from the loader/model diff --git a/src/winml/modelkit/config/build.py b/src/winml/modelkit/config/build.py index f458ac27f..49a96f641 100644 --- a/src/winml/modelkit/config/build.py +++ b/src/winml/modelkit/config/build.py @@ -1441,6 +1441,7 @@ def _assemble_config( optim=optim_config, quant=quant_config, compile=compile_config, + skip_optimize=registered.skip_optimize if registered else False, ) diff --git a/tests/unit/commands/test_build.py b/tests/unit/commands/test_build.py index 46eb62f9f..af27b01d6 100644 --- a/tests/unit/commands/test_build.py +++ b/tests/unit/commands/test_build.py @@ -2094,6 +2094,48 @@ def test_explicit_device_preserves_quant_for_raw_skip_optimize_config( assert result.exit_code == 0, result.output assert mock_run_single_build.call_args.kwargs["config"].quant is not None + def test_explicit_device_populates_quant_for_raw_skip_optimize_config( + self, tmp_path: Path, mock_run_single_build: MagicMock + ) -> None: + from winml.modelkit.config import WinMLBuildConfig + + cfg = WinMLBuildConfig.from_dict( + { + "loader": {"task": "text-generation"}, + "export": {"opset_version": 18}, + "optim": {}, + "quant": None, + "compile": None, + "skip_optimize": True, + } + ) + + with ( + patch("winml.modelkit.config.generate_build_config", return_value=cfg), + patch( + "winml.modelkit.commands.build._validate_loader_tasks_for_model", + return_value=None, + ), + ): + result = _invoke( + [ + "-m", + "Qwen/Qwen3-1.7B", + "-o", + str(tmp_path / "out"), + "--device", + "npu", + "--no-compile", + ] + ) + + assert result.exit_code == 0, result.output + patched = mock_run_single_build.call_args.kwargs["config"] + assert patched.skip_optimize is True + assert patched.quant is not None + assert patched.quant.task == "text-generation" + assert patched.quant.model_id == "Qwen/Qwen3-1.7B" + def test_shape_config_rejected_for_onnx_input( self, tmp_path: Path, mock_run_single_build: MagicMock ): diff --git a/tests/unit/config/test_build.py b/tests/unit/config/test_build.py index 300f99662..ac1345878 100644 --- a/tests/unit/config/test_build.py +++ b/tests/unit/config/test_build.py @@ -926,6 +926,47 @@ def test_registry_without_export_falls_through_to_optimum( # Optimum SHOULD have been called mock_optimum.assert_called_once() + def test_registry_skip_optimize_survives_assembly( + self, + mock_hf_config: MagicMock, + mock_model_class: MagicMock, + mock_export_config: WinMLExportConfig, + ) -> None: + """Registered skip_optimize is carried into generated configs.""" + registry_config = WinMLBuildConfig( + export=WinMLExportConfig(dynamo=False, opset_version=18), + optim=WinMLOptimizationConfig(), + skip_optimize=True, + ) + loader_config = WinMLLoaderConfig( + task="text-generation", + model_class="Qwen3ForCausalLM", + model_type="qwen3_transformer_only", + ) + mock_hf_config.model_type = "qwen3" + + with ( + patch( + "winml.modelkit.config.build.resolve_loader_config", + return_value=(loader_config, mock_hf_config, mock_model_class, MagicMock()), + ), + patch( + "winml.modelkit.config.build._resolve_export_config_from_specs", + return_value=mock_export_config, + ), + patch( + "winml.modelkit.models.hf.MODEL_BUILD_CONFIGS", + {"qwen3-transformer-only": registry_config}, + ), + ): + result = generate_hf_build_config( + "Qwen/Qwen3-1.7B", + model_type="qwen3_transformer_only", + task="text-generation", + ) + + assert result.skip_optimize is True + def test_registry_with_none_input_tensors_falls_through( self, mock_hf_config: MagicMock, From 73ae442402ad158bbac560d1eae2aa47ba042eac Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 30 Jul 2026 12:01:04 +0800 Subject: [PATCH 4/5] fix: release exported pytorch model reference --- src/winml/modelkit/build/hf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/winml/modelkit/build/hf.py b/src/winml/modelkit/build/hf.py index 6e64e3d34..e2f624735 100644 --- a/src/winml/modelkit/build/hf.py +++ b/src/winml/modelkit/build/hf.py @@ -233,7 +233,7 @@ def _name(base: str) -> str: stage_timings["export"] = time.monotonic() - t0 stages_completed.append("export") logger.info("Export done (%.1fs) -> %s", stage_timings["export"], export_path) - pytorch_model = None + del pytorch_model gc.collect() # ========================================================================= From 749f024eb46291bdbf7083657b6c0b1e1a7969ef Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 30 Jul 2026 12:23:54 +0800 Subject: [PATCH 5/5] fix: disarm WinML catalog at shutdown --- src/winml/modelkit/ep_path.py | 10 ++++++++++ .../unit/ep_path/test_winml_catalog_source.py | 19 ++++++++++++++----- 2 files changed, 24 insertions(+), 5 deletions(-) diff --git a/src/winml/modelkit/ep_path.py b/src/winml/modelkit/ep_path.py index 66e92ec7d..d5a0c3699 100644 --- a/src/winml/modelkit/ep_path.py +++ b/src/winml/modelkit/ep_path.py @@ -40,6 +40,7 @@ from __future__ import annotations +import atexit import dataclasses import functools import logging @@ -716,6 +717,14 @@ def iter_eps(self) -> Iterable[str]: _winml_catalog_warned_keys: set[str] = set() +def _disarm_winml_catalog(catalog: Any) -> None: + """Prevent ``EpCatalog.__del__`` from releasing native state during shutdown.""" + if not hasattr(catalog, "_handle"): + logger.debug("Windows ML catalog cleanup skipped: catalog has no _handle") + return + catalog._handle = None + + @functools.cache def _get_catalog() -> Any | None: """Return the cached ``windowsml.EpCatalog`` or ``None``. @@ -747,6 +756,7 @@ def _get_catalog() -> Any | None: logger.warning("WinMLCatalogSource: EpCatalog() failed: %s", e) return None + atexit.register(_disarm_winml_catalog, catalog) return catalog diff --git a/tests/unit/ep_path/test_winml_catalog_source.py b/tests/unit/ep_path/test_winml_catalog_source.py index 13acfe9a7..29771fad1 100644 --- a/tests/unit/ep_path/test_winml_catalog_source.py +++ b/tests/unit/ep_path/test_winml_catalog_source.py @@ -15,8 +15,8 @@ Also covers: - The default EP source list includes the 5 ``WinMLCatalogSource`` rows with the canonical EP names from the design doc. - - ``_get_catalog()`` does not register native catalog cleanup at process - exit. + - ``_get_catalog()`` disarms the native catalog handle at process exit + without calling into native release during interpreter shutdown. """ from __future__ import annotations @@ -90,6 +90,7 @@ def __init__( self._providers = providers self._find_raises = find_raises self.closed = False + self._handle: object | None = object() def find_all_providers(self) -> list[_FakeProvider]: if self._find_raises is not None: @@ -548,9 +549,9 @@ def test_is_ready(value: Any, expected: bool) -> None: class TestCatalogLifecycle: - """The catalog singleton avoids process-exit native cleanup hooks.""" + """The catalog singleton avoids process-exit native release calls.""" - def test_get_catalog_does_not_register_atexit_cleanup( + def test_get_catalog_registers_atexit_handle_disarm_without_close( self, reset_catalog_singleton: None, monkeypatch: pytest.MonkeyPatch, @@ -586,4 +587,12 @@ def fake_register(func: Any, *args: Any, **kwargs: Any) -> Any: assert c1 is not None assert c2 is c1 assert c3 is c1 - assert registered == [] + assert len(registered) == 1 + + cleanup, args, kwargs = registered[0] + assert args == (catalog,) + assert kwargs == {} + cleanup(*args, **kwargs) + + assert catalog._handle is None + assert catalog.closed is False