refactor(warmup): drop the orphan purge and join the warm-up once before the engine imports

strix/interface/main.py

@@ -136,8 +136,6 @@ def _subscription_error_hint(exc: BaseException) -> str | None:
 
 
 async def warm_up_llm(show_model_warning: bool = True) -> None:
-    wait_for_import_warmup()
-
     from agents.models.interface import ModelTracing
 
     from strix.config.models import (
@@ -467,6 +465,9 @@ def main() -> None:
     pull_docker_image()
     validate_environment()
 
+    # Everything below imports the scan engine; do not race the warm-up thread.
+    wait_for_import_warmup()
+
     if args.non_interactive:
         _bootstrap_scan(args)
 

strix/interface/scan_setup.py

@@ -31,7 +31,6 @@ from strix.interface.utils import (
     stage_api_specs,
     write_fetched_collection,
 )
-from strix.llm.warmup import wait_for_import_warmup
 from strix.telemetry import posthog, scarf
 from strix.utils.api_spec import (
     SpecParseError,
@@ -65,8 +64,6 @@ async def preflight_model_connection(
     settings: Settings | None = None,
 ) -> None:
     """Verify the configured model route before starting a scan."""
-    wait_for_import_warmup()
-
     from agents.models.interface import ModelTracing
 
     from strix.config.models import StrixProvider, configure_sdk_model_defaults

strix/llm/warmup.py

@@ -4,17 +4,18 @@ The scan engine's import graph (the agents SDK, OpenAI client, LiteLLM, the
 Caido SDK, the Docker SDK) costs seconds to import cold, but none of it is
 needed until a scan actually starts. Importing it on a daemon thread at CLI
 entry overlaps that cost with the I/O-bound startup work that always precedes
-a scan (argument parsing, Docker checks, image pull, TUI setup), so by the
-time the scan begins the modules are already in ``sys.modules``. Any thread
-that needs one of them before the warm-up finishes just blocks on the normal
-import lock, so behaviour is unchanged either way.
+a scan (argument parsing, Docker checks, image pull, TUI setup).
+
+The main thread must call :func:`wait_for_import_warmup` before its first
+import from that graph. Two threads that enter the same package graph from
+different modules hold each other's import locks, and CPython breaks the cycle
+by failing one of the imports.
 """
 
 from __future__ import annotations
 
 import importlib
 import logging
-import sys
 import threading
 
 
@@ -27,42 +28,15 @@ WARMUP_MODULES = (
     "docker",
 )
 
-_lock = threading.Lock()
 _thread: threading.Thread | None = None
 
 
-def _purge_orphaned_modules(before: frozenset[str]) -> None:
-    """Remove submodules stranded by an import attempt that just failed.
-
-    When a package import fails partway (for example CPython's import-lock
-    deadlock avoidance breaking a cross-thread cycle), the failed package is
-    removed from ``sys.modules`` but submodules it already finished stay
-    behind. A later import of one of those submodules then short-circuits on
-    the cached entry without re-importing its parent, and re-entering the
-    parent from inside a submodule crashes with "partially initialized
-    module". Dropping the orphans (cached submodules whose ancestor package is
-    gone) restores a clean slate, and touches nothing another thread imported
-    successfully.
-    """
-    added = set(sys.modules) - before
-    for name in added:
-        parent = name.rpartition(".")[0]
-        while parent:
-            if parent not in sys.modules:
-                sys.modules.pop(name, None)
-                logger.debug("Import warm-up purged orphaned module %r", name)
-                break
-            parent = parent.rpartition(".")[0]
-
-
 def _warm(modules: tuple[str, ...]) -> None:
     for name in modules:
-        before = frozenset(sys.modules)
         try:
             importlib.import_module(name)
         except Exception:  # noqa: BLE001 - a failed warm-up must never fail the run.
             logger.debug("Import warm-up for %r failed", name, exc_info=True)
-            _purge_orphaned_modules(before)
 
 
 def start_import_warmup(modules: tuple[str, ...] = WARMUP_MODULES) -> threading.Thread:
@@ -72,25 +46,15 @@ def start_import_warmup(modules: tuple[str, ...] = WARMUP_MODULES) -> threading.
     runtime that has no local Docker) warm a narrower set.
     """
     global _thread  # noqa: PLW0603
-    with _lock:
-        if _thread is not None:
-            return _thread
+    if _thread is None:
         _thread = threading.Thread(
             target=_warm, args=(modules,), name="strix-import-warmup", daemon=True
         )
         _thread.start()
