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

Skip to content

Commit 2762144

Browse files
vstinnervstinner
committed
gh-109649: Add affinity parameter to os.cpu_count()
Implement cpu_count(affinity=True) with sched_getaffinity(). Changes: * Fix test_posix.test_sched_getaffinity(): restore the old CPU mask when the test completes! * Doc: Specify that os.cpu_count() counts *logicial* CPUs. Mention that Linux cgroups are ignored.
1 parent 3f5c564 commit 2762144

11 files changed

+191
-50
lines changed

‎Doc/library/os.rst

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5153,14 +5153,21 @@ Miscellaneous System Information
51535153
.. availability:: Unix.
51545154

51555155

5156-
.. function:: cpu_count()
5156+
.. function:: cpu_count(*, affinity=False)
51575157

5158-
Return the number of CPUs in the system. Returns ``None`` if undetermined.
5158+
Return the number of logical CPUs in the system. Returns ``None`` if
5159+
undetermined.
51595160

5160-
This number is not equivalent to the number of CPUs the current process can
5161-
use. The number of usable CPUs can be obtained with
5162-
``len(os.sched_getaffinity(0))``
5161+
If *affinity* is true, return the number of logical CPUs the current process
5162+
can use. If the :func:`sched_getaffinity` function is available,
5163+
return ``len(os.sched_getaffinity(0))``. Otherwise, return
5164+
``cpu_count(affinity=False)``.
51635165

5166+
Linux control groups, *cgroups*, are not taken in account to get the number
5167+
of logical CPUs.
5168+
5169+
.. versionchanged:: 3.13
5170+
Add *affinity* parameter.
51645171

51655172
.. versionadded:: 3.4
51665173

‎Doc/whatsnew/3.13.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,13 @@ opcode
163163
documented or exposed through ``dis``, and were not intended to be
164164
used externally.
165165

166+
os
167+
--
168+
169+
* Add *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs
170+
the current process can use.
171+
(Contributed by Victor Stinner in :gh:`109649`.)
172+
166173
pathlib
167174
-------
168175

