Portable way of detecting number of *usable* CPUs in Python
I don't think you will get any truly portable answers, so I will give a correct one.
The correct* answer for Linux is len(os.sched_getaffinity(pid))
, where pid
may be 0
for the current process. This function is exposed in Python 3.3 and later; if you need it in earlier, you'll have to do some fancy cffi
coding.
Edit: you might try to see if you can use a function int omp_get_num_procs();
if it exists, it is the only meaningful answer I found on this question but I haven't tried it from Python.
Use psutil:
from the doc https://psutil.readthedocs.io/en/latest/:
>>> import psutil
>>> psutil.cpu_count()
4
>>> psutil.cpu_count(logical=False) # Ignoring virtual cores
2
This is portable
Here's an approach that gets the number of available CPU cores for the current process on systems that implement sched_getaffinity
, and Windows:
import ctypes
import ctypes.wintypes
import os
from platform import system
def num_available_cores() -> int:
if hasattr(os, 'sched_getaffinity'):
return len(os.sched_getaffinity(0))
elif system() == 'Windows':
kernel32 = ctypes.WinDLL('kernel32')
DWORD_PTR = ctypes.wintypes.WPARAM
PDWORD_PTR = ctypes.POINTER(DWORD_PTR)
GetCurrentProcess = kernel32.GetCurrentProcess
GetCurrentProcess.restype = ctypes.wintypes.HANDLE
GetProcessAffinityMask = kernel32.GetProcessAffinityMask
GetProcessAffinityMask.argtypes = (ctypes.wintypes.HANDLE, PDWORD_PTR, PDWORD_PTR)
mask = DWORD_PTR()
if not GetProcessAffinityMask(GetCurrentProcess(), ctypes.byref(mask), ctypes.byref(DWORD_PTR())):
raise Exception("Call to 'GetProcessAffinityMask' failed")
return bin(mask.value).count('1')
else:
raise Exception('Cannot determine the number of available cores')
On Linux and any other systems that implement sched_getaffinity
, we use Python's built-in wrapper for it.
On Windows we use ctypes
to call GetProcessAffinityMask
.
As far as I know there are no user APIs or tools to get/set the CPU affinity on macOS. In most cases os.cpu_count()
will work fine, but if you truly need the number of available cores you may be out of luck.