diff --git a/Doc/library/os.rst b/Doc/library/os.rst index c67b966f777db8..c8f963255fc37b 100644 --- a/Doc/library/os.rst +++ b/Doc/library/os.rst @@ -5153,14 +5153,21 @@ Miscellaneous System Information .. availability:: Unix. -.. function:: cpu_count() +.. function:: cpu_count(*, affinity=False) - Return the number of CPUs in the system. Returns ``None`` if undetermined. + Return the number of logical 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))`` + If *affinity* is true, return the number of logical CPUs the current process + can use. If the :func:`sched_getaffinity` function is available, + return ``len(os.sched_getaffinity(0))``. Otherwise, return + ``cpu_count(affinity=False)``. + Linux control groups, *cgroups*, are not taken in account to get the number + of logical CPUs. + + .. versionchanged:: 3.13 + Add *affinity* parameter. .. versionadded:: 3.4 diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst index b4411a587c8bc7..606d15c8c7b9c0 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 *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs + the current process can use. + (Contributed by Victor Stinner in :gh:`109649`.) + pathlib ------- diff --git a/Include/internal/pycore_global_objects_fini_generated.h b/Include/internal/pycore_global_objects_fini_generated.h index 6361f5d1100231..2fc7b05c511050 100644 --- a/Include/internal/pycore_global_objects_fini_generated.h +++ b/Include/internal/pycore_global_objects_fini_generated.h @@ -784,6 +784,7 @@ _PyStaticObjects_CheckRefcnt(PyInterpreterState *interp) { _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(aclose)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(add)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(add_done_callback)); + _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(affinity)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(after_in_child)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(after_in_parent)); _PyStaticObject_CheckRefcnt((PyObject *)&_Py_ID(aggregate_class)); diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h index 504008d67fe9cd..eb8e65e34db2ad 100644 --- a/Include/internal/pycore_global_strings.h +++ b/Include/internal/pycore_global_strings.h @@ -273,6 +273,7 @@ struct _Py_global_strings { STRUCT_FOR_ID(aclose) STRUCT_FOR_ID(add) STRUCT_FOR_ID(add_done_callback) + STRUCT_FOR_ID(affinity) STRUCT_FOR_ID(after_in_child) STRUCT_FOR_ID(after_in_parent) STRUCT_FOR_ID(aggregate_class) diff --git a/Include/internal/pycore_runtime_init_generated.h b/Include/internal/pycore_runtime_init_generated.h index 8cc3287ce35e5b..a30d5e65d2e2be 100644 --- a/Include/internal/pycore_runtime_init_generated.h +++ b/Include/internal/pycore_runtime_init_generated.h @@ -782,6 +782,7 @@ extern "C" { INIT_ID(aclose), \ INIT_ID(add), \ INIT_ID(add_done_callback), \ + INIT_ID(affinity), \ INIT_ID(after_in_child), \ INIT_ID(after_in_parent), \ INIT_ID(aggregate_class), \ diff --git a/Include/internal/pycore_unicodeobject_generated.h b/Include/internal/pycore_unicodeobject_generated.h index 50400db2919a73..90945a222dd043 100644 --- a/Include/internal/pycore_unicodeobject_generated.h +++ b/Include/internal/pycore_unicodeobject_generated.h @@ -660,6 +660,9 @@ _PyUnicode_InitStaticStrings(PyInterpreterState *interp) { string = &_Py_ID(add_done_callback); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); + string = &_Py_ID(affinity); + assert(_PyUnicode_CheckConsistency(string, 1)); + _PyUnicode_InternInPlace(interp, &string); string = &_Py_ID(after_in_child); assert(_PyUnicode_CheckConsistency(string, 1)); _PyUnicode_InternInPlace(interp, &string); diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py index 66aece2c4b3eb9..f41a5fb33c2aeb 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_cpu_count_affinity(self): + cpus = os.cpu_count(affinity=True) + 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_cpu_count_affinity_setaffinity(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 cpu_count(affinity=True) + affinity = os.cpu_count(affinity=True) + 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 444f8abe4607b7..9d72dba159c6be 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 00000000000000..dc53c6a106d254 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst @@ -0,0 +1,2 @@ +Add *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs the +current process can use. Patch by Victor Stinner. diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h index c3e7f86b3e33f1..ed71815af82fbe 100644 --- a/Modules/clinic/posixmodule.c.h +++ b/Modules/clinic/posixmodule.c.h @@ -10422,25 +10422,68 @@ os_get_terminal_size(PyObject *module, PyObject *const *args, Py_ssize_t nargs) #endif /* (defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL)) */ PyDoc_STRVAR(os_cpu_count__doc__, -"cpu_count($module, /)\n" +"cpu_count($module, /, *, affinity=False)\n" "--\n" "\n" "Return the number of CPUs in the system; return None if indeterminable.\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))``"); +"If \'affinity\' is true, return the number of CPUs the current process can use."); #define OS_CPU_COUNT_METHODDEF \ - {"cpu_count", (PyCFunction)os_cpu_count, METH_NOARGS, os_cpu_count__doc__}, + {"cpu_count", _PyCFunction_CAST(os_cpu_count), METH_FASTCALL|METH_KEYWORDS, os_cpu_count__doc__}, static PyObject * -os_cpu_count_impl(PyObject *module); +os_cpu_count_impl(PyObject *module, int affinity); static PyObject * -os_cpu_count(PyObject *module, PyObject *Py_UNUSED(ignored)) +os_cpu_count(PyObject *module, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames) { - return os_cpu_count_impl(module); + PyObject *return_value = NULL; + #if defined(Py_BUILD_CORE) && !defined(Py_BUILD_CORE_MODULE) + + #define NUM_KEYWORDS 1 + static struct { + PyGC_Head _this_is_not_used; + PyObject_VAR_HEAD + PyObject *ob_item[NUM_KEYWORDS]; + } _kwtuple = { + .ob_base = PyVarObject_HEAD_INIT(&PyTuple_Type, NUM_KEYWORDS) + .ob_item = { &_Py_ID(affinity), }, + }; + #undef NUM_KEYWORDS + #define KWTUPLE (&_kwtuple.ob_base.ob_base) + + #else // !Py_BUILD_CORE + # define KWTUPLE NULL + #endif // !Py_BUILD_CORE + + static const char * const _keywords[] = {"affinity", NULL}; + static _PyArg_Parser _parser = { + .keywords = _keywords, + .fname = "cpu_count", + .kwtuple = KWTUPLE, + }; + #undef KWTUPLE + PyObject *argsbuf[1]; + Py_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - 0; + int affinity = 0; + + args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, 0, 0, 0, argsbuf); + if (!args) { + goto exit; + } + if (!noptargs) { + goto skip_optional_kwonly; + } + affinity = PyObject_IsTrue(args[0]); + if (affinity < 0) { + goto exit; + } +skip_optional_kwonly: + return_value = os_cpu_count_impl(module, affinity); + +exit: + return return_value; } PyDoc_STRVAR(os_get_inheritable__doc__, @@ -11988,4 +12031,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=c176e219ae3343de input=a9049054013a1b77]*/ diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index 096aa043514c85..d8a1075ba5be0d 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,85 @@ os_get_terminal_size_impl(PyObject *module, int fd) /*[clinic input] os.cpu_count + * + affinity: bool = False + Return the number of CPUs in the system; return None if indeterminable. -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))`` +If 'affinity' is true, return the number of CPUs the current process can use. [clinic start generated code]*/ static PyObject * -os_cpu_count_impl(PyObject *module) -/*[clinic end generated code: output=5fc29463c3936a9c input=e7c8f4ba6dbbadd3]*/ -{ - int ncpu = 0; +os_cpu_count_impl(PyObject *module, int affinity) +/*[clinic end generated code: output=0cd2ead51703a781 input=8b61eda84766f638]*/ +{ + int ncpu; + if (affinity) { +#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; + } + + ncpu = CPU_COUNT_S(setsize, mask); + CPU_FREE(mask); + return PyLong_FromLong(ncpu); +#endif + } + #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); }