diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index ee49c2a..4fa8b6d 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -14,6 +14,15 @@ 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). +* 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 5fdeebb..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 @@ -1403,6 +1404,55 @@ 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_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 9e45fbe..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: @@ -171,7 +244,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: