Key EPContext cache by compile identity (#1295)

## Summary

- replace direct `WinMLSession` mtime-only EPContext reuse with a strict
compile identity
- compile cache misses into immutable UUID generations selected by an
atomic marker
- serialize same-identity compilation across threads/processes while
isolating different identities
- hash source ONNX/external-data bytes, EP DLL, hardware, effective
options, embed mode, ORT version, and generated artifacts
- publish stable CLI outputs atomically with content-addressed binary
references while retaining the legacy stable binary alias

## Problem

`WinMLSession.compile()` previously reused `<model>_<device>_ctx.onnx`
whenever it was newer than the source ONNX. The cache key omitted the
EP/source/version, hardware, provider/session options, embed mode, ORT
version, external-data sidecars, and context binary integrity. Different
configurations could therefore reuse or overwrite one shared EPContext
artifact.

## Behavior

- legacy or malformed caches without an identity marker rebuild once
- exact identities reuse an immutable generation
- source/external data, EP DLL, hardware, provider/session options,
embed mode, or ORT changes produce a different cache namespace
- source changes observed while waiting or compiling are retried and
never published under a stale identity
- custom opaque `SessionOptions` and uncertain source identity use
unique non-cacheable generations
- marker/artifact failures fail closed without discarding a successfully
compiled generation
- missing, modified, or timestamp-preserving substituted context
binaries invalidate reuse
- compiler finalization consumes the exact session generation and never
scans stale siblings after raw fallback
- public ONNX output remains stable, references an immutable
content-addressed binary, and still publishes the documented stable
binary alias

## Validation

- affected unit suites: 125 passed, 6 existing skips
- focused cache/concurrency suite: 16 passed
- Ruff check and format check on all changed Python files
- mypy on all changed production modules
- real QNN/NPU keen_hominy compile:
  - first compile produced one identity marker and immutable generation
  - second compile logged `Using cached EPContext`
  - second compile reduced total time from about 13.0s to 1.75s
- public `model_qnn_ctx.onnx` references
`model_qnn_ctx_qnn.<sha256>.bin`
- legacy `model_qnn_ctx_qnn.bin` alias remains present with matching
size

`uv lock --offline --check` could not resolve uncached Linux `onnx`
metadata on this ARM64 host. The lockfile TOML parses and the added unit
test verifies that every direct project dependency is represented in
both the editable package dependencies and `requires-dist` metadata; CI
runs on the supported AMD64 environment.

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: DingmaomaoBJTU

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

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. |

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",

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)
 

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"),

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 = {}