diff options
| -rw-r--r-- | Doc/library/os.rst | 31 | ||||
| -rw-r--r-- | Doc/whatsnew/3.13.rst | 7 | ||||
| -rw-r--r-- | Lib/os.py | 14 | ||||
| -rw-r--r-- | Lib/test/test_os.py | 36 | ||||
| -rw-r--r-- | Lib/test/test_posix.py | 1 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst | 2 | ||||
| -rw-r--r-- | Modules/clinic/posixmodule.c.h | 8 | ||||
| -rw-r--r-- | Modules/posixmodule.c | 73 |
8 files changed, 125 insertions, 47 deletions
diff --git a/Doc/library/os.rst b/Doc/library/os.rst index 4ffd520..141ab0b 100644 --- a/Doc/library/os.rst +++ b/Doc/library/os.rst @@ -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: @@ -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 @@ -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 diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst index 56618b9..484443a 100644 --- a/Doc/whatsnew/3.13.rst +++ b/Doc/whatsnew/3.13.rst @@ -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 ------- @@ -1136,3 +1136,17 @@ if name == 'nt': 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 diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py index 66aece2..c1a78a7 100644 --- a/Lib/test/test_os.py +++ b/Lib/test/test_os.py @@ -3996,14 +3996,42 @@ class OSErrorTests(unittest.TestCase): 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() diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py index 444f8ab..9d72dba 100644 --- a/Lib/test/test_posix.py +++ b/Lib/test/test_posix.py @@ -1205,6 +1205,7 @@ class PosixTester(unittest.TestCase): @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() diff --git a/Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst b/Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst new file mode 100644 index 0000000..ab708e6 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst @@ -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. diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h index e77a31b..fc39ab7 100644 --- a/Modules/clinic/posixmodule.c.h +++ b/Modules/clinic/posixmodule.c.h @@ -10425,11 +10425,9 @@ PyDoc_STRVAR(os_cpu_count__doc__, "cpu_count($module, /)\n" "--\n" "\n" -"Return the number of CPUs in the system; return None if indeterminable.\n" +"Return the number of logical CPUs in the system.\n" "\n" -"This number is not equivalent to the number of CPUs the current process can\n" -"use. The number of usable CPUs can be obtained with\n" -"``len(os.sched_getaffinity(0))``"); +"Return None if indeterminable."); #define OS_CPU_COUNT_METHODDEF \ {"cpu_count", (PyCFunction)os_cpu_count, METH_NOARGS, os_cpu_count__doc__}, @@ -11988,4 +11986,4 @@ exit: #ifndef OS_WAITSTATUS_TO_EXITCODE_METHODDEF #define OS_WAITSTATUS_TO_EXITCODE_METHODDEF #endif /* !defined(OS_WAITSTATUS_TO_EXITCODE_METHODDEF) */ -/*[clinic end generated code: output=51aa26bc6a41e1da input=a9049054013a1b77]*/ +/*[clinic end generated code: output=8b60de6ddb925bc3 input=a9049054013a1b77]*/ diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index 0b25209..d3c0aa6 100644 --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -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; @@ -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 */ @@ -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); } |