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/compiler/stages/compile.py b/src/winml/modelkit/compiler/stages/compile.py index d877867ea..7f8780828 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.""" @@ -107,8 +122,10 @@ def _compile_single_model(self, context: CompileContext) -> None: ep_device=ep_device, ep_config=ep_config, ) + running_model_path = model_path try: winml_session.compile() + running_model_path = winml_session.running_model_path session = winml_session._session context.session = session @@ -122,11 +139,15 @@ def _compile_single_model(self, context: CompileContext) -> None: winml_session.reset() 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, ) def _compile_shared_context(self, context: CompileContext) -> None: @@ -306,6 +327,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 +365,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 +384,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 +434,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 +466,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 +527,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..96844789c 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,10 @@ import numpy as np import onnxruntime as ort +from filelock import FileLock 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 +55,16 @@ logger = logging.getLogger(__name__) +_EPCONTEXT_THREAD_LOCKS_GUARD = threading.Lock() +_EPCONTEXT_THREAD_LOCKS: dict[Path, threading.Lock] = {} + + +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()) + @contextmanager def _suppress_native_output(log_path: str | Path | None = None) -> Iterator[None]: @@ -472,47 +488,17 @@ 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 # 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) - - 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 = self._compile_epcontext_with_stable_source(compile_log) try: # Create the runtime InferenceSession against the (possibly compiled) model. @@ -551,6 +537,321 @@ def compile(self) -> None: self._running_model_path = model_path self._state = SessionState.COMPILED + def _compile_epcontext_with_stable_source(self, compile_log: Path) -> Path: + """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") + with _epcontext_thread_lock(lock_path), FileLock(lock_path): + prepared = self._prepare_epcontext_model( + cache_path, + compile_log, + None, + ) + return prepared or self._onnx_path + + cache_path = self._epcontext_cache_path(expected_identity) + lock_path = cache_path.with_name(f"{cache_path.name}.lock") + with _epcontext_thread_lock(lock_path), FileLock(lock_path): + 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): + continue + if final_identity == locked_identity: + return prepared + + logger.warning("ONNX source changed during EPContext compilation; using original model") + return self._onnx_path + + def _prepare_epcontext_model( + self, + cache_path: Path, + compile_log: Path, + cache_identity: dict[str, object] | None, + ) -> Path | 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 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: + logger.warning("ModelCompiler failed, using original: %s", exc) + return self._onnx_path + + if not generation_path.exists(): + return self._onnx_path + 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, + ) + logger.info("Compiled to EPContext: %s", generation_path) + return generation_path + + @classmethod + def _discard_epcontext_generation(cls, generation_path: Path) -> None: + """Best-effort removal of an unpublished generation and its sidecars.""" + try: + sidecars = cls._epcontext_external_sidecars(generation_path) + except (OSError, ValueError): + sidecars = () + for path in (*sidecars, generation_path): + cls._unlink_generation_file(path) + + @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())), + "session_options": dict(sorted(self._active_session_option_entries.items())), + "embed_context": self._embed_context, + "ort_version": ort.__version__, + } + + @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, inputs: dict[str, Any], diff --git a/tests/unit/compiler/test_compiler_stages.py b/tests/unit/compiler/test_compiler_stages.py index 8cc7d3745..bcd112ff7 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) @@ -451,6 +455,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 +498,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 +737,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 +894,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_winml_session.py b/tests/unit/session/test_winml_session.py index 1632b742b..f1b393672 100644 --- a/tests/unit/session/test_winml_session.py +++ b/tests/unit/session/test_winml_session.py @@ -17,15 +17,18 @@ 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 onnx import pytest +from onnx import TensorProto, helper, numpy_helper from winml.modelkit.compiler import EPConfig from winml.modelkit.session import ( @@ -51,6 +54,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 + onnx.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 +250,617 @@ 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_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_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)]) + onnx.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 = onnx.load(simple_matmul_onnx) + model.producer_name = "changed-during-compile" + onnx.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_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 + + 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=lambda: 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 + + 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" },