Merge pull request #1180 from projectsyn/fix/respect-sched-affinity

Fix overcommitted thread counts on Linux

Author: simu

commodore/cli/__init__.py

@@ -12,6 +12,7 @@
 
 from commodore import __git_version__, __version__, tools
 from commodore.config import Config
+from commodore.helpers import cpu_count
 from commodore.version import version_info
 
 import commodore.cli.options as options
@@ -85,7 +86,10 @@ def commodore(ctx, working_dir, verbose, request_timeout):
 
 def main():
     multiprocessing.set_start_method("spawn")
-    Reclass.set_thread_count(0)
+    # `Reclass.set_thread_count()` wants 0 if it should use its internal detection logic, so we call
+    # the `cpu_count()` helper with `fallback=0` to return 0 if the number of CPUs can't be
+    # determined.
+    Reclass.set_thread_count(cpu_count(fallback=0))
     tools.setup_path()
 
     load_dotenv(dotenv_path=find_dotenv(usecwd=True))

commodore/helpers.py

@@ -271,7 +271,7 @@ def kapitan_compile(
     cached.args.verbose = config.trace
     cached.args.output_path = output_dir
     cached.args.targets = targets
-    cached.args.parallelism = None
+    cached.args.parallelism = cpu_count()
     cached.args.labels = None
     cached.args.prune = False
     cached.args.indent = 2
@@ -385,3 +385,15 @@ def python3_executable() -> str:
         fg="yellow",
     )
     return path_python
+
+
+def cpu_count(fallback: Optional[int] = None) -> Optional[int]:
+    """Return the number of available CPU cores or `fallback` if CPU count can't be determined.
+
+    On Linux this respects the process's scheduling affinity.
+
+    NOTE(sg): This can be replaced by `os.process_cpu_count() on Python >= 3.13.
+    """
+    if hasattr(os, "sched_getaffinity"):
+        return len(os.sched_getaffinity(0))
+    return os.cpu_count() or fallback

tests/test_helpers.py

@@ -503,3 +503,42 @@ def test_python3_executable_err(mock_shutil, mock_sys, capsys):
         "[ERROR] No python3 in $PATH, components which use `${_python3}` will fail."
         in out
     )
+
+
+class MockOS:
+    _cpu_count: Optional[int]
+    _sched_affinity: list[int]
+
+    def __init__(self):
+        self._cpu_count = None
+        self._sched_affinitiy = []
+
+    def cpu_count(self) -> Optional[int]:
+        return self._cpu_count
+
+
+class MockOSAff(MockOS):
+    def sched_getaffinity(self, pid: int) -> list[int]:
+        assert pid == 0
+        return self._sched_affinity
+
+
+@pytest.mark.parametrize(
+    "cpu_count,fallback", [(None, None), (None, 0), (4, None), (4, 0)]
+)
+@patch.object(helpers, "os", new_callable=MockOS)
+def test_cpu_count(mock_os, cpu_count, fallback):
+    mock_os._cpu_count = cpu_count
+    r = helpers.cpu_count(fallback=fallback)
+    if cpu_count is None:
+        assert r == fallback
+    else:
+        assert r == cpu_count
+
+
+@pytest.mark.parametrize("sched_affinity", [[0, 3], [0, 1, 2, 3, 4, 5, 6]])
+@patch.object(helpers, "os", new_callable=MockOSAff)
+def test_cpu_count_affinity(mock_os, sched_affinity):
+    mock_os._cpu_count = 4
+    mock_os._sched_affinity = sched_affinity
+    assert helpers.cpu_count() == len(sched_affinity)