Thanks to visit codestin.com
Credit goes to github.com

Skip to content

gh-109649: Add os.process_cpu_count() function #109907

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Sep 30, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 24 additions & 7 deletions Doc/library/os.rst
Original file line number Diff line number Diff line change
Expand Up @@ -5141,8 +5141,12 @@ operating system.

.. function:: sched_getaffinity(pid, /)

Return the set of CPUs the process with PID *pid* (or the current process
if zero) is restricted to.
Return the set of CPUs the process with PID *pid* is restricted to.

If *pid* is zero, return the set of CPUs the calling thread of the current
process is restricted to.

See also the :func:`process_cpu_count` function.


.. _os-path:
Expand Down Expand Up @@ -5183,12 +5187,11 @@ Miscellaneous System Information

.. function:: cpu_count()

Return the number of CPUs in the system. Returns ``None`` if undetermined.

This number is not equivalent to the number of CPUs the current process can
use. The number of usable CPUs can be obtained with
``len(os.sched_getaffinity(0))``
Return the number of logical CPUs in the **system**. Returns ``None`` if
undetermined.

The :func:`process_cpu_count` function can be used to get the number of
logical CPUs usable by the calling thread of the **current process**.

.. versionadded:: 3.4

Expand All @@ -5202,6 +5205,20 @@ Miscellaneous System Information
.. availability:: Unix.


.. function:: process_cpu_count()

Get the number of logical CPUs usable by the calling thread of the **current
process**. Returns ``None`` if undetermined. It can be less than
:func:`cpu_count` depending on the CPU affinity.

The :func:`cpu_count` function can be used to get the number of logical CPUs
in the **system**.

See also the :func:`sched_getaffinity` functions.

.. versionadded:: 3.13


.. function:: sysconf(name, /)

Return integer-valued system configuration values. If the configuration value
Expand Down
7 changes: 7 additions & 0 deletions Doc/whatsnew/3.13.rst
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,13 @@ opcode
documented or exposed through ``dis``, and were not intended to be
used externally.

os
--

* Add :func:`os.process_cpu_count` function to get the number of logical CPUs
usable by the calling thread of the current process.
(Contributed by Victor Stinner in :gh:`109649`.)

pathlib
-------

Expand Down
14 changes: 14 additions & 0 deletions Lib/os.py
Original file line number Diff line number Diff line change
Expand Up @@ -1136,3 +1136,17 @@ def add_dll_directory(path):
cookie,
nt._remove_dll_directory
)


if _exists('sched_getaffinity'):
def process_cpu_count():
"""
Get the number of CPUs of the current process.

Return the number of logical CPUs usable by the calling thread of the
current process. Return None if indeterminable.
"""
return len(sched_getaffinity(0))
else:
# Just an alias to cpu_count() (same docstring)
process_cpu_count = cpu_count
36 changes: 32 additions & 4 deletions Lib/test/test_os.py
Original file line number Diff line number Diff line change
Expand Up @@ -3996,14 +3996,42 @@ def test_oserror_filename(self):
self.fail(f"No exception thrown by {func}")

class CPUCountTests(unittest.TestCase):
def check_cpu_count(self, cpus):
if cpus is None:
self.skipTest("Could not determine the number of CPUs")

self.assertIsInstance(cpus, int)
self.assertGreater(cpus, 0)

def test_cpu_count(self):
cpus = os.cpu_count()
if cpus is not None:
self.assertIsInstance(cpus, int)
self.assertGreater(cpus, 0)
else:
self.check_cpu_count(cpus)

def test_process_cpu_count(self):
cpus = os.process_cpu_count()
self.assertLessEqual(cpus, os.cpu_count())
self.check_cpu_count(cpus)

@unittest.skipUnless(hasattr(os, 'sched_setaffinity'),
"don't have sched affinity support")
def test_process_cpu_count_affinity(self):
ncpu = os.cpu_count()
if ncpu is None:
self.skipTest("Could not determine the number of CPUs")

# Disable one CPU
mask = os.sched_getaffinity(0)
if len(mask) <= 1:
self.skipTest(f"sched_getaffinity() returns less than "
f"2 CPUs: {sorted(mask)}")
self.addCleanup(os.sched_setaffinity, 0, list(mask))
mask.pop()
os.sched_setaffinity(0, mask)

# test process_cpu_count()
affinity = os.process_cpu_count()
self.assertEqual(affinity, ncpu - 1)


# FD inheritance check is only useful for systems with process support.
@support.requires_subprocess()
Expand Down
1 change: 1 addition & 0 deletions Lib/test/test_posix.py
Original file line number Diff line number Diff line change
Expand Up @@ -1205,6 +1205,7 @@ def test_sched_getaffinity(self):
@requires_sched_affinity
def test_sched_setaffinity(self):
mask = posix.sched_getaffinity(0)
self.addCleanup(posix.sched_setaffinity, 0, list(mask))
if len(mask) > 1:
# Empty masks are forbidden
mask.pop()
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Add :func:`os.process_cpu_count` function to get the number of logical CPUs
usable by the calling thread of the current process. Patch by Victor Stinner.
8 changes: 3 additions & 5 deletions Modules/clinic/posixmodule.c.h

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

