expand unit tests of util.misc.available_cpu_count()

expand unit tests of util.misc.available_cpu_count() to include tests where cgroup v2 is used, where cgroup v1 is used, where limits are imposed (or not) on either, or where multiprocessing.cpu_count() is used as the fallback in the event a hex bitmask cannot be found in the usual fallback of /proc/self/status

docker/calc_mem.py

@@ -4,6 +4,7 @@
 Print result to stdout.
 """
 
+import math
 import argparse
 import sys
 import os
@@ -14,6 +15,8 @@
 
 import psutil
 
+#from util.misc import available_cpu_count # use the version of available_cpu_count() from viral-core/util/misc.py
+
 log = logging.getLogger(__name__)
 
 parser = argparse.ArgumentParser('Calculated memory allocated to the process')
@@ -48,9 +51,11 @@ def _load(path, encoding="utf-8"):
 
         # cgroup CPU count determination (w/ v2) adapted from:
         #   https://github.com/conan-io/conan/blob/2.9.2/conan/tools/build/cpu.py#L31-L54
+        #
+        # see also:
+        #   https://docs.kernel.org/scheduler/sched-bwc.html
+
         # This is necessary to determine docker cpu_count
-        cfs_quota_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"))
-        cfs_period_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_period_us"))
         cfs_quota_us = cfs_period_us = 0
         # cgroup v2
         if os.path.exists("/sys/fs/cgroup/cgroup.controllers"):

test/unit/test_util_misc.py

@@ -5,6 +5,7 @@
 import os, random, collections
 import unittest
 import subprocess
+import multiprocessing
 import util.misc
 import util.file
 import pytest
@@ -284,15 +285,76 @@ def test_chk():
 
 def test_available_cpu_count(monkeypatch_function_result):
     reported_cpu_count = util.misc.available_cpu_count()
-
     assert reported_cpu_count >= int(os.environ.get('PYTEST_XDIST_WORKER_COUNT', '1'))
+    assert util.misc.available_cpu_count() == reported_cpu_count
 
-    with monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='1'), \
-         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'):
+    # cgroup v2 limited to 1 cpu
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=True), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu.max', patch_result="100000 100000"):
         assert util.misc.available_cpu_count() == 1
 
-    assert util.misc.available_cpu_count() == reported_cpu_count
+    # cgroup v2 limited to 2 cpu
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=True), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu.max', patch_result="200000 100000"):
+        assert util.misc.available_cpu_count() == 2
+
+    # cgroup v2 with no CPU limit imposed on cgroup
+    # (fall back to /proc/self/status method, with limit imposed there):
+    #   'Cpus_allowed:  d' = 0b1101 bitmask (meaning execution allowed on 3 CPUs)
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=True), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu.max', patch_result="max 100000"), \
+         monkeypatch_function_result(util.file.slurp_file, '/proc/self/status', patch_result='Cpus_allowed:  d'):
+        assert util.misc.available_cpu_count() == 3
+
+    # cgroup v1 limited to 2 CPUs
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='200000'), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='100000'):
+
+        assert util.misc.available_cpu_count() == 2
+
+    # cgroup v1 limited to 1 CPU
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'):
+
+        assert util.misc.available_cpu_count() == 1
 
-    with monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='-1'), \
+    # cgroup v1 with no limit imposed on the cgroup
+    # (fall back to /proc/self/status method, with limit imposed there):
+    #   'Cpus_allowed:  c' = 0b1100 bitmask (meaning execution allowed on 2 CPUs)
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='-1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/proc/self/status', patch_result='Cpus_allowed:  c'):
+
+        assert util.misc.available_cpu_count() == 2
+
+    # cgroup v1 with no limit imposed on the cgoup or via /proc/self/status
+    # (fall back to /proc/self/status method, with no limit imposed there)
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='-1'), \
          monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'):
+
+        assert util.misc.available_cpu_count() == reported_cpu_count
+
+    # cgroup v1 with no limit imposed on the cgoup
+    # with 'Cpus_allowed' not present in /proc/self/status
+    # (fall back to multiprocessing.cpu_count() method)
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='-1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/proc/self/status', patch_result='unexpected_key:  1'):
+
         assert util.misc.available_cpu_count() == reported_cpu_count
+
+    # cgroup v1 with no limit imposed on the cgoup
+    # with 'Cpus_allowed' not present in /proc/self/status
+    # (fall back to multiprocessing.cpu_count() method with CPU count of 2 reported)
+    with monkeypatch_function_result(os.path.exists, "/sys/fs/cgroup/cgroup.controllers", patch_result=False, patch_module=os.path), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_quota_us', patch_result='-1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/sys/fs/cgroup/cpu/cpu.cfs_period_us', patch_result='1'), \
+         monkeypatch_function_result(util.file.slurp_file, '/proc/self/status', patch_result='unexpected_key:  1'), \
+         monkeypatch_function_result(multiprocessing.cpu_count, patch_result=2, patch_module=multiprocessing):
+
+        assert util.misc.available_cpu_count() == 2

util/misc.py

@@ -1,4 +1,5 @@
 '''A few miscellaneous tools. '''
+import math
 import collections
 import contextlib
 import itertools, functools, operator
@@ -334,30 +335,28 @@ def available_cpu_count():
 
     cgroup_cpus = MAX_INT32
     try:
-        def _load(path, encoding="utf-8"):
-            """ Loads a file content """
-            with open(path, 'r', encoding=encoding, newline="") as handle:
-                tmp = handle.read()
-            return tmp
-
         # cgroup CPU count determination (w/ v2) adapted from:
         #   https://github.com/conan-io/conan/blob/2.9.2/conan/tools/build/cpu.py#L31-L54
+        #
+        # see also:
+        #   https://docs.kernel.org/scheduler/sched-bwc.html
+
         # This is necessary to determine docker cpu_count
-        cfs_quota_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"))
-        cfs_period_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_period_us"))
         cfs_quota_us = cfs_period_us = 0
         # cgroup v2
         if os.path.exists("/sys/fs/cgroup/cgroup.controllers"):
-            cpu_max = _load("/sys/fs/cgroup/cpu.max").split()
+            log.debug("cgroup v2 detected")
+            cpu_max = util.file.slurp_file("/sys/fs/cgroup/cpu.max").split()
             if cpu_max[0] != "max":
                 if len(cpu_max) == 1:
                     cfs_quota_us, cfs_period_us = int(cpu_max[0]), 100_000
                 else:
                     cfs_quota_us, cfs_period_us = map(int, cpu_max)
         # cgroup v1
         else:
-            cfs_quota_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"))
-            cfs_period_us = int(_load("/sys/fs/cgroup/cpu/cpu.cfs_period_us"))
+            log.debug("cgroup v1 detected")
+            cfs_quota_us = int(util.file.slurp_file("/sys/fs/cgroup/cpu/cpu.cfs_quota_us"))
+            cfs_period_us = int(util.file.slurp_file("/sys/fs/cgroup/cpu/cpu.cfs_period_us"))
 
         log.debug('cfs_quota_us %s, cfs_period_us %s', cfs_quota_us, cfs_period_us)
         if cfs_quota_us > 0 and cfs_period_us > 0:
@@ -367,7 +366,7 @@ def _load(path, encoding="utf-8"):
 
     proc_cpus = MAX_INT32
     try:
-        m = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', _load('/proc/self/status'))
+        m = re.search(r'(?m)^Cpus_allowed:\s*(.*)$', util.file.slurp_file('/proc/self/status'))
         if m:
             res = bin(int(m.group(1).replace(',', ''), 16)).count('1')
             if res > 0: