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.

Tags:

Python

Cpu