Calculate linter.config.jobs in cgroupsv2 environments (#10089)

Co-authored-by: Pierre Sassoulas <[email protected]>
(cherry picked from commit 6456374)

custom_dict.txt

@@ -68,6 +68,7 @@ contextlib
 contextmanager
 contravariance
 contravariant
+cgroup
 CPython
 cpython
 csv

doc/whatsnew/fragments/10103.bugfix

@@ -0,0 +1,3 @@
+Fixes a crash that occurred when pylint was run in a container on a host with cgroupsv2 and restrictions on CPU usage.
+
+Closes #10103

pylint/lint/run.py

@@ -43,8 +43,30 @@ def _query_cpu() -> int | None:
     This is based on discussion and copied from suggestions in
     https://bugs.python.org/issue36054.
     """
-    cpu_quota, avail_cpu = None, None
+    if Path("/sys/fs/cgroup/cpu.max").is_file():
+        avail_cpu = _query_cpu_cgroupv2()
+    else:
+        avail_cpu = _query_cpu_cgroupsv1()
+    return _query_cpu_handle_k8s_pods(avail_cpu)
+
+
+def _query_cpu_cgroupv2() -> int | None:
+    avail_cpu = None
+    with open("/sys/fs/cgroup/cpu.max", encoding="utf-8") as file:
+        line = file.read().rstrip()
+        fields = line.split()
+        if len(fields) == 2:
+            str_cpu_quota = fields[0]
+            cpu_period = int(fields[1])
+            # Make sure this is not in an unconstrained cgroup
+            if str_cpu_quota != "max":
+                cpu_quota = int(str_cpu_quota)
+                avail_cpu = int(cpu_quota / cpu_period)
+    return avail_cpu
+
 
+def _query_cpu_cgroupsv1() -> int | None:
+    cpu_quota, avail_cpu = None, None
     if Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").is_file():
         with open("/sys/fs/cgroup/cpu/cpu.cfs_quota_us", encoding="utf-8") as file:
             # Not useful for AWS Batch based jobs as result is -1, but works on local linux systems
@@ -65,7 +87,10 @@ def _query_cpu() -> int | None:
             cpu_shares = int(file.read().rstrip())
         # For AWS, gives correct value * 1024.
         avail_cpu = int(cpu_shares / 1024)
+    return avail_cpu
+
 
+def _query_cpu_handle_k8s_pods(avail_cpu: int | None) -> int | None:
     # In K8s Pods also a fraction of a single core could be available
     # As multiprocessing is not able to run only a "fraction" of process
     # assume we have 1 CPU available

tests/test_pylint_runners.py

@@ -18,6 +18,7 @@
 import pytest
 
 from pylint import run_pylint, run_pyreverse, run_symilar
+from pylint.lint.run import _query_cpu
 from pylint.testutils import GenericTestReporter as Reporter
 from pylint.testutils._run import _Run as Run
 from pylint.testutils.utils import _test_cwd
@@ -70,33 +71,157 @@ def test_pylint_argument_deduplication(
     assert err.value.code == 0
 
 
-def test_pylint_run_jobs_equal_zero_dont_crash_with_cpu_fraction(
+@pytest.mark.parametrize(
+    "quota,shares,period",
+    [
+        # Shares path
+        ("-1", "2", ""),
+        ("-1", "1023", ""),
+        ("-1", "1024", ""),
+        # Periods path
+        ("100", "", "200"),
+        ("999", "", "1000"),
+        ("1000", "", "1000"),
+    ],
+)
+def test_pylint_run_dont_crash_with_cgroupv1(
     tmp_path: pathlib.Path,
+    quota: str,
+    shares: str,
+    period: str,
 ) -> None:
     """Check that the pylint runner does not crash if `pylint.lint.run._query_cpu`
     determines only a fraction of a CPU core to be available.
     """
-    builtin_open = open
+    filepath = os.path.abspath(__file__)
+    testargs = [filepath, "--jobs=0"]
 
-    def _mock_open(*args: Any, **kwargs: Any) -> BufferedReader:
-        if args[0] == "/sys/fs/cgroup/cpu/cpu.cfs_quota_us":
-            return mock_open(read_data=b"-1")(*args, **kwargs)  # type: ignore[no-any-return]
-        if args[0] == "/sys/fs/cgroup/cpu/cpu.shares":
-            return mock_open(read_data=b"2")(*args, **kwargs)  # type: ignore[no-any-return]
-        return builtin_open(*args, **kwargs)  # type: ignore[no-any-return]
+    with _test_cwd(tmp_path):
+        with pytest.raises(SystemExit) as err:
+            with patch(
+                "builtins.open",
+                mock_cgroup_fs(quota=quota, shares=shares, period=period),
+            ):
+                with patch("pylint.lint.run.Path", mock_cgroup_path(v2=False)):
+                    Run(testargs, reporter=Reporter())
+        assert err.value.code == 0
 
-    pathlib_path = pathlib.Path
-
-    def _mock_path(*args: str, **kwargs: Any) -> pathlib.Path:
-        if args[0] == "/sys/fs/cgroup/cpu/cpu.shares":
-            return MagicMock(is_file=lambda: True)
-        return pathlib_path(*args, **kwargs)
 
+@pytest.mark.parametrize(
+    "contents",
+    [
+        "1 2",
+        "max 100000",
+    ],
+)
+def test_pylint_run_dont_crash_with_cgroupv2(
+    tmp_path: pathlib.Path,
+    contents: str,
+) -> None:
+    """Check that the pylint runner does not crash if `pylint.lint.run._query_cpu`
+    determines only a fraction of a CPU core to be available.
+    """
     filepath = os.path.abspath(__file__)
     testargs = [filepath, "--jobs=0"]
+
     with _test_cwd(tmp_path):
         with pytest.raises(SystemExit) as err:
-            with patch("builtins.open", _mock_open):
-                with patch("pylint.lint.run.Path", _mock_path):
+            with patch("builtins.open", mock_cgroup_fs(max_v2=contents)):
+                with patch("pylint.lint.run.Path", mock_cgroup_path(v2=True)):
                     Run(testargs, reporter=Reporter())
         assert err.value.code == 0
+
+
+@pytest.mark.parametrize(
+    "contents,expected",
+    [
+        ("50000 100000", 1),
+        ("100000 100000", 1),
+        ("200000 100000", 2),
+        ("299999 100000", 2),
+        ("300000 100000", 3),
+        # Unconstrained cgroup
+        ("max 100000", None),
+    ],
+)
+def test_query_cpu_cgroupv2(
+    tmp_path: pathlib.Path,
+    contents: str,
+    expected: int,
+) -> None:
+    """Check that `pylint.lint.run._query_cpu` generates realistic values in cgroupsv2
+    systems.
+    """
+    with _test_cwd(tmp_path):
+        with patch("builtins.open", mock_cgroup_fs(max_v2=contents)):
+            with patch("pylint.lint.run.Path", mock_cgroup_path(v2=True)):
+                cpus = _query_cpu()
+                assert cpus == expected
+
+
+@pytest.mark.parametrize(
+    "quota,shares,period,expected",
+    [
+        # Shares path
+        ("-1", "2", "", 1),
+        ("-1", "1023", "", 1),
+        ("-1", "1024", "", 1),
+        ("-1", "2048", "", 2),
+        # Periods path
+        ("100", "", "200", 1),
+        ("999", "", "1000", 1),
+        ("1000", "", "1000", 1),
+        ("2000", "", "1000", 2),
+    ],
+)
+def test_query_cpu_cgroupv1(
+    tmp_path: pathlib.Path,
+    quota: str,
+    shares: str,
+    period: str,
+    expected: int,
+) -> None:
+    """Check that `pylint.lint.run._query_cpu` generates realistic values in cgroupsv1
+    systems.
+    """
+    with _test_cwd(tmp_path):
+        with patch(
+            "builtins.open", mock_cgroup_fs(quota=quota, shares=shares, period=period)
+        ):
+            with patch("pylint.lint.run.Path", mock_cgroup_path(v2=False)):
+                cpus = _query_cpu()
+                assert cpus == expected
+
+
+def mock_cgroup_path(v2: bool) -> Any:
+    def _mock_path(*args: str, **kwargs: Any) -> pathlib.Path:
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.cfs_period_us":
+            return MagicMock(is_file=lambda: not v2)
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.shares":
+            return MagicMock(is_file=lambda: not v2)
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.cfs_quota_us":
+            return MagicMock(is_file=lambda: not v2)
+        if args[0] == "/sys/fs/cgroup/cpu.max":
+            return MagicMock(is_file=lambda: v2)
+        return pathlib.Path(*args, **kwargs)
+
+    return _mock_path
+
+
+def mock_cgroup_fs(
+    quota: str = "", shares: str = "", period: str = "", max_v2: str = ""
+) -> Any:
+    builtin_open = open
+
+    def _mock_open(*args: Any, **kwargs: Any) -> BufferedReader:
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.cfs_quota_us":
+            return mock_open(read_data=quota)(*args, **kwargs)  # type: ignore[no-any-return]
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.shares":
+            return mock_open(read_data=shares)(*args, **kwargs)  # type: ignore[no-any-return]
+        if args[0] == "/sys/fs/cgroup/cpu/cpu.cfs_period_us":
+            return mock_open(read_data=period)(*args, **kwargs)  # type: ignore[no-any-return]
+        if args[0] == "/sys/fs/cgroup/cpu.max":
+            return mock_open(read_data=max_v2)(*args, **kwargs)  # type: ignore[no-any-return]
+        return builtin_open(*args, **kwargs)  # type: ignore[no-any-return]
+
+    return _mock_open