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_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 diff --git a/test/test_thread_defaults.py b/test/test_thread_defaults.py new file mode 100644 index 0000000000..4c94f9dc0a --- /dev/null +++ b/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"