From 9bb91506bea79689867147ca9440aae25de3752c Mon Sep 17 00:00:00 2001 From: parashardhapola Date: Mon, 26 Oct 2026 22:58:17 +0100 Subject: [PATCH] Keep compiling and garbage collection off the worker threads of mixed models The mixed-model fits and Kenward and Roger's adjustment run their compiled routines on Python threads. A garbage collection that starts on one of those threads runs the finalizers of what it frees there, and some objects may only be finalized on the thread that made them: Tk's, which matplotlib's interactive backend makes, raised errors in the test suite, and a cold compile on a worker thread once crashed the interpreter in the collector. The shared helper by_runs now runs the first run on the calling thread, so the routine compiles there, and pauses the collector while the workers run. The tests draw figures with matplotlib's Agg backend, as CI does. --- cytepool/_jit.py | 44 ++++++++++++++- cytepool/stats/_fixed_effect_inference.py | 27 +++------- cytepool/stats/mixed_models.py | 28 +++------- tests/conftest.py | 8 +++ tests/test_jit.py | 66 +++++++++++++++++++++++ 5 files changed, 130 insertions(+), 43 deletions(-) create mode 100644 tests/conftest.py diff --git a/cytepool/_jit.py b/cytepool/_jit.py index 5eb8731..23630f5 100644 --- a/cytepool/_jit.py +++ b/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. diff --git a/cytepool/stats/_fixed_effect_inference.py b/cytepool/stats/_fixed_effect_inference.py index 5f8af8e..f9dece7 100644 --- a/cytepool/stats/_fixed_effect_inference.py +++ b/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( diff --git a/cytepool/stats/mixed_models.py b/cytepool/stats/mixed_models.py index acfb62d..8477213 100644 --- a/cytepool/stats/mixed_models.py +++ b/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, diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..f96565e --- /dev/null +++ b/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") diff --git a/tests/test_jit.py b/tests/test_jit.py index 9114a08..5f7620d 100644 --- a/tests/test_jit.py +++ b/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()