diff --git a/strix/interface/main.py b/strix/interface/main.py index 8f9dae7a..e5633ce6 100644 --- a/strix/interface/main.py +++ b/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) diff --git a/strix/interface/scan_setup.py b/strix/interface/scan_setup.py index d87d41c8..ae7caf2f 100644 --- a/strix/interface/scan_setup.py +++ b/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 diff --git a/strix/llm/warmup.py b/strix/llm/warmup.py index 25d284a0..e4c7357b 100644 --- a/strix/llm/warmup.py +++ b/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() diff --git a/tests/test_import_warmup.py b/tests/test_import_warmup.py index 53cfc296..c2ea6798 100644 --- a/tests/test_import_warmup.py +++ b/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()