Respect CPU affinity and cgroup CPU quota in the default thread count

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
 -----------------------------

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: "<quota> <period>" 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:

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: