diff --git a/continuous_integration/run_tests.sh b/continuous_integration/run_tests.sh index f26f0fdc..078bc89a 100755 --- a/continuous_integration/run_tests.sh +++ b/continuous_integration/run_tests.sh @@ -18,4 +18,8 @@ fi # launching the tests: python -m threadpoolctl -i numpy scipy.linalg tests._openmp_test_helper.openmp_helpers_inner +# Introspect API scope: +PYTHONPATH=. python tests/empirical_scope_observation.py blas +PYTHONPATH=. python tests/empirical_scope_observation.py openmp + pytest -vlrxXs -W error -k "$TESTS" --junitxml=test_result.xml --cov=threadpoolctl --cov-report xml diff --git a/dev-requirements.txt b/dev-requirements.txt index 04901b66..d98aa93a 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -1,5 +1,6 @@ flit coverage +psutil pytest pytest-cov cython diff --git a/tests/empirical_scope_observation.py b/tests/empirical_scope_observation.py new file mode 100644 index 00000000..8624820a --- /dev/null +++ b/tests/empirical_scope_observation.py @@ -0,0 +1,123 @@ +""" +Run this script to empirically determine if OpenMP and BLAS limiting API set +the size of a shared process-wide thread pool or a per-thread limit. +""" + +import sys +from concurrent.futures import ThreadPoolExecutor +from os import cpu_count +from pprint import pprint +from threading import Thread +from typing import Callable + +import threadpoolctl +import psutil + + +def start_counting_threads() -> Callable[[], int]: + num_threads = 0 + stop = False + + def in_thread(): + nonlocal num_threads, stop + process = psutil.Process() + while not stop: + num_threads = max(num_threads, process.num_threads()) + + thread = Thread(target=in_thread) + thread.start() + + def finish(): + nonlocal stop, num_threads + stop = True + thread.join() + return num_threads + + return finish + + +def set_limits(): + threadpoolctl.threadpool_limits(2) + print("== Library info ==") + pprint(threadpoolctl.threadpool_info()) + print() + + +def blas(num_python_threads): + try: + import numpy as np + except ImportError: + print("BLAS not available") + return False + + set_limits() + A = np.ones((10_000_000,)) + + def blas_math(_): + A.dot(A) + + with ThreadPoolExecutor(num_python_threads) as pool: + for _ in pool.map(blas_math, range(400)): + pass + + return True + + +def openmp(num_python_threads): + try: + from tests._openmp_test_helper.openmp_helpers_inner import ( + check_openmp_num_threads, + ) + except ImportError: + print("OpenMP not available") + return False + + set_limits() + + def run_openmp(_): + check_openmp_num_threads(1000) + + with ThreadPoolExecutor(num_python_threads) as pool: + for _ in pool.map(run_openmp, range(400)): + pass + + return True + + +def run(which: str) -> None: + if which == "blas": + func = blas + else: + func = openmp + + num_python_threads = 10 * cpu_count() + count_threads = start_counting_threads() + if not func(num_python_threads): + count_threads() + return + max_num_threads = count_threads() + + # We're creating os.cpu_count() * 10 Python threads, and asking for 2 + # native threads. If number is close to number of Python threads, we'll + # assume process-wide shared thread pool. If the number is close to double + # the number of Python threads, we'll assume a thread pool per thread. + if max_num_threads < num_python_threads * 1.25: + scope = "process-wide shared thread pool" + elif max_num_threads > num_python_threads * 1.75: + scope = "thread pool per thread" + else: + scope = "not sure" + + print(f"== Observed behavior: {which} ==") + print("Maximum number of observed threads:", max_num_threads) + print("Presumed API scope:", scope) + print() + + +if __name__ == "__main__": + if sys.argv[1] == "openmp": + run("openmp") + elif sys.argv[1] == "blas": + run("blas") + else: + raise SystemExit("Bad argument")