Improve CPU count detection in cgroup environments and fix CI (#6098)

Co-authored-by: Jacob Walls <jacobtylerwalls@gmail.com>

7 files changed, 69 insertions, 5 deletions

diff --git a/pylint/lint/run.py b/pylint/lint/run.py
index 5ee387ead..0a90272bf 100644
--- a/pylint/lint/run.py
+++ b/pylint/lint/run.py
@@ -7,6 +7,7 @@ from __future__ import annotations
 import os
 import sys
 import warnings
+from pathlib import Path
 
 from pylint import config
 from pylint.config.callback_actions import (
@@ -36,15 +37,51 @@ except ImportError:
     multiprocessing = None  # type: ignore[assignment]
 
 
+def _query_cpu() -> int | None:
+    """Try to determine number of CPUs allotted in a docker container.
+
+    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/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
+            cpu_quota = int(file.read().rstrip())
+
+    if (
+        cpu_quota
+        and cpu_quota != -1
+        and Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us").is_file()
+    ):
+        with open("/sys/fs/cgroup/cpu/cpu.cfs_period_us", encoding="utf-8") as file:
+            cpu_period = int(file.read().rstrip())
+        # Divide quota by period and you should get num of allotted CPU to the container, rounded down if fractional.
+        avail_cpu = int(cpu_quota / cpu_period)
+    elif Path("/sys/fs/cgroup/cpu/cpu.shares").is_file():
+        with open("/sys/fs/cgroup/cpu/cpu.shares", encoding="utf-8") as file:
+            cpu_shares = int(file.read().rstrip())
+        # For AWS, gives correct value * 1024.
+        avail_cpu = int(cpu_shares / 1024)
+    return avail_cpu
+
+
 def _cpu_count() -> int:
     """Use sched_affinity if available for virtualized or containerized environments."""
+    cpu_share = _query_cpu()
+    cpu_count = None
     sched_getaffinity = getattr(os, "sched_getaffinity", None)
     # pylint: disable=not-callable,using-constant-test,useless-suppression
     if sched_getaffinity:
-        return len(sched_getaffinity(0))
-    if multiprocessing:
-        return multiprocessing.cpu_count()
-    return 1
+        cpu_count = len(sched_getaffinity(0))
+    elif multiprocessing:
+        cpu_count = multiprocessing.cpu_count()
+    else:
+        cpu_count = 1
+    if cpu_share is not None:
+        return min(cpu_share, cpu_count)
+    return cpu_count
 
 
 UNUSED_PARAM_SENTINEL = object()
diff --git a/setup.cfg b/setup.cfg
index 94a5f197f..49da28993 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -86,6 +86,7 @@ markers =
     primer_external_batch_two: Checks for crashes and errors when running pylint on external libs (batch two)
     benchmark: Baseline of pylint performance, if this regress something serious happened
     timeout: Marks from pytest-timeout.
+    needs_two_cores: Checks that need 2 or more cores to be meaningful
 
 [isort]
 profile = black
diff --git a/tests/benchmark/test_baseline_benchmarks.py b/tests/benchmark/test_baseline_benchmarks.py
index f9e997ef3..b020711b4 100644
--- a/tests/benchmark/test_baseline_benchmarks.py
+++ b/tests/benchmark/test_baseline_benchmarks.py
@@ -136,6 +136,7 @@ class TestEstablishBaselineBenchmarks:
             linter.msg_status == 0
         ), f"Expected no errors to be thrown: {pprint.pformat(linter.reporter.messages)}"
 
+    @pytest.mark.needs_two_cores
     def test_baseline_benchmark_j2(self, benchmark):
         """Establish a baseline of pylint performance with no work across threads.
 
@@ -158,6 +159,7 @@ class TestEstablishBaselineBenchmarks:
             linter.msg_status == 0
         ), f"Expected no errors to be thrown: {pprint.pformat(linter.reporter.messages)}"
 
+    @pytest.mark.needs_two_cores
     def test_baseline_benchmark_check_parallel_j2(self, benchmark):
         """Should demonstrate times very close to `test_baseline_benchmark_j2`."""
         linter = PyLinter(reporter=Reporter())
@@ -190,6 +192,7 @@ class TestEstablishBaselineBenchmarks:
             linter.msg_status == 0
         ), f"Expected no errors to be thrown: {pprint.pformat(linter.reporter.messages)}"
 
+    @pytest.mark.needs_two_cores
     def test_baseline_lots_of_files_j2(self, benchmark):
         """Establish a baseline with only 'master' checker being run in -j2.
 
@@ -230,6 +233,7 @@ class TestEstablishBaselineBenchmarks:
             linter.msg_status == 0
         ), f"Expected no errors to be thrown: {pprint.pformat(linter.reporter.messages)}"
 
+    @pytest.mark.needs_two_cores
     def test_baseline_lots_of_files_j2_empty_checker(self, benchmark):
         """Baselines pylint for a single extra checker being run in -j2, for N-files.
 
