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

Skip to content

Commit 0362cbf

Browse files
corona10vstinnergpshead
authored
gh-109595: Add -Xcpu_count=<n> cmdline for container users (#109667)
 --------- Co-authored-by: Victor Stinner <[email protected]> Co-authored-by: Gregory P. Smith [Google LLC] <[email protected]>
1 parent 5aa62a8 commit 0362cbf

File tree

15 files changed

+192
-11
lines changed

15 files changed

+192
-11
lines changed

‎Doc/c-api/init_config.rst

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -878,6 +878,19 @@ PyConfig
878878
879879
.. versionadded:: 3.12
880880
881+
.. c:member:: int cpu_count
882+
883+
If the value of :c:member:`~PyConfig.cpu_count` is not ``-1`` then it will
884+
override the return values of :func:`os.cpu_count`,
885+
:func:`os.process_cpu_count`, and :func:`multiprocessing.cpu_count`.
886+
887+
Configured by the :samp:`-X cpu_count={n|default}` command line
888+
flag or the :envvar:`PYTHON_CPU_COUNT` environment variable.
889+
890+
Default: ``-1``.
891+
892+
.. versionadded:: 3.13
893+
881894
.. c:member:: int isolated
882895
883896
If greater than ``0``, enable isolated mode:

‎Doc/library/multiprocessing.rst

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -996,13 +996,20 @@ Miscellaneous
996996

997997
This number is not equivalent to the number of CPUs the current process can
998998
use. The number of usable CPUs can be obtained with
999-
:func:`os.process_cpu_count`.
999+
:func:`os.process_cpu_count` (or ``len(os.sched_getaffinity(0))``).
10001000

10011001
When the number of CPUs cannot be determined a :exc:`NotImplementedError`
10021002
is raised.
10031003

10041004
.. seealso::
1005-
:func:`os.cpu_count` and :func:`os.process_cpu_count`
1005+
:func:`os.cpu_count`
1006+
:func:`os.process_cpu_count`
1007+
1008+
.. versionchanged:: 3.13
1009+
1010+
The return value can also be overridden using the
1011+
:option:`-X cpu_count <-X>` flag or :envvar:`PYTHON_CPU_COUNT` as this is
1012+
merely a wrapper around the :mod:`os` cpu count APIs.
10061013

10071014
.. function:: current_process()
10081015

‎Doc/library/os.rst

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5406,6 +5406,10 @@ Miscellaneous System Information
54065406

54075407
.. versionadded:: 3.4
54085408

5409+
.. versionchanged:: 3.13
5410+
If :option:`-X cpu_count <-X>` is given or :envvar:`PYTHON_CPU_COUNT` is set,
5411+
:func:`cpu_count` returns the overridden value *n*.
5412+
54095413

54105414
.. function:: getloadavg()
54115415

@@ -5425,6 +5429,9 @@ Miscellaneous System Information
54255429
The :func:`cpu_count` function can be used to get the number of logical CPUs
54265430
in the **system**.
54275431

5432+
If :option:`-X cpu_count <-X>` is given or :envvar:`PYTHON_CPU_COUNT` is set,
5433+
:func:`process_cpu_count` returns the overridden value *n*.
5434+
54285435
See also the :func:`sched_getaffinity` functions.
54295436

54305437
.. versionadded:: 3.13

‎Doc/using/cmdline.rst

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,12 @@ Miscellaneous options
546546
report Python calls. This option is only available on some platforms and
547547
will do nothing if is not supported on the current system. The default value
548548
is "off". See also :envvar:`PYTHONPERFSUPPORT` and :ref:`perf_profiling`.
549+
* :samp:`-X cpu_count={n}` overrides :func:`os.cpu_count`,
550+
:func:`os.process_cpu_count`, and :func:`multiprocessing.cpu_count`.
551+
*n* must be greater than or equal to 1.
552+
This option may be useful for users who need to limit CPU resources of a
553+
container system. See also :envvar:`PYTHON_CPU_COUNT`.
554+
If *n* is ``default``, nothing is overridden.
549555

550556
It also allows passing arbitrary values and retrieving them through the
551557
:data:`sys._xoptions` dictionary.
@@ -593,6 +599,9 @@ Miscellaneous options
593599
.. versionadded:: 3.12
594600
The ``-X perf`` option.
595601

602+
.. versionadded:: 3.13
603+
The ``-X cpu_count`` option.
604+
596605

597606
Options you shouldn't use
598607
~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1063,6 +1072,15 @@ conflict.
10631072

10641073
.. versionadded:: 3.12
10651074

1075+
.. envvar:: PYTHON_CPU_COUNT
1076+
1077+
If this variable is set to a positive integer, it overrides the return
1078+
values of :func:`os.cpu_count` and :func:`os.process_cpu_count`.
1079+
1080+
See also the :option:`-X cpu_count <-X>` command-line option.
1081+
1082+
.. versionadded:: 3.13
1083+
10661084

10671085
Debug-mode variables
10681086
~~~~~~~~~~~~~~~~~~~~

‎Doc/whatsnew/3.13.rst

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,12 @@ os
188188
:const:`os.TFD_TIMER_ABSTIME`, and :const:`os.TFD_TIMER_CANCEL_ON_SET`
189189
(Contributed by Masaru Tsuchiyama in :gh:`108277`.)
190190

191+
* :func:`os.cpu_count` and :func:`os.process_cpu_count` can be overridden through
192+
the new environment variable :envvar:`PYTHON_CPU_COUNT` or the new command-line option
193+
:option:`-X cpu_count <-X>`. This option is useful for users who need to limit
194+
CPU resources of a container system without having to modify the container (application code).
195+
(Contributed by Donghee Na in :gh:`109595`)
196+
191197
pathlib
192198
-------
193199

‎Include/cpython/initconfig.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,8 @@ typedef struct PyConfig {
180180
int safe_path;
181181
int int_max_str_digits;
182182

183+
int cpu_count;
184+
183185
/* --- Path configuration inputs ------------ */
184186
int pathconfig_warnings;
185187
wchar_t *program_name;

‎Lib/os.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1138,7 +1138,7 @@ def add_dll_directory(path):
11381138
)
11391139

11401140

1141-
if _exists('sched_getaffinity'):
1141+
if _exists('sched_getaffinity') and sys._get_cpu_count_config() < 0:
11421142
def process_cpu_count():
11431143
"""
11441144
Get the number of CPUs of the current process.

‎Lib/test/test_cmd_line.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -878,11 +878,8 @@ def test_int_max_str_digits(self):
878878
assert_python_failure('-c', code, PYTHONINTMAXSTRDIGITS='foo')
879879
assert_python_failure('-c', code, PYTHONINTMAXSTRDIGITS='100')
880880

881-
def res2int(res):
882-
out = res.out.strip().decode("utf-8")
883-
return tuple(int(i) for i in out.split())
884-
885881
res = assert_python_ok('-c', code)
882+
res2int = self.res2int
886883
current_max = sys.get_int_max_str_digits()
887884
self.assertEqual(res2int(res), (current_max, current_max))
888885
res = assert_python_ok('-X', 'int_max_str_digits=0', '-c', code)
@@ -902,6 +899,26 @@ def res2int(res):
902899
)
903900
self.assertEqual(res2int(res), (6000, 6000))
904901

902+
def test_cpu_count(self):
903+
code = "import os; print(os.cpu_count(), os.process_cpu_count())"
904+
res = assert_python_ok('-X', 'cpu_count=4321', '-c', code)
905+
self.assertEqual(self.res2int(res), (4321, 4321))
906+
res = assert_python_ok('-c', code, PYTHON_CPU_COUNT='1234')
907+
self.assertEqual(self.res2int(res), (1234, 1234))
908+
909+
def test_cpu_count_default(self):
910+
code = "import os; print(os.cpu_count(), os.process_cpu_count())"
911+
res = assert_python_ok('-X', 'cpu_count=default', '-c', code)
912+
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
913+
res = assert_python_ok('-X', 'cpu_count=default', '-c', code, PYTHON_CPU_COUNT='1234')
914+
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
915+
es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='default')
916+
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
917+
918+
def res2int(self, res):
919+
out = res.out.strip().decode("utf-8")
920+
return tuple(int(i) for i in out.split())
921+
905922

906923
@unittest.skipIf(interpreter_requires_environment(),
907924
'Cannot run -I tests when PYTHON env vars are required.')

‎Lib/test/test_embed.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,6 +445,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
445445
'use_hash_seed': 0,
446446
'hash_seed': 0,
447447
'int_max_str_digits': sys.int_info.default_max_str_digits,
448+
'cpu_count': -1,
448449
'faulthandler': 0,
449450
'tracemalloc': 0,
450451
'perf_profiling': 0,
@@ -893,6 +894,7 @@ def test_init_from_config(self):
893894
'module_search_paths': self.IGNORE_CONFIG,
894895
'safe_path': 1,
895896
'int_max_str_digits': 31337,
897+
'cpu_count': 4321,
896898

897899
'check_hash_pycs_mode': 'always',
898900
'pathconfig_warnings': 0,

‎Misc/NEWS.d/next/Core and Builtins/2023-09-22-01-44-53.gh-issue-109595.fVINgD.rst

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Add :option:`-X cpu_count <-X>` command line option to override return results of
2+
:func:`os.cpu_count` and :func:`os.process_cpu_count`.
3+
This option is useful for users who need to limit CPU resources of a container system
4+
without having to modify the container (application code).
5+
Patch by Donghee Na.

‎Modules/posixmodule.c

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14592,7 +14592,6 @@ os_get_terminal_size_impl(PyObject *module, int fd)
1459214592
}
1459314593
#endif /* defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL) */
1459414594

14595-
1459614595
/*[clinic input]
1459714596
os.cpu_count
1459814597
@@ -14605,7 +14604,12 @@ static PyObject *
1460514604
os_cpu_count_impl(PyObject *module)
1460614605
/*[clinic end generated code: output=5fc29463c3936a9c input=ba2f6f8980a0e2eb]*/
1460714606
{
14608-
int ncpu;
14607+
const PyConfig *config = _Py_GetConfig();
14608+
if (config->cpu_count > 0) {
14609+
return PyLong_FromLong(config->cpu_count);
14610+
}
14611+
14612+
int ncpu = 0;
1460914613
#ifdef MS_WINDOWS
1461014614
# ifdef MS_WINDOWS_DESKTOP
1461114615
ncpu = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);

‎Programs/_testembed.c

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -715,6 +715,7 @@ static int test_init_from_config(void)
715715

716716
putenv("PYTHONINTMAXSTRDIGITS=6666");
717717
config.int_max_str_digits = 31337;
718+
config.cpu_count = 4321;
718719

719720
init_from_config_clear(&config);
720721

‎Python/clinic/sysmodule.c.h

Lines changed: 29 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Python/initconfig.c

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,7 @@ static const PyConfigSpec PYCONFIG_SPEC[] = {
9292
SPEC(use_frozen_modules, UINT),
9393
SPEC(safe_path, UINT),
9494
SPEC(int_max_str_digits, INT),
95+
SPEC(cpu_count, INT),
9596
SPEC(pathconfig_warnings, UINT),
9697
SPEC(program_name, WSTR),
9798
SPEC(pythonpath_env, WSTR_OPT),
@@ -229,7 +230,11 @@ The following implementation-specific options are available:\n\
229230
\n\
230231
-X int_max_str_digits=number: limit the size of int<->str conversions.\n\
231232
This helps avoid denial of service attacks when parsing untrusted data.\n\
232-
The default is sys.int_info.default_max_str_digits. 0 disables."
233+
The default is sys.int_info.default_max_str_digits. 0 disables.\n\
234+
\n\
235+
-X cpu_count=[n|default]: Override the return value of os.cpu_count(),\n\
236+
os.process_cpu_count(), and multiprocessing.cpu_count(). This can help users who need\n\
237+
to limit resources in a container."
233238

234239
#ifdef Py_STATS
235240
"\n\
@@ -267,6 +272,8 @@ static const char usage_envvars[] =
267272
" locale coercion and locale compatibility warnings on stderr.\n"
268273
"PYTHONBREAKPOINT: if this variable is set to 0, it disables the default\n"
269274
" debugger. It can be set to the callable of your debugger of choice.\n"
275+
"PYTHON_CPU_COUNT: Overrides the return value of os.process_cpu_count(),\n"
276+
" os.cpu_count(), and multiprocessing.cpu_count() if set to a positive integer.\n"
270277
"PYTHONDEVMODE: enable the development mode.\n"
271278
"PYTHONPYCACHEPREFIX: root directory for bytecode cache (pyc) files.\n"
272279
"PYTHONWARNDEFAULTENCODING: enable opt-in EncodingWarning for 'encoding=None'.\n"
@@ -732,6 +739,8 @@ config_check_consistency(const PyConfig *config)
732739
assert(config->_is_python_build >= 0);
733740
assert(config->safe_path >= 0);
734741
assert(config->int_max_str_digits >= 0);
742+
// cpu_count can be -1 if the user doesn't override it.
743+
assert(config->cpu_count != 0);
735744
// config->use_frozen_modules is initialized later
736745
// by _PyConfig_InitImportConfig().
737746
#ifdef Py_STATS
@@ -832,6 +841,7 @@ _PyConfig_InitCompatConfig(PyConfig *config)
832841
config->int_max_str_digits = -1;
833842
config->_is_python_build = 0;
834843
config->code_debug_ranges = 1;
844+
config->cpu_count = -1;
835845
}
836846

837847

@@ -1617,6 +1627,45 @@ config_read_env_vars(PyConfig *config)
16171627
return _PyStatus_OK();
16181628
}
16191629

1630+
static PyStatus
1631+
config_init_cpu_count(PyConfig *config)
1632+
{
1633+
const char *env = config_get_env(config, "PYTHON_CPU_COUNT");
1634+
if (env) {
1635+
int cpu_count = -1;
1636+
if (strcmp(env, "default") == 0) {
1637+
cpu_count = -1;
1638+
}
1639+
else if (_Py_str_to_int(env, &cpu_count) < 0 || cpu_count < 1) {
1640+
goto error;
1641+
}
1642+
config->cpu_count = cpu_count;
1643+
}
1644+
1645+
const wchar_t *xoption = config_get_xoption(config, L"cpu_count");
1646+
if (xoption) {
1647+
int cpu_count = -1;
1648+
const wchar_t *sep = wcschr(xoption, L'=');
1649+
if (sep) {
1650+
if (wcscmp(sep + 1, L"default") == 0) {
1651+
cpu_count = -1;
1652+
}
1653+
else if (config_wstr_to_int(sep + 1, &cpu_count) < 0 || cpu_count < 1) {
1654+
goto error;
1655+
}
1656+
}
1657+
else {
1658+
goto error;
1659+
}
1660+
config->cpu_count = cpu_count;
1661+
}
1662+
return _PyStatus_OK();
1663+
1664+
error:
1665+
return _PyStatus_ERR("-X cpu_count=n option: n is missing or an invalid number, "
1666+
"n must be greater than 0");
1667+
}
1668+
16201669
static PyStatus
16211670
config_init_perf_profiling(PyConfig *config)
16221671
{
@@ -1799,6 +1848,13 @@ config_read_complex_options(PyConfig *config)
17991848
}
18001849
}
18011850

1851+
if (config->cpu_count < 0) {
1852+
status = config_init_cpu_count(config);
1853+
if (_PyStatus_EXCEPTION(status)) {
1854+
return status;
1855+
}
1856+
}
1857+
18021858
if (config->pycache_prefix == NULL) {
18031859
status = config_init_pycache_prefix(config);
18041860
if (_PyStatus_EXCEPTION(status)) {

0 commit comments

Comments
 (0)