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()