One thread per MPI rank by default (OMP/OpenBLAS/MKL/Accelerate), threads untouched without MPI

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

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

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

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

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

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

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.")

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

test/test_thread_defaults.py

@@ -0,0 +1,174 @@
+"""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 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_COMM_WORLD")}
+    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"