73 changes: 42 additions & 31 deletions Modules/posixmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -8133,39 +8133,45 @@ static PyObject *
os_sched_getaffinity_impl(PyObject *module, pid_t pid)
/*[clinic end generated code: output=f726f2c193c17a4f input=983ce7cb4a565980]*/
{
int cpu, ncpus, count;
int ncpus = NCPUS_START;
size_t setsize;
cpu_set_t *mask = NULL;
PyObject *res = NULL;
cpu_set_t *mask;

ncpus = NCPUS_START;
while (1) {
setsize = CPU_ALLOC_SIZE(ncpus);
mask = CPU_ALLOC(ncpus);
if (mask == NULL)
if (mask == NULL) {
return PyErr_NoMemory();
if (sched_getaffinity(pid, setsize, mask) == 0)
}
if (sched_getaffinity(pid, setsize, mask) == 0) {
break;
}
CPU_FREE(mask);
if (errno != EINVAL)
if (errno != EINVAL) {
return posix_error();
}
if (ncpus > INT_MAX / 2) {
PyErr_SetString(PyExc_OverflowError, "could not allocate "
"a large enough CPU set");
PyErr_SetString(PyExc_OverflowError,
"could not allocate a large enough CPU set");
return NULL;
}
ncpus = ncpus * 2;
ncpus *= 2;
}

res = PySet_New(NULL);
if (res == NULL)
PyObject *res = PySet_New(NULL);
if (res == NULL) {
goto error;
for (cpu = 0, count = CPU_COUNT_S(setsize, mask); count; cpu++) {
}

int cpu = 0;
int count = CPU_COUNT_S(setsize, mask);
for (; count; cpu++) {
if (CPU_ISSET_S(cpu, setsize, mask)) {
PyObject *cpu_num = PyLong_FromLong(cpu);
--count;
if (cpu_num == NULL)
if (cpu_num == NULL) {
goto error;
}
if (PySet_Add(res, cpu_num)) {
Py_DECREF(cpu_num);
goto error;
Expand All @@ -8177,12 +8183,12 @@ os_sched_getaffinity_impl(PyObject *module, pid_t pid)
return res;

error:
if (mask)
if (mask) {
CPU_FREE(mask);
}
Py_XDECREF(res);
return NULL;
}

#endif /* HAVE_SCHED_SETAFFINITY */

#endif /* HAVE_SCHED_H */
Expand Down Expand Up @@ -14333,44 +14339,49 @@ os_get_terminal_size_impl(PyObject *module, int fd)
/*[clinic input]
os.cpu_count

Return the number of CPUs in the system; return None if indeterminable.
Return the number of logical CPUs in the system.

This number is not equivalent to the number of CPUs the current process can
use. The number of usable CPUs can be obtained with
``len(os.sched_getaffinity(0))``
Return None if indeterminable.
[clinic start generated code]*/

static PyObject *
os_cpu_count_impl(PyObject *module)
/*[clinic end generated code: output=5fc29463c3936a9c input=e7c8f4ba6dbbadd3]*/
/*[clinic end generated code: output=5fc29463c3936a9c input=ba2f6f8980a0e2eb]*/
{
int ncpu = 0;
int ncpu;
#ifdef MS_WINDOWS
#ifdef MS_WINDOWS_DESKTOP
# ifdef MS_WINDOWS_DESKTOP
ncpu = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
#endif
# else
ncpu = 0;
# endif

#elif defined(__hpux)
ncpu = mpctl(MPC_GETNUMSPUS, NULL, NULL);

#elif defined(HAVE_SYSCONF) && defined(_SC_NPROCESSORS_ONLN)
ncpu = sysconf(_SC_NPROCESSORS_ONLN);

#elif defined(__VXWORKS__)
ncpu = _Py_popcount32(vxCpuEnabledGet());

#elif defined(__DragonFly__) || \
defined(__OpenBSD__) || \
defined(__FreeBSD__) || \
defined(__NetBSD__) || \
defined(__APPLE__)
int mib[2];
ncpu = 0;
size_t len = sizeof(ncpu);
mib[0] = CTL_HW;
mib[1] = HW_NCPU;
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) != 0)
int mib[2] = {CTL_HW, HW_NCPU};
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) != 0) {
ncpu = 0;
}
#endif
if (ncpu >= 1)
return PyLong_FromLong(ncpu);
else

if (ncpu < 1) {
Py_RETURN_NONE;
}
return PyLong_FromLong(ncpu);
}


Expand Down