-        return _thread
+    return _thread
 
 
-def wait_for_import_warmup(timeout: float | None = None) -> None:
-    """Block until the warm-up thread has finished, if one was started.
-
-    Call this before the first import of a warmed module on another thread.
-    Two threads walking the same package graph hold each other's import locks,
-    CPython breaks the cycle by failing one side, and the failed side's
-    orphan purge can remove a package the other thread is still importing.
-    """
-    with _lock:
-        thread = _thread
-    if thread is not None and thread is not threading.current_thread():
-        thread.join(timeout)
+def wait_for_import_warmup() -> None:
+    """Block until the warm-up thread has finished, if one was started."""
+    if _thread is not None:
+        _thread.join()

tests/test_import_warmup.py

@@ -1,12 +1,10 @@
-"""The import warm-up thread must never leave the import system poisoned.
+"""The import warm-up thread must never race the main thread into the engine.
 
-Field failure: the warm-up thread's ``strix.core.runner`` import and the main
-thread's ``strix.report`` import both walked the agents SDK graph, and the two
-held each other's import locks (report -> dedupe -> agents while runner ->
-hooks -> report.state). CPython's deadlock avoidance breaks such a cycle by
-failing one import, which strands finished submodules in ``sys.modules`` with
-their parent package gone — and the next import of one of those submodules
-crashes with "partially initialized module".
+Two threads that enter the same package graph from different modules hold
+each other's import locks (warm-up: ``strix.core.runner`` -> ``agents``;
+main: ``agents.models.interface``). CPython breaks such a cycle by failing one
+of the imports, so the main thread waits for the warm-up before its first
+engine import.
 """
 
 from __future__ import annotations
@@ -56,42 +54,6 @@ def test_check_duplicate_resolves_lazily() -> None:
     assert result.returncode == 0, result.stderr
 
 
-def test_failed_warm_import_purges_orphaned_submodules() -> None:
-    result = _run(
-        """
-        import sys
-
-        from strix.llm.warmup import _warm
-
-        # A package whose import fails after a submodule already completed:
-        # CPython removes the package but leaves the submodule stranded.
-        import pathlib
-        import tempfile
-
-        root = pathlib.Path(tempfile.mkdtemp())
-        pkg = root / "stranded_pkg"
-        pkg.mkdir()
-        (pkg / "ok.py").write_text("VALUE = 1")
-        (pkg / "__init__.py").write_text("from . import ok\\nraise RuntimeError('boom')")
-        sys.path.insert(0, str(root))
-
-        _warm(("stranded_pkg",))
-
-        assert "stranded_pkg" not in sys.modules
-        assert "stranded_pkg.ok" not in sys.modules, "orphan survived the purge"
-
-        # And the subtree imports cleanly afterwards up to the real error.
-        try:
-            import stranded_pkg  # noqa: F401
-        except RuntimeError:
-            pass
-        else:
-            raise AssertionError("expected the package's own error")
-        """
-    )
-    assert result.returncode == 0, result.stderr
-
-
 def test_wait_for_import_warmup_lets_main_thread_import_the_agents_graph() -> None:
     result = _run(
         """
@@ -100,7 +62,7 @@ def test_wait_for_import_warmup_lets_main_thread_import_the_agents_graph() -> No
         from strix.llm.warmup import start_import_warmup, wait_for_import_warmup
 
         # Same shape as the CLI: warm-up starts, then the main thread needs a
-        # module from the middle of the agents graph before it has finished.
+        # module from the middle of the agents graph.
         start_import_warmup()
         wait_for_import_warmup()
 
@@ -114,12 +76,9 @@ def test_wait_for_import_warmup_lets_main_thread_import_the_agents_graph() -> No
     assert result.returncode == 0, result.stderr
 
 
+def test_failed_warm_import_does_not_raise() -> None:
+    warmup._warm(("strix_no_such_module_for_warmup_test",))
+
+
 def test_wait_for_import_warmup_is_a_no_op_without_a_thread() -> None:
-    warmup.wait_for_import_warmup(timeout=0)
-
-
-def test_purge_does_not_touch_preexisting_or_healthy_modules() -> None:
-    before = frozenset(sys.modules) - {"strix.llm.warmup"}
-    warmup._purge_orphaned_modules(before)
-    assert "strix.llm.warmup" in sys.modules  # parent chain intact -> kept
-    assert "strix" in sys.modules
+    warmup.wait_for_import_warmup()