‎Include/internal/pycore_global_objects_fini_generated.h

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_global_strings.h

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,7 @@ struct _Py_global_strings {
273273
STRUCT_FOR_ID(aclose)
274274
STRUCT_FOR_ID(add)
275275
STRUCT_FOR_ID(add_done_callback)
276+
STRUCT_FOR_ID(affinity)
276277
STRUCT_FOR_ID(after_in_child)
277278
STRUCT_FOR_ID(after_in_parent)
278279
STRUCT_FOR_ID(aggregate_class)

‎Include/internal/pycore_runtime_init_generated.h

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Include/internal/pycore_unicodeobject_generated.h

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Lib/test/test_os.py

Lines changed: 32 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3996,14 +3996,42 @@ def test_oserror_filename(self):
39963996
self.fail(f"No exception thrown by {func}")
39973997

39983998
class CPUCountTests(unittest.TestCase):
3999+
def check_cpu_count(self, cpus):
4000+
if cpus is None:
4001+
self.skipTest("Could not determine the number of CPUs")
4002+
4003+
self.assertIsInstance(cpus, int)
4004+
self.assertGreater(cpus, 0)
4005+
39994006
def test_cpu_count(self):
40004007
cpus = os.cpu_count()
4001-
if cpus is not None:
4002-
self.assertIsInstance(cpus, int)
4003-
self.assertGreater(cpus, 0)
4004-
else:
4008+
self.check_cpu_count(cpus)
4009+
4010+
def test_cpu_count_affinity(self):
4011+
cpus = os.cpu_count(affinity=True)
4012+
self.assertLessEqual(cpus, os.cpu_count())
4013+
self.check_cpu_count(cpus)
4014+
4015+
@unittest.skipUnless(hasattr(os, 'sched_setaffinity'),
4016+
"don't have sched affinity support")
4017+
def test_cpu_count_affinity_setaffinity(self):
4018+
ncpu = os.cpu_count()
4019+
if ncpu is None:
40054020
self.skipTest("Could not determine the number of CPUs")
40064021

4022+
# Disable one CPU
4023+
mask = os.sched_getaffinity(0)
4024+
if len(mask) <= 1:
4025+
self.skipTest(f"sched_getaffinity() returns less than "
4026+
f"2 CPUs: {sorted(mask)}")
4027+
self.addCleanup(os.sched_setaffinity, 0, list(mask))
4028+
mask.pop()
4029+
os.sched_setaffinity(0, mask)
4030+
4031+
# test cpu_count(affinity=True)
4032+
affinity = os.cpu_count(affinity=True)
4033+
self.assertEqual(affinity, ncpu - 1)
4034+
40074035

40084036
# FD inheritance check is only useful for systems with process support.
40094037
@support.requires_subprocess()

‎Lib/test/test_posix.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1205,6 +1205,7 @@ def test_sched_getaffinity(self):
12051205
@requires_sched_affinity
12061206
def test_sched_setaffinity(self):
12071207
mask = posix.sched_getaffinity(0)
1208+
self.addCleanup(posix.sched_setaffinity, 0, list(mask))
12081209
if len(mask) > 1:
12091210
# Empty masks are forbidden
12101211
mask.pop()

‎Misc/NEWS.d/next/Library/2023-09-21-16-21-19.gh-issue-109649.YYCjAF.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Add *affinity* parameter to :func:`os.cpu_count` to get the number of CPUs the
2+
current process can use. Patch by Victor Stinner.

‎Modules/clinic/posixmodule.c.h

Lines changed: 52 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Modules/posixmodule.c

Lines changed: 79 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8138,39 +8138,45 @@ static PyObject *
81388138
os_sched_getaffinity_impl(PyObject *module, pid_t pid)
81398139
/*[clinic end generated code: output=f726f2c193c17a4f input=983ce7cb4a565980]*/
81408140
{
8141-
int cpu, ncpus, count;
8141+
int ncpus = NCPUS_START;
81428142
size_t setsize;
8143-
cpu_set_t *mask = NULL;
8144-
PyObject *res = NULL;
8143+
cpu_set_t *mask;
81458144

8146-
ncpus = NCPUS_START;
81478145
while (1) {
81488146
setsize = CPU_ALLOC_SIZE(ncpus);
81498147
mask = CPU_ALLOC(ncpus);
8150-
if (mask == NULL)
8148+
if (mask == NULL) {
81518149
return PyErr_NoMemory();
8152-
if (sched_getaffinity(pid, setsize, mask) == 0)
8150+
}
8151+
if (sched_getaffinity(pid, setsize, mask) == 0) {
81538152
break;
8153+
}
81548154
CPU_FREE(mask);
8155-
if (errno != EINVAL)
8155+
if (errno != EINVAL) {
81568156
return posix_error();
8157+
}
81578158
if (ncpus > INT_MAX / 2) {
8158-
PyErr_SetString(PyExc_OverflowError, "could not allocate "
8159-
"a large enough CPU set");
8159+
PyErr_SetString(PyExc_OverflowError,
8160+
"could not allocate a large enough CPU set");
81608161
return NULL;
81618162
}
8162-
ncpus = ncpus * 2;
8163+
ncpus *= 2;
81638164
}
81648165

8165-
res = PySet_New(NULL);
8166-
if (res == NULL)
8166+
PyObject *res = PySet_New(NULL);
8167+
if (res == NULL) {
81678168
goto error;
8168-
for (cpu = 0, count = CPU_COUNT_S(setsize, mask); count; cpu++) {
8169+
}
8170+
8171+
int cpu = 0;
8172+
int count = CPU_COUNT_S(setsize, mask);
8173+
for (; count; cpu++) {
81698174
if (CPU_ISSET_S(cpu, setsize, mask)) {
81708175
PyObject *cpu_num = PyLong_FromLong(cpu);
81718176
--count;
8172-
if (cpu_num == NULL)
8177+
if (cpu_num == NULL) {
81738178
goto error;
8179+
}
81748180
if (PySet_Add(res, cpu_num)) {
81758181
Py_DECREF(cpu_num);
81768182
goto error;
@@ -8182,12 +8188,12 @@ os_sched_getaffinity_impl(PyObject *module, pid_t pid)
81828188
return res;
81838189

81848190
error:
8185-
if (mask)
8191+
if (mask) {
81868192
CPU_FREE(mask);
8193+
}
81878194
Py_XDECREF(res);
81888195
return NULL;
81898196
}
8190-
81918197
#endif /* HAVE_SCHED_SETAFFINITY */
81928198

81938199
#endif /* HAVE_SCHED_H */
@@ -14338,44 +14344,85 @@ os_get_terminal_size_impl(PyObject *module, int fd)
1433814344
/*[clinic input]
1433914345
os.cpu_count
1434014346
14347+
*
14348+
affinity: bool = False
14349+
1434114350
Return the number of CPUs in the system; return None if indeterminable.
1434214351
14343-
This number is not equivalent to the number of CPUs the current process can
14344-
use. The number of usable CPUs can be obtained with
14345-
``len(os.sched_getaffinity(0))``
14352+
If 'affinity' is true, return the number of CPUs the current process can use.
1434614353
[clinic start generated code]*/
1434714354

1434814355
static PyObject *
14349-
os_cpu_count_impl(PyObject *module)
14350-
/*[clinic end generated code: output=5fc29463c3936a9c input=e7c8f4ba6dbbadd3]*/
14351-
{
14352-
int ncpu = 0;
14356+
os_cpu_count_impl(PyObject *module, int affinity)
14357+
/*[clinic end generated code: output=0cd2ead51703a781 input=8b61eda84766f638]*/
14358+
{
14359+
int ncpu;
14360+
if (affinity) {
14361+
#if defined(HAVE_SCHED_H) && defined(HAVE_SCHED_SETAFFINITY)
14362+
int ncpus = NCPUS_START;
14363+
cpu_set_t *mask;
14364+
size_t setsize;
14365+
14366+
while (1) {
14367+
setsize = CPU_ALLOC_SIZE(ncpus);
14368+
mask = CPU_ALLOC(ncpus);
14369+
if (mask == NULL) {
14370+
return PyErr_NoMemory();
14371+
}
14372+
if (sched_getaffinity(0, setsize, mask) == 0) {
14373+
break;
14374+
}
14375+
CPU_FREE(mask);
14376+
if (errno != EINVAL) {
14377+
return posix_error();
14378+
}
14379+
if (ncpus > INT_MAX / 2) {
14380+
PyErr_SetString(PyExc_OverflowError,
14381+
"could not allocate a large enough CPU set");
14382+
return NULL;
14383+
}
14384+
ncpus *= 2;
14385+
}
14386+
14387+
ncpu = CPU_COUNT_S(setsize, mask);
14388+
CPU_FREE(mask);
14389+
return PyLong_FromLong(ncpu);
14390+
#endif
14391+
}
14392+
1435314393
#ifdef MS_WINDOWS
14354-
#ifdef MS_WINDOWS_DESKTOP
14394+
# ifdef MS_WINDOWS_DESKTOP
1435514395
ncpu = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
14356-
#endif
14396+
# else
14397+
ncpu = 0;
14398+
# endif
14399+
1435714400
#elif defined(__hpux)
1435814401
ncpu = mpctl(MPC_GETNUMSPUS, NULL, NULL);
14402+
1435914403
#elif defined(HAVE_SYSCONF) && defined(_SC_NPROCESSORS_ONLN)
1436014404
ncpu = sysconf(_SC_NPROCESSORS_ONLN);
14405+
1436114406
#elif defined(__VXWORKS__)
1436214407
ncpu = _Py_popcount32(vxCpuEnabledGet());
14408+
1436314409
#elif defined(__DragonFly__) || \
1436414410
defined(__OpenBSD__) || \
1436514411
defined(__FreeBSD__) || \
1436614412
defined(__NetBSD__) || \
1436714413
defined(__APPLE__)
14368-
int mib[2];
14414+
ncpu = 0;
1436914415
size_t len = sizeof(ncpu);
14370-
mib[0] = CTL_HW;
14371-
mib[1] = HW_NCPU;
14372-
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) != 0)
14416+
int mib[2] = {CTL_HW, HW_NCPU};
14417+
if (sysctl(mib, 2, &ncpu, &len, NULL, 0) != 0) {
1437314418
ncpu = 0;
14419+
}
1437414420
#endif
14375-
if (ncpu >= 1)
14376-
return PyLong_FromLong(ncpu);
14377-
else
14421+
14422+
if (ncpu < 1) {
1437814423
Py_RETURN_NONE;
14424+
}
14425+
return PyLong_FromLong(ncpu);
1437914426
}
1438014427

1438114428

0 commit comments

Comments
 (0)