@@ -276,6 +280,7 @@ class TestEstablishBaselineBenchmarks:
             linter.msg_status == 0
         ), f"Expected no errors to be thrown: {pprint.pformat(linter.reporter.messages)}"
 
+    @pytest.mark.needs_two_cores
     def test_baseline_benchmark_j2_single_working_checker(self, benchmark):
         """Establishes baseline of multi-worker performance for PyLinter/check_parallel.
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 71d37d9c5..829b9bacc 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,6 +3,9 @@
 # Copyright (c) https://github.com/PyCQA/pylint/blob/main/CONTRIBUTORS.txt
 
 # pylint: disable=redefined-outer-name
+
+from __future__ import annotations
+
 import os
 from pathlib import Path
 
@@ -10,6 +13,7 @@ import pytest
 
 from pylint import checkers
 from pylint.lint import PyLinter
+from pylint.lint.run import _cpu_count
 from pylint.testutils import MinimalTestReporter
 
 
@@ -88,7 +92,9 @@ def pytest_addoption(parser) -> None:
     )
 
 
-def pytest_collection_modifyitems(config, items) -> None:
+def pytest_collection_modifyitems(
+    config: pytest.Config, items: list[pytest.Function]
+) -> None:
     """Convert command line options to markers."""
     # Add skip_primer_external mark
     if not config.getoption("--primer-external"):
@@ -109,3 +115,12 @@ def pytest_collection_modifyitems(config, items) -> None:
         for item in items:
             if "primer_stdlib" in item.keywords:
                 item.add_marker(skip_primer_stdlib)
+
+    # Add skip_cpu_cores mark
+    if _cpu_count() < 2:
+        skip_cpu_cores = pytest.mark.skip(
+            reason="Need 2 or more cores for test to be meaningful"
+        )
+        for item in items:
+            if "needs_two_cores" in item.keywords:
+                item.add_marker(skip_cpu_cores)
diff --git a/tests/lint/unittest_lint.py b/tests/lint/unittest_lint.py
index 82ed9d3d5..97a9220cd 100644
--- a/tests/lint/unittest_lint.py
+++ b/tests/lint/unittest_lint.py
@@ -723,6 +723,7 @@ class _CustomPyLinter(PyLinter):
         )
 
 
+@pytest.mark.needs_two_cores
 def test_custom_should_analyze_file() -> None:
     """Check that we can write custom should_analyze_file that work
     even for arguments.
@@ -751,6 +752,7 @@ def test_custom_should_analyze_file() -> None:
 
 # we do the check with jobs=1 as well, so that we are sure that the duplicates
 # are created by the multiprocessing problem.
+@pytest.mark.needs_two_cores
 @pytest.mark.parametrize("jobs", [1, 2])
 def test_multiprocessing(jobs: int) -> None:
     """Check that multiprocessing does not create duplicates."""
diff --git a/tests/test_check_parallel.py b/tests/test_check_parallel.py
index 839f75c71..b642eadf7 100644
--- a/tests/test_check_parallel.py
+++ b/tests/test_check_parallel.py
@@ -176,6 +176,7 @@ class TestCheckParallelFramework:
         worker_initialize(linter=dill.dumps(linter))
         assert isinstance(pylint.lint.parallel._worker_linter, type(linter))
 
+    @pytest.mark.needs_two_cores
     def test_worker_initialize_pickling(self) -> None:
         """Test that we can pickle objects that standard pickling in multiprocessing can't.
 
@@ -394,6 +395,7 @@ class TestCheckParallel:
         assert linter.stats.warning == 0
         assert linter.msg_status == 0, "We expect a single-file check to exit cleanly"
 
+    @pytest.mark.needs_two_cores
     @pytest.mark.parametrize(
         "num_files,num_jobs,num_checkers",
         [
@@ -491,6 +493,7 @@ class TestCheckParallel:
             expected_stats
         ), "The lint is returning unexpected results, has something changed?"
 
+    @pytest.mark.needs_two_cores
     @pytest.mark.parametrize(
         "num_files,num_jobs,num_checkers",
         [
diff --git a/tests/test_self.py b/tests/test_self.py
index f7c1ec0e7..18dcf7d6b 100644
--- a/tests/test_self.py
+++ b/tests/test_self.py
@@ -1043,6 +1043,7 @@ class TestRunTC:
             stderr=subprocess.PIPE,
         )
 
+    @pytest.mark.needs_two_cores
     def test_jobs_score(self) -> None:
         path = join(HERE, "regrtest_data", "unused_variable.py")
         expected = "Your code has been rated at 7.50/10"