From 66116459218aa2e45df7a3c566f9b4007d771b77 Mon Sep 17 00:00:00 2001 From: cekees <613654+cekees@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:19:43 +0200 Subject: [PATCH 1/3] One thread per MPI rank by default: OMP, OpenBLAS, MKL and Accelerate thread counts are 1 under MPI Under MPI with more than one rank, import proteus (and proteus.Comm, which knows the real rank count) sets OMP_NUM_THREADS, OPENBLAS_NUM_THREADS, MKL_NUM_THREADS and VECLIB_MAXIMUM_THREADS to 1 unless the user set them. A serial run, or a single rank, is untouched and may use threads. OPENBLAS_NUM_THREADS alone is not enough: conda-forge's OpenBLAS is built on OpenMP and ignores it (13.8 cores busy in one process with only that set; one with OMP_NUM_THREADS=1). The variables are read at library load, so when numpy was imported before proteus, Comm limits the loaded libraries through threadpoolctl if it is installed (added to the dev environments) and logs a warning if not. Tests cover the rules on a fake environment, import proteus under each launcher variable, and real mpiexec ranks measuring BLAS CPU time per wall time with proteus first and with numpy first. Co-Authored-By: Claude Sonnet 5.5 --- docs/source/index.rst | 12 +++ environment-dev-up.yml | 1 + environment-mpich-dev.yml | 1 + environment-openmpi-dev.yml | 1 + proteus/Comm.py | 7 ++ proteus/__init__.py | 6 ++ proteus/_threads.py | 113 +++++++++++++++++++++++ test/test_thread_defaults.py | 170 +++++++++++++++++++++++++++++++++++ 8 files changed, 311 insertions(+) create mode 100644 proteus/_threads.py create mode 100644 test/test_thread_defaults.py diff --git a/docs/source/index.rst b/docs/source/index.rst index cd6e5c1758..35e86ec4ba 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -86,6 +86,18 @@ XDMF view of it first:: and open the resulting .xmf. +Threads and MPI +--------------- + +Run under MPI with more than one rank, Proteus gives each rank one thread: it sets +``OMP_NUM_THREADS``, ``OPENBLAS_NUM_THREADS``, ``MKL_NUM_THREADS`` and ``VECLIB_MAXIMUM_THREADS`` +to 1 (``proteus/_threads.py``), because one rank per core is how a node is filled and +OpenBLAS would otherwise start a thread per core in every rank. A value you set yourself +is kept, so ``OMP_NUM_THREADS=4 mpiexec -n 8 parun ...`` runs eight ranks of four threads. +A serial run is not touched and uses whatever threads the libraries choose. The variables +are read when numpy is imported: import ``proteus`` before numpy, or install +``threadpoolctl``, which lets Proteus limit libraries that are already loaded. + .. _capabilities-sec: Capabilities diff --git a/environment-dev-up.yml b/environment-dev-up.yml index 5389f3754c..7b7b66e2cc 100644 --- a/environment-dev-up.yml +++ b/environment-dev-up.yml @@ -45,6 +45,7 @@ dependencies: - pytest-xdist >=3.6.1 - scipy >=1.18.0 - tetgen >=1.6.0 + - threadpoolctl >=3.5.0 - python >=3.13.15 - ipyparallel >=9.2 - git-lfs >=3.7 diff --git a/environment-mpich-dev.yml b/environment-mpich-dev.yml index 0244e8d6eb..321a7b904c 100644 --- a/environment-mpich-dev.yml +++ b/environment-mpich-dev.yml @@ -68,6 +68,7 @@ dependencies: - pytest-forked =1.7.5 - scipy =1.18.0 - tetgen =1.6.0 + - threadpoolctl >=3.5.0 - python =3.13.15 - pybind11 =3.1.0 - eigen =5.0.1 diff --git a/environment-openmpi-dev.yml b/environment-openmpi-dev.yml index 22a35ed525..35b0b3906d 100644 --- a/environment-openmpi-dev.yml +++ b/environment-openmpi-dev.yml @@ -53,6 +53,7 @@ dependencies: - pytest-forked =1.7.5 - scipy =1.18.0 - tetgen =1.6.0 + - threadpoolctl >=3.5.0 - python =3.13.15 - pybind11 =3.1.0 - eigen =5.0.1 diff --git a/proteus/Comm.py b/proteus/Comm.py index 915ac06638..f6b1bda084 100644 --- a/proteus/Comm.py +++ b/proteus/Comm.py @@ -22,10 +22,17 @@ # Special workaround for broken MPI on certain Cray systems from . import config +from . import _threads mpi_preload_libs=[] for lib in config.PROTEUS_PRELOAD_LIBS: mpi_preload_libs.append(ctypes.CDLL(lib,mode=ctypes.RTLD_GLOBAL)) +# One thread per rank under MPI. import proteus already did this when the launcher announces +# its size; this covers a launcher that does not, and numpy imported before proteus. +_thread_warning = _threads.after_mpi_init(MPI.COMM_WORLD.Get_size()) +if _thread_warning: + logEvent(_thread_warning, 0) + comm = None argv = sys.argv petscInitialized = False diff --git a/proteus/__init__.py b/proteus/__init__.py index e3f2634cfd..e0cbcd8b27 100644 --- a/proteus/__init__.py +++ b/proteus/__init__.py @@ -35,6 +35,12 @@ if _sys.platform == "darwin": _os.environ.setdefault("FI_PROVIDER", "tcp") +# Under MPI, one thread per rank for the BLAS and OpenMP runtimes numpy brings in, +# unless the user set the variables (proteus/_threads.py). Without MPI nothing +# changes. The variables are read when numpy is imported, so this comes first. +from . import _threads +_threads.apply_defaults(_threads.launched_size()) + __version__ = '2.0.0.dev' #: What ``from proteus import *`` binds. It is that surface, not an diff --git a/proteus/_threads.py b/proteus/_threads.py new file mode 100644 index 0000000000..e8f1378197 --- /dev/null +++ b/proteus/_threads.py @@ -0,0 +1,113 @@ +"""One thread per MPI rank, unless the user asked for something else. + +Proteus runs MPI-only on a node: its legacy kernels are limited by +instruction issue on one core, so a rank cannot use a second core +(ON_NODE_PERFORMANCE_BASELINE.md, section 6.2). The numerical libraries +that sit under it do not know that. OpenBLAS starts one thread per core +in every process (16 per rank on a 16-core Mac), so a node running one rank +per core is oversubscribed by the number of cores, and a rank waiting in a +BLAS call is slowed by every other rank's threads. + +Under MPI with more than one rank, the thread-count variables of the +libraries numpy and scipy bring in are therefore given the value 1, +**unless the user already set them**. Without MPI (or with one rank) nothing +is touched: a serial run is free to use the node's threads. + +``OPENBLAS_NUM_THREADS`` alone is not enough: an OpenBLAS built on OpenMP +(conda-forge's, on macOS and on Linux) takes its thread count from +``OMP_NUM_THREADS`` and ignores ``OPENBLAS_NUM_THREADS=1`` (measured: 13.8 +cores busy in one process with only the latter set, one with the former). All +the variables are set. + +The variables are read when the library is loaded, which is when numpy is +imported, so this has to happen before that. ``import proteus`` does it +for a launch whose size the launcher announces in the environment. If numpy +was imported first, or the launcher announces nothing, ``proteus.Comm`` -- which +knows the real size -- sets the variables and, when threadpoolctl is installed, +also limits the libraries that are already loaded. +""" +import os +import sys + +#: variable -> threadpoolctl ``user_api`` of the library it controls +THREAD_VARIABLES = { + "OMP_NUM_THREADS": "openmp", + "OPENBLAS_NUM_THREADS": "blas", + "MKL_NUM_THREADS": "blas", + "VECLIB_MAXIMUM_THREADS": "blas", +} + +#: variables in which launchers announce the number of ranks in the job +LAUNCHER_SIZE_VARIABLES = ( + "OMPI_COMM_WORLD_SIZE", # Open MPI (mpiexec, mpirun) + "PMI_SIZE", # MPICH/Hydra, Slurm pmi2, Cray + "PMIX_SIZE", + "SLURM_STEP_NUM_TASKS", # srun +) + +#: the variables this process has defaulted (the user set none of them) +defaulted = [] +#: whether numpy was already loaded when they were defaulted +numpy_was_loaded = False + + +def launched_size(environ=None): + """The number of ranks the launcher announces, or 1 if it announces none.""" + environ = os.environ if environ is None else environ + size = 1 + for name in LAUNCHER_SIZE_VARIABLES: + try: + size = max(size, int(environ.get(name, "1"))) + except ValueError: + pass + return size + + +def apply_defaults(size, environ=None): + """Give every thread variable the user left unset the value 1, if size > 1. + + Returns the names set by this call. + """ + global numpy_was_loaded + environ = os.environ if environ is None else environ + if size <= 1: + return [] + names = [n for n in THREAD_VARIABLES if n not in environ] + for name in names: + environ[name] = "1" + if names: + defaulted.extend(names) + numpy_was_loaded = numpy_was_loaded or "numpy" in sys.modules + return names + + +def limit_loaded_libraries(): + """Apply the defaulted variables to libraries that were loaded before them. + + Returns the ``user_api`` kinds limited; empty if there was nothing to do or + threadpoolctl is not installed. + """ + if not (defaulted and numpy_was_loaded): + return [] + try: + import threadpoolctl + except ImportError: + return [] + apis = sorted({THREAD_VARIABLES[n] for n in defaulted}) + for api in apis: + threadpoolctl.threadpool_limits(limits=1, user_api=api) + return apis + + +def after_mpi_init(size): + """Called by proteus.Comm once the number of ranks is known. + + Returns a message to log when libraries already loaded could not be limited. + """ + apply_defaults(size) + if not (defaulted and numpy_was_loaded): + return None + if limit_loaded_libraries(): + return None + return ("numpy was imported before proteus, so its BLAS may use more than one thread per MPI rank. " + "Import proteus first, set OMP_NUM_THREADS=1 and OPENBLAS_NUM_THREADS=1, or install threadpoolctl.") diff --git a/test/test_thread_defaults.py b/test/test_thread_defaults.py new file mode 100644 index 0000000000..3c3dc04ddc --- /dev/null +++ b/test/test_thread_defaults.py @@ -0,0 +1,170 @@ +"""One thread per MPI rank by default; threads left alone without MPI. + +proteus/_threads.py gives the thread-count variables of numpy's BLAS and +OpenMP runtimes the value 1 when the job has more than one rank and the user +did not set them. The tests check the rules on a fake environment, then the +effect on real processes: that ``import proteus`` sets the variables for a +launcher that announces its size, that a user's value survives, that a single +process is untouched, and -- under a real ``mpiexec`` -- that a rank's BLAS +runs on one core, whether numpy or proteus was imported first. +""" +import json +import os +import shutil +import subprocess +import sys +import textwrap + +import pytest + +from proteus import _threads + +VARIABLES = tuple(_threads.THREAD_VARIABLES) +LAUNCHER = _threads.LAUNCHER_SIZE_VARIABLES + + +def clean_environment(**extra): + """This process's environment without thread settings or launcher announcements.""" + env = {k: v for k, v in os.environ.items() + if k not in VARIABLES and k not in LAUNCHER and not k.startswith(("OMPI_", "PMIX_", "PMI_"))} + env.update(extra) + return env + + +# --- the rules, on a fake environment -------------------------------------------------------- + +def test_launched_size_reads_the_launchers_variables(): + assert _threads.launched_size({}) == 1 + assert _threads.launched_size({"OMPI_COMM_WORLD_SIZE": "4"}) == 4 + assert _threads.launched_size({"PMI_SIZE": "8"}) == 8 + assert _threads.launched_size({"SLURM_STEP_NUM_TASKS": "16"}) == 16 + assert _threads.launched_size({"PMI_SIZE": "not a number"}) == 1 + + +@pytest.fixture +def fresh_state(monkeypatch): + monkeypatch.setattr(_threads, "defaulted", []) + monkeypatch.setattr(_threads, "numpy_was_loaded", False) + + +def test_one_rank_changes_nothing(fresh_state): + env = {} + assert _threads.apply_defaults(1, env) == [] + assert env == {} + + +def test_several_ranks_set_every_unset_variable_to_one(fresh_state): + env = {} + assert sorted(_threads.apply_defaults(4, env)) == sorted(VARIABLES) + assert env == {name: "1" for name in VARIABLES} + + +def test_a_value_the_user_set_is_kept(fresh_state): + env = {"OMP_NUM_THREADS": "4", "OPENBLAS_NUM_THREADS": "2"} + names = _threads.apply_defaults(4, env) + assert env["OMP_NUM_THREADS"] == "4" and env["OPENBLAS_NUM_THREADS"] == "2" + assert "OMP_NUM_THREADS" not in names and "OPENBLAS_NUM_THREADS" not in names + assert env["MKL_NUM_THREADS"] == "1" + + +def test_nothing_to_warn_about_when_the_user_set_everything(fresh_state): + env = {name: "2" for name in VARIABLES} + assert _threads.after_mpi_init(4) is None + assert _threads.apply_defaults(4, env) == [] + + +# --- import proteus, no MPI ------------------------------------------------------------------ + +def report_environment(**extra): + out = subprocess.run( + [sys.executable, "-c", + "import json, os, proteus; print(json.dumps({k: os.environ.get(k) for k in %r}))" % (VARIABLES,)], + env=clean_environment(**extra), capture_output=True, text=True, timeout=120) + assert out.returncode == 0, out.stderr + return json.loads(out.stdout.splitlines()[-1]) + + +def test_import_proteus_without_mpi_leaves_threads_alone(): + assert set(report_environment().values()) == {None} + + +@pytest.mark.parametrize("launcher", ["OMPI_COMM_WORLD_SIZE", "PMI_SIZE"]) +def test_import_proteus_under_a_launcher_with_several_ranks_sets_one_thread(launcher): + assert report_environment(**{launcher: "4"}) == {name: "1" for name in VARIABLES} + + +def test_import_proteus_under_a_launcher_with_one_rank_leaves_threads_alone(): + assert set(report_environment(OMPI_COMM_WORLD_SIZE="1").values()) == {None} + + +def test_import_proteus_keeps_the_users_value(): + got = report_environment(OMPI_COMM_WORLD_SIZE="4", OMP_NUM_THREADS="3") + assert got["OMP_NUM_THREADS"] == "3" and got["OPENBLAS_NUM_THREADS"] == "1" + + +# --- real ranks ------------------------------------------------------------------------------ + +def _find_mpiexec(): + beside = os.path.join(os.path.dirname(sys.executable), "mpiexec") + return beside if os.path.exists(beside) else shutil.which("mpiexec") + + +MPIEXEC = _find_mpiexec() + +RANK_SCRIPT = textwrap.dedent(''' + import json, os, sys, time + if sys.argv[1] == "numpy-first": + import numpy as np + from proteus import Comm + import numpy as np + size = Comm.get().size() + a = np.random.rand(1200, 1200) + a @ a + wall0, cpu0 = time.perf_counter(), time.process_time() + for _ in range(4): + a @ a + wall, cpu = time.perf_counter() - wall0, time.process_time() - cpu0 + print("RANK " + json.dumps({"size": size, "cpu_per_wall": cpu / wall, + "env": {k: os.environ.get(k) for k in %r}}), flush=True) +''' % (VARIABLES,)) + + +def run_ranks(tmp_path, ranks, order, **extra): + pytest.importorskip("proteus.Comm") + if MPIEXEC is None: + pytest.skip("no mpiexec") + script = tmp_path / "rank.py" + script.write_text(RANK_SCRIPT) + out = subprocess.run([MPIEXEC, "-n", str(ranks), sys.executable, str(script), order], + env=clean_environment(**extra), capture_output=True, text=True, timeout=300) + assert out.returncode == 0, out.stdout + out.stderr + reports = [json.loads(line[5:]) for line in out.stdout.splitlines() if line.startswith("RANK ")] + assert len(reports) == ranks, out.stdout + out.stderr + return reports + + +def test_ranks_run_their_blas_on_one_core(tmp_path): + """proteus imported first: the variables are set before numpy loads its BLAS.""" + for report in run_ranks(tmp_path, 2, "proteus-first"): + assert report["size"] == 2 + assert report["env"] == {name: "1" for name in VARIABLES} + assert report["cpu_per_wall"] < 1.5, report + + +def test_ranks_run_their_blas_on_one_core_when_numpy_was_imported_first(tmp_path): + """The variables are read at load, so too late; threadpoolctl limits the loaded libraries.""" + pytest.importorskip("threadpoolctl") + for report in run_ranks(tmp_path, 2, "numpy-first"): + assert report["cpu_per_wall"] < 1.5, report + + +def test_a_single_rank_keeps_its_threads(tmp_path): + for report in run_ranks(tmp_path, 1, "proteus-first"): + assert report["size"] == 1 + assert set(report["env"].values()) == {None} + + +def test_a_value_the_user_set_survives_under_mpi(tmp_path): + for report in run_ranks(tmp_path, 2, "proteus-first", OMP_NUM_THREADS="2"): + assert report["env"]["OMP_NUM_THREADS"] == "2" + assert report["env"]["OPENBLAS_NUM_THREADS"] == "1" From 53b35619cbbbc5f6f917ea4790be064a43f8b8d2 Mon Sep 17 00:00:00 2001 From: cekees <613654+cekees@users.noreply.github.com> Date: Sat, 24 Oct 2026 09:33:10 +0200 Subject: [PATCH 2/3] test: the thread-default tests keep the MPI runtime's own settings in the ranks' environment clean_environment dropped every OMPI_*, PMIX_* and PMI_* variable, including CI's OMPI_MCA_* settings that keep Open MPI off UCX; two-rank launches then failed to create a UCP worker and hung until the timeout on linux-64 and linux-aarch64. Only the launcher's size announcements are removed now. Co-Authored-By: Claude Sonnet 5.5 --- test/test_thread_defaults.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/test/test_thread_defaults.py b/test/test_thread_defaults.py index 3c3dc04ddc..4c94f9dc0a 100644 --- a/test/test_thread_defaults.py +++ b/test/test_thread_defaults.py @@ -24,9 +24,13 @@ def clean_environment(**extra): - """This process's environment without thread settings or launcher announcements.""" + """This process's environment without thread settings or launcher size announcements. + + Only the announcements: the MPI runtime's own configuration (OMPI_MCA_*, which CI uses to keep Open MPI off UCX) + has to reach the ranks, or they fail to start. + """ env = {k: v for k, v in os.environ.items() - if k not in VARIABLES and k not in LAUNCHER and not k.startswith(("OMPI_", "PMIX_", "PMI_"))} + if k not in VARIABLES and k not in LAUNCHER and not k.startswith("OMPI_COMM_WORLD")} env.update(extra) return env From 21052224d1f3bcf0188a0f087942c457cfec3ede Mon Sep 17 00:00:00 2001 From: cekees <613654+cekees@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:21:49 +0200 Subject: [PATCH 3/3] test: the Dirac advection-diffusion check draws from its own generator until it has enough points It used the module's shared generator and 40 draws. Where that generator stood depended on which tests the xdist worker had run first, and for n=3 the peak is narrow enough that one linux-aarch64 run got 4 points above the cutoff and failed 'assert checked >= 5'. A generator of its own, and draws until 8 points are above the cutoff, make it independent of order. Co-Authored-By: Claude Sonnet 5.5 --- test/test_analyticalSolutions_verify.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/test/test_analyticalSolutions_verify.py b/test/test_analyticalSolutions_verify.py index f583551bbd..fca59f8d36 100644 --- a/test/test_analyticalSolutions_verify.py +++ b/test/test_analyticalSolutions_verify.py @@ -77,9 +77,14 @@ def test_poissons_equation_source_is_the_laplacian(nd): def test_dirac_ic_satisfies_the_advection_diffusion_equation(n, b): a = 0.02 s = AS.LinearAD_DiracIC(n=n, b=np.array(b), a=a, tStart=0.3, u0=0.7, x0=np.array([0.1, 0.0, 0.0])) + # its own generator: the module's is shared, so where it stands depends on which tests this xdist worker + # ran before, and with 40 draws a 3-D peak this narrow left only 4 above the cutoff on one CI run + local = np.random.default_rng(1000 + n) checked = 0 - for _ in range(40): - x = np.zeros(3); x[:n] = rng.uniform(-0.2, 1.0, n); T = rng.uniform(0.1, 0.4) + for _ in range(4000): + if checked == 8: + break + x = np.zeros(3); x[:n] = local.uniform(-0.2, 1.0, n); T = local.uniform(0.1, 0.4) u = s.uOfXT(x, T) if u < 1e-3: continue