diff --git a/Doc/library/os.rst b/Doc/library/os.rst index 4ffd520f9ecd8bf..9ab4bda2c642e25 100644 --- a/Doc/library/os.rst +++ b/Doc/library/os.rst @@ -5183,12 +5183,10 @@ 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. + See also the :func:`process_cpu_count` function. .. versionadded:: 3.4 @@ -5202,6 +5200,17 @@ Miscellaneous System Information .. availability:: Unix. +.. function:: process_cpu_count() + + Get the number of logical CPUs usable by the current process. Returns + ``None`` if undetermined. It can be less than :func:`cpu_count` depending on + the process affinity. + + See also the :func:`cpu_count` function. + + .. 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 c9e6ca8bf88866c..d907ad20add28f8 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:`process_cpu_count` function to get the number of logical CPUs + usable by the current process. + (Contributed by Victor Stinner in :gh:`109649`.) + pathlib ------- diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py index 66aece2c4b3eb99..c1a78a70c094411 100644 --- a/Lib/test/test_os.py +++ b/Lib/test/test_os.py @@ -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() diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py index 444f8abe4607b71..9d72dba159c6bee 100644 --- a/Lib/test/test_posix.py +++ b/Lib/test/test_posix.py @@ -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() 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 000000000000000..93ade8e346a9acf --- /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:`process_cpu_count` function to get the number of logical CPUs usable +by the current process. Patch by Victor Stinner. diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h index c3e7f86b3e33f18..0b65793fd98d76b 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__}, @@ -10443,6 +10441,26 @@ os_cpu_count(PyObject *module, PyObject *Py_UNUSED(ignored)) return os_cpu_count_impl(module); } +PyDoc_STRVAR(os_process_cpu_count__doc__, +"process_cpu_count($module, /)\n" +"--\n" +"\n" +"Get the number of logical CPUs usable by the current process.\n" +"\n" +"Return None if indeterminable."); + +#define OS_PROCESS_CPU_COUNT_METHODDEF \ + {"process_cpu_count", (PyCFunction)os_process_cpu_count, METH_NOARGS, os_process_cpu_count__doc__}, + +static PyObject * +os_process_cpu_count_impl(PyObject *module); + +static PyObject * +os_process_cpu_count(PyObject *module, PyObject *Py_UNUSED(ignored)) +{ + return os_process_cpu_count_impl(module); +} + PyDoc_STRVAR(os_get_inheritable__doc__, "get_inheritable($module, fd, /)\n" "--\n" @@ -11988,4 +12006,4 @@ os_waitstatus_to_exitcode(PyObject *module, PyObject *const *args, Py_ssize_t na #ifndef OS_WAITSTATUS_TO_EXITCODE_METHODDEF #define OS_WAITSTATUS_TO_EXITCODE_METHODDEF #endif /* !defined(OS_WAITSTATUS_TO_EXITCODE_METHODDEF) */ -/*[clinic end generated code: output=1dd5aa7495cd6e3a input=a9049054013a1b77]*/ +/*[clinic end generated code: output=883f5bdcbc36cd7b input=a9049054013a1b77]*/ diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index 096aa043514c85c..ad3aab320e9cfc8 100644 --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -8138,39 +8138,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; @@ -8182,12 +8188,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 */ @@ -14338,44 +14344,96 @@ 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); +} + + +/*[clinic input] +os.process_cpu_count + +Get the number of logical CPUs usable by the current process. + +Return None if indeterminable. +[clinic start generated code]*/ + +static PyObject * +os_process_cpu_count_impl(PyObject *module) +/*[clinic end generated code: output=dc750a336e010b50 input=b9b63acbae0dbe45]*/ +{ +#if defined(HAVE_SCHED_H) && defined(HAVE_SCHED_SETAFFINITY) + int ncpus = NCPUS_START; + cpu_set_t *mask; + size_t setsize; + + while (1) { + setsize = CPU_ALLOC_SIZE(ncpus); + mask = CPU_ALLOC(ncpus); + if (mask == NULL) { + return PyErr_NoMemory(); + } + if (sched_getaffinity(0, setsize, mask) == 0) { + break; + } + CPU_FREE(mask); + if (errno != EINVAL) { + return posix_error(); + } + if (ncpus > INT_MAX / 2) { + PyErr_SetString(PyExc_OverflowError, + "could not allocate a large enough CPU set"); + return NULL; + } + ncpus *= 2; + } + + int ncpu = CPU_COUNT_S(setsize, mask); + CPU_FREE(mask); + return PyLong_FromLong(ncpu); +#else + return os_cpu_count_impl(NULL); +#endif } @@ -15972,6 +16030,7 @@ static PyMethodDef posix_methods[] = { OS_GET_TERMINAL_SIZE_METHODDEF OS_CPU_COUNT_METHODDEF + OS_PROCESS_CPU_COUNT_METHODDEF OS_GET_INHERITABLE_METHODDEF OS_SET_INHERITABLE_METHODDEF OS_GET_HANDLE_INHERITABLE_METHODDEF