Keep compiling and garbage collection off the worker threads of mixed models

cytepool/_jit.py

@@ -12,15 +12,18 @@
 compiled in every process instead.
 
 Parallel routines use numba's threads; ``thread_limit`` caps how many of them one
-call uses, and ``by_chunks`` runs a loop over items in parallel or on the calling
-thread.
+call uses, ``by_chunks`` runs a loop over items in parallel or on the calling
+thread, and ``by_runs`` runs a routine that releases the GIL over runs of items on
+Python threads.
 """
 
+import gc
 import hashlib
 import logging
 import os
 import tempfile
 from collections.abc import Callable, Iterator
+from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from functools import cache
 from pathlib import Path
@@ -132,6 +135,43 @@ def by_chunks(
     return serial(*arguments)
 
 
+def by_runs(run: Callable[[int, int], Any], n_items: int) -> None:
+    """Run a compiled routine over runs of consecutive items on Python threads.
+
+    ``run(start, stop)`` computes the items ``start:stop`` and releases the GIL while
+    it works, as a routine compiled with ``nogil=True`` does. With more than one of
+    numba's threads and more than one item, the items are split into about four runs
+    per thread. The first run goes on the calling thread, so that the routine is
+    compiled, or loaded from the cache, there and never on a worker thread; the
+    others go to a pool of Python threads, with the garbage collector paused until
+    they finish. A collection that started on a worker thread would run the
+    finalizers of the objects it frees on that thread, and some objects may only be
+    finalized on the thread that made them: those of Tk, which matplotlib's
+    interactive backend makes, raise an error there, and a cold compile on a worker
+    thread once crashed the interpreter in the collector. Each item is computed by
+    the same operations whatever its run, so the result does not depend on the
+    number of threads.
+    """
+    threads = min(numba.get_num_threads(), n_items)
+    if threads <= 1:
+        run(0, n_items)
+        return
+    n_runs = min(n_items, 4 * threads)
+    bounds = [
+        (k * n_items // n_runs, (k + 1) * n_items // n_runs) for k in range(n_runs)
+    ]
+    run(*bounds[0])
+    collecting = gc.isenabled()
+    gc.disable()
+    try:
+        with ThreadPoolExecutor(max_workers=threads) as pool:
+            for _ in pool.map(lambda bound: run(*bound), bounds[1:]):
+                pass
+    finally:
+        if collecting:
+            gc.enable()
+
+
 @cache
 def cache_key() -> str:
     """Return the key that names this installation's cache directory.

cytepool/stats/_fixed_effect_inference.py

