gh-109649: Add affinity parameter to os.cpu_count()

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
 

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
 -------
 

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)

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()

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()

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.

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);
 }