diff --git a/packages/flow/openadapt_flow/__main__.py b/packages/flow/openadapt_flow/__main__.py index efba5f140..ae84308c4 100644 --- a/packages/flow/openadapt_flow/__main__.py +++ b/packages/flow/openadapt_flow/__main__.py @@ -10204,10 +10204,15 @@ def _qualify_default_form(argv: list[str]) -> list[str]: return argv +# Flags that print and exit before any work, so loading OCR would only slow +# the exit down. +_PREWARM_SKIP_FLAGS = frozenset({"-h", "--help", "--version"}) + + def main(argv: Sequence[str] | None = None) -> int: """CLI entry point. Returns a process exit code.""" raw = list(sys.argv[1:] if argv is None else argv) - if raw[:1] in (["demo"], ["tutorial"]): + if raw[:1] in (["demo"], ["tutorial"]) and not _PREWARM_SKIP_FLAGS & set(raw): # Building the parser imports for a second or two on a new install; # start loading OCR and checking the browser driver meanwhile. from openadapt_flow import _prewarm diff --git a/packages/flow/openadapt_flow/_prewarm.py b/packages/flow/openadapt_flow/_prewarm.py index 4e58325eb..ef03f5a7b 100644 --- a/packages/flow/openadapt_flow/_prewarm.py +++ b/packages/flow/openadapt_flow/_prewarm.py @@ -14,6 +14,7 @@ from __future__ import annotations +import atexit import threading _started = False @@ -80,6 +81,16 @@ def _touch_driver() -> None: pass +# onnxruntime starts native threads while it loads. If the interpreter exits +# during that load, the C++ runtime aborts the process ("terminate called +# without an active exception", exit -6), so exit waits up to this long for it. +_EXIT_WAIT_S = 30.0 + + +def _wait_for_load(thread: threading.Thread) -> None: + thread.join(timeout=_EXIT_WAIT_S) + + def start() -> None: """Start the background work once per process and return at once.""" global _started @@ -88,4 +99,6 @@ def start() -> None: return _started = True threading.Thread(target=_touch_driver, name="driver-prewarm", daemon=True).start() - threading.Thread(target=_load, name="ocr-prewarm", daemon=True).start() + ocr = threading.Thread(target=_load, name="ocr-prewarm", daemon=True) + ocr.start() + atexit.register(_wait_for_load, ocr) diff --git a/packages/flow/tests/test_ocr_backends.py b/packages/flow/tests/test_ocr_backends.py index 2de98a9bc..1ee901200 100644 --- a/packages/flow/tests/test_ocr_backends.py +++ b/packages/flow/tests/test_ocr_backends.py @@ -143,3 +143,46 @@ def test_prewarm_image_is_a_valid_png() -> None: with Image.open(io.BytesIO(_prewarm._blank_png())) as image: image.load() assert image.size == (64, 16) + + +@pytest.mark.parametrize("flag", ["-h", "--help", "--version"]) +@pytest.mark.parametrize("command", ["demo", "tutorial"]) +def test_help_does_not_start_the_prewarm( + monkeypatch: pytest.MonkeyPatch, command: str, flag: str +) -> None: + """`tutorial --help` exited while onnxruntime was still loading on the + prewarm thread, and the native runtime aborted the process (exit -6).""" + from openadapt_flow import __main__ as cli + from openadapt_flow import _prewarm + + started: list[bool] = [] + monkeypatch.setattr(_prewarm, "start", lambda: started.append(True)) + with pytest.raises(SystemExit): + cli.main([command, flag]) + assert started == [] + + +def test_exit_waits_for_an_ocr_load_in_progress( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import threading + + from openadapt_flow import _prewarm + + release = threading.Event() + registered: list[Any] = [] + monkeypatch.setattr(_prewarm, "_started", False) + monkeypatch.setattr(_prewarm, "_load", release.wait) + monkeypatch.setattr(_prewarm, "_touch_driver", lambda: None) + monkeypatch.setattr( + _prewarm.atexit, "register", lambda fn, *args: registered.append((fn, args)) + ) + _prewarm.start() + + assert len(registered) == 1 + fn, (thread,) = registered[0] + assert fn is _prewarm._wait_for_load and thread.name == "ocr-prewarm" + assert thread.is_alive() + release.set() + fn(thread) + assert not thread.is_alive()