|
| 1 | +"""Single-thread executor for the daemon's PyTorch/MPS model inference. |
| 2 | +
|
| 3 | +The sentence-transformer embedder and the cross-encoder reranker run on MPS |
| 4 | +(Apple-Silicon GPU). PyTorch MPS is NOT safe to drive from multiple threads: |
| 5 | +concurrent forward passes from different threads share the model's internal |
| 6 | +buffers, which both |
| 7 | +
|
| 8 | + 1. **deadlocks** after the first cross-thread call (the priming RPC served |
| 9 | + exactly one request, then every later request hung), and |
| 10 | + 2. **corrupts** intermediate tensors — observed as |
| 11 | + ``RuntimeError: size of tensor a (45) must match tensor b (35) at |
| 12 | + non-singleton dimension 1`` (one query's sequence length colliding with |
| 13 | + another's attention mask). |
| 14 | +
|
| 15 | +Several daemon threads want to run inference: the priming RPC handler pool, the |
| 16 | +Scholar/Librarian agent-research workers (embedding search via |
| 17 | +``library_tools``), and the startup warmup. Funnel EVERY model forward pass |
| 18 | +through the one dedicated thread here, so MPS is only ever driven from a single |
| 19 | +thread. Connection I/O and BM25 (CPU) stay concurrent; only the model calls |
| 20 | +serialise. |
| 21 | +
|
| 22 | +Wrap LEAF model calls only (``index.search`` / ``rerank`` / ``encode``), never a |
| 23 | +whole pipeline that itself calls :func:`run` — submitting to this 1-worker pool |
| 24 | +from its own worker thread would deadlock. The startup warmup must also go |
| 25 | +through :func:`run` so the MPS kernels JIT-compile on the serving thread. |
| 26 | +""" |
| 27 | + |
| 28 | +from __future__ import annotations |
| 29 | + |
| 30 | +import threading |
| 31 | +from concurrent.futures import ThreadPoolExecutor |
| 32 | +from typing import Callable, Optional, ParamSpec, TypeVar |
| 33 | + |
| 34 | +_P = ParamSpec("_P") |
| 35 | +_T = TypeVar("_T") |
| 36 | + |
| 37 | +# Lazily-created single-worker executor. Re-creatable: :func:`shutdown` resets it |
| 38 | +# so a fresh one is built on the next :func:`run` — without this, a shutdown |
| 39 | +# (daemon exit, or a test exercising it) would poison every later call with |
| 40 | +# "cannot schedule new futures after shutdown". |
| 41 | +_lock = threading.Lock() |
| 42 | +_executor: Optional[ThreadPoolExecutor] = None |
| 43 | + |
| 44 | + |
| 45 | +def _get_executor() -> ThreadPoolExecutor: |
| 46 | + global _executor |
| 47 | + with _lock: |
| 48 | + if _executor is None: |
| 49 | + _executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mps-inference") |
| 50 | + return _executor |
| 51 | + |
| 52 | + |
| 53 | +def run(fn: Callable[_P, _T], *args: _P.args, **kwargs: _P.kwargs) -> _T: |
| 54 | + """Run a leaf model-inference callable on the single inference thread and |
| 55 | + return its result (blocking the caller until it completes). |
| 56 | +
|
| 57 | + MUST wrap only leaf calls (``encode`` / ``predict`` / ``index.search`` / |
| 58 | + ``rerank``); never a pipeline that re-enters :func:`run`, which would submit |
| 59 | + to this 1-worker pool from its own worker thread and deadlock. |
| 60 | + """ |
| 61 | + return _get_executor().submit(fn, *args, **kwargs).result() |
| 62 | + |
| 63 | + |
| 64 | +def shutdown() -> None: |
| 65 | + """Stop the inference thread (daemon shutdown) so the executor's at-exit |
| 66 | + join doesn't delay process exit. Idempotent; a later :func:`run` lazily |
| 67 | + recreates the executor.""" |
| 68 | + global _executor |
| 69 | + with _lock: |
| 70 | + ex, _executor = _executor, None |
| 71 | + if ex is not None: |
| 72 | + ex.shutdown(wait=False, cancel_futures=True) |
0 commit comments