Fix Qwen3 genai bundle component builds

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:

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)
+    del pytorch_model
+    gc.collect()
 
     # =========================================================================
     # [3]-[6] OPTIMIZE -> QUANTIZE -> COMPILE -> FINALIZE

src/winml/modelkit/commands/build.py

@@ -1040,14 +1040,22 @@ 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
 
                 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 is_pre_quantized_onnx_input:
                     cfg.quant = None
                 elif cfg.quant is None:
                     # Populate calibration identifiers from the loader/model
@@ -1641,8 +1649,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 +1803,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 +2179,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 +2191,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))

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
@@ -1449,6 +1441,7 @@ def _assemble_config(
         optim=optim_config,
         quant=quant_config,
         compile=compile_config,
+        skip_optimize=registered.skip_optimize if registered else False,
     )
 
 

src/winml/modelkit/ep_path.py

@@ -713,22 +713,16 @@ 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)
+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
@@ -762,7 +756,7 @@ def _get_catalog() -> Any | None:
         logger.warning("WinMLCatalogSource: EpCatalog() failed: %s", e)
         return None
 
-    atexit.register(_release_winml_catalog, catalog)
+    atexit.register(_disarm_winml_catalog, catalog)
     return catalog