FIX cpu_count on linux and windows (#425)

Author: jeremiedbb

.github/workflows/test.yml

@@ -52,6 +52,7 @@ jobs:
           - name: windows-py39
             os: windows-latest
             PYTHON_VERSION: "3.9"
+            LOKY_TEST_NO_CIM: "true"
 
           - name: macos-py311
             os: macos-latest
@@ -63,6 +64,7 @@ jobs:
           - name: linux-py311
             os: ubuntu-latest
             PYTHON_VERSION: "3.11"
+            LOKY_TEST_NO_LSCPU: "true"
           - name: linux-py39-joblib-tests
             os: ubuntu-latest
             PYTHON_VERSION: "3.9"

CHANGES.md

@@ -1,5 +1,10 @@
 ### 3.5.0 - in development
 
+
+- Fix detection of the number of physical cores in
+  `cpu_count(only_physical_cores=True)` on some Linux systems and recent
+  Windows versions. (#425)
+
 - Drop support for Python 3.7 and Python 3.8. (#409)
 
 - Drop support for PyPy. (#427)

loky/backend/context.py

@@ -242,33 +242,11 @@ def _count_physical_cores():
     # Not cached yet, find it
     try:
         if sys.platform == "linux":
-            cpu_info = subprocess.run(
-                "lscpu --parse=core".split(), capture_output=True, text=True
-            )
-            cpu_info = cpu_info.stdout.splitlines()
-            cpu_info = {line for line in cpu_info if not line.startswith("#")}
-            cpu_count_physical = len(cpu_info)
+            cpu_count_physical = _count_physical_cores_linux()
         elif sys.platform == "win32":
-            cpu_info = subprocess.run(
-                "wmic CPU Get NumberOfCores /Format:csv".split(),
-                capture_output=True,
-                text=True,
-            )
-            cpu_info = cpu_info.stdout.splitlines()
-            cpu_info = [
-                l.split(",")[1]
-                for l in cpu_info
-                if (l and l != "Node,NumberOfCores")
-            ]
-            cpu_count_physical = sum(map(int, cpu_info))
+            cpu_count_physical = _count_physical_cores_win32()
         elif sys.platform == "darwin":
-            cpu_info = subprocess.run(
-                "sysctl -n hw.physicalcpu".split(),
-                capture_output=True,
-                text=True,
-            )
-            cpu_info = cpu_info.stdout
-            cpu_count_physical = int(cpu_info)
+            cpu_count_physical = _count_physical_cores_darwin()
         else:
             raise NotImplementedError(f"unsupported platform: {sys.platform}")
 
@@ -286,6 +264,60 @@ def _count_physical_cores():
     return cpu_count_physical, exception
 
 
+def _count_physical_cores_linux():
+    try:
+        cpu_info = subprocess.run(
+            "lscpu --parse=core".split(), capture_output=True, text=True
+        )
+        cpu_info = cpu_info.stdout.splitlines()
+        cpu_info = {line for line in cpu_info if not line.startswith("#")}
+        return len(cpu_info)
+    except:
+        pass  # fallback to /proc/cpuinfo
+
+    cpu_info = subprocess.run(
+        "cat /proc/cpuinfo".split(), capture_output=True, text=True
+    )
+    cpu_info = cpu_info.stdout.splitlines()
+    cpu_info = {line for line in cpu_info if line.startswith("core id")}
+    return len(cpu_info)
+
+
+def _count_physical_cores_win32():
+    try:
+        cmd = "-Command (Get-CimInstance -ClassName Win32_Processor).NumberOfCores"
+        cpu_info = subprocess.run(
+            f"powershell.exe {cmd}".split(),
+            capture_output=True,
+            text=True,
+        )
+        cpu_info = cpu_info.stdout.splitlines()
+        return int(cpu_info[0])
+    except:
+        pass  # fallback to wmic (older Windows versions; deprecated now)
+
+    cpu_info = subprocess.run(
+        "wmic CPU Get NumberOfCores /Format:csv".split(),
+        capture_output=True,
+        text=True,
+    )
+    cpu_info = cpu_info.stdout.splitlines()
+    cpu_info = [
+        l.split(",")[1] for l in cpu_info if (l and l != "Node,NumberOfCores")
+    ]
+    return sum(map(int, cpu_info))
+
+
+def _count_physical_cores_darwin():
+    cpu_info = subprocess.run(
+        "sysctl -n hw.physicalcpu".split(),
+        capture_output=True,
+        text=True,
+    )
+    cpu_info = cpu_info.stdout
+    return int(cpu_info)
+
+
 class LokyContext(BaseContext):
     """Context relying on the LokyProcess."""
 

tests/test_loky_module.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import shutil
+import subprocess
 import tempfile
 import warnings
 from subprocess import check_output
@@ -19,7 +20,29 @@ def test_version():
     ), "There are no __version__ argument on the loky module"
 
 
-def test_cpu_count():
+def test_cpu_count(monkeypatch):
+
+    # Monkeypatch subprocess.run to simulate the absence of lscpu on linux or CIM on
+    # windows to test the different code paths in _cpu_count_physical.
+    old_run = subprocess.run
+
+    def mock_run(*args, **kwargs):
+        if (
+            "lscpu" in args[0]
+            and os.environ.get("LOKY_TEST_NO_LSCPU") == "true"
+        ):
+            raise RuntimeError("lscpu not available")
+
+        if (
+            "powershell.exe" in args[0]
+            and os.environ.get("LOKY_TEST_NO_CIM") == "true"
+        ):
+            raise RuntimeError("Cim not available")
+
+        return old_run(*args, **kwargs)
+
+    monkeypatch.setattr(subprocess, "run", mock_run)
+
     cpus = cpu_count()
     assert type(cpus) is int
     assert cpus >= 1