diff --git a/loky/backend/context.py b/loky/backend/context.py index 62795d65..d8cc59ba 100644 --- a/loky/backend/context.py +++ b/loky/backend/context.py @@ -145,22 +145,34 @@ def _cpu_count_cgroup(os_cpu_count): cpu_max_fname = "/sys/fs/cgroup/cpu.max" cfs_quota_fname = "/sys/fs/cgroup/cpu/cpu.cfs_quota_us" cfs_period_fname = "/sys/fs/cgroup/cpu/cpu.cfs_period_us" + + cpu_quota_us = None + cpu_period_us = None + if os.path.exists(cpu_max_fname): # cgroup v2 # https://www.kernel.org/doc/html/latest/admin-guide/cgroup-v2.html with open(cpu_max_fname) as fh: - cpu_quota_us, cpu_period_us = fh.read().strip().split() - elif os.path.exists(cfs_quota_fname) and os.path.exists(cfs_period_fname): - # cgroup v1 - # https://www.kernel.org/doc/html/latest/scheduler/sched-bwc.html#management - with open(cfs_quota_fname) as fh: - cpu_quota_us = fh.read().strip() - with open(cfs_period_fname) as fh: - cpu_period_us = fh.read().strip() - else: - # No Cgroup CPU bandwidth limit (e.g. non-Linux platform) - cpu_quota_us = "max" - cpu_period_us = 100_000 # unused, for consistency with default values + # Parse the quota and period values + parts = fh.read().strip().split() + if len(parts) == 2: + cpu_quota_us, cpu_period_us = parts + # If len(parts) != 2, leave as None and fall back to v1 + + # If we didn't get values from cgroup v2, try cgroup v1 + if cpu_quota_us is None or cpu_period_us is None: + if os.path.exists(cfs_quota_fname) and os.path.exists( + cfs_period_fname + ): + # cgroup v1 + # https://www.kernel.org/doc/html/latest/scheduler/sched-bwc.html#management + with open(cfs_quota_fname) as fh: + cpu_quota_us = fh.read().strip() + with open(cfs_period_fname) as fh: + cpu_period_us = fh.read().strip() + else: + # No Cgroup CPU bandwidth limit (e.g. non-Linux platform) + cpu_quota_us = "max" if cpu_quota_us == "max": # No active Cgroup quota on a Cgroup-capable platform @@ -172,7 +184,7 @@ def _cpu_count_cgroup(os_cpu_count): return math.ceil(cpu_quota_us / cpu_period_us) else: # pragma: no cover # Setting a negative cpu_quota_us value is a valid way to disable - # cgroup CPU bandwith limits + # cgroup CPU bandwidth limits return os_cpu_count diff --git a/tests/test_loky_module.py b/tests/test_loky_module.py index 3fd1b542..ef1194c3 100644 --- a/tests/test_loky_module.py +++ b/tests/test_loky_module.py @@ -6,6 +6,7 @@ import tempfile import warnings from subprocess import check_output +from unittest.mock import patch, mock_open import pytest @@ -142,7 +143,27 @@ def test_cpu_count_cgroup_limit(): os.path.join(loky_module_path, os.pardir) ) - # The following will always run using the Python 3.7 docker image. + # Check if Docker can actually set cgroup CPU limits in this environment + # by verifying that --cpus flag writes to cgroup files + cgroup_check = check_output( + f'{docker_bin} run --rm --cpus 0.5 python:3.10 python3 -c "' + "import os; " + "v2 = '/sys/fs/cgroup/cpu.max'; " + "v1_quota = '/sys/fs/cgroup/cpu/cpu.cfs_quota_us'; " + "v2_content = open(v2).read().strip() if os.path.exists(v2) else ''; " + "v1_content = open(v1_quota).read().strip() if os.path.exists(v1_quota) else ''; " + "print('ok' if (v2_content and v2_content != 'max') or (v1_content and v1_content != '-1') else 'skip')" + '"', + shell=True, + text=True, + ).strip() + + if cgroup_check != "ok": + pytest.skip( + "Docker doesn't properly set cgroup CPU limits in this environment" + ) + + # The following will always run using the Python 3.10 docker image. # We mount the loky source as /loky inside the container, # so it can be imported when running commands under / @@ -238,3 +259,35 @@ def test_only_physical_cores_with_user_limitation(): if cpu_count_user < cpu_count_mp: assert cpu_count() == cpu_count_user assert cpu_count(only_physical_cores=True) == cpu_count_user + + +@pytest.mark.parametrize( + "read_data,description", + [ + ("", "empty file"), + ("max\n", "max value"), + ], +) +def test_cpu_count_cgroup_invalid_content(read_data, description): + # Test that invalid cgroup cpu.max file content is handled gracefully + # and doesn't cause a ValueError when trying to unpack values + if sys.platform != "linux": + pytest.skip() + + from loky.backend.context import _cpu_count_cgroup + + os_cpu_count = mp.cpu_count() + + # Mock the file with the provided read_data + with patch("builtins.open", mock_open(read_data=read_data)): + with patch("os.path.exists") as mock_exists: + # cpu.max exists, but other files don't + mock_exists.side_effect = ( + lambda path: path == "/sys/fs/cgroup/cpu.max" + ) + + # This should not raise ValueError and return os_cpu_count + result = _cpu_count_cgroup(os_cpu_count) + assert ( + result == os_cpu_count + ), f"cpu.max with {description} should return os_cpu_count"