diff --git a/docker/calc_mem.py b/docker/calc_mem.py index d963adbf..4fa72a09 100755 --- a/docker/calc_mem.py +++ b/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"): diff --git a/test/unit/test_util_misc.py b/test/unit/test_util_misc.py index aa9a3d98..acb4c87a 100644 --- a/test/unit/test_util_misc.py +++ b/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 \ No newline at end of file diff --git a/util/misc.py b/util/misc.py index 46aa3987..452ed989 100644 --- a/util/misc.py +++ b/util/misc.py @@ -1,4 +1,5 @@ '''A few miscellaneous tools. ''' +import math import collections import contextlib import itertools, functools, operator @@ -334,21 +335,18 @@ 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 @@ -356,8 +354,9 @@ def _load(path, encoding="utf-8"): 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: