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

Skip to content

Commit 38ba1c3

Browse files
vstinnervstinner
committed
gh-109649: Add affinity parameter to os.cpu_count()
Implement cpu_count(affinity=True) with sched_getaffinity() on Unix and GetProcessAffinityMask() on Windows. 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 and mention that Linux cgroups are ignored. * _Py_popcount32() uses UINT32_C() for M1, M2 and M4 constants. * Add _Py_popcount64(). Add tests on _Py_popcount64().
1 parent 869f177 commit 38ba1c3

13 files changed

+302
-68
lines changed

‎Doc/library/os.rst

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5153,14 +5153,19 @@ 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. On error, raise an exception, instead of returning ``None``.
51635163

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

51655170
.. versionadded:: 3.4
51665171

‎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_bitutils.h

Lines changed: 43 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -110,18 +110,19 @@ _Py_popcount32(uint32_t x)
110110
// 32-bit SWAR (SIMD Within A Register) popcount
111111

112112
// Binary: 0 1 0 1 ...
113-
const uint32_t M1 = 0x55555555;
113+
const uint32_t M1 = UINT32_C(0x55555555);
114114
// Binary: 00 11 00 11. ..
115-
const uint32_t M2 = 0x33333333;
115+
const uint32_t M2 = UINT32_C(0x33333333);
116116
// Binary: 0000 1111 0000 1111 ...
117-
const uint32_t M4 = 0x0F0F0F0F;
117+
const uint32_t M4 = UINT32_C(0x0F0F0F0F);
118118

119119
// Put count of each 2 bits into those 2 bits
120120
x = x - ((x >> 1) & M1);
121121
// Put count of each 4 bits into those 4 bits
122122
x = (x & M2) + ((x >> 2) & M2);
123123
// Put count of each 8 bits into those 8 bits
124124
x = (x + (x >> 4)) & M4;
125+
125126
// Sum of the 4 byte counts.
126127
// Take care when considering changes to the next line. Portability and
127128
// correctness are delicate here, thanks to C's "integer promotions" (C99
@@ -140,6 +141,45 @@ _Py_popcount32(uint32_t x)
140141
}
141142

142143

144+
static inline int
145+
_Py_popcount64(uint64_t x)
146+
{
147+
#if (defined(__clang__) || defined(__GNUC__))
148+
149+
#if SIZEOF_INT >= 8
150+
Py_BUILD_ASSERT(sizeof(x) <= sizeof(unsigned int));
151+
return __builtin_popcount(x);
152+
#else
153+
// The C standard guarantees that unsigned long will always be big enough
154+
// to hold a uint32_t value without losing information.
155+
Py_BUILD_ASSERT(sizeof(x) <= sizeof(unsigned long));
156+
return __builtin_popcountl(x);
157+
#endif
158+
159+
#else
160+
// 64-bit SWAR (SIMD Within A Register) popcount
161+
162+
// Binary: 0 1 0 1 ...
163+
const uint64_t M1 = UINT64_C(0x5555555555555555);
164+
// Binary: 00 11 00 11. ..
165+
const uint64_t M2 = UINT64_C(0x3333333333333333);
166+
// Binary: 0000 1111 0000 1111 ...
167+
const uint64_t M4 = UINT64_C(0xF0F0F0F0F0F0F0F);
168+
169+
// Put count of each 2 bits into those 2 bits
170+
x = x - ((x >> 1) & M1);
171+
// Put count of each 4 bits into those 4 bits
172+
x = (x & M2) + ((x >> 2) & M2);
173+
// Put count of each 8 bits into those 8 bits
174+
x = (x + (x >> 4)) & M4;
175+
176+
// Sum of the 8 byte counts.
177+
// See _Py_popcount32() for the rationale on the final cast.
178+
return (uint64_t)(x * UINT64_C(0x0101010101010101)) >> 56;
179+
#endif
180+
}
181+
182+
143183
// Return the index of the most significant 1 bit in 'x'. This is the smallest
144184
// integer k such that x < 2**k. Equivalent to floor(log2(x)) + 1 for x != 0.
145185
static inline int

‎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: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3996,14 +3996,43 @@ 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+
# None is valid: os.cpu_count() cannot determine
4002+
# the number of CPUs
4003+
return
4004+
4005+
self.assertIsInstance(cpus, int)
4006+
self.assertGreater(cpus, 0)
4007+
39994008
def test_cpu_count(self):
40004009
cpus = os.cpu_count()
4001-
if cpus is not None:
4002-
self.assertIsInstance(cpus, int)
4003-
self.assertGreater(cpus, 0)
4004-
else:
4010+
self.check_cpu_count(cpus)
4011+
4012+
def test_cpu_count_affinity(self):
4013+
cpus = os.cpu_count(affinity=True)
4014+
self.check_cpu_count(cpus)
4015+
4016+
@unittest.skipUnless(hasattr(os, 'sched_setaffinity'),
4017+
"don't have sched affinity support")
4018+
def test_cpu_count_affinity_setaffinity(self):
4019+
ncpu = os.cpu_count()
4020+
if ncpu is None:
40054021
self.skipTest("Could not determine the number of CPUs")
40064022

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

40084037
# FD inheritance check is only useful for systems with process support.
40094038
@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/_testinternalcapi.c

Lines changed: 54 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -138,14 +138,14 @@ test_bswap(PyObject *self, PyObject *Py_UNUSED(args))
138138

139139

140140
static int
141-
check_popcount(uint32_t x, int expected)
141+
check_popcount32(uint32_t x, int expected)
142142
{
143143
// Use volatile to prevent the compiler to optimize out the whole test
144144
volatile uint32_t u = x;
145145
int bits = _Py_popcount32(u);
146146
if (bits != expected) {
147147
PyErr_Format(PyExc_AssertionError,
148-
"_Py_popcount32(%lu) returns %i, expected %i",
148+
"_Py_popcount32(0x%08lx) returns %i, expected %i",
149149
(unsigned long)x, bits, expected);
150150
return -1;
151151
}
@@ -154,23 +154,63 @@ check_popcount(uint32_t x, int expected)
154154

155155

156156
static PyObject*
157-
test_popcount(PyObject *self, PyObject *Py_UNUSED(args))
157+
test_popcount32(PyObject *self, PyObject *Py_UNUSED(args))
158158
{
159159
#define CHECK(X, RESULT) \
160160
do { \
161-
if (check_popcount(X, RESULT) < 0) { \
161+
if (check_popcount32(X, RESULT) < 0) { \
162162
return NULL; \
163163
} \
164164
} while (0)
165165

166-
CHECK(0, 0);
167-
CHECK(1, 1);
168-
CHECK(0x08080808, 4);
169-
CHECK(0x10000001, 2);
170-
CHECK(0x10101010, 4);
171-
CHECK(0x10204080, 4);
172-
CHECK(0xDEADCAFE, 22);
173-
CHECK(0xFFFFFFFF, 32);
166+
CHECK(UINT32_C(0), 0);
167+
CHECK(UINT32_C(1), 1);
168+
CHECK(UINT32_C(0x08080808), 4);
169+
CHECK(UINT32_C(0x10000001), 2);
170+
CHECK(UINT32_C(0x10101010), 4);
171+
CHECK(UINT32_C(0x10204080), 4);
172+
CHECK(UINT32_C(0xDEADCAFE), 22);
173+
CHECK(UINT32_C(0xFFFFFFFF), 32);
174+
Py_RETURN_NONE;
175+
176+
#undef CHECK
177+
}
178+
179+
180+
static int
181+
check_popcount64(uint64_t x, int expected)
182+
{
183+
// Use volatile to prevent the compiler to optimize out the whole test
184+
volatile uint64_t u = x;
185+
int bits = _Py_popcount64(u);
186+
if (bits != expected) {
187+
PyErr_Format(PyExc_AssertionError,
188+
"_Py_popcount64(0x%016llx) returns %i, expected %i",
189+
(unsigned long long)x, bits, expected);
190+
return -1;
191+
}
192+
return 0;
193+
}
194+
195+
196+
static PyObject*
197+
test_popcount64(PyObject *self, PyObject *Py_UNUSED(args))
198+
{
199+
#define CHECK(X, RESULT) \
200+
do { \
201+
if (check_popcount64(X, RESULT) < 0) { \
202+
return NULL; \
203+
} \
204+
} while (0)
205+
206+
CHECK(UINT64_C(0), 0);
207+
CHECK(UINT64_C(1), 1);
208+
CHECK(UINT64_C(0x0808080808080808), 8);
209+
CHECK(UINT64_C(0x1000000000000001), 2);
210+
CHECK(UINT64_C(0x1010101010101010), 8);
211+
CHECK(UINT64_C(0x1020408080402010), 8);
212+
CHECK(UINT64_C(0xDEADCAFEFEE4ABED), 44);
213+
CHECK(UINT64_C(0xFFFFFFFFFFFFFFFF), 64);
174214
Py_RETURN_NONE;
175215

176216
#undef CHECK
@@ -1482,7 +1522,8 @@ static PyMethodDef module_functions[] = {
14821522
{"get_configs", get_configs, METH_NOARGS},
14831523
{"get_recursion_depth", get_recursion_depth, METH_NOARGS},
14841524
{"test_bswap", test_bswap, METH_NOARGS},
1485-
{"test_popcount", test_popcount, METH_NOARGS},
1525+
{"test_popcount32", test_popcount32, METH_NOARGS},
1526+
{"test_popcount64", test_popcount64, METH_NOARGS},
14861527
{"test_bit_length", test_bit_length, METH_NOARGS},
14871528
{"test_hashtable", test_hashtable, METH_NOARGS},
14881529
{"get_config", test_get_config, METH_NOARGS},

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

0 commit comments

Comments
 (0)