@@ -66,14 +66,12 @@
 """
 
 import math
-from concurrent.futures import ThreadPoolExecutor
 from typing import Any, NamedTuple
 
-import numba
 import numpy as np
 import numpy.typing as npt
 
-from .._jit import jit
+from .._jit import by_runs, jit
 from .._linear_algebra import cholesky, column_space, symmetric
 from .distributions import f_sf, t_sf
 
@@ -175,23 +173,12 @@ def kenward_roger_parts(fit: Any, average_weight: bool) -> InferenceParts:
         covariance,
     )
     outputs = (adjusted, derivatives, parameter_covariance)
-    threads = min(numba.get_num_threads(), n_features)
-    if threads > 1:
-        # Runs of consecutive features on Python threads; the compiled routine
-        # releases the GIL, and each feature is computed by the same operations.
-        n_runs = min(n_features, 4 * threads)
-        bounds = [
-            (run * n_features // n_runs, (run + 1) * n_features // n_runs)
-            for run in range(n_runs)
-        ]
-        with ThreadPoolExecutor(max_workers=threads) as pool:
-            for _ in pool.map(
-                lambda run: _kenward_roger_features(*arguments, *run, *outputs),
-                bounds,
-            ):
-                pass
-    else:
-        _kenward_roger_features(*arguments, 0, n_features, *outputs)
+    # Runs of consecutive features on Python threads; the compiled routine releases
+    # the GIL, and each feature is computed by the same operations.
+    by_runs(
+        lambda start, stop: _kenward_roger_features(*arguments, start, stop, *outputs),
+        n_features,
+    )
     # Back to the coefficients of x: beta = r^-1 beta_q.
     adjusted = np.einsum("ij,gjk,lk->gil", inverse, adjusted, inverse, optimize=True)
     derivatives = np.einsum(

cytepool/stats/mixed_models.py

@@ -220,15 +220,13 @@
 import math
 from collections import namedtuple
 from collections.abc import Sequence
-from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
 from typing import Any, Literal
 
-import numba
 import numpy as np
 import numpy.typing as npt
 
-from cytepool._jit import jit, thread_limit
+from cytepool._jit import by_runs, jit, thread_limit
 from cytepool._linear_algebra import RANK_TOLERANCE, cholesky, symmetric
 from cytepool._masked_arrays import real_array, reject_masked, unmasked_array
 from cytepool._options import checked_choice
@@ -605,24 +603,12 @@ def _fit(
         level_starts,
         ml,
     )
-    threads = min(numba.get_num_threads(), n_features)
-    if threads > 1:
-        # Runs of consecutive features on Python threads, several per thread to
-        # balance the work; the compiled fits release the GIL. Each feature is
-        # fitted by the same operations whatever its run, so the result does not
-        # depend on the number of threads.
-        n_runs = min(n_features, 4 * threads)
-        bounds = [
-            (run * n_features // n_runs, (run + 1) * n_features // n_runs)
-            for run in range(n_runs)
-        ]
-        with ThreadPoolExecutor(max_workers=threads) as pool:
-            for _ in pool.map(
-                lambda run: _fit_features(*arguments, *run, *outputs), bounds
-            ):
-                pass
-    else:
-        _fit_features(*arguments, 0, n_features, *outputs)
+    # Runs of consecutive features on Python threads; the compiled fits release the
+    # GIL, and each feature is fitted by the same operations whatever its run.
+    by_runs(
+        lambda start, stop: _fit_features(*arguments, start, stop, *outputs),
+        n_features,
+    )
     (
         coefficients,
         unscaled_covariance,

tests/conftest.py

@@ -0,0 +1,8 @@
+"""Settings shared by the whole test suite."""
+
+import os
+
+# Figures are drawn without windows, as on a machine without a display. Where a
+# display is available, matplotlib chooses an interactive backend, whose Tk objects
+# raise errors when the garbage collector finalizes them on another thread.
+os.environ.setdefault("MPLBACKEND", "Agg")

tests/test_jit.py

@@ -1,3 +1,4 @@
+import gc
 import importlib.util
 import os
 import sys
@@ -229,3 +230,68 @@ def test_thread_limit_rejects_invalid_counts(
         with _jit.thread_limit(n_threads):  # type: ignore[arg-type]
             pass
     assert numba.get_num_threads() == before
+
+
+@pytest.mark.skipif(
+    numba.config.NUMBA_NUM_THREADS < 2, reason="needs two threads to use a pool"
+)
+def test_by_runs_covers_every_item_once_and_starts_on_the_calling_thread() -> None:
+    calls: list[tuple[int, int, int]] = []
+    lock = threading.Lock()
+
+    def run(start: int, stop: int) -> None:
+        with lock:
+            calls.append((start, stop, threading.get_ident()))
+
+    with _jit.thread_limit(2):
+        _jit.by_runs(run, 23)
+    items = sorted(i for start, stop, _ in calls for i in range(start, stop))
+    assert items == list(range(23))
+    # Two threads give eight runs, the first of them on this thread, where the
+    # routine compiles before any worker uses it.
+    assert len(calls) == 8
+    assert calls[0][:2] == (0, 2) and calls[0][2] == threading.get_ident()
+
+
+def test_by_runs_on_one_thread_runs_everything_here() -> None:
+    calls: list[tuple[int, int, int]] = []
+    with _jit.thread_limit(1):
+        _jit.by_runs(
+            lambda start, stop: calls.append((start, stop, threading.get_ident())), 5
+        )
+    assert calls == [(0, 5, threading.get_ident())]
+
+
+@pytest.mark.skipif(
+    numba.config.NUMBA_NUM_THREADS < 2, reason="needs two threads to use a pool"
+)
+def test_by_runs_finalizes_the_garbage_of_its_workers_on_the_calling_thread() -> None:
+    """Objects that only the collector frees are finalized where a collection runs.
+    With the collector running at almost every allocation, the workers' garbage
+    would be finalized on the workers; by_runs pauses it, so they are finalized
+    here, where objects such as Tk's must be."""
+    finalized_on: list[int] = []
+
+    class Cycle:
+        def __init__(self) -> None:
+            self.itself = self
+
+        def __del__(self) -> None:
+            finalized_on.append(threading.get_ident())
+
+    def run(start: int, stop: int) -> None:
+        for _ in range(200):
+            Cycle()
+
+    threshold = gc.get_threshold()
+    gc.collect()
+    gc.set_threshold(1)
+    try:
+        with _jit.thread_limit(2):
+            _jit.by_runs(run, 16)
+    finally:
+        gc.set_threshold(*threshold)
+    gc.collect()
+    assert len(finalized_on) == 8 * 200
+    assert set(finalized_on) == {threading.get_ident()}
+    assert gc.isenabled()