diff --git a/docs/commands/compile.md b/docs/commands/compile.md index c8aa3d214..c57fa63f9 100644 --- a/docs/commands/compile.md +++ b/docs/commands/compile.md @@ -43,6 +43,11 @@ eliminating graph partitioning at load time. An optional post-compilation validation pass runs a forward pass through the target EP; skip it with `--no-validate` when the target hardware is absent. +If a provider option names a compiler input file, list its key in +`compile.provider_option_file_keys` in the JSON config. The CLI canonicalizes +that option path and fingerprints its contents for EPContext cache identity; +other provider-option strings are always passed through unchanged. + ## Examples ```bash diff --git a/docs/reference/index.md b/docs/reference/index.md index 30541b579..ebb3e99e0 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -122,6 +122,7 @@ Set to `null` to skip compilation. | `ep_config.embed_context` | `bool` | `false` | Embed binary in ONNX (true) or external .bin (false). | | `ep_config.compiler` | `str` | `"ort"` | Compiler backend: `ort` or `qairt`. | | `ep_config.provider_options` | `dict` | `{}` | EP-specific options. | +| `ep_config.provider_option_file_keys` | `list[str]` | `[]` | Keys in `provider_options` whose values are input files. Declared paths are canonicalized and content-fingerprinted for EPContext cache identity. | | `ep_config.qnn_sdk_root` | `str \| null` | `null` | QNN SDK path for QAIRT compiler backend. | | `validate` | `bool` | `true` | Validate compiled model. | | `verbose` | `bool` | `false` | Verbose compilation logging. | diff --git a/pyproject.toml b/pyproject.toml index 65a07ab62..63f1b7925 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,6 +37,7 @@ dependencies = [ "diffusers>=0.36", "evaluate>=0.4.6", "fastapi>=0.135.3", + "filelock>=3.20", "hf_xet>=1.1.10", "httpx>=0.24.0", "jsonschema>=4.23", diff --git a/src/winml/modelkit/commands/compile.py b/src/winml/modelkit/commands/compile.py index 42e69d304..2c566e92e 100644 --- a/src/winml/modelkit/commands/compile.py +++ b/src/winml/modelkit/commands/compile.py @@ -183,6 +183,7 @@ def compile( # Apply build config defaults (CLI explicit options take precedence). # Read raw JSON so missing keys are distinguishable from dataclass defaults. config_provider_options: dict[str, str] = {} + config_provider_option_file_keys: set[str] = set() if config_file is not None: try: build_cfg, raw_cfg = cli_utils.load_build_config(config_file) @@ -197,6 +198,8 @@ def compile( # EP provider options (e.g. QNN htp_arch/soc_model/vtcm_mb) for the compile session. if "provider_options" in cc: config_provider_options = dict(cc["provider_options"]) + if "provider_option_file_keys" in cc: + config_provider_option_file_keys = set(cc["provider_option_file_keys"]) if not cli_utils.is_cli_provided(ctx, "device"): if configured_target is not None: device = configured_target.device @@ -309,6 +312,8 @@ def compile( # for duplicate keys. if config_provider_options: config.ep_config.provider_options.update(config_provider_options) + if config_provider_option_file_keys: + config.ep_config.provider_option_file_keys.update(config_provider_option_file_keys) if cli_provider_options: config.ep_config.provider_options.update(cli_provider_options) diff --git a/src/winml/modelkit/compiler/configs.py b/src/winml/modelkit/compiler/configs.py index dfe4db9e5..85df35909 100644 --- a/src/winml/modelkit/compiler/configs.py +++ b/src/winml/modelkit/compiler/configs.py @@ -36,6 +36,7 @@ class EPConfig: Attributes: provider: Target execution provider (qnn, cpu, cuda, dml) provider_options: EP-specific options as key=value dict + provider_option_file_keys: Provider option keys whose values are file paths enable_ep_context: Generate EPContext model with pre-compiled graph embed_context: Embed context in ONNX (True) or external .bin file (False) compiler: Compiler backend ("ort", "ort_session", or "qairt"). @@ -51,6 +52,7 @@ class EPConfig: compiler: CompilerName = "ort" qnn_sdk_root: Path | None = None device: str = "auto" + provider_option_file_keys: set[str] = field(default_factory=set) @dataclass @@ -274,18 +276,21 @@ def for_vitisai(cls, device: str | None = None) -> WinMLCompileConfig: from pathlib import Path as _Path provider_options: dict[str, str] = {} + provider_option_file_keys: set[str] = set() ryzen_ai = os.environ.get("RYZEN_AI_INSTALLATION_PATH") if ryzen_ai: xclbin = _Path(ryzen_ai) / "voe-4.0-win_amd64" / "xclbins" / "phoenix" / "4x4.xclbin" if xclbin.exists(): provider_options["target"] = "X1" provider_options["xclbin"] = str(xclbin) + provider_option_file_keys.add("xclbin") provider_options["xlnx_enable_py3_round"] = "0" ep_cfg = EPConfig( provider="vitisai", enable_ep_context=True, provider_options=provider_options, device=device or "auto", + provider_option_file_keys=provider_option_file_keys, ) return cls(ep_config=ep_cfg) @@ -305,6 +310,7 @@ def to_dict(self) -> dict[str, Any]: return { "execution_provider": self.ep_config.provider, "provider_options": self.ep_config.provider_options, + "provider_option_file_keys": sorted(self.ep_config.provider_option_file_keys), "enable_ep_context": self.ep_config.enable_ep_context, "embed_context": self.ep_config.embed_context, "compiler": self.ep_config.compiler, @@ -324,6 +330,7 @@ def from_dict(cls, data: dict[str, Any]) -> WinMLCompileConfig: ep_config = EPConfig( provider=data.get("execution_provider"), provider_options=data.get("provider_options", {}), + provider_option_file_keys=set(data.get("provider_option_file_keys", [])), enable_ep_context=data.get("enable_ep_context", True), embed_context=data.get("embed_context", False), compiler=data.get("compiler", "ort"), diff --git a/src/winml/modelkit/compiler/stages/compile.py b/src/winml/modelkit/compiler/stages/compile.py index d877867ea..c912c2b22 100644 --- a/src/winml/modelkit/compiler/stages/compile.py +++ b/src/winml/modelkit/compiler/stages/compile.py @@ -6,13 +6,17 @@ from __future__ import annotations +import hashlib +import os import shutil import tempfile +import threading import time from pathlib import Path from typing import TYPE_CHECKING, ClassVar, cast import numpy as np +from filelock import FileLock from onnx import AttributeProto from ...onnx import load_onnx, save_onnx @@ -31,6 +35,7 @@ if TYPE_CHECKING: import onnxruntime as ort + from onnx import ModelProto from ...utils.constants import EPAlias from ..context import CompileContext @@ -42,6 +47,16 @@ "qairt": WinMLQairtSession, } +_FINALIZE_THREAD_LOCKS_GUARD = threading.Lock() +_FINALIZE_THREAD_LOCKS: dict[Path, threading.Lock] = {} + + +def _finalize_thread_lock(lock_path: Path) -> threading.Lock: + """Return the process-local lock paired with one public output path.""" + resolved = lock_path.resolve(strict=False) + with _FINALIZE_THREAD_LOCKS_GUARD: + return _FINALIZE_THREAD_LOCKS.setdefault(resolved, threading.Lock()) + class CompileStage(BaseStage): """Compile model.""" @@ -109,6 +124,7 @@ def _compile_single_model(self, context: CompileContext) -> None: ) try: winml_session.compile() + running_model_path = winml_session.running_model_path session = winml_session._session context.session = session @@ -117,18 +133,22 @@ def _compile_single_model(self, context: CompileContext) -> None: if context.validate: self._validate_model(session, context) self._collect_model_info(session, context) + + if ep_config.enable_ep_context: + if running_model_path == model_path: + context.add_warning(f"No EPContext produced for {model_path.name}") + return + self._finalize_output( + context, + model_path, + output_dir, + device=ep_device.device.device_type.lower(), + src_ctx_path=running_model_path, + ) finally: context.session = None winml_session.reset() - if ep_config.enable_ep_context: - self._finalize_output( - context, - model_path, - output_dir, - device=ep_device.device.device_type.lower(), - ) - def _compile_shared_context(self, context: CompileContext) -> None: """Compile through shared SessionOptions for multi-model and ORT-session flows.""" import onnxruntime as ort @@ -306,6 +326,7 @@ def _finalize_output( output_dir: Path, *, device: str | None = None, + src_ctx_path: Path | None = None, ) -> None: """Find EPContext files and copy to output directory. @@ -343,11 +364,11 @@ def _finalize_output( ] ) - src_ctx_path = None - for pattern in ctx_patterns: - if pattern.exists(): - src_ctx_path = pattern - break + if src_ctx_path is None: + for pattern in ctx_patterns: + if pattern.exists(): + src_ctx_path = pattern + break if src_ctx_path is None: context.add_warning("EPContext model not found in work directory") @@ -362,7 +383,23 @@ def _finalize_output( else: final_ctx_path = output_dir / f"{original_stem}_{output_suffix}_ctx.onnx" - # Ensure output directory exists + publish_lock = final_ctx_path.with_name(f"{final_ctx_path.name}.publish.lock") + with _finalize_thread_lock(publish_lock), FileLock(publish_lock): + self._publish_finalized_output( + context, + src_ctx_path, + final_ctx_path, + output_dir, + ) + + def _publish_finalized_output( + self, + context: CompileContext, + src_ctx_path: Path, + final_ctx_path: Path, + output_dir: Path, + ) -> None: + """Publish a self-consistent EPContext bundle while holding its output lock.""" output_dir.mkdir(parents=True, exist_ok=True) # Validate every external context reference before publishing the final @@ -396,7 +433,7 @@ def _finalize_output( source_root = src_ctx_path.parent.resolve() output_root = output_dir.resolve() - binary_exports: list[tuple[Path, Path, bytes, list[AttributeProto]]] = [] + binary_exports: list[tuple[Path, Path, Path, bytes, list[AttributeProto]]] = [] sources_by_final_binary: dict[Path, Path] = {} for raw_ref, cache_attrs in external_refs.items(): try: @@ -428,51 +465,57 @@ def _finalize_output( suffix = relative_ref.name[len(src_ctx_path.stem) :] final_relative_ref = relative_ref.with_name(f"{final_ctx_path.stem}{suffix}") - final_binary = (output_root / final_relative_ref).resolve() + stable_binary = (output_root / final_relative_ref).resolve() try: - final_binary.relative_to(output_root) + stable_binary.relative_to(output_root) except ValueError as exc: raise ValueError(f"unsafe EPContext binary reference: {cache_ref!r}") from exc - existing_source = sources_by_final_binary.get(final_binary) + existing_source = sources_by_final_binary.get(stable_binary) if existing_source is not None and existing_source != source_binary: raise ValueError( "Distinct EPContext binaries map to the same output path: " - f"{existing_source}, {source_binary} -> {final_binary}" + f"{existing_source}, {source_binary} -> {stable_binary}" ) - sources_by_final_binary[final_binary] = source_binary + sources_by_final_binary[stable_binary] = source_binary + content_token = self._file_sha256(source_binary)[:16] + unique_relative_ref = final_relative_ref.with_name( + f"{final_relative_ref.stem}.{content_token}{final_relative_ref.suffix}" + ) + unique_binary = (output_root / unique_relative_ref).resolve() + try: + unique_binary.relative_to(output_root) + except ValueError as exc: + raise ValueError(f"unsafe EPContext binary reference: {cache_ref!r}") from exc binary_exports.append( ( source_binary, - final_binary, - final_relative_ref.as_posix().encode("utf-8"), + unique_binary, + stable_binary, + unique_relative_ref.as_posix().encode("utf-8"), cache_attrs, ) ) - cache_refs_updated = False first_final_binary: Path | None = None - for source_binary, final_binary, final_ref_bytes, cache_attrs in binary_exports: - final_binary.parent.mkdir(parents=True, exist_ok=True) - if source_binary != final_binary: - shutil.copy2(source_binary, final_binary) - context.log(f"Copied binary to: {final_binary}") + for ( + source_binary, + unique_binary, + stable_binary, + final_ref_bytes, + cache_attrs, + ) in binary_exports: + self._atomic_copy(source_binary, unique_binary) + self._atomic_copy(source_binary, stable_binary) + context.log(f"Published binary generation: {unique_binary}") if first_final_binary is None: - first_final_binary = final_binary + first_final_binary = unique_binary for cache_attr in cache_attrs: - if cache_attr.s != final_ref_bytes: - cache_attr.s = final_ref_bytes - cache_refs_updated = True - - if cache_refs_updated: - save_onnx(model, final_ctx_path) - context.log("Updated external EPContext binary references") - elif src_ctx_path != final_ctx_path: - shutil.copy2(src_ctx_path, final_ctx_path) - context.log(f"Copied EPContext to: {final_ctx_path}") - else: - context.log(f"EPContext already at: {final_ctx_path}") + cache_attr.s = final_ref_bytes + + self._atomic_save_onnx(model, final_ctx_path) + context.log(f"Published EPContext: {final_ctx_path}") context.output_path = final_ctx_path context.context_binary_path = first_final_binary @@ -483,9 +526,49 @@ def _finalize_output( src_schematic = src_ctx_path.parent / schematic_name final_schematic = output_dir / schematic_name if src_schematic.is_file() and src_schematic != final_schematic: - shutil.copy2(src_schematic, final_schematic) + self._atomic_copy(src_schematic, final_schematic) context.log(f"Copied schematic to: {final_schematic}") + @staticmethod + def _file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source_file: + for chunk in iter(lambda: source_file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + @staticmethod + def _atomic_copy(source: Path, destination: Path) -> None: + """Copy one file through a same-directory temporary and atomic replace.""" + if source.resolve() == destination.resolve(strict=False): + return + destination.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=f".{destination.name}.", suffix=".tmp", dir=destination.parent + ) + os.close(fd) + temporary_path = Path(temporary_name) + try: + shutil.copy2(source, temporary_path) + temporary_path.replace(destination) + finally: + temporary_path.unlink(missing_ok=True) + + @staticmethod + def _atomic_save_onnx(model: ModelProto, destination: Path) -> None: + """Save an ONNX model beside its destination and atomically replace it.""" + fd, temporary_name = tempfile.mkstemp( + prefix=f".{destination.stem}.", suffix=destination.suffix, dir=destination.parent + ) + os.close(fd) + temporary_path = Path(temporary_name) + temporary_path.unlink(missing_ok=True) + try: + save_onnx(model, temporary_path) + temporary_path.replace(destination) + finally: + temporary_path.unlink(missing_ok=True) + def _collect_model_info(self, session: ort.InferenceSession, context: CompileContext) -> None: """Collect model input/output information.""" input_shapes = {} diff --git a/src/winml/modelkit/onnx/external_data.py b/src/winml/modelkit/onnx/external_data.py index 5363a1985..97d4d52d3 100644 --- a/src/winml/modelkit/onnx/external_data.py +++ b/src/winml/modelkit/onnx/external_data.py @@ -166,15 +166,37 @@ def _update_hash_from_path_metadata(hash_obj: Any, path: Path) -> None: hash_obj.update(str(stat.st_mtime_ns).encode("ascii")) -def get_onnx_model_hash(model_path: str | Path) -> str: - """Compute a lightweight metadata hash for an ONNX model and external data.""" +def _update_hash_from_path_content(hash_obj: Any, path: Path) -> None: + """Update *hash_obj* with a resolved path and the file's complete bytes.""" + resolved = path.resolve(strict=True) + hash_obj.update(str(resolved).encode("utf-8", "surrogatepass")) + hash_obj.update(b"\0") + with resolved.open("rb") as source_file: + for chunk in iter(lambda: source_file.read(1024 * 1024), b""): + hash_obj.update(chunk) + + +def get_onnx_model_hash(model_path: str | Path, *, strict: bool = False) -> str: + """Compute a lightweight metadata hash for an ONNX model and external data. + + Args: + model_path: ONNX graph whose source metadata participates in the hash. + strict: Raise when external-data references cannot be inspected or a + referenced sidecar is unavailable. Cache identities should use + strict mode so uncertainty forces a rebuild instead of a false hit. + """ model_path = Path(model_path).resolve() hash_obj = hashlib.sha256() - _update_hash_from_path_metadata(hash_obj, model_path) + if strict: + _update_hash_from_path_content(hash_obj, model_path) + else: + _update_hash_from_path_metadata(hash_obj, model_path) try: external_files = get_external_data_files(model_path) except Exception: + if strict: + raise logger.debug("Could not inspect ONNX external data for hashing: %s", model_path) external_files = [] @@ -186,8 +208,13 @@ def get_onnx_model_hash(model_path: str | Path) -> str: hash_obj.update(location.replace("\\", "/").encode("utf-8")) hash_obj.update(b"\0") try: - _update_hash_from_path_metadata(hash_obj, data_path) + if strict: + _update_hash_from_path_content(hash_obj, data_path) + else: + _update_hash_from_path_metadata(hash_obj, data_path) except FileNotFoundError: + if strict: + raise logger.debug( "ONNX external data file referenced by %s is missing: %s", model_path, diff --git a/src/winml/modelkit/session/session.py b/src/winml/modelkit/session/session.py index 648b85058..d62a74fe5 100644 --- a/src/winml/modelkit/session/session.py +++ b/src/winml/modelkit/session/session.py @@ -6,8 +6,13 @@ from __future__ import annotations +import hashlib +import json import logging import os +import tempfile +import threading +import uuid from contextlib import contextmanager from dataclasses import dataclass, replace from enum import Enum @@ -16,9 +21,11 @@ import numpy as np import onnxruntime as ort +from filelock import FileLock, Timeout +from google.protobuf.message import DecodeError from ..core.onnx_utils import get_io_config -from ..onnx import is_compiled_onnx +from ..onnx import get_onnx_model_hash, is_compiled_onnx from ..utils.native_stderr import ( get_win32_fd_handle, get_win32_std_handle, @@ -49,6 +56,94 @@ logger = logging.getLogger(__name__) +_EPCONTEXT_THREAD_LOCKS_GUARD = threading.Lock() +_EPCONTEXT_THREAD_LOCKS: dict[Path, threading.Lock] = {} +_EPCONTEXT_CACHE_MAX_GENERATIONS = 4 + + +def _epcontext_thread_lock(lock_path: Path) -> threading.Lock: + """Return the process-local lock paired with one EPContext lockfile.""" + resolved = lock_path.resolve(strict=False) + with _EPCONTEXT_THREAD_LOCKS_GUARD: + return _EPCONTEXT_THREAD_LOCKS.setdefault(resolved, threading.Lock()) + + +@dataclass +class _EPContextCacheLease: + """Held cache lock that can be transferred from compile selection to runtime open.""" + + lock_path: Path + thread_lock: threading.Lock + file_lock: FileLock + _released: bool = False + + @classmethod + def acquire(cls, lock_path: Path, *, blocking: bool = True) -> _EPContextCacheLease | None: + thread_lock = _epcontext_thread_lock(lock_path) + if not thread_lock.acquire(blocking=blocking): + return None + file_lock = FileLock(lock_path) + try: + if blocking: + file_lock.acquire() + else: + file_lock.acquire(timeout=0) + except Timeout: + thread_lock.release() + return None + except Exception: + thread_lock.release() + raise + return cls(lock_path=lock_path, thread_lock=thread_lock, file_lock=file_lock) + + def release(self) -> None: + if self._released: + return + try: + self.file_lock.release() + finally: + self.thread_lock.release() + self._released = True + + def __enter__(self) -> _EPContextCacheLease: + return self + + def __exit__( + self, + _exc_type: type[BaseException] | None, + _exc: BaseException | None, + _tb: TracebackType | None, + ) -> None: + self.release() + + +@dataclass +class _PreparedEPContextModel: + """Model path selected for runtime loading plus any lock held until ORT opens it.""" + + path: Path + markerless: bool = False + lease: _EPContextCacheLease | None = None + + @property + def lock_path(self) -> Path | None: + return self.lease.lock_path if self.lease is not None else None + + def release(self) -> None: + if self.lease is not None: + self.lease.release() + + def __enter__(self) -> _PreparedEPContextModel: + return self + + def __exit__( + self, + _exc_type: type[BaseException] | None, + _exc: BaseException | None, + _tb: TracebackType | None, + ) -> None: + self.release() + @contextmanager def _suppress_native_output(log_path: str | Path | None = None) -> Iterator[None]: @@ -366,10 +461,15 @@ def __init__( ) # Snapshots preserved across perf()/reset()/compile() entry/exit (see perf()). - self._provider_options: dict[str, str] = _build_provider_options( - ep_device, ep_config, ep_monitor + self._provider_option_file_keys: frozenset[str] = frozenset( + ep_config.provider_option_file_keys if ep_config is not None else () + ) + self._provider_options: dict[str, str] = self._canonicalize_option_files( + _build_provider_options(ep_device, ep_config, ep_monitor), + self._provider_option_file_keys, ) self._active_session_option_entries: dict[str, str] = dict(initial_session_option_entries) + self._markerless_epcontext_generations: dict[Path, Path | None] = {} # Convenience: the canonical EP name from the chosen handle. self._ep: str = ep_device.device.ep_name @@ -472,47 +572,19 @@ def compile(self) -> None: self._state = SessionState.COMPILED return - # Derive the output ctx path from the original model path. - ctx_path = self._onnx_path.parent / f"{self._onnx_path.stem}_{target_device}_ctx.onnx" - model_path = self._onnx_path + prepared_model = _PreparedEPContextModel(self._onnx_path) # Native QNN SDK compiler writes progress to stdout/stderr; # redirect to log file to keep the console clean. compile_log = self._onnx_path.parent / "compile.log" - # Check for existing fresh EPContext (skip re-compile if cache is fresh). - if ctx_path.exists() and ctx_path.stat().st_mtime >= self._onnx_path.stat().st_mtime: - model_path = ctx_path - logger.info("Using cached EPContext: %s", ctx_path) - elif is_compiled_onnx(self._onnx_path): + if is_compiled_onnx(self._onnx_path): # Input model is already an EPContext — use it directly. logger.info("Model already compiled (EPContext), skipping ModelCompiler") else: - # AOT compile to .ctx.onnx via ort.ModelCompiler. - try: - so = _build_session_options( - self._ep_device, - self._ep_config, - None, # no monitor at compile time - self._session_options_factory, - session_option_entries=self._active_session_option_entries, - provider_options=self._provider_options, - ) - model_compiler = ort.ModelCompiler( - so, - str(self._onnx_path), - embed_compiled_data_into_model=self._embed_context, - ) - with _suppress_native_output(compile_log): - model_compiler.compile_to_file(str(ctx_path)) - - if ctx_path.exists(): - model_path = ctx_path - logger.info("Compiled to EPContext: %s", ctx_path) + prepared_model = self._compile_epcontext_with_stable_source(compile_log) - except Exception as e: - # Some EPs don't support compilation — fall back to original model. - logger.warning("ModelCompiler failed, using original: %s", e) + model_path = prepared_model.path try: # Create the runtime InferenceSession against the (possibly compiled) model. @@ -524,7 +596,7 @@ def compile(self) -> None: session_option_entries=self._active_session_option_entries, provider_options=self._provider_options, ) - with _suppress_native_output(compile_log): + with prepared_model, _suppress_native_output(compile_log): session = ort.InferenceSession(str(model_path), sess_options=runtime_so) actual_providers = session.get_providers() @@ -535,6 +607,9 @@ def compile(self) -> None: ) except Exception as e: + prepared_model.release() + if prepared_model.markerless and model_path != self._onnx_path: + self._discard_epcontext_generation(model_path) self._state = SessionState.ERROR self._last_error = e raise CompilationError( @@ -550,6 +625,541 @@ def compile(self) -> None: self._session = session self._running_model_path = model_path self._state = SessionState.COMPILED + if prepared_model.markerless and model_path != self._onnx_path: + self._markerless_epcontext_generations[model_path] = prepared_model.lock_path + + def _compile_epcontext_with_stable_source(self, compile_log: Path) -> _PreparedEPContextModel: + """Prepare an EPContext whose marker matches a stable source snapshot.""" + for _attempt in range(3): + try: + expected_identity = self._epcontext_cache_identity() + except (OSError, ValueError) as exc: + logger.warning( + "Could not establish EPContext source identity; cache reuse disabled: %s", + exc, + ) + cache_path = self._epcontext_cache_path(None) + lock_path = cache_path.with_name(f"{cache_path.name}.lock") + lease = _EPContextCacheLease.acquire(lock_path) + assert lease is not None + try: + prepared = self._prepare_epcontext_model( + cache_path, + compile_log, + None, + ) + if prepared is not None and prepared.path != self._onnx_path: + prepared.lease = lease + lease = None + return prepared or _PreparedEPContextModel(self._onnx_path) + finally: + if lease is not None: + lease.release() + + cache_path = self._epcontext_cache_path(expected_identity) + lock_path = cache_path.with_name(f"{cache_path.name}.lock") + lease = _EPContextCacheLease.acquire(lock_path) + assert lease is not None + try: + try: + locked_identity = self._epcontext_cache_identity() + except (OSError, ValueError): + continue + if locked_identity != expected_identity: + continue + prepared = self._prepare_epcontext_model( + cache_path, + compile_log, + locked_identity, + ) + if prepared is None: + continue + try: + final_identity = self._epcontext_cache_identity() + except (OSError, ValueError): + if prepared.markerless and prepared.path != self._onnx_path: + self._discard_epcontext_generation(prepared.path) + continue + if final_identity == locked_identity: + if prepared.path != self._onnx_path: + prepared.lease = lease + lease = None + return prepared + if prepared.markerless and prepared.path != self._onnx_path: + self._discard_epcontext_generation(prepared.path) + finally: + if lease is not None: + lease.release() + + logger.warning("ONNX source changed during EPContext compilation; using original model") + return _PreparedEPContextModel(self._onnx_path) + + def _prepare_epcontext_model( + self, + cache_path: Path, + compile_log: Path, + cache_identity: dict[str, object] | None, + ) -> _PreparedEPContextModel | None: + """Reuse or compile one EPContext while the caller holds its file lock.""" + if cache_identity is not None: + cached_generation = self._epcontext_cached_generation(cache_path, cache_identity) + if cached_generation is not None: + logger.info("Using cached EPContext: %s", cached_generation) + return _PreparedEPContextModel(cached_generation) + + generation_path = cache_path.with_name( + f"{cache_path.stem}_{uuid.uuid4().hex[:16]}{cache_path.suffix}" + ) + try: + so = _build_session_options( + self._ep_device, + self._ep_config, + None, + self._session_options_factory, + session_option_entries=self._active_session_option_entries, + provider_options=self._provider_options, + ) + model_compiler = ort.ModelCompiler( + so, + str(self._onnx_path), + embed_compiled_data_into_model=self._embed_context, + ) + with _suppress_native_output(compile_log): + model_compiler.compile_to_file(str(generation_path)) + except Exception as exc: + self._discard_epcontext_generation(generation_path) + logger.warning("ModelCompiler failed, using original: %s", exc) + return _PreparedEPContextModel(self._onnx_path) + + if not generation_path.exists(): + self._discard_epcontext_generation(generation_path) + return _PreparedEPContextModel(self._onnx_path) + markerless = cache_identity is None + if cache_identity is not None: + try: + current_identity = self._epcontext_cache_identity() + except (OSError, ValueError): + self._discard_epcontext_generation(generation_path) + return None + if current_identity != cache_identity: + self._discard_epcontext_generation(generation_path) + return None + try: + self._write_epcontext_cache_marker( + cache_path, + generation_path, + cache_identity, + ) + except (OSError, TypeError, ValueError) as exc: + logger.warning( + "Compiled EPContext but could not write cache marker %s: %s", + self._epcontext_cache_marker_path(cache_path), + exc, + ) + markerless = True + else: + self._prune_epcontext_cache(cache_path) + logger.info("Compiled to EPContext: %s", generation_path) + return _PreparedEPContextModel(generation_path, markerless=markerless) + + @classmethod + def _discard_epcontext_generation(cls, generation_path: Path) -> None: + """Best-effort removal of an unpublished generation and its sidecars.""" + paths: dict[Path, None] = {} + try: + for sidecar in cls._epcontext_external_sidecars(generation_path): + paths.setdefault(sidecar, None) + except (OSError, ValueError, DecodeError): + logger.debug( + "Could not parse EPContext sidecars for %s; using prefix cleanup fallback", + generation_path, + exc_info=True, + ) + paths.setdefault(generation_path, None) + try: + for candidate in generation_path.parent.iterdir(): + if candidate.name.startswith(generation_path.stem) and candidate.is_file(): + paths.setdefault(candidate, None) + except OSError: + logger.debug( + "Could not enumerate EPContext generation sidecars for %s", + generation_path, + ) + for path in paths: + cls._unlink_generation_file(path) + + def _prune_epcontext_cache(self, current_cache_path: Path) -> None: + """Bound successful EPContext cache entries for this source model and device.""" + keep_count = max(1, _EPCONTEXT_CACHE_MAX_GENERATIONS) + marker_suffix = ".meta.json" + cache_name_suffix = "_ctx.onnx" + marker_name_suffix = f"{cache_name_suffix}{marker_suffix}" + marker_prefix = f"{self._onnx_path.stem}_{self._device}_" + current_marker = self._epcontext_cache_marker_path(current_cache_path).resolve(strict=False) + entries: list[tuple[bool, int, str, Path, Path, Path | None, bytes]] = [] + for marker_path in current_cache_path.parent.iterdir(): + try: + if not marker_path.is_file(): + continue + if not marker_path.name.startswith(marker_prefix) or not marker_path.name.endswith( + marker_name_suffix + ): + continue + cache_name = marker_path.name.removesuffix(marker_suffix) + if not cache_name.startswith(marker_prefix) or not cache_name.endswith( + cache_name_suffix + ): + continue + cache_path = marker_path.with_name(cache_name) + marker_contents = marker_path.read_bytes() + generation_path = self._epcontext_marker_generation( + marker_path, + marker_contents, + ) + lock_path = cache_path.with_name(f"{cache_path.name}.lock") + marker_stat = marker_path.stat() + except OSError: + continue + is_current = marker_path.resolve(strict=False) == current_marker + entries.append( + ( + is_current, + marker_stat.st_mtime_ns, + marker_path.name, + marker_path, + lock_path, + generation_path, + marker_contents, + ) + ) + entries.sort(key=lambda entry: (entry[0], entry[1], entry[2]), reverse=True) + for ( + is_current, + *_unused, + marker_path, + lock_path, + generation_path, + marker_contents, + ) in entries[keep_count:]: + if is_current: + continue + lease = _EPContextCacheLease.acquire(lock_path, blocking=False) + if lease is None: + continue + try: + try: + locked_contents = marker_path.read_bytes() + except OSError: + continue + locked_generation = self._epcontext_marker_generation( + marker_path, + locked_contents, + ) + if locked_contents != marker_contents: + if generation_path is not None and generation_path != locked_generation: + self._discard_epcontext_generation(generation_path) + continue + if locked_generation is not None: + self._discard_epcontext_generation(locked_generation) + self._unlink_generation_file(marker_path) + finally: + lease.release() + self._unlink_generation_file(lock_path) + + @staticmethod + def _epcontext_marker_generation(marker_path: Path, contents: bytes) -> Path | None: + """Return a marker's validated sibling generation path.""" + try: + recorded = json.loads(contents) + except (json.JSONDecodeError, UnicodeDecodeError): + return None + if not isinstance(recorded, dict): + return None + generation_name = recorded.get("generation") + if not isinstance(generation_name, str) or Path(generation_name).name != generation_name: + return None + return marker_path.parent / generation_name + + def _cleanup_markerless_epcontext_generations(self) -> None: + """Best-effort cleanup for non-reusable EPContext generations owned by this session.""" + remaining: dict[Path, Path | None] = {} + for generation_path, lock_path in self._markerless_epcontext_generations.items(): + self._discard_epcontext_generation(generation_path) + if self._epcontext_generation_artifacts_exist(generation_path): + remaining[generation_path] = lock_path + continue + if lock_path is not None: + self._unlink_generation_file(lock_path) + self._markerless_epcontext_generations = remaining + + @staticmethod + def _epcontext_generation_artifacts_exist(generation_path: Path) -> bool: + try: + return any( + candidate.name.startswith(generation_path.stem) and candidate.is_file() + for candidate in generation_path.parent.iterdir() + ) + except OSError: + return generation_path.exists() + + @staticmethod + def _unlink_generation_file(path: Path) -> None: + try: + path.unlink(missing_ok=True) + except OSError: + logger.debug("Could not remove stale EPContext generation file %s", path) + + @staticmethod + def _epcontext_cache_marker_path(ctx_path: Path) -> Path: + """Return the sidecar that records a direct-session cache identity.""" + return ctx_path.with_name(f"{ctx_path.name}.meta.json") + + def _epcontext_cache_path(self, identity: dict[str, object] | None) -> Path: + """Return an immutable identity path, or a unique non-cacheable path.""" + if identity is None: + identity_token = uuid.uuid4().hex[:16] + else: + encoded_identity = json.dumps(identity, sort_keys=True, separators=(",", ":")) + identity_token = hashlib.sha256(encoded_identity.encode("utf-8")).hexdigest()[:16] + return self._onnx_path.with_name( + f"{self._onnx_path.stem}_{self._device}_{identity_token}_ctx.onnx" + ) + + def _epcontext_cache_identity(self) -> dict[str, object]: + """Return all inputs that affect a direct-session EPContext artifact.""" + if self._session_options_factory is not None: + raise ValueError( + "custom SessionOptions factory cannot be represented in cache identity" + ) + hardware = self._ep_device.device.ort_handle.device + + def _optional_text(value: object) -> str | None: + return value if isinstance(value, str) and value else None + + dll_fingerprint = ( + None if self._ep_device.is_builtin else self._file_fingerprint(self._ep_device.dll_path) + ) + + return { + "schema_version": 1, + "source_model_hash": get_onnx_model_hash(self._onnx_path, strict=True), + "ep": self._ep_device.device.ep_name, + "ep_source": self._ep_device.source_tag, + "ep_version": self._ep_device.version, + "ep_dll": dll_fingerprint, + "device": self._device, + "hardware": { + "vendor_id": hardware.vendor_id, + "device_id": hardware.device_id, + "name": _optional_text(self._ep_device.device.hardware_name), + "driver_version": _optional_text(self._ep_device.device.driver_version), + "compiler_version": _optional_text(self._ep_device.device.compiler_version), + }, + "provider_options": dict(sorted(self._provider_options.items())), + "provider_option_files": self._option_file_fingerprints( + self._provider_options, + self._provider_option_file_keys, + ), + "session_options": dict(sorted(self._active_session_option_entries.items())), + "session_option_files": {}, + "embed_context": self._embed_context, + "ort_version": ort.__version__, + } + + def _option_file_fingerprints( + self, + options: dict[str, str], + file_keys: frozenset[str], + ) -> dict[str, dict[str, object]]: + """Fingerprint provider options explicitly declared as file-backed.""" + fingerprints = {} + for key in sorted(file_keys): + value = options.get(key) + if value is None: + continue + option_path = self._resolve_option_file(value) + if option_path is None: + raise ValueError(f"file-backed provider option {key!r} does not resolve to a file") + fingerprints[key] = self._file_fingerprint(option_path) + return fingerprints + + def _canonicalize_option_files( + self, + options: dict[str, str], + file_keys: frozenset[str], + ) -> dict[str, str]: + """Resolve only provider options explicitly declared as file-backed.""" + canonicalized = dict(options) + for key in file_keys: + value = options.get(key) + if value is None: + raise ValueError(f"file-backed provider option {key!r} is not configured") + option_path = self._resolve_option_file(value) + if option_path is None: + raise ValueError(f"file-backed provider option {key!r} does not resolve to a file") + canonicalized[key] = str(option_path) + return canonicalized + + def _resolve_option_file(self, value: object) -> Path | None: + if not isinstance(value, str) or not value: + return None + try: + raw_path = Path(value).expanduser() + except (RuntimeError, ValueError): + return None + candidates: tuple[Path, ...] = ( + (raw_path,) if raw_path.is_absolute() else (Path.cwd() / raw_path,) + ) + seen: set[Path] = set() + for candidate in candidates: + try: + resolved = candidate.resolve(strict=True) + except (OSError, RuntimeError, ValueError): + continue + if resolved in seen: + continue + seen.add(resolved) + if resolved.is_file(): + return resolved + return None + + @staticmethod + def _file_fingerprint(path: Path) -> dict[str, object]: + """Return a strict metadata and content fingerprint for one file.""" + resolved = path.resolve(strict=True) + stat = resolved.stat() + digest = hashlib.sha256() + with resolved.open("rb") as source_file: + for chunk in iter(lambda: source_file.read(1024 * 1024), b""): + digest.update(chunk) + return { + "path": str(resolved), + "size": stat.st_size, + "mtime_ns": stat.st_mtime_ns, + "sha256": digest.hexdigest(), + } + + @classmethod + def _epcontext_cached_generation( + cls, + cache_path: Path, + expected_identity: dict[str, object], + ) -> Path | None: + """Return the immutable generation selected by a valid identity marker.""" + marker_path = cls._epcontext_cache_marker_path(cache_path) + try: + recorded = json.loads(marker_path.read_text(encoding="utf-8")) + if not isinstance(recorded, dict) or recorded.get("identity") != expected_identity: + return None + generation_name = recorded.get("generation") + if ( + not isinstance(generation_name, str) + or Path(generation_name).name != generation_name + ): + return None + generation_path = cache_path.parent / generation_name + current_artifacts = cls._epcontext_artifact_fingerprint(generation_path) + except Exception: + return None + if recorded.get("artifacts") != current_artifacts: + return None + return generation_path + + @staticmethod + def _epcontext_external_sidecars(ctx_path: Path) -> tuple[Path, ...]: + """Return validated external binaries referenced by an EPContext graph.""" + from onnx import AttributeProto + + from ..onnx import load_onnx + + model = load_onnx(ctx_path, load_weights=False, validate=False) + source_root = ctx_path.parent.resolve() + sidecars: dict[Path, None] = {} + for node in model.graph.node: + if node.op_type != "EPContext": + continue + attrs = {attr.name: attr for attr in node.attribute} + embed_mode = attrs.get("embed_mode") + if embed_mode is None or (embed_mode.type == AttributeProto.INT and embed_mode.i != 0): + continue + if embed_mode.type != AttributeProto.INT: + raise ValueError("EPContext embed_mode must be an integer") + main_context = attrs.get("main_context") + cache_attr = attrs.get("ep_cache_context") + is_secondary = ( + main_context is not None + and main_context.type == AttributeProto.INT + and main_context.i == 0 + ) + if cache_attr is None and is_secondary: + continue + if cache_attr is None or cache_attr.type != AttributeProto.STRING: + raise ValueError("External EPContext node must have a string ep_cache_context") + try: + cache_ref = cache_attr.s.decode("utf-8") + except UnicodeDecodeError as exc: + raise ValueError("EPContext cache reference must be valid UTF-8") from exc + relative_ref = Path(cache_ref) + if not cache_ref or relative_ref.is_absolute() or relative_ref.drive: + raise ValueError(f"unsafe EPContext cache reference: {cache_ref!r}") + try: + sidecar = (source_root / relative_ref).resolve() + sidecar.relative_to(source_root) + except (OSError, ValueError) as exc: + raise ValueError(f"unsafe EPContext cache reference: {cache_ref!r}") from exc + if not sidecar.is_file() or sidecar.stat().st_size == 0: + raise FileNotFoundError(f"EPContext sidecar is unavailable: {sidecar}") + sidecars.setdefault(sidecar, None) + return tuple(sidecars) + + @classmethod + def _epcontext_artifact_fingerprint(cls, ctx_path: Path) -> list[dict[str, object]]: + """Return metadata fingerprints for the graph and external binaries.""" + source_root = ctx_path.parent.resolve() + paths = (ctx_path.resolve(), *cls._epcontext_external_sidecars(ctx_path)) + fingerprints = [] + for path in paths: + fingerprint = cls._file_fingerprint(path) + fingerprints.append( + { + "path": path.relative_to(source_root).as_posix(), + "size": fingerprint["size"], + "mtime_ns": fingerprint["mtime_ns"], + "sha256": fingerprint["sha256"], + } + ) + return fingerprints + + @classmethod + def _write_epcontext_cache_marker( + cls, + cache_path: Path, + generation_path: Path, + identity: dict[str, object], + ) -> None: + """Atomically publish the identity for a successfully compiled context.""" + marker_path = cls._epcontext_cache_marker_path(cache_path) + marker_path.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=f".{marker_path.name}.", suffix=".tmp", dir=marker_path.parent + ) + temporary_path = Path(temporary_name) + try: + with os.fdopen(fd, "w", encoding="utf-8") as marker_file: + json.dump( + { + "identity": identity, + "generation": generation_path.name, + "artifacts": cls._epcontext_artifact_fingerprint(generation_path), + }, + marker_file, + indent=2, + sort_keys=True, + ) + temporary_path.replace(marker_path) + except Exception: + temporary_path.unlink(missing_ok=True) + raise def run( self, @@ -629,7 +1239,13 @@ def reset(self) -> None: Clears compiled session and error state. """ + self._reset_runtime_state(cleanup_markerless=True) + + def _reset_runtime_state(self, *, cleanup_markerless: bool) -> None: + """Release the native session and optionally its private EPContext artifacts.""" self._session = None + if cleanup_markerless: + self._cleanup_markerless_epcontext_generations() self._running_model_path = None self._state = SessionState.INITIALIZED self._last_error = None @@ -639,6 +1255,7 @@ def __del__(self) -> None: """Clean up resources on deletion.""" try: self._session = None + self._cleanup_markerless_epcontext_generations() except Exception: pass # Suppress errors during interpreter shutdown @@ -875,9 +1492,12 @@ def perf( saved_sess_entries, dict(monitor.get_session_options()) if monitor is not None else None, ) - new_prov = _overlay_options( - saved_prov, - dict(monitor.get_provider_options()) if monitor is not None else None, + new_prov = self._canonicalize_option_files( + _overlay_options( + saved_prov, + dict(monitor.get_provider_options()) if monitor is not None else None, + ), + self._provider_option_file_keys, ) # Rebuild InferenceSession only when monitor-contributed provider/session @@ -891,7 +1511,7 @@ def perf( ) if had_baseline_session and _session_rebuilt: logger.info("auto-resetting compiled session to apply monitor session/provider options") - self.reset() + self._reset_runtime_state(cleanup_markerless=False) stats = PerfStats(warmup=warmup) restore_baseline = _session_rebuilt or getattr( @@ -997,10 +1617,10 @@ def _restore_baseline() -> Exception | None: if exc_info[1] is None: monitor_error = error - # C-2: for monitors that require session teardown, reset() BEFORE - # monitor.__exit__ so the flushed data is available in __exit__. + # C-2: for monitors that require session teardown, release the + # native session BEFORE monitor.__exit__ so flushed data is available. if getattr(effective_monitor, "requires_session_teardown", False): - self.reset() + self._reset_runtime_state(cleanup_markerless=False) # Call monitor.__exit__ — propagate exc_info so monitor sees the # exception (exception transparency contract). diff --git a/tests/unit/commands/test_compile_cli.py b/tests/unit/commands/test_compile_cli.py index b77b02021..643e6e474 100644 --- a/tests/unit/commands/test_compile_cli.py +++ b/tests/unit/commands/test_compile_cli.py @@ -89,6 +89,7 @@ def compile_cli_mocks() -> CompileCliMocks: compile_config.ep_config.qnn_sdk_root = None compile_config.ep_config.embed_context = False compile_config.ep_config.provider_options = {} + compile_config.ep_config.provider_option_file_keys = set() compile_config.ep_config.enable_ep_context = False def _build_compile_config(ep_device: EPDeviceTarget) -> MagicMock: @@ -349,6 +350,37 @@ def test_cli_ep_options_override_matching_config_keys( "soc_model": "57", } + def test_config_provider_option_file_keys_reach_compile_config( + self, + runner: CliRunner, + fake_onnx: Path, + compile_cli_mocks: CompileCliMocks, + ) -> None: + """Compile config preserves explicit file-backed provider option metadata.""" + dependency = fake_onnx.parent / "compiler-input.bin" + dependency.write_bytes(b"compiler input") + config_path = fake_onnx.parent / "compile.json" + config_path.write_text( + json.dumps( + { + "compile": { + "provider_options": {"compiler_input": str(dependency)}, + "provider_option_file_keys": ["compiler_input"], + } + } + ) + ) + + result = runner.invoke( + compile, + ["-m", str(fake_onnx), "--config", str(config_path)], + ) + + _assert_successful_compile_call(result, compile_cli_mocks, fake_onnx) + assert compile_cli_mocks.compile_config.ep_config.provider_option_file_keys == { + "compiler_input" + } + def test_invalid_ep_option_is_rejected( self, runner: CliRunner, fake_onnx: Path, compile_cli_mocks: CompileCliMocks ) -> None: diff --git a/tests/unit/compiler/test_compiler_configs.py b/tests/unit/compiler/test_compiler_configs.py index ccfc24ac2..ff96f3d45 100644 --- a/tests/unit/compiler/test_compiler_configs.py +++ b/tests/unit/compiler/test_compiler_configs.py @@ -41,6 +41,10 @@ def test_custom_values(self): assert config.enable_ep_context is False assert config.embed_context is True + def test_file_backed_provider_options_are_explicit(self) -> None: + """EP configs expose architecture-agnostic file dependency metadata.""" + assert "provider_option_file_keys" in EPConfig.__dataclass_fields__ + class TestCompileConfig: """Test WinMLCompileConfig dataclass.""" @@ -105,6 +109,21 @@ def test_for_vitisai(self): assert config.ep_config.provider == "vitisai" assert config.ep_config.enable_ep_context is True + def test_for_vitisai_declares_discovered_xclbin_as_file_backed( + self, + tmp_path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """The discovered compiler input participates in EPContext identity.""" + xclbin = tmp_path / "voe-4.0-win_amd64" / "xclbins" / "phoenix" / "4x4.xclbin" + xclbin.parent.mkdir(parents=True) + xclbin.write_bytes(b"xclbin") + monkeypatch.setenv("RYZEN_AI_INSTALLATION_PATH", str(tmp_path)) + + config = WinMLCompileConfig.for_vitisai() + + assert config.ep_config.provider_option_file_keys == {"xclbin"} + def test_for_migraphx(self): """Test MIGraphX factory method.""" config = WinMLCompileConfig.for_migraphx() @@ -173,6 +192,22 @@ def test_roundtrip(self): assert restored.ep_config.enable_ep_context == original.ep_config.enable_ep_context assert restored.validate == original.validate + def test_roundtrip_preserves_provider_option_file_keys(self) -> None: + """Explicit file-backed option keys survive config serialization.""" + original = WinMLCompileConfig( + ep_config=EPConfig( + provider="qnn", + provider_options={"compiler_input": "inputs.bin"}, + provider_option_file_keys={"compiler_input"}, + ) + ) + + serialized = original.to_dict() + + assert serialized.get("provider_option_file_keys") == ["compiler_input"] + restored = WinMLCompileConfig.from_dict(serialized) + assert restored.ep_config.provider_option_file_keys == {"compiler_input"} + def test_roundtrip_preserves_ep_device(self) -> None: """Round-trip retains the resolved EP/device/source binding.""" from winml.modelkit.session import EPDeviceTarget diff --git a/tests/unit/compiler/test_compiler_stages.py b/tests/unit/compiler/test_compiler_stages.py index 8cc7d3745..2b8db62c9 100644 --- a/tests/unit/compiler/test_compiler_stages.py +++ b/tests/unit/compiler/test_compiler_stages.py @@ -410,6 +410,7 @@ def test_process_preserves_trtrtx_provider_options(self, tmp_path): fake_winml_session = MagicMock() fake_winml_session._session = fake_session + fake_winml_session.running_model_path = tmp_path / "model_trtrtx_ctx.onnx" context = CompileContext( model_path=model_path, @@ -422,10 +423,13 @@ def test_process_preserves_trtrtx_provider_options(self, tmp_path): ) mock_session_cls = MagicMock(return_value=fake_winml_session) - with patch.dict( - "winml.modelkit.compiler.stages.compile.COMPILER_SESSION_MAPPING", - {"ort": mock_session_cls}, - clear=False, + with ( + patch.dict( + "winml.modelkit.compiler.stages.compile.COMPILER_SESSION_MAPPING", + {"ort": mock_session_cls}, + clear=False, + ), + patch.object(CompileStage, "_finalize_output"), ): stage = CompileStage() stage.process(context) @@ -435,6 +439,57 @@ def test_process_preserves_trtrtx_provider_options(self, tmp_path): fake_winml_session.reset.assert_called_once() assert context.session is None + def test_process_finalizes_markerless_context_before_reset(self, tmp_path): + """Session-owned EPContext artifacts remain present until publication.""" + from unittest.mock import MagicMock, patch + + from winml.modelkit.compiler import CompileContext, CompileStage + from winml.modelkit.session import EPDeviceTarget + + model_path = tmp_path / "model.onnx" + create_simple_model(model_path) + markerless_path = tmp_path / "model_npu_ctx_private.onnx" + markerless_path.write_bytes(b"markerless context") + events: list[str] = [] + fake_winml_session = MagicMock() + fake_winml_session._session = None + fake_winml_session.running_model_path = markerless_path + + def _reset() -> None: + events.append("reset") + markerless_path.unlink() + + def _finalize(*_args, src_ctx_path: Path, **_kwargs) -> None: + assert src_ctx_path.is_file() + events.append("finalize") + + fake_winml_session.reset.side_effect = _reset + ep_device = MagicMock() + ep_device.device.device_type = "NPU" + context = CompileContext( + model_path=model_path, + config={ + "execution_provider": "qnn", + "enable_ep_context": True, + "validate": False, + "ep_device": EPDeviceTarget(ep="qnn", device="npu").to_dict(), + }, + ) + + with ( + patch.dict( + "winml.modelkit.compiler.stages.compile.COMPILER_SESSION_MAPPING", + {"ort": MagicMock(return_value=fake_winml_session)}, + clear=False, + ), + patch("winml.modelkit.compiler.stages.compile.WinMLEPRegistry.instance") as registry, + patch.object(CompileStage, "_finalize_output", side_effect=_finalize), + ): + registry.return_value.auto_device.return_value = ep_device + CompileStage().process(context) + + assert events == ["finalize", "reset"] + def test_process_reconstructs_explicit_serialized_device(self, tmp_path): from unittest.mock import MagicMock, patch @@ -451,6 +506,8 @@ def test_process_reconstructs_explicit_serialized_device(self, tmp_path): fake_winml_session = MagicMock() fake_winml_session._session = fake_session + identity_ctx_path = tmp_path / "model_1234567890abcdef_ctx.onnx" + fake_winml_session.running_model_path = identity_ctx_path context = CompileContext( model_path=model_path, @@ -492,6 +549,53 @@ def test_process_reconstructs_explicit_serialized_device(self, tmp_path): mock_resolve_device.assert_called_once_with(EPDeviceTarget(ep="qnn", device="gpu")) assert mock_finalize_output.call_args.kwargs["device"] == "gpu" + assert mock_finalize_output.call_args.kwargs["src_ctx_path"] == identity_ctx_path + + def test_single_model_compile_fallback_does_not_publish_stale_context(self, tmp_path): + """A raw running path wins over stale identity artifacts in the same directory.""" + from unittest.mock import MagicMock, patch + + from winml.modelkit.compiler import CompileContext, CompileStage + + model_path = tmp_path / "model.onnx" + create_simple_model(model_path) + stale_context = tmp_path / "model_npu_staleidentity_ctx.onnx" + create_epcontext_onnx(stale_context, "embedded", embed_mode=1) + fake_session = MagicMock() + fake_session.get_providers.return_value = ["QNNExecutionProvider"] + fake_session.get_inputs.return_value = [] + fake_session.get_outputs.return_value = [] + fake_winml_session = MagicMock() + fake_winml_session._session = fake_session + fake_winml_session.running_model_path = model_path + context = CompileContext( + model_path=model_path, + config={ + "execution_provider": "qnn", + "device": "npu", + "enable_ep_context": True, + "validate": False, + }, + ) + stage = CompileStage() + ep_device = MagicMock() + ep_device.device.device_type = "NPU" + + with ( + patch.dict( + "winml.modelkit.compiler.stages.compile.COMPILER_SESSION_MAPPING", + {"ort": MagicMock(return_value=fake_winml_session)}, + clear=False, + ), + patch("winml.modelkit.compiler.stages.compile.WinMLEPRegistry.instance") as registry, + patch.object(stage, "_finalize_output") as finalize_output, + ): + registry.return_value.auto_device.return_value = ep_device + stage.process(context) + + finalize_output.assert_not_called() + assert context.output_path is None + assert context.warnings == [f"No EPContext produced for {model_path.name}"] def test_multi_model_sequence_shares_options_and_closes_context(self, tmp_path): """First, intermediate, and final models share one EP context in sequence.""" @@ -684,6 +788,64 @@ def test_updates_ep_cache_context_in_external_mode(self, tmp_path): assert b"mymodel_qnn_ctx" in attr.s, f"Expected updated name, got {attr.s}" break + def test_finalize_output_preserves_previous_referenced_binary_generation(self, tmp_path): + """A later publication cannot overwrite a binary used by an older ONNX.""" + from winml.modelkit.compiler import CompileContext, CompileStage + + work_dir = tmp_path / "work" + output_dir = tmp_path / "output" + work_dir.mkdir() + output_dir.mkdir() + original_model_path = tmp_path / "mymodel.onnx" + create_simple_model(original_model_path) + context = CompileContext( + model_path=original_model_path, + config={"execution_provider": "qnn", "output_path": str(output_dir)}, + work_dir=work_dir, + ) + stage = CompileStage() + + first_ctx = work_dir / "first_identity_ctx.onnx" + create_epcontext_onnx(first_ctx, "first_identity_ctx_qnn.bin", embed_mode=0) + (work_dir / "first_identity_ctx_qnn.bin").write_bytes(b"first binary") + stage._finalize_output( + context, + work_dir / "model_to_compile.onnx", + output_dir, + src_ctx_path=first_ctx, + ) + first_published_model = onnx.load(str(context.output_path), load_external_data=False) + first_ref = next( + attr.s.decode("utf-8") + for node in first_published_model.graph.node + for attr in node.attribute + if node.op_type == "EPContext" and attr.name == "ep_cache_context" + ) + first_published_binary = output_dir / first_ref + assert first_published_binary.read_bytes() == b"first binary" + + second_ctx = work_dir / "second_identity_ctx.onnx" + create_epcontext_onnx(second_ctx, "second_identity_ctx_qnn.bin", embed_mode=0) + (work_dir / "second_identity_ctx_qnn.bin").write_bytes(b"second binary") + stage._finalize_output( + context, + work_dir / "model_to_compile.onnx", + output_dir, + src_ctx_path=second_ctx, + ) + second_published_model = onnx.load(str(context.output_path), load_external_data=False) + second_ref = next( + attr.s.decode("utf-8") + for node in second_published_model.graph.node + for attr in node.attribute + if node.op_type == "EPContext" and attr.name == "ep_cache_context" + ) + + assert second_ref != first_ref + assert first_published_binary.read_bytes() == b"first binary" + assert (output_dir / second_ref).read_bytes() == b"second binary" + assert (output_dir / "mymodel_qnn_ctx_qnn.bin").read_bytes() == b"second binary" + def test_updates_matching_cache_reference_with_malformed_main_context(self, tmp_path): """Binary renames follow the referenced file, not malformed main metadata.""" from winml.modelkit.compiler import CompileContext, CompileStage @@ -783,9 +945,17 @@ def test_finalize_output_copies_all_referenced_context_binaries(self, tmp_path): for attr in node.attribute if node.op_type == "EPContext" and attr.name == "ep_cache_context" } - assert cache_refs == { - "mymodel_qnn_ctx.bin", - "mymodel_qnn_ctx_partition_1.bin", + assert len(cache_refs) == 2 + assert any( + ref.startswith("mymodel_qnn_ctx.") and ref.endswith(".bin") for ref in cache_refs + ) + assert any( + ref.startswith("mymodel_qnn_ctx_partition_1.") and ref.endswith(".bin") + for ref in cache_refs + ) + assert {((output_dir / ref).read_bytes()) for ref in cache_refs} == { + b"main binary", + b"secondary binary", } assert (output_dir / "mymodel_qnn_ctx.bin").read_bytes() == b"main binary" assert (output_dir / "mymodel_qnn_ctx_partition_1.bin").read_bytes() == b"secondary binary" diff --git a/tests/unit/onnx/test_external_data.py b/tests/unit/onnx/test_external_data.py index 2acb81a82..02c5b132c 100644 --- a/tests/unit/onnx/test_external_data.py +++ b/tests/unit/onnx/test_external_data.py @@ -19,6 +19,7 @@ from winml.modelkit.onnx.external_data import ( copy_onnx_model, get_external_data_files, + get_onnx_model_hash, has_external_data, ) @@ -91,6 +92,25 @@ def test_with_external_data(self, tmp_path: Path) -> None: assert get_external_data_files(path) == ["ext.onnx.data"] +def test_strict_model_hash_detects_content_change_with_preserved_metadata(tmp_path: Path) -> None: + """Strict cache identity hashes bytes, not only path/size/mtime metadata.""" + path = tmp_path / "model.onnx" + model = _make_filled_model(1.0, (4, 4)) + onnx.save(model, path) + original_stat = path.stat() + original_hash = get_onnx_model_hash(path, strict=True) + + replacement = _make_filled_model(2.0, (4, 4)) + onnx.save(replacement, path) + assert path.stat().st_size == original_stat.st_size + path.touch() + import os + + os.utime(path, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns)) + + assert get_onnx_model_hash(path, strict=True) != original_hash + + class TestHasExternalData: """Tests for has_external_data().""" diff --git a/tests/unit/session/conftest.py b/tests/unit/session/conftest.py index f445bf996..c58dbbccb 100644 --- a/tests/unit/session/conftest.py +++ b/tests/unit/session/conftest.py @@ -42,17 +42,18 @@ def test_qnn_inference(self, simple_matmul_onnx, qnn_npu_ep_device, fake_ort_npu def _stub_ep_entry(ep_name: str) -> EPEntry: """Build a minimal EPEntry suitable for wrapping a mocked OrtEpDevice. - The dll_path is fictional — tests never load the DLL because they - construct WinMLEP/WinMLEPDevice directly. + This fixture file stands in for an existing DLL so cache-identity tests can + fingerprint a stable path without loading any native library. """ return EPEntry( ep_name=ep_name, - dll_path=Path(f"C:/fake/{ep_name}.dll"), + dll_path=Path(__file__), source=PyPISource( distribution="fake-dist", relative_dll="fake.dll", eps=(ep_name,), ), + version="test-version", ) diff --git a/tests/unit/session/test_perf_auto_reset.py b/tests/unit/session/test_perf_auto_reset.py index ee4f0770c..3b7722e0f 100644 --- a/tests/unit/session/test_perf_auto_reset.py +++ b/tests/unit/session/test_perf_auto_reset.py @@ -287,7 +287,7 @@ def _fake_inference_session(*_args, **_kwargs): assert "useful error" in stderr -def test_requires_teardown_reset_preserves_warning_stderr(monkeypatch, capfd): +def test_requires_teardown_preserves_warning_stderr(monkeypatch, capfd): """Library perf teardown leaves stderr untouched; CLI owns warning filtering.""" from winml.modelkit.session.monitor.ep_monitor import WinMLEPMonitor @@ -309,16 +309,21 @@ def to_dict(self): session, _cpu_dev, _cpu_ep = _make_cpu_session(get_minimal_onnx_model_path()) session.compile() + cleanup_modes: list[bool] = [] + original_reset = session._reset_runtime_state - def noisy_reset() -> None: + def noisy_reset(*, cleanup_markerless: bool) -> None: os.write(2, b"2026 [W:custom-native:, file.cc:1 ResetWarn] reset warning\n") + cleanup_modes.append(cleanup_markerless) + original_reset(cleanup_markerless=cleanup_markerless) - monkeypatch.setattr(session, "reset", noisy_reset) + monkeypatch.setattr(session, "_reset_runtime_state", noisy_reset) with session.perf(monitor=_TeardownMonitor()): pass assert "reset warning" in capfd.readouterr().err + assert cleanup_modes == [False] def test_no_auto_reset_when_monitor_empty(): diff --git a/tests/unit/session/test_winml_session.py b/tests/unit/session/test_winml_session.py index 1632b742b..903042363 100644 --- a/tests/unit/session/test_winml_session.py +++ b/tests/unit/session/test_winml_session.py @@ -17,15 +17,17 @@ from __future__ import annotations -from typing import TYPE_CHECKING +import os +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import numpy as np - - -if TYPE_CHECKING: - from pathlib import Path import pytest +from onnx import TensorProto, helper, load, numpy_helper, save, save_model from winml.modelkit.compiler import EPConfig from winml.modelkit.session import ( @@ -51,6 +53,75 @@ def _stub_registry(monkeypatch: pytest.MonkeyPatch, ep_device: object) -> MagicM return registry +def _write_fake_epcontext(session: WinMLSession, path: str) -> None: + """Write a valid EPContext graph and its optional external binary.""" + ctx_path = Path(path) + if session._embed_context: + cache_value = b"embedded context" + else: + binary_path = ctx_path.with_name(f"{ctx_path.stem}_qnn.bin") + binary_path.write_bytes(b"external context") + cache_value = binary_path.name + node = helper.make_node( + "EPContext", + inputs=[], + outputs=["output"], + name="ep_context_0", + domain="com.microsoft", + embed_mode=1 if session._embed_context else 0, + ep_cache_context=cache_value, + main_context=1, + ) + output = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1]) + graph = helper.make_graph([node], "epcontext_graph", [], [output]) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 17), + helper.make_opsetid("com.microsoft", 1), + ], + ) + model.ir_version = 9 + save(model, ctx_path) + + +def _compile_with_fake_ort(session: WinMLSession) -> MagicMock: + """Compile through mocked ORT while preserving its file-output contract.""" + inference_session = MagicMock() + inference_session.get_providers.return_value = ["QNNExecutionProvider"] + model_compiler = MagicMock() + + model_compiler.return_value.compile_to_file.side_effect = lambda path: _write_fake_epcontext( + session, path + ) + with ( + patch("winml.modelkit.session.session._build_session_options", return_value=MagicMock()), + patch("winml.modelkit.session.session.ort.ModelCompiler", model_compiler), + patch( + "winml.modelkit.session.session.ort.InferenceSession", + return_value=inference_session, + ), + ): + session.compile() + return model_compiler + + +def _cache_path(session: WinMLSession) -> Path: + """Return the deterministic EPContext path for this test session.""" + return session._epcontext_cache_path(session._epcontext_cache_identity()) + + +def _compiled_generation(session: WinMLSession, model_compiler: MagicMock) -> Path: + """Return the single compiled generation and validate its identity namespace.""" + model_compiler.return_value.compile_to_file.assert_called_once() + generation = Path(model_compiler.return_value.compile_to_file.call_args.args[0]) + cache_path = _cache_path(session) + assert generation.parent == cache_path.parent + assert generation.name.startswith(f"{cache_path.stem}_") + assert generation.suffix == cache_path.suffix + return generation + + class TestWinMLSessionInstantiation: """Test WinMLSession instantiation with EPDeviceTarget-based selection.""" @@ -178,6 +249,1159 @@ def test_compile_is_idempotent(self, cpu_winml_session: WinMLSession): session.compile() assert session._session is first_session + def test_compile_rebuilds_legacy_cache_without_identity_marker( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """A newer sibling context without cache identity is not reusable.""" + ctx_path = simple_matmul_onnx.with_name(f"{simple_matmul_onnx.stem}_npu_ctx.onnx") + ctx_path.write_bytes(b"legacy context") + source_mtime = simple_matmul_onnx.stat().st_mtime_ns + os.utime(ctx_path, ns=(source_mtime + 1_000_000, source_mtime + 1_000_000)) + + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + model_compiler = _compile_with_fake_ort(session) + compiled_path = _compiled_generation(session, model_compiler) + + assert session.running_model_path == compiled_path + assert ctx_path.read_bytes() == b"legacy context" + + def test_compile_rebuilds_cache_with_non_object_marker( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """A valid JSON marker with the wrong shape is a cache miss, not an error.""" + legacy_path = simple_matmul_onnx.with_name(f"{simple_matmul_onnx.stem}_npu_ctx.onnx") + legacy_path.write_bytes(b"legacy context") + marker_path = WinMLSession._epcontext_cache_marker_path(legacy_path) + marker_path.write_text("[]", encoding="utf-8") + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(session) + compiled_path = _compiled_generation(session, model_compiler) + + assert session.running_model_path == compiled_path + + def test_compile_reuses_cache_with_matching_identity( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """An unchanged source and compile identity reuse the sibling context.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"htp_performance_mode": "burst"}, + enable_ep_context=True, + ), + ) + first_compiler = _compile_with_fake_ort(first_session) + first_compiler.return_value.compile_to_file.assert_called_once() + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"htp_performance_mode": "burst"}, + enable_ep_context=True, + ), + ) + + model_compiler = _compile_with_fake_ort(session) + + model_compiler.assert_not_called() + + def test_compile_rebuilds_cache_when_provider_options_change( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """Effective provider options are part of the direct-session cache key.""" + old_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"htp_performance_mode": "default"}, + enable_ep_context=True, + ), + ) + _compile_with_fake_ort(old_session) + new_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"htp_performance_mode": "burst"}, + enable_ep_context=True, + ), + ) + + model_compiler = _compile_with_fake_ort(new_session) + ctx_path = _compiled_generation(new_session, model_compiler) + + assert new_session.running_model_path == ctx_path + + def test_compile_rebuilds_cache_when_provider_option_file_content_changes( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + tmp_path: Path, + ) -> None: + """Existing file-valued provider options are fingerprinted by content.""" + option_file = tmp_path / "compiler-input.bin" + option_file.write_bytes(b"alpha") + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"compiler_input": str(option_file)}, + provider_option_file_keys={"compiler_input"}, + enable_ep_context=True, + ), + ) + _compile_with_fake_ort(first_session) + original_stat = option_file.stat() + option_file.write_bytes(b"bravo") + os.utime( + option_file, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"compiler_input": str(option_file)}, + provider_option_file_keys={"compiler_input"}, + enable_ep_context=True, + ), + ) + + model_compiler = _compile_with_fake_ort(second_session) + ctx_path = _compiled_generation(second_session, model_compiler) + + assert second_session.running_model_path == ctx_path + + def test_relative_provider_option_file_is_canonicalized_for_identity_and_ort( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """File-valued relative options are resolved once before hashing and ORT binding.""" + option_name = "compiler-input.bin" + cwd_dir = tmp_path / "cwd" + cwd_dir.mkdir() + cwd_option_file = cwd_dir / option_name + cwd_option_file.write_bytes(b"cwd option") + model_dir_option_file = simple_matmul_onnx.parent / option_name + model_dir_option_file.write_bytes(b"model option") + monkeypatch.chdir(cwd_dir) + captured_provider_options: list[dict[str, str]] = [] + + def _session_options(*_args, provider_options, **_kwargs): + captured_provider_options.append(dict(provider_options)) + return MagicMock() + + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"compiler_input": option_name}, + provider_option_file_keys={"compiler_input"}, + enable_ep_context=True, + ), + ) + expected_path = str(cwd_option_file.resolve()) + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + _session_options, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + MagicMock( + return_value=SimpleNamespace( + compile_to_file=lambda path: _write_fake_epcontext(session, path) + ) + ), + ) + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + lambda *_args, **_kwargs: runtime_session, + ) + + session.compile() + + identity = session._epcontext_cache_identity() + assert session._provider_options["compiler_input"] == expected_path + assert identity["provider_option_files"]["compiler_input"]["path"] == expected_path + assert captured_provider_options + assert all( + options["compiler_input"] == expected_path for options in captured_provider_options + ) + + def test_plain_provider_option_matching_cwd_file_is_not_rewritten( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Existing files do not make unrelated provider-option values file-backed.""" + (tmp_path / "default").write_bytes(b"unrelated") + monkeypatch.chdir(tmp_path) + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"htp_performance_mode": "default"}, + enable_ep_context=True, + ), + ) + + identity = session._epcontext_cache_identity() + + assert session._provider_options["htp_performance_mode"] == "default" + assert "htp_performance_mode" not in identity["provider_option_files"] + + def test_declared_provider_option_file_must_resolve( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """An explicitly file-backed option cannot retain an ambiguous relative value.""" + with pytest.raises(ValueError, match="compiler_input"): + WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"compiler_input": "missing.bin"}, + provider_option_file_keys={"compiler_input"}, + enable_ep_context=True, + ), + ) + + def test_different_compile_identities_use_distinct_context_paths( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """One identity cannot overwrite artifacts loaded by another session.""" + sessions = [ + WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": mode}, + enable_ep_context=True, + ), + ) + for mode in ("first", "second") + ] + + paths = [ + session._epcontext_cache_path(session._epcontext_cache_identity()) + for session in sessions + ] + + assert paths[0] != paths[1] + assert all(path.parent == simple_matmul_onnx.parent for path in paths) + assert all(path.name.startswith(f"{simple_matmul_onnx.stem}_npu_") for path in paths) + assert all(path.name.endswith("_ctx.onnx") for path in paths) + + def test_compile_rebuilds_cache_when_embed_mode_changes( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """Embedded and external EPContext artifacts never share a cache entry.""" + old_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True, embed_context=False), + ) + _compile_with_fake_ort(old_session) + new_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True, embed_context=True), + ) + + model_compiler = _compile_with_fake_ort(new_session) + ctx_path = _compiled_generation(new_session, model_compiler) + + assert new_session.running_model_path == ctx_path + + def test_compile_rebuilds_cache_when_external_context_binary_is_missing( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """A marker cannot make an EPContext with a missing binary reusable.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + ctx_path = first_session.running_model_path + binary_path = ctx_path.with_name(f"{ctx_path.stem}_qnn.bin") + binary_path.unlink() + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(second_session) + + _compiled_generation(second_session, model_compiler) + + def test_compile_rebuilds_cache_when_external_context_binary_changes( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """A replaced EPContext binary invalidates an otherwise matching marker.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + ctx_path = first_session.running_model_path + binary_path = ctx_path.with_name(f"{ctx_path.stem}_qnn.bin") + binary_path.write_bytes(b"replaced external context") + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(second_session) + + _compiled_generation(second_session, model_compiler) + + def test_compile_rebuilds_cache_when_binary_content_changes_with_same_metadata( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """Content digests catch replacements that preserve size and mtime.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + ctx_path = first_session.running_model_path + binary_path = ctx_path.with_name(f"{ctx_path.stem}_qnn.bin") + original_stat = binary_path.stat() + binary_path.write_bytes(b"tampered content") + os.utime( + binary_path, + ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns), + ) + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(second_session) + + _compiled_generation(second_session, model_compiler) + + def test_compile_failure_preserves_other_identity_cache_and_uses_source( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """A failed identity compile leaves other caches intact and uses source.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "old"}, + enable_ep_context=True, + ), + ) + _compile_with_fake_ort(first_session) + ctx_path = _cache_path(first_session) + marker_path = first_session._epcontext_cache_marker_path(ctx_path) + assert marker_path.is_file() + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "new"}, + enable_ep_context=True, + ), + ) + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + inference_session = MagicMock(return_value=runtime_session) + model_compiler = MagicMock() + model_compiler.return_value.compile_to_file.side_effect = RuntimeError("compile failed") + with ( + patch( + "winml.modelkit.session.session._build_session_options", + return_value=MagicMock(), + ), + patch("winml.modelkit.session.session.ort.ModelCompiler", model_compiler), + patch( + "winml.modelkit.session.session.ort.InferenceSession", + inference_session, + ), + ): + second_session.compile() + + assert marker_path.exists() + assert inference_session.call_count == 1 + assert inference_session.call_args.args[0] == str(simple_matmul_onnx) + + def test_compile_failure_removes_private_generation_and_sidecar( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Failed private generations are removed with sidecars before fallback.""" + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + written_paths: list[Path] = [] + + class _FailingCompiler: + def __init__(self, *_args, **_kwargs): + pass + + def compile_to_file(self, path: str) -> None: + generation_path = Path(path) + sidecar_path = generation_path.with_name(f"{generation_path.stem}_partial.bin") + generation_path.write_bytes(b"partial context") + sidecar_path.write_bytes(b"partial sidecar") + written_paths.extend([generation_path, sidecar_path]) + raise RuntimeError("compile failed") + + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + _FailingCompiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + lambda *_args, **_kwargs: runtime_session, + ) + + session.compile() + + assert session.running_model_path == simple_matmul_onnx + assert written_paths + assert all(not path.exists() for path in written_paths) + + def test_compile_rebuilds_cache_when_source_external_data_changes( + self, + tmp_path: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """Referenced source weight sidecars participate in cache identity.""" + source_path = tmp_path / "external_model.onnx" + weight = numpy_helper.from_array(np.ones((4, 4), dtype=np.float32), name="weight") + input_info = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 4]) + output_info = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 4]) + node = helper.make_node("MatMul", ["input", "weight"], ["output"]) + graph = helper.make_graph([node], "external_graph", [input_info], [output_info], [weight]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + save_model( + model, + source_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location="external_model.onnx.data", + size_threshold=0, + ) + data_path = tmp_path / "external_model.onnx.data" + first_session = WinMLSession( + onnx_path=source_path, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + data_path.write_bytes(data_path.read_bytes() + b"changed") + second_session = WinMLSession( + onnx_path=source_path, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(second_session) + + ctx_path = _compiled_generation(second_session, model_compiler) + assert second_session.running_model_path == ctx_path + + def test_source_external_data_introspection_failure_disables_cache( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Uncertain source identity recompiles instead of falling back to ONNX-only cache keys.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + ctx_path = _cache_path(first_session) + marker_path = first_session._epcontext_cache_marker_path(ctx_path) + assert marker_path.is_file() + + def _fail_external_data(_model_path: Path) -> list[str]: + raise PermissionError("cannot inspect external data") + + monkeypatch.setattr( + "winml.modelkit.onnx.external_data.get_external_data_files", + _fail_external_data, + ) + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + + model_compiler = _compile_with_fake_ort(second_session) + + model_compiler.return_value.compile_to_file.assert_called_once() + assert marker_path.exists() + + def test_source_change_during_compile_retries_with_new_identity( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A generation is published only under the source identity it compiled.""" + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + compile_calls = 0 + + class _SourceChangingCompiler: + def __init__(self, *_args, **_kwargs): + pass + + def compile_to_file(self, path: str) -> None: + nonlocal compile_calls + compile_calls += 1 + _write_fake_epcontext(session, path) + if compile_calls == 1: + model = load(simple_matmul_onnx) + model.producer_name = "changed-during-compile" + save(model, simple_matmul_onnx) + + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + _SourceChangingCompiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + lambda *_args, **_kwargs: runtime_session, + ) + + session.compile() + + current_identity = session._epcontext_cache_identity() + cache_path = session._epcontext_cache_path(current_identity) + assert compile_calls == 2 + assert session.running_model_path == session._epcontext_cached_generation( + cache_path, + current_identity, + ) + + def test_final_identity_failure_discards_unpublished_generation( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A retry cannot abandon a markerless generation after preparation.""" + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + stable_identity = session._epcontext_cache_identity() + identity_calls = 0 + compiled_paths: list[Path] = [] + + def _identity() -> dict[str, object]: + nonlocal identity_calls + identity_calls += 1 + if identity_calls == 4: + raise OSError("final identity unavailable") + if identity_calls == 5: + raise ValueError("cache identity unavailable") + return stable_identity + + def _compile(path: str) -> None: + compiled_path = Path(path) + compiled_paths.append(compiled_path) + _write_fake_epcontext(session, path) + + def _fail_marker(*_args, **_kwargs) -> None: + raise PermissionError("marker publication failed") + + model_compiler = MagicMock() + model_compiler.return_value.compile_to_file.side_effect = _compile + monkeypatch.setattr(session, "_epcontext_cache_identity", _identity) + monkeypatch.setattr(session, "_write_epcontext_cache_marker", _fail_marker) + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + model_compiler, + ) + + prepared = session._compile_epcontext_with_stable_source( + simple_matmul_onnx.parent / "compile.log" + ) + try: + assert len(compiled_paths) == 2 + assert not compiled_paths[0].exists() + assert not compiled_paths[0].with_name(f"{compiled_paths[0].stem}_qnn.bin").exists() + assert prepared.path == compiled_paths[1] + assert prepared.path.exists() + finally: + prepared.release() + for compiled_path in compiled_paths: + session._discard_epcontext_generation(compiled_path) + + def test_marker_write_failure_keeps_compiled_context_usable( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Cache metadata failure does not discard a successful compilation.""" + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + first_compiler = _compile_with_fake_ort(session) + first_compiler.return_value.compile_to_file.assert_called_once() + cached_path = session.running_model_path + cached_bytes = cached_path.read_bytes() + marker_path = session._epcontext_cache_marker_path(_cache_path(session)) + marker_path.unlink() + session.reset() + + def _fail_marker(*_args, **_kwargs): + raise PermissionError("marker directory is read-only") + + monkeypatch.setattr(session, "_write_epcontext_cache_marker", _fail_marker) + model_compiler = _compile_with_fake_ort(session) + + model_compiler.return_value.compile_to_file.assert_called_once() + assert session.running_model_path != cached_path + assert cached_path.read_bytes() == cached_bytes + assert session._session is not None + generation_path = Path(model_compiler.return_value.compile_to_file.call_args.args[0]) + sidecar_path = generation_path.with_name(f"{generation_path.stem}_qnn.bin") + assert generation_path.is_file() + assert sidecar_path.is_file() + + session.reset() + + assert not generation_path.exists() + assert not sidecar_path.exists() + + def test_custom_session_options_factory_disables_cache_reuse( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + ) -> None: + """Opaque SessionOptions factory state is never represented as a reusable cache key.""" + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + ) + _compile_with_fake_ort(first_session) + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + session_options=MagicMock, + ) + + model_compiler = _compile_with_fake_ort(second_session) + + model_compiler.return_value.compile_to_file.assert_called_once() + assert second_session.running_model_path != first_session.running_model_path + generation_path = Path(model_compiler.return_value.compile_to_file.call_args.args[0]) + sidecar_path = generation_path.with_name(f"{generation_path.stem}_qnn.bin") + assert generation_path.is_file() + assert sidecar_path.is_file() + + second_session.reset() + + assert not generation_path.exists() + assert not sidecar_path.exists() + + def test_markerless_generation_survives_perf_rebuild_until_reset( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Temporary perf rebuilds can reopen a session-owned markerless model.""" + session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig(provider="qnn", enable_ep_context=True), + session_options=MagicMock, + ) + opened_paths: list[Path] = [] + + def _runtime_session(path: str | Path, *_args, **_kwargs) -> MagicMock: + model_path = Path(path) + if model_path != simple_matmul_onnx: + assert model_path.is_file() + opened_paths.append(model_path) + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + return runtime_session + + model_compiler = MagicMock() + model_compiler.return_value.compile_to_file.side_effect = lambda path: ( + _write_fake_epcontext(session, path) + ) + monitor = MagicMock() + monitor.ep_name = "qnn" + monitor.requires_session_teardown = False + monitor.get_provider_options.return_value = {"profiling_level": "detailed"} + monitor.get_session_options.return_value = {} + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + model_compiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + _runtime_session, + ) + + session.compile() + generation_path = session.running_model_path + sidecar_path = generation_path.with_name(f"{generation_path.stem}_qnn.bin") + + assert generation_path.is_file() + with session.perf(monitor=monitor): + assert session.running_model_path == generation_path + assert generation_path.is_file() + assert opened_paths.count(generation_path) == 3 + + session.reset() + + assert not generation_path.exists() + assert not sidecar_path.exists() + + def test_successful_epcontext_cache_prunes_old_identity_artifacts( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A bounded successful cache removes old generations, markers, locks, and sidecars.""" + monkeypatch.setattr( + "winml.modelkit.session.session._EPCONTEXT_CACHE_MAX_GENERATIONS", + 1, + raising=False, + ) + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "old"}, + enable_ep_context=True, + ), + ) + first_compiler = _compile_with_fake_ort(first_session) + first_generation = _compiled_generation(first_session, first_compiler) + first_sidecar = first_generation.with_name(f"{first_generation.stem}_qnn.bin") + first_cache_path = _cache_path(first_session) + first_marker = first_session._epcontext_cache_marker_path(first_cache_path) + first_lock = first_cache_path.with_name(f"{first_cache_path.name}.lock") + assert first_generation.is_file() + assert first_sidecar.is_file() + assert first_marker.is_file() + first_lock.write_text("stale lock", encoding="utf-8") + assert first_lock.is_file() + first_session.reset() + + second_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "new"}, + enable_ep_context=True, + ), + ) + second_compiler = _compile_with_fake_ort(second_session) + second_generation = _compiled_generation(second_session, second_compiler) + second_cache_path = _cache_path(second_session) + second_marker = second_session._epcontext_cache_marker_path(second_cache_path) + + assert not first_generation.exists() + assert not first_sidecar.exists() + assert not first_marker.exists() + assert not first_lock.exists() + assert second_generation.is_file() + assert second_marker.is_file() + + def test_cache_prune_skips_marker_replaced_before_lease( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A stale scan cannot unlink a newly published generation marker.""" + from winml.modelkit.session.session import _EPContextCacheLease + + old_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "old"}, + enable_ep_context=True, + ), + ) + old_compiler = _compile_with_fake_ort(old_session) + old_generation = _compiled_generation(old_session, old_compiler) + old_identity = old_session._epcontext_cache_identity() + old_cache_path = old_session._epcontext_cache_path(old_identity) + old_lock_path = old_cache_path.with_name(f"{old_cache_path.name}.lock") + current_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "current"}, + enable_ep_context=True, + ), + ) + current_compiler = _compile_with_fake_ort(current_session) + _compiled_generation(current_session, current_compiler) + current_cache_path = _cache_path(current_session) + replacement_generation = old_cache_path.with_name( + f"{old_cache_path.stem}_replacement{old_cache_path.suffix}" + ) + real_acquire = _EPContextCacheLease.acquire + replaced = False + + def _acquire(lock_path: Path, *, blocking: bool = True): + nonlocal replaced + if ( + not blocking + and not replaced + and lock_path.resolve(strict=False) == old_lock_path.resolve(strict=False) + ): + replaced = True + _write_fake_epcontext(old_session, str(replacement_generation)) + old_session._write_epcontext_cache_marker( + old_cache_path, + replacement_generation, + old_identity, + ) + return real_acquire(lock_path, blocking=blocking) + + monkeypatch.setattr( + "winml.modelkit.session.session._EPCONTEXT_CACHE_MAX_GENERATIONS", + 1, + ) + monkeypatch.setattr( + _EPContextCacheLease, + "acquire", + staticmethod(_acquire), + ) + + current_session._prune_epcontext_cache(current_cache_path) + + assert replaced is True + assert replacement_generation.is_file() + assert ( + old_session._epcontext_cached_generation( + old_cache_path, + old_identity, + ) + == replacement_generation + ) + assert not old_generation.exists() + + def test_cache_hit_generation_is_pinned_until_runtime_session_opens( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Pruning a different identity cannot delete a cache hit before ORT opens it.""" + from winml.modelkit.session.session import _epcontext_thread_lock + + checked = False + old_generation: Path | None = None + old_lock: Path | None = None + + def _session_options(*_args, provider_options, **_kwargs): + return SimpleNamespace(mode=provider_options["mode"]) + + class _ModeCompiler: + def __init__(self, session_options, *_args, **_kwargs): + self.mode = session_options.mode + + def compile_to_file(self, path: str) -> None: + _write_fake_epcontext(first_session, path) + + def _runtime_session(path: str, *_args, **_kwargs): + nonlocal checked + model_path = Path(path) + if old_generation is not None and old_lock is not None and model_path == old_generation: + checked = True + assert _epcontext_thread_lock(old_lock).locked() + runtime_session = MagicMock() + runtime_session.get_providers.return_value = ["QNNExecutionProvider"] + return runtime_session + + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + _session_options, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + _ModeCompiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + _runtime_session, + ) + first_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "old"}, + enable_ep_context=True, + ), + ) + first_session.compile() + old_generation = first_session.running_model_path + old_cache_path = _cache_path(first_session) + old_lock = old_cache_path.with_name(f"{old_cache_path.name}.lock") + first_session.reset() + hit_session = WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "old"}, + enable_ep_context=True, + ), + ) + + hit_session.compile() + + assert checked is True + assert hit_session.running_model_path == old_generation + + def test_concurrent_different_identities_use_distinct_artifacts( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """Concurrent different identities never write one shared artifact.""" + sessions = [ + WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": mode}, + enable_ep_context=True, + ), + ) + for mode in ("first", "second") + ] + state_lock = threading.Lock() + first_entered = threading.Event() + second_entered = threading.Event() + state = {"active": 0, "max_active": 0} + compiled_paths: dict[str, Path] = {} + + def _session_options(*_args, provider_options, **_kwargs): + return SimpleNamespace(mode=provider_options["mode"]) + + class _ConcurrentCompiler: + def __init__(self, session_options, *_args, **_kwargs): + self.mode = session_options.mode + + def compile_to_file(self, path: str) -> None: + with state_lock: + state["active"] += 1 + state["max_active"] = max(state["max_active"], state["active"]) + if self.mode == "first": + first_entered.set() + assert second_entered.wait(timeout=1) + else: + assert first_entered.wait(timeout=1) + second_entered.set() + compiled_paths[self.mode] = Path(path) + session = sessions[0] if self.mode == "first" else sessions[1] + _write_fake_epcontext(session, path) + with state_lock: + state["active"] -= 1 + + inference_session = MagicMock() + inference_session.get_providers.return_value = ["QNNExecutionProvider"] + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + _session_options, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + _ConcurrentCompiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + lambda *_args, **_kwargs: inference_session, + ) + monkeypatch.setattr( + "winml.modelkit.session.session._suppress_native_output", + lambda *_args, **_kwargs: nullcontext(), + ) + with ThreadPoolExecutor(max_workers=2) as executor: + first_future = executor.submit(sessions[0].compile) + assert first_entered.wait(timeout=1) + second_future = executor.submit(sessions[1].compile) + first_future.result(timeout=5) + second_future.result(timeout=5) + + assert state["max_active"] == 2 + assert compiled_paths["first"] != compiled_paths["second"] + assert sessions[0].running_model_path == compiled_paths["first"] + assert sessions[1].running_model_path == compiled_paths["second"] + + def test_concurrent_matching_identity_compiles_once( + self, + simple_matmul_onnx: Path, + qnn_npu_ep_device: WinMLEPDevice, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + """A waiter rechecks the marker and reuses the first identity artifact.""" + sessions = [ + WinMLSession( + onnx_path=simple_matmul_onnx, + ep_device=qnn_npu_ep_device, + ep_config=EPConfig( + provider="qnn", + provider_options={"mode": "shared"}, + enable_ep_context=True, + ), + ) + for _ in range(2) + ] + first_entered = threading.Event() + release_first = threading.Event() + compile_calls = 0 + + class _SingleCompiler: + def __init__(self, *_args, **_kwargs): + pass + + def compile_to_file(self, path: str) -> None: + nonlocal compile_calls + compile_calls += 1 + first_entered.set() + assert release_first.wait(timeout=1) + _write_fake_epcontext(sessions[0], path) + + inference_session = MagicMock() + inference_session.get_providers.return_value = ["QNNExecutionProvider"] + monkeypatch.setattr( + "winml.modelkit.session.session._build_session_options", + lambda *_args, **_kwargs: MagicMock(), + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.ModelCompiler", + _SingleCompiler, + ) + monkeypatch.setattr( + "winml.modelkit.session.session.ort.InferenceSession", + lambda *_args, **_kwargs: inference_session, + ) + monkeypatch.setattr( + "winml.modelkit.session.session._suppress_native_output", + lambda *_args, **_kwargs: nullcontext(), + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + first_future = executor.submit(sessions[0].compile) + assert first_entered.wait(timeout=1) + second_future = executor.submit(sessions[1].compile) + release_first.set() + first_future.result(timeout=5) + second_future.result(timeout=5) + + assert compile_calls == 1 + assert sessions[0].running_model_path == sessions[1].running_model_path + def test_runtime_compile_bypasses_model_compiler( self, simple_matmul_onnx: Path, diff --git a/tests/unit/test_uv_lock.py b/tests/unit/test_uv_lock.py index 734ac208c..9e928cf0b 100644 --- a/tests/unit/test_uv_lock.py +++ b/tests/unit/test_uv_lock.py @@ -38,3 +38,37 @@ def test_uv_lock_does_not_include_cuda_accelerator_packages() -> None: assert not disallowed_refs, "Unexpected CUDA/NVIDIA lock entries: " + ", ".join( sorted(disallowed_refs) ) + + +def test_uv_lock_records_direct_project_dependencies() -> None: + """The editable lock entry must retain every direct project dependency.""" + repo_root = Path(__file__).resolve().parents[2] + project_data = tomllib.loads((repo_root / "pyproject.toml").read_text(encoding="utf-8")) + lock_data = tomllib.loads((repo_root / "uv.lock").read_text(encoding="utf-8")) + direct_names = { + dependency.split(";", 1)[0] + .split("[", 1)[0] + .split("=", 1)[0] + .split("<", 1)[0] + .split(">", 1)[0] + .strip() + .lower() + .replace("_", "-") + for dependency in project_data["project"]["dependencies"] + } + root_package = next( + package + for package in lock_data["package"] + if package["name"] == "winml-cli" and package.get("source", {}).get("editable") == "." + ) + locked_dependencies = { + _dependency_name(dependency).lower().replace("_", "-") + for dependency in root_package["dependencies"] + } + locked_requirements = { + str(requirement["name"]).lower().replace("_", "-") + for requirement in root_package["metadata"]["requires-dist"] + } + + assert direct_names <= locked_dependencies + assert direct_names <= locked_requirements diff --git a/uv.lock b/uv.lock index 8a529bff1..637397f63 100644 --- a/uv.lock +++ b/uv.lock @@ -3228,6 +3228,7 @@ dependencies = [ { name = "diffusers", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "evaluate", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "fastapi", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, + { name = "filelock", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "hf-xet", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "httpx", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, { name = "jsonschema", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'AMD64' and sys_platform == 'win32')" }, @@ -3323,6 +3324,7 @@ requires-dist = [ { name = "diffusers", specifier = ">=0.36" }, { name = "evaluate", specifier = ">=0.4.6" }, { name = "fastapi", specifier = ">=0.135.3" }, + { name = "filelock", specifier = ">=3.20" }, { name = "hf-xet", specifier = ">=1.1.10" }, { name = "httpx", specifier = ">=0.24.0" }, { name = "jsonschema", specifier = ">=4.23" },