mirror of
https://hubproxy.babadafafafafa.cn/https://github.com/usestrix/strix.git
synced 2026-09-20 08:03:42 +08:00
refactor(warmup): drop the orphan purge and join the warm-up once before the engine imports
This commit is contained in:
@@ -136,8 +136,6 @@ def _subscription_error_hint(exc: BaseException) -> str | None:
|
|||||||
|
|
||||||
|
|
||||||
async def warm_up_llm(show_model_warning: bool = True) -> None:
|
async def warm_up_llm(show_model_warning: bool = True) -> None:
|
||||||
wait_for_import_warmup()
|
|
||||||
|
|
||||||
from agents.models.interface import ModelTracing
|
from agents.models.interface import ModelTracing
|
||||||
|
|
||||||
from strix.config.models import (
|
from strix.config.models import (
|
||||||
@@ -467,6 +465,9 @@ def main() -> None:
|
|||||||
pull_docker_image()
|
pull_docker_image()
|
||||||
validate_environment()
|
validate_environment()
|
||||||
|
|
||||||
|
# Everything below imports the scan engine; do not race the warm-up thread.
|
||||||
|
wait_for_import_warmup()
|
||||||
|
|
||||||
if args.non_interactive:
|
if args.non_interactive:
|
||||||
_bootstrap_scan(args)
|
_bootstrap_scan(args)
|
||||||
|
|
||||||
|
|||||||
@@ -31,7 +31,6 @@ from strix.interface.utils import (
|
|||||||
stage_api_specs,
|
stage_api_specs,
|
||||||
write_fetched_collection,
|
write_fetched_collection,
|
||||||
)
|
)
|
||||||
from strix.llm.warmup import wait_for_import_warmup
|
|
||||||
from strix.telemetry import posthog, scarf
|
from strix.telemetry import posthog, scarf
|
||||||
from strix.utils.api_spec import (
|
from strix.utils.api_spec import (
|
||||||
SpecParseError,
|
SpecParseError,
|
||||||
@@ -65,8 +64,6 @@ async def preflight_model_connection(
|
|||||||
settings: Settings | None = None,
|
settings: Settings | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
"""Verify the configured model route before starting a scan."""
|
"""Verify the configured model route before starting a scan."""
|
||||||
wait_for_import_warmup()
|
|
||||||
|
|
||||||
from agents.models.interface import ModelTracing
|
from agents.models.interface import ModelTracing
|
||||||
|
|
||||||
from strix.config.models import StrixProvider, configure_sdk_model_defaults
|
from strix.config.models import StrixProvider, configure_sdk_model_defaults
|
||||||
|
|||||||
@@ -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
|
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
|
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
|
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
|
a scan (argument parsing, Docker checks, image pull, TUI setup).
|
||||||
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
|
The main thread must call :func:`wait_for_import_warmup` before its first
|
||||||
import lock, so behaviour is unchanged either way.
|
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
|
from __future__ import annotations
|
||||||
|
|
||||||
import importlib
|
import importlib
|
||||||
import logging
|
import logging
|
||||||
import sys
|
|
||||||
import threading
|
import threading
|
||||||
|
|
||||||
|
|
||||||
@@ -27,42 +28,15 @@ WARMUP_MODULES = (
|
|||||||
"docker",
|
"docker",
|
||||||
)
|
)
|
||||||
|
|
||||||
_lock = threading.Lock()
|
|
||||||
_thread: threading.Thread | None = None
|
_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:
|
def _warm(modules: tuple[str, ...]) -> None:
|
||||||
for name in modules:
|
for name in modules:
|
||||||
before = frozenset(sys.modules)
|
|
||||||
try:
|
try:
|
||||||
importlib.import_module(name)
|
importlib.import_module(name)
|
||||||
except Exception: # noqa: BLE001 - a failed warm-up must never fail the run.
|
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)
|
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:
|
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.
|
runtime that has no local Docker) warm a narrower set.
|
||||||
"""
|
"""
|
||||||
global _thread # noqa: PLW0603
|
global _thread # noqa: PLW0603
|
||||||
with _lock:
|
if _thread is None:
|
||||||
if _thread is not None:
|
|
||||||
return _thread
|
|
||||||
_thread = threading.Thread(
|
_thread = threading.Thread(
|
||||||
target=_warm, args=(modules,), name="strix-import-warmup", daemon=True
|
target=_warm, args=(modules,), name="strix-import-warmup", daemon=True
|
||||||
)
|
)
|
||||||
_thread.start()
|
_thread.start()
|
||||||
return _thread
|
return _thread
|
||||||
|
|
||||||
|
|
||||||
def wait_for_import_warmup(timeout: float | None = None) -> None:
|
def wait_for_import_warmup() -> None:
|
||||||
"""Block until the warm-up thread has finished, if one was started.
|
"""Block until the warm-up thread has finished, if one was started."""
|
||||||
|
if _thread is not None:
|
||||||
Call this before the first import of a warmed module on another thread.
|
_thread.join()
|
||||||
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)
|
|
||||||
|
|||||||
@@ -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
|
Two threads that enter the same package graph from different modules hold
|
||||||
thread's ``strix.report`` import both walked the agents SDK graph, and the two
|
each other's import locks (warm-up: ``strix.core.runner`` -> ``agents``;
|
||||||
held each other's import locks (report -> dedupe -> agents while runner ->
|
main: ``agents.models.interface``). CPython breaks such a cycle by failing one
|
||||||
hooks -> report.state). CPython's deadlock avoidance breaks such a cycle by
|
of the imports, so the main thread waits for the warm-up before its first
|
||||||
failing one import, which strands finished submodules in ``sys.modules`` with
|
engine import.
|
||||||
their parent package gone — and the next import of one of those submodules
|
|
||||||
crashes with "partially initialized module".
|
|
||||||
"""
|
"""
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
@@ -56,42 +54,6 @@ def test_check_duplicate_resolves_lazily() -> None:
|
|||||||
assert result.returncode == 0, result.stderr
|
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:
|
def test_wait_for_import_warmup_lets_main_thread_import_the_agents_graph() -> None:
|
||||||
result = _run(
|
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
|
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
|
# 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()
|
start_import_warmup()
|
||||||
wait_for_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
|
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:
|
def test_wait_for_import_warmup_is_a_no_op_without_a_thread() -> None:
|
||||||
warmup.wait_for_import_warmup(timeout=0)
|
warmup.wait_for_import_warmup()
|
||||||
|
|
||||||
|
|
||||||
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
|
|
||||||
|
|||||||
Reference in New Issue
Block a user