From adb4f9e56ed08ef72cc660b6d66f270c9ff66f24 Mon Sep 17 00:00:00 2001 From: Manu Nicholas Jacob Date: Fri, 30 Oct 2026 16:45:05 -0500 Subject: [PATCH 1/2] Respect CPU affinity in detect_number_of_cores() detect_number_of_cores() returned SC_NPROCESSORS_ONLN, the machine's online core count, ignoring any CPU affinity mask the process runs under. A process pinned to a subset of cores by taskset, by `docker run --cpuset-cpus`, or by the cpuset a SLURM or Kubernetes job is placed in was still told the full core count, so the default thread pool oversubscribed the cores it was actually allowed to use. This is the common failure mode when numexpr is pulled in as an indirect dependency and nobody has set NUMEXPR_MAX_THREADS (gh-328): pinned to 2 of 20 cores the default pool still started 16 threads. Prefer len(os.sched_getaffinity(0)) where the OS exposes it (Linux), falling back to the existing sysconf / sysctl / NUMBER_OF_PROCESSORS probes everywhere else. The returned count is only ever reduced, and only when the process is genuinely restricted; the unconstrained default and every explicitly configured path (NUMEXPR_MAX_THREADS, NUMEXPR_NUM_THREADS, OMP_NUM_THREADS) are unchanged. Verified on x86-64 (20 cores, pinned to 2 -> 2 threads) and on aarch64 (Raspberry Pi 5, 4 cores, pinned to 2 -> 2 threads). Adds a regression test that pins the process and asserts the detected count matches the affinity mask. Signed-off-by: Manu Nicholas Jacob --- RELEASE_NOTES.rst | 5 +++++ numexpr/tests/test_numexpr.py | 17 +++++++++++++++++ numexpr/utils.py | 16 ++++++++++++++++ 3 files changed, 38 insertions(+) diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index ee49c2a..8cec4b5 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -14,6 +14,11 @@ Changes from 2.14.2 to 2.14.3 by the sanitizer raise ``ValueError``; and unknown functions raise ``TypeError``. Sanitization can still be explicitly disabled with ``sanitize=False`` or ``NUMEXPR_SANITIZE=0``. +* ``detect_number_of_cores()`` now honours a restricted CPU affinity mask + (``taskset``, ``docker run --cpuset-cpus``, or the cpuset a cluster manager + such as SLURM pins a job to) when choosing the default thread count, so the + default pool no longer oversubscribes the cores a pinned process may run on + (gh-328). Changes from 2.14.1 to 2.14.2 ----------------------------- diff --git a/numexpr/tests/test_numexpr.py b/numexpr/tests/test_numexpr.py index 5fdeebb..af2424e 100644 --- a/numexpr/tests/test_numexpr.py +++ b/numexpr/tests/test_numexpr.py @@ -1403,6 +1403,23 @@ def test_numexpr_max_threads_empty_string(self): else: self.assertEqual(min(detect_number_of_cores(), MAX_THREADS), numexpr._init_num_threads()) + def test_detect_cores_respects_affinity(self): + # A process pinned to a subset of the machine's cores (taskset, + # docker --cpuset-cpus, a SLURM cpuset) must not be told there are + # more cores than it may run on; otherwise the default pool + # oversubscribes them (gh-328). + if not (hasattr(os, 'sched_getaffinity') and hasattr(os, 'sched_setaffinity')): + self.skipTest('CPU affinity control not available on this platform') + original = os.sched_getaffinity(0) + if len(original) < 2: + self.skipTest('need at least two allowed CPUs to restrict the mask') + subset = {sorted(original)[0]} + try: + os.sched_setaffinity(0, subset) + self.assertEqual(detect_number_of_cores(), len(subset)) + finally: + os.sched_setaffinity(0, original) + def test_vml_threads_round_trip(self): n_threads = 3 if use_vml: diff --git a/numexpr/utils.py b/numexpr/utils.py index 9e45fbe..c8c544d 100644 --- a/numexpr/utils.py +++ b/numexpr/utils.py @@ -171,7 +171,23 @@ def _init_num_threads(): def detect_number_of_cores(): """ Detects the number of cores on a system. Cribbed from pp. + + When the process is restricted to a subset of the machine's cores -- by + ``taskset``, ``docker run --cpuset-cpus``, or the cpuset that cluster + managers such as SLURM pin a job to -- the CPU affinity mask reports fewer + cores than are physically online. Prefer that mask so the default thread + pool does not oversubscribe the cores the process may actually run on, + which is the common failure mode when numexpr is pulled in as an indirect + dependency (see gh-328). """ + # Respect a restricted CPU affinity mask where the OS exposes one (Linux). + if hasattr(os, "sched_getaffinity"): + try: + ncpus = len(os.sched_getaffinity(0)) + if ncpus > 0: + return ncpus + except OSError: + pass # Linux, Unix and MacOS: if hasattr(os, "sysconf"): if "SC_NPROCESSORS_ONLN" in os.sysconf_names: From d17942c496ff00fd1ec0faa90e44da8d3c347a08 Mon Sep 17 00:00:00 2001 From: Manu Nicholas Jacob Date: Fri, 30 Oct 2026 17:31:55 -0500 Subject: [PATCH 2/2] Honour a cgroup CPU quota in the default thread count The affinity fix covers cpuset-style restriction (taskset, --cpuset-cpus, a SLURM cpuset). A CPU *quota* set without a cpuset does not change the affinity mask, so it was still invisible: `docker run --cpus=2` or a Kubernetes CPU limit on a 20-core host left the default pool at 16 threads, oversubscribing the 2-CPU share eightfold. This is the most common containerised case. When numexpr is left to auto-detect (no NUMEXPR_MAX_THREADS / NUMEXPR_NUM_THREADS / OMP_NUM_THREADS), cap the default by the cgroup CPU quota in addition to the affinity mask. The quota is read from cgroup v2 `cpu.max` or cgroup v1 `cpu.cfs_quota_us` / `cpu.cfs_period_us` for the process cgroup and its ancestors, taking the tightest, rounded up. No quota, or a non-Linux platform, leaves the count unchanged. Verified end to end: inside a `cpu.max = 200000 100000` (2.0 CPUs) cgroup with the affinity mask left at all 20 cores, numexpr.nthreads drops from 16 to 2; with no quota it stays 16. Adds a root-free regression test that parses v2 and v1 quota fixtures. Signed-off-by: Manu Nicholas Jacob --- RELEASE_NOTES.rst | 4 ++ numexpr/tests/test_numexpr.py | 35 ++++++++++++++++- numexpr/utils.py | 73 +++++++++++++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 1 deletion(-) diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index 8cec4b5..4fa8b6d 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -19,6 +19,10 @@ Changes from 2.14.2 to 2.14.3 such as SLURM pins a job to) when choosing the default thread count, so the default pool no longer oversubscribes the cores a pinned process may run on (gh-328). +* The default thread count also honours a cgroup CPU quota (``docker run + --cpus``, a Kubernetes CPU limit, a SLURM cpu bandwidth limit) read from + cgroup v2 ``cpu.max`` or cgroup v1 ``cpu.cfs_quota_us``, rounded up, so a + quota-limited container does not oversubscribe its share (gh-328). Changes from 2.14.1 to 2.14.2 ----------------------------- diff --git a/numexpr/tests/test_numexpr.py b/numexpr/tests/test_numexpr.py index af2424e..09ac04d 100644 --- a/numexpr/tests/test_numexpr.py +++ b/numexpr/tests/test_numexpr.py @@ -14,6 +14,7 @@ import platform import subprocess import sys +import tempfile import unittest import warnings import weakref @@ -38,7 +39,7 @@ from numexpr import (E, NumExpr, disassemble, evaluate, re_evaluate, use_vml, validate) from numexpr.expressions import ConstantNode -from numexpr.utils import detect_number_of_cores +from numexpr.utils import _cgroup_cpu_limit, detect_number_of_cores try: import pytest @@ -1420,6 +1421,38 @@ def test_detect_cores_respects_affinity(self): finally: os.sched_setaffinity(0, original) + def test_cgroup_cpu_quota_parsing(self): + # A cgroup CPU quota (docker --cpus, a Kubernetes CPU limit) must be + # read so the default pool does not oversubscribe it (gh-328). Parsing + # is exercised against fixture files so the test needs no root and no + # real cgroup. + with tempfile.TemporaryDirectory() as rootdir: + proc = os.path.join(rootdir, 'proc_cgroup') + # cgroup v2: " " in a leaf cpu.max == 1.5 CPUs + with open(proc, 'w') as f: + f.write('0::/mygroup\n') + leaf = os.path.join(rootdir, 'mygroup') + os.makedirs(leaf) + with open(os.path.join(leaf, 'cpu.max'), 'w') as f: + f.write('150000 100000\n') + self.assertAlmostEqual(_cgroup_cpu_limit(proc, rootdir), 1.5) + # "max" means no quota -> None + with open(os.path.join(leaf, 'cpu.max'), 'w') as f: + f.write('max 100000\n') + self.assertIsNone(_cgroup_cpu_limit(proc, rootdir)) + with tempfile.TemporaryDirectory() as rootdir: + # cgroup v1: cpu.cfs_quota_us / cpu.cfs_period_us == 2.0 CPUs + proc = os.path.join(rootdir, 'proc_cgroup') + with open(proc, 'w') as f: + f.write('5:cpu,cpuacct:/mygroup\n') + leaf = os.path.join(rootdir, 'cpu,cpuacct', 'mygroup') + os.makedirs(leaf) + with open(os.path.join(leaf, 'cpu.cfs_quota_us'), 'w') as f: + f.write('200000\n') + with open(os.path.join(leaf, 'cpu.cfs_period_us'), 'w') as f: + f.write('100000\n') + self.assertAlmostEqual(_cgroup_cpu_limit(proc, rootdir), 2.0) + def test_vml_threads_round_trip(self): n_threads = 3 if use_vml: diff --git a/numexpr/utils.py b/numexpr/utils.py index c8c544d..dba2667 100644 --- a/numexpr/utils.py +++ b/numexpr/utils.py @@ -13,6 +13,7 @@ log = logging.getLogger(__name__) import contextvars +import math import os import subprocess @@ -118,6 +119,69 @@ def get_num_threads(): """ return _get_num_threads() + +def _read_cgroup_cpu_limit(cgroup_dir, v2): + """Return the CPU limit (in CPUs) set on one cgroup directory, or None.""" + try: + if v2: + with open(os.path.join(cgroup_dir, "cpu.max")) as f: + quota, period = f.read().split()[:2] + if quota == "max": + return None + else: + with open(os.path.join(cgroup_dir, "cpu.cfs_quota_us")) as f: + quota = f.read().strip() + with open(os.path.join(cgroup_dir, "cpu.cfs_period_us")) as f: + period = f.read().strip() + if int(quota) <= 0 or int(period) <= 0: + return None + return int(quota) / int(period) + except (OSError, ValueError): + return None + + +def _cgroup_cpu_limit(proc_cgroup="/proc/self/cgroup", cgroup_root="/sys/fs/cgroup"): + """Return the tightest CPU quota (in CPUs) on this process's cgroup, or None. + + Reads cgroup v2 ``cpu.max`` or cgroup v1 ``cpu.cfs_quota_us`` / + ``cpu.cfs_period_us`` for the process's own cgroup and its ancestors, for + example a ``docker run --cpus`` limit or a Kubernetes CPU limit. Returns + None when no quota is set, or on non-Linux platforms where the files are + absent. + """ + try: + with open(proc_cgroup) as f: + lines = f.read().splitlines() + except OSError: + return None + limits = [] + for line in lines: + parts = line.split(":", 2) + if len(parts) != 3: + continue + controllers, path = parts[1], parts[2] + if controllers == "": + mounts, v2 = [cgroup_root], True + elif "cpu" in controllers.split(","): + names = ("cpu", "cpu,cpuacct", "cpuacct,cpu") + mounts, v2 = [os.path.join(cgroup_root, n) for n in names], False + else: + continue + for mount in mounts: + # Walk from the process cgroup up to the mount root. Inside a + # container without a cgroup namespace the mount root is also the + # process cgroup, where ``path`` does not exist under the mount. + rel = path.strip("/") + while True: + limit = _read_cgroup_cpu_limit(os.path.join(mount, rel), v2) + if limit is not None: + limits.append(limit) + if not rel: + break + rel = os.path.dirname(rel) + return min(limits) if limits else None + + def _init_num_threads(): """ Detects the environment variable 'NUMEXPR_MAX_THREADS' to set the threadpool @@ -144,6 +208,15 @@ def _init_num_threads(): else: # The use has not set 'NUMEXPR_MAX_THREADS', so likely they have not # configured NumExpr as desired, so we emit info logs. + # A cgroup CPU quota (`docker run --cpus`, a Kubernetes CPU limit, a + # SLURM cpu bandwidth limit) caps how much CPU time the process gets + # without restricting which cores it runs on, so it is invisible to + # the affinity mask detect_number_of_cores() reads. Honour it too, + # rounding up, so the default pool does not oversubscribe the quota + # (gh-328). + cpu_quota = _cgroup_cpu_limit() + if cpu_quota is not None: + n_cores = max(min(n_cores, math.ceil(cpu_quota)), 1) if n_cores > MAX_THREADS: log.info('Note: detected %d virtual cores but NumExpr set to maximum of %d, check "NUMEXPR_MAX_THREADS" environment variable.'%(n_cores, MAX_THREADS)) if n_cores > 16: