diff --git a/docs/samples/qwen3-genai-bundle.md b/docs/samples/qwen3-genai-bundle.md index accf848b4..a3d416fb8 100644 --- a/docs/samples/qwen3-genai-bundle.md +++ b/docs/samples/qwen3-genai-bundle.md @@ -117,31 +117,6 @@ time-to-first-token (prefill) and decode throughput, and writes a results JSON u generated tokens and the saved perf metrics are unaffected — and originates in the native runtime below winml-cli, not in the bundle or the build. -!!! tip "More reliable on the NPU: compile once, then run the compiled bundle directly" - Doing the EPContext compilation and the generation load in the **same** `--compile` - process can fault natively at model load (before the first token) on some setups — the - stage compilation can leave the process in a fragile native state. The `--compile` run - still writes the compiled bundle to `out/qwen3-bundle/_compiled/` (its `genai_config.json` - points at the compiled `context_ctx.onnx` / `iterator_ctx.onnx`), so run generation as a - **second, fresh** command pointed straight at that directory — without `--compile`: - - ```bash - # 1. Compile once (produces out/qwen3-bundle/_compiled/; re-run if a stage is still missing) - winml perf -m out/qwen3-bundle --runtime winml-genai --device npu --compile \ - --compile-timeout 600 --max-new-tokens 20 --prompt "What is the capital of France?" - - # 2. Run the compiled bundle in a fresh process (loads EPContext directly, no re-compile) - winml perf -m out/qwen3-bundle/_compiled --runtime winml-genai --device npu \ - --max-new-tokens 20 --prompt "What is the capital of France?" - ``` - - Dropping `--compile` is safe **only** because `-m` points at the already-compiled - `_compiled/` directory, whose stages are EPContext graphs loaded as-is — nothing is - JIT-compiled. (Dropping `--compile` on the original bundle would JIT-compile the source - ONNX and fault, per the warning above.) A root-cause fix that isolates model load into its - own process is tracked in - [issue #1087](https://github.com/microsoft/winml-cli/issues/1087). - ## How it maps to the composite system The bundle reuses winml-cli's existing composite-model machinery — it does not add diff --git a/src/winml/modelkit/session/genai_session.py b/src/winml/modelkit/session/genai_session.py index 329b2ec23..712c4e721 100644 --- a/src/winml/modelkit/session/genai_session.py +++ b/src/winml/modelkit/session/genai_session.py @@ -129,6 +129,50 @@ def _compile_stage_worker(src: str, dst: str, ep_alias: str, provider_options: d raise RuntimeError(f"Compilation failed: {result.errors}") +def _prepare_derived_bundle_worker( + result_queue: Any, + bundle_dir: str, + compile_timeout: int, + effective_cfg: dict[str, Any], + overridden: bool, +) -> None: + """Run the EPContext compile orchestration in an isolated subprocess. + + Executed via ``multiprocessing`` (spawn) by + :meth:`GenaiSession._prepare_derived_bundle_isolated`. Running the whole + orchestration here — rather than in the process that later creates + ``og.Model`` — is the fix for issue #1087: per-stage accelerator compile + workers can native-fault during interpreter/driver teardown (QNN on the + Hexagon NPU in particular), and orchestrating those crash-prone subprocesses + leaves the *orchestrating* process's native accelerator state fragile. When + that orchestrating process is also the one that loads ``og.Model`` for + generation, the load native-crashes (``0xC0000005`` / ``0xC0000374``) before + the first token. Isolating the orchestration in this throwaway process keeps + the parent pristine for the model load. + + The compiled artifacts are written to disk (``bundle_dir/_compiled/``) by + :meth:`GenaiSession._prepare_derived_bundle`; only the resolved load + directory is handed back to the parent through *result_queue* as a + ``(status, payload)`` tuple — ``("ok", )`` on success or + ``("error", )`` if the orchestration raised. + + Args: + result_queue: A ``multiprocessing`` queue the parent drains for the + single ``(status, payload)`` result. + bundle_dir: The genai bundle directory (as a string path). + compile_timeout: Per-stage compile timeout forwarded to the reconstructed + session so the nested per-stage workers keep the same bound. + effective_cfg: The post-override config to compile/load from. + overridden: Whether *effective_cfg* differs from the on-disk config. + """ + try: + session = GenaiSession(bundle_dir, compile=True, compile_timeout=compile_timeout) + load_dir = session._prepare_derived_bundle(effective_cfg, overridden=overridden) + result_queue.put(("ok", str(load_dir))) + except Exception as exc: + result_queue.put(("error", repr(exc))) + + # --------------------------------------------------------------------------- # Data classes # --------------------------------------------------------------------------- @@ -323,6 +367,15 @@ class GenaiSession: # Sub-directory within the bundle that holds pre-compiled EPContext ONNX files. _COMPILED_SUBDIR: str = "_compiled" + # Sentinel dropped into ``_compiled/`` as the final step of a successful + # :meth:`_prepare_derived_bundle`. Its presence proves the derived bundle was + # written in full for the current effective config and stage inputs. The + # isolated-compile parent removes any stale copy before spawning the child and, + # on the post-failure fallback, trusts an existing ``_compiled/`` only when this + # marker is present — so it never serves stale EPContext artifacts from an + # earlier run (issue #1087). + _BUNDLE_COMPLETE_MARKER: str = ".winml_bundle_complete" + # Standalone chat-template sidecar written next to genai_config.json (the # conventional onnxruntime-genai / Hugging Face filename). Read to format # prompts with the model's own template; absent for bundles that ship none. @@ -458,9 +511,17 @@ def load(self) -> None: # Materialize the directory og.Model loads from. A derived ``_compiled/`` # bundle is written when the override rewrote the config and/or stages # were pre-compiled; otherwise the original bundle dir is used as-is. + # + # When compiling, the EPContext orchestration runs in an isolated + # subprocess so its crash-prone accelerator teardown cannot poison THIS + # process, which must stay pristine to load og.Model (issue #1087). An + # override-only pass (no compilation) merely rewrites JSON + mirrors + # files, touches no native accelerator state, and stays in-process. load_dir = self._bundle_dir - if self._compile or overridden: - load_dir = self._prepare_compiled_bundle(effective_cfg, overridden=overridden) + if self._compile: + load_dir = self._prepare_derived_bundle_isolated(effective_cfg, overridden=overridden) + elif overridden: + load_dir = self._prepare_derived_bundle(effective_cfg, overridden=overridden) try: config = og.Config(str(load_dir)) @@ -995,7 +1056,126 @@ def _existing_opts_for_ep(provider_options: Any, target: EPName) -> dict: return dict(opts) return {} - def _prepare_compiled_bundle( + def _prepare_derived_bundle_isolated( + self, effective_cfg: dict[str, Any], *, overridden: bool + ) -> Path: + """Run :meth:`_prepare_derived_bundle` in a throwaway subprocess. + + This is the root-cause fix for issue #1087. The EPContext compile + orchestration spawns and reaps per-stage accelerator compile workers that + can native-fault during interpreter / driver teardown (QNN on the Hexagon + NPU). Doing that orchestration in the *same* process that then creates + ``og.Model`` leaves the process's native accelerator state fragile, so the + subsequent model load native-crashes (``0xC0000005`` / ``0xC0000374``) + before the first token — even when every stage is already compiled and + cached. + + Delegating the orchestration to a disposable subprocess (which writes the + compiled artifacts to ``bundle_dir/_compiled/`` and exits) keeps *this* + process — the one that loads and runs ``og.Model`` — pristine. The child + performs exactly the work today's parent does up to, but not including, + the model load, so it completes and reports the load directory back before + exiting; only the resolved path crosses the process boundary. + + Args: + effective_cfg: The post-override config to compile/load from. + overridden: Whether *effective_cfg* differs from the on-disk config. + + Returns: + Path to the directory ``og.Model`` should load from — the derived + ``_compiled/`` bundle when the child produced one (reported directly, + or proven current-run fresh by its completion marker on the fallback + path), otherwise the original bundle directory (a clean JIT load in + this un-poisoned process). + """ + import multiprocessing + import queue as queue_mod + + ctx = multiprocessing.get_context("spawn") + result_queue = ctx.Queue() + proc = ctx.Process( + target=_prepare_derived_bundle_worker, + args=( + result_queue, + str(self._bundle_dir), + self._compile_timeout, + effective_cfg, + overridden, + ), + ) + logger.info( + "Compiling bundle in an isolated subprocess so the compile orchestration " + "cannot poison this model-load process (issue #1087)" + ) + # Drop any completion marker left by a previous run before spawning the + # child. The post-failure fallback below trusts an existing _compiled/ + # only when THIS run's child re-creates the marker, so a stale one must + # never survive into that check (issue #1087 stale-bundle guard). + compiled_dir = self._bundle_dir / self._COMPILED_SUBDIR + (compiled_dir / self._BUNDLE_COMPLETE_MARKER).unlink(missing_ok=True) + proc.start() + + # Drain the single (status, payload) the worker posts when the compile + # orchestration finishes. Poll rather than block indefinitely so a native + # crash in the worker (which would post nothing) is noticed promptly + # instead of hanging; the worker's runtime is already bounded by the + # per-stage compile timeouts it enforces internally. + status: str | None = None + payload: str | None = None + while proc.is_alive(): + try: + status, payload = result_queue.get(timeout=1.0) + break + except queue_mod.Empty: + continue + if status is None: + try: + status, payload = result_queue.get_nowait() + except queue_mod.Empty: + # The worker exited without ever posting a result (it crashed or + # was killed before finishing). Leave status as None: the failure + # path below reports the error and never falls back to an + # in-process compile (issue #1087). + pass + + # The result (if any) is already in hand, so the child's remaining work is + # just process teardown. Bound the wait and kill a child whose native + # accelerator teardown hangs instead of exiting: a throwaway compile child + # must never be able to wedge this model-load process (issue #1087), and + # driver teardown can hang as readily as it can fault. + proc.join(timeout=self._compile_timeout) + if proc.is_alive(): + logger.warning( + "Isolated compile subprocess did not exit within %ds after reporting; " + "terminating it (its compiled artifacts are already on disk)", + self._compile_timeout, + ) + proc.kill() + proc.join() + + if status == "ok" and payload: + return Path(payload) + + # The child failed or crashed before reporting. Never fall back to an + # in-process compile (that would reintroduce the poisoning this method + # exists to prevent). A previously written _compiled/ is reused only when + # this run's child left the completion marker: it was removed before spawn + # and rewritten only after a full prepare, so its presence proves the + # bundle matches the current effective config + stage inputs. Without it + # we cannot rule out a stale bundle from an earlier run, so we surface the + # failure by loading the original bundle and letting og.Model report. + logger.warning( + "Isolated compile subprocess did not report success (status=%s, detail=%s); " + "falling back to on-disk state", + status, + payload, + ) + marker_present = (compiled_dir / self._BUNDLE_COMPLETE_MARKER).exists() + if marker_present and (compiled_dir / "genai_config.json").exists(): + return compiled_dir + return self._bundle_dir + + def _prepare_derived_bundle( self, effective_cfg: dict[str, Any] | None = None, *, overridden: bool = False ) -> Path: """Create (or reuse) a derived bundle directory og.Model can load from. @@ -1153,6 +1333,14 @@ def _prepare_compiled_bundle( self._mirror_non_onnx_files(compiled_dir, skip_filenames=compiled_stage_filenames) logger.info("Compiled bundle prepared at %s", compiled_dir) + # Final step: record that the derived bundle was written in full. The + # isolated-compile parent removes this marker before spawning the child and + # trusts an existing _compiled/ on its post-failure fallback only when the + # marker is present, so writing it last (after the config + mirrored files) + # makes its presence prove a complete, current-run bundle (issue #1087). + (compiled_dir / self._BUNDLE_COMPLETE_MARKER).write_text( + "winml genai derived bundle completed\n", encoding="utf-8" + ) return compiled_dir @staticmethod diff --git a/tests/unit/session/test_genai_session.py b/tests/unit/session/test_genai_session.py index d1ccc2f71..0c1cfcbfd 100644 --- a/tests/unit/session/test_genai_session.py +++ b/tests/unit/session/test_genai_session.py @@ -13,6 +13,7 @@ import json import logging import os +import queue import sys import time from pathlib import Path @@ -1148,7 +1149,7 @@ def test_compile_timeout_passed_to_compile_stage( patch("multiprocessing.get_context", return_value=ctx_mock), _patch_og(mock_og), ): - session._prepare_compiled_bundle() + session._prepare_derived_bundle() # proc.join was called with timeout=42 for each stage for call in proc_mock.join.call_args_list: @@ -1611,15 +1612,354 @@ def test_raises_when_compile_unsuccessful(self) -> None: # --------------------------------------------------------------------------- -# Tests: _prepare_compiled_bundle +# Tests: _prepare_derived_bundle_worker (isolated compile subprocess target) # --------------------------------------------------------------------------- -class TestPrepareCompiledBundle: +class TestPrepareDerivedBundleWorker: + """The module-level target executed inside the isolated compile subprocess.""" + + def test_posts_ok_with_resolved_load_dir( + self, bundle_dir_with_pipeline: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """On success the worker posts ("ok", ) to the result queue.""" + from winml.modelkit.session.genai_session import _prepare_derived_bundle_worker + + expected = bundle_dir_with_pipeline / "_compiled" + monkeypatch.setattr( + GenaiSession, + "_prepare_derived_bundle", + lambda self, cfg, *, overridden: expected, + ) + result_queue = MagicMock() + _prepare_derived_bundle_worker( + result_queue, str(bundle_dir_with_pipeline), 42, {"model": {}}, False + ) + result_queue.put.assert_called_once_with(("ok", str(expected))) + + def test_forwards_effective_cfg_and_overridden( + self, bundle_dir_with_pipeline: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The worker passes the effective config and overridden flag through.""" + from winml.modelkit.session.genai_session import _prepare_derived_bundle_worker + + seen: dict = {} + + def _capture(self, cfg, *, overridden): + seen["cfg"] = cfg + seen["overridden"] = overridden + return bundle_dir_with_pipeline + + monkeypatch.setattr(GenaiSession, "_prepare_derived_bundle", _capture) + _prepare_derived_bundle_worker( + MagicMock(), str(bundle_dir_with_pipeline), 7, {"model": {"x": 1}}, True + ) + assert seen == {"cfg": {"model": {"x": 1}}, "overridden": True} + + def test_reconstructs_compile_enabled_session_with_timeout( + self, bundle_dir_with_pipeline: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """The reconstructed session enables compile and preserves the timeout.""" + from winml.modelkit.session.genai_session import _prepare_derived_bundle_worker + + captured: dict = {} + + def _capture(self, cfg, *, overridden): + captured["compile"] = self._compile + captured["timeout"] = self._compile_timeout + return bundle_dir_with_pipeline + + monkeypatch.setattr(GenaiSession, "_prepare_derived_bundle", _capture) + _prepare_derived_bundle_worker( + MagicMock(), str(bundle_dir_with_pipeline), 123, {"model": {}}, False + ) + assert captured == {"compile": True, "timeout": 123} + + def test_posts_error_when_orchestration_raises( + self, bundle_dir_with_pipeline: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """A failure inside the orchestration is reported as ("error", ).""" + from winml.modelkit.session.genai_session import _prepare_derived_bundle_worker + + def _boom(self, cfg, *, overridden): + raise RuntimeError("compile blew up") + + monkeypatch.setattr(GenaiSession, "_prepare_derived_bundle", _boom) + result_queue = MagicMock() + _prepare_derived_bundle_worker( + result_queue, str(bundle_dir_with_pipeline), 42, {"model": {}}, False + ) + status, payload = result_queue.put.call_args.args[0] + assert status == "error" + assert "compile blew up" in payload + + +# --------------------------------------------------------------------------- +# Tests: _prepare_derived_bundle_isolated (spawns + drains the worker) +# --------------------------------------------------------------------------- + + +class TestPrepareDerivedBundleIsolated: + """Isolation wrapper: spawn the compile worker, drain its single result.""" + + @staticmethod + def _mock_ctx(proc: MagicMock, result_queue: MagicMock) -> MagicMock: + ctx = MagicMock() + ctx.Queue.return_value = result_queue + ctx.Process.return_value = proc + return ctx + + def test_returns_load_dir_reported_by_worker(self, bundle_dir_with_pipeline: Path) -> None: + """The path the worker posts becomes the og.Model load dir.""" + compiled = bundle_dir_with_pipeline / "_compiled" + proc = MagicMock() + # Alive during the drain (get() breaks the loop), exited by the post-join + # liveness check so the hang-kill path is not taken. + proc.is_alive.side_effect = [True, False] + result_queue = MagicMock() + result_queue.get.return_value = ("ok", str(compiled)) + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == compiled + proc.start.assert_called_once() + proc.join.assert_called_once() + proc.kill.assert_not_called() + # The load itself never happens in this (parent) process's subprocess call. + ctx.Process.assert_called_once() + + def test_drains_result_posted_as_worker_exits(self, bundle_dir_with_pipeline: Path) -> None: + """A result posted just as the worker exits is still drained via get_nowait.""" + compiled = bundle_dir_with_pipeline / "_compiled" + proc = MagicMock() + proc.is_alive.return_value = False # already exited when first polled + result_queue = MagicMock() + result_queue.get_nowait.return_value = ("ok", str(compiled)) + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == compiled + + def test_reuses_compiled_dir_when_marker_proves_current_run( + self, bundle_dir_with_pipeline: Path + ) -> None: + """A silent crash after the child wrote _compiled/ + its marker still reuses it. + + Models a child that finished compiling (leaving a complete, current-run + _compiled/ and its completion marker) but died before its "ok" result was + drained; the fallback safely reuses the freshly written bundle. + """ + compiled = bundle_dir_with_pipeline / "_compiled" + compiled.mkdir() + (compiled / "genai_config.json").write_text("{}", encoding="utf-8") + marker = compiled / GenaiSession._BUNDLE_COMPLETE_MARKER + + proc = MagicMock() + proc.is_alive.return_value = False + # The parent clears any stale marker before spawning; simulate the child + # writing a fresh one as it completes the compile, just before dying. + proc.start.side_effect = lambda: marker.write_text("done", encoding="utf-8") + result_queue = MagicMock() + result_queue.get_nowait.side_effect = queue.Empty + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == compiled + + def test_ignores_compiled_dir_without_completion_marker( + self, bundle_dir_with_pipeline: Path + ) -> None: + """A _compiled/ with no current-run marker is refused after a silent failure. + + Guards issue #1087: the parent must not load stale EPContext artifacts/config + from an earlier run that may predate a config or stage-input change; it + surfaces the failure by loading the original bundle instead. + """ + compiled = bundle_dir_with_pipeline / "_compiled" + compiled.mkdir() + (compiled / "genai_config.json").write_text("{}", encoding="utf-8") + # No completion marker, and the mocked child never writes one. + + proc = MagicMock() + proc.is_alive.return_value = False + result_queue = MagicMock() + result_queue.get_nowait.side_effect = queue.Empty + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == bundle_dir_with_pipeline + + def test_clears_stale_marker_before_spawning_child( + self, bundle_dir_with_pipeline: Path + ) -> None: + """A marker left by a previous run is removed before spawn so it can't be trusted. + + The mocked child never re-writes the marker, so after a silent failure the + stale marker must be gone and the fallback must decline the stale _compiled/ + (issue #1087 stale-bundle guard). + """ + compiled = bundle_dir_with_pipeline / "_compiled" + compiled.mkdir() + (compiled / "genai_config.json").write_text("{}", encoding="utf-8") + stale_marker = compiled / GenaiSession._BUNDLE_COMPLETE_MARKER + stale_marker.write_text("from a previous run", encoding="utf-8") + + proc = MagicMock() + proc.is_alive.return_value = False + result_queue = MagicMock() + result_queue.get_nowait.side_effect = queue.Empty + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert not stale_marker.exists() + assert result == bundle_dir_with_pipeline + + def test_falls_back_to_bundle_dir_without_compiled_output( + self, bundle_dir_with_pipeline: Path + ) -> None: + """With neither a reported path nor a _compiled/ on disk, use the bundle dir.""" + proc = MagicMock() + proc.is_alive.return_value = False + result_queue = MagicMock() + result_queue.get_nowait.side_effect = queue.Empty + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == bundle_dir_with_pipeline + + def test_error_status_never_reused_as_load_dir( + self, bundle_dir_with_pipeline: Path, caplog: pytest.LogCaptureFixture + ) -> None: + """An ("error", ...) result is not treated as a load dir; it warns and falls back.""" + proc = MagicMock() + # Alive during the drain (get() returns the error), exited by the post-join + # check so the hang-kill path is not taken. + proc.is_alive.side_effect = [True, False] + result_queue = MagicMock() + result_queue.get.return_value = ("error", "RuntimeError('boom')") + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with ( + patch("multiprocessing.get_context", return_value=ctx), + caplog.at_level(logging.WARNING), + ): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == bundle_dir_with_pipeline + assert "did not report success" in caplog.text + + def test_kills_child_that_hangs_after_reporting(self, bundle_dir_with_pipeline: Path) -> None: + """A child that reports its result but then hangs in teardown is killed, not awaited.""" + compiled = bundle_dir_with_pipeline / "_compiled" + proc = MagicMock() + # Always alive: the drain loop still breaks because get() returns a result, + # and the post-join check then sees it stuck (teardown hang) and kills it. + proc.is_alive.return_value = True + result_queue = MagicMock() + result_queue.get.return_value = ("ok", str(compiled)) + ctx = self._mock_ctx(proc, result_queue) + + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + with patch("multiprocessing.get_context", return_value=ctx): + result = session._prepare_derived_bundle_isolated({"model": {}}, overridden=False) + + assert result == compiled # the already-reported result is still returned + proc.kill.assert_called_once() + assert proc.join.call_count == 2 # bounded join, then join after kill + + +# --------------------------------------------------------------------------- +# Tests: load() routes compilation through the isolated subprocess (issue #1087) +# --------------------------------------------------------------------------- + + +class TestLoadCompileIsolation: + """``load()`` must never run the compile orchestration in the model-load process.""" + + def test_compile_load_uses_isolated_subprocess( + self, bundle_dir_with_pipeline: Path, mock_og: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + """compile=True loads from the isolated compile dir, never the in-process one.""" + session = GenaiSession(bundle_dir_with_pipeline, ep="qnn", compile=True) + compiled = bundle_dir_with_pipeline / "_compiled" + isolated = MagicMock(return_value=compiled) + in_process = MagicMock() + monkeypatch.setattr(session, "_prepare_derived_bundle_isolated", isolated) + monkeypatch.setattr(session, "_prepare_derived_bundle", in_process) + monkeypatch.setattr(session, "_register_eps", lambda: None) + + with _patch_og(mock_og): + session.load() + + isolated.assert_called_once() + in_process.assert_not_called() + mock_og.Config.assert_called_once_with(str(compiled)) + + def test_override_only_load_stays_in_process( + self, bundle_dir_with_pipeline: Path, mock_og: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + """An ep override without --compile keeps the cheap in-process derived bundle.""" + session = GenaiSession(bundle_dir_with_pipeline, ep="cpu") # compile=False + compiled = bundle_dir_with_pipeline / "_compiled" + isolated = MagicMock() + in_process = MagicMock(return_value=compiled) + monkeypatch.setattr(session, "_prepare_derived_bundle_isolated", isolated) + monkeypatch.setattr(session, "_prepare_derived_bundle", in_process) + monkeypatch.setattr(session, "_register_eps", lambda: None) + + with _patch_og(mock_og): + session.load() + + in_process.assert_called_once() + isolated.assert_not_called() + + def test_plain_load_prepares_nothing( + self, bundle_dir: Path, mock_og: MagicMock, monkeypatch: pytest.MonkeyPatch + ) -> None: + """No compile and no override loads the bundle dir directly (no derived bundle).""" + session = GenaiSession(bundle_dir) # compile=False, no override + isolated = MagicMock() + in_process = MagicMock() + monkeypatch.setattr(session, "_prepare_derived_bundle_isolated", isolated) + monkeypatch.setattr(session, "_prepare_derived_bundle", in_process) + + with _patch_og(mock_og): + session.load() + + isolated.assert_not_called() + in_process.assert_not_called() + mock_og.Config.assert_called_once_with(str(bundle_dir)) + + +# --------------------------------------------------------------------------- +# Tests: _prepare_derived_bundle +# --------------------------------------------------------------------------- + + +class TestPrepareDerivedBundle: def test_no_compilable_stages_returns_original_bundle_dir(self, bundle_dir: Path) -> None: """When no EPContext-capable stages exist, bundle_dir is returned unchanged.""" session = GenaiSession(bundle_dir, ep="qnn", compile=True) - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() assert result == bundle_dir assert not (bundle_dir / "_compiled").exists() @@ -1646,7 +1986,7 @@ def test_non_epcontext_stage_is_skipped(self, tmp_path: Path) -> None: (tmp_path / "ctx.onnx").write_bytes(b"fake") session = GenaiSession(tmp_path, ep="dml", compile=True) - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() assert result == tmp_path assert not (tmp_path / "_compiled").exists() @@ -1667,7 +2007,7 @@ def test_writes_modified_genai_config_to_compiled_dir( compiled_dir = bundle_dir_with_pipeline / "_compiled" with patch("multiprocessing.get_context", return_value=ctx_mock): - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() assert result == compiled_dir config_out = compiled_dir / "genai_config.json" @@ -1691,7 +2031,7 @@ def test_override_only_writes_derived_bundle_without_compile( assert overridden is True compiled_dir = bundle_dir_with_pipeline / "_compiled" - result = session._prepare_compiled_bundle(effective, overridden=True) + result = session._prepare_derived_bundle(effective, overridden=True) assert result == compiled_dir written = json.loads((compiled_dir / "genai_config.json").read_text(encoding="utf-8")) @@ -1722,7 +2062,7 @@ def test_reuses_cached_epcontext_when_fresh( spy = MagicMock(return_value=True) monkeypatch.setattr(session, "_compile_stage", spy) - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() spy.assert_not_called() assert result == compiled_dir @@ -1738,7 +2078,7 @@ def test_recompiles_when_provider_options_change( spy = MagicMock(return_value=True) monkeypatch.setattr(session, "_compile_stage", spy) - session._prepare_compiled_bundle() + session._prepare_derived_bundle() assert spy.call_count == 2 @@ -1758,7 +2098,7 @@ def test_recompiles_when_data_sidecar_is_newer( spy = MagicMock(return_value=True) monkeypatch.setattr(session, "_compile_stage", spy) - session._prepare_compiled_bundle() + session._prepare_derived_bundle() # Only the context stage (whose .data changed) is recompiled. assert spy.call_count == 1 @@ -1786,7 +2126,7 @@ def _fake_compile(src, ctx, stage_key, ep_alias, ep_opts): return True monkeypatch.setattr(session, "_compile_stage", _fake_compile) - session._prepare_compiled_bundle(effective, overridden=overridden) + session._prepare_derived_bundle(effective, overridden=overridden) return next(p for p in captured if p.name.startswith("context")) ov_ctx = _context_ctx_path_for("openvino") @@ -1820,7 +2160,7 @@ def test_different_ep_does_not_reuse_cached_epcontext( effective, overridden = session._apply_ep_override(session._read_genai_config()) spy = MagicMock(return_value=True) monkeypatch.setattr(session, "_compile_stage", spy) - session._prepare_compiled_bundle(effective, overridden=overridden) + session._prepare_derived_bundle(effective, overridden=overridden) # Both stages recompiled for VitisAI; the OpenVINO cache is untouched. assert spy.call_count == 2 @@ -1857,7 +2197,7 @@ def fake_compile(src_onnx, ctx_out, stage_key, ep_alias, ep_opts=None): monkeypatch.setattr(session, "_compile_stage", fake_compile) compiled_dir = bundle_dir_with_pipeline / "_compiled" - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() assert result == compiled_dir written = json.loads((compiled_dir / "genai_config.json").read_text(encoding="utf-8")) @@ -1898,7 +2238,7 @@ def fake_compile(src_onnx, ctx_out, stage_key, ep_alias, ep_opts=None): monkeypatch.setattr(session, "_compile_stage", fake_compile) compiled_dir = bundle_dir_with_pipeline / "_compiled" - result = session._prepare_compiled_bundle() + result = session._prepare_derived_bundle() assert result == compiled_dir written = json.loads((compiled_dir / "genai_config.json").read_text(encoding="utf-8"))