mirror of https://github.com/python/cpython
gh-109595: Add -Xcpu_count=<n> cmdline for container users (#109667)
--------- Co-authored-by: Victor Stinner <vstinner@python.org> Co-authored-by: Gregory P. Smith [Google LLC] <greg@krypto.org>
This commit is contained in:
parent
5aa62a8de1
commit
0362cbf908
|
@ -878,6 +878,19 @@ PyConfig
|
|||
|
||||
.. versionadded:: 3.12
|
||||
|
||||
.. c:member:: int cpu_count
|
||||
|
||||
If the value of :c:member:`~PyConfig.cpu_count` is not ``-1`` then it will
|
||||
override the return values of :func:`os.cpu_count`,
|
||||
:func:`os.process_cpu_count`, and :func:`multiprocessing.cpu_count`.
|
||||
|
||||
Configured by the :samp:`-X cpu_count={n|default}` command line
|
||||
flag or the :envvar:`PYTHON_CPU_COUNT` environment variable.
|
||||
|
||||
Default: ``-1``.
|
||||
|
||||
.. versionadded:: 3.13
|
||||
|
||||
.. c:member:: int isolated
|
||||
|
||||
If greater than ``0``, enable isolated mode:
|
||||
|
|
|
@ -996,13 +996,20 @@ Miscellaneous
|
|||
|
||||
This number is not equivalent to the number of CPUs the current process can
|
||||
use. The number of usable CPUs can be obtained with
|
||||
:func:`os.process_cpu_count`.
|
||||
:func:`os.process_cpu_count` (or ``len(os.sched_getaffinity(0))``).
|
||||
|
||||
When the number of CPUs cannot be determined a :exc:`NotImplementedError`
|
||||
is raised.
|
||||
|
||||
.. seealso::
|
||||
:func:`os.cpu_count` and :func:`os.process_cpu_count`
|
||||
:func:`os.cpu_count`
|
||||
:func:`os.process_cpu_count`
|
||||
|
||||
.. versionchanged:: 3.13
|
||||
|
||||
The return value can also be overridden using the
|
||||
:option:`-X cpu_count <-X>` flag or :envvar:`PYTHON_CPU_COUNT` as this is
|
||||
merely a wrapper around the :mod:`os` cpu count APIs.
|
||||
|
||||
.. function:: current_process()
|
||||
|
||||
|
|
|
@ -5406,6 +5406,10 @@ Miscellaneous System Information
|
|||
|
||||
.. versionadded:: 3.4
|
||||
|
||||
.. versionchanged:: 3.13
|
||||
If :option:`-X cpu_count <-X>` is given or :envvar:`PYTHON_CPU_COUNT` is set,
|
||||
:func:`cpu_count` returns the overridden value *n*.
|
||||
|
||||
|
||||
.. function:: getloadavg()
|
||||
|
||||
|
@ -5425,6 +5429,9 @@ Miscellaneous System Information
|
|||
The :func:`cpu_count` function can be used to get the number of logical CPUs
|
||||
in the **system**.
|
||||
|
||||
If :option:`-X cpu_count <-X>` is given or :envvar:`PYTHON_CPU_COUNT` is set,
|
||||
:func:`process_cpu_count` returns the overridden value *n*.
|
||||
|
||||
See also the :func:`sched_getaffinity` functions.
|
||||
|
||||
.. versionadded:: 3.13
|
||||
|
|
|
@ -546,6 +546,12 @@ Miscellaneous options
|
|||
report Python calls. This option is only available on some platforms and
|
||||
will do nothing if is not supported on the current system. The default value
|
||||
is "off". See also :envvar:`PYTHONPERFSUPPORT` and :ref:`perf_profiling`.
|
||||
* :samp:`-X cpu_count={n}` overrides :func:`os.cpu_count`,
|
||||
:func:`os.process_cpu_count`, and :func:`multiprocessing.cpu_count`.
|
||||
*n* must be greater than or equal to 1.
|
||||
This option may be useful for users who need to limit CPU resources of a
|
||||
container system. See also :envvar:`PYTHON_CPU_COUNT`.
|
||||
If *n* is ``default``, nothing is overridden.
|
||||
|
||||
It also allows passing arbitrary values and retrieving them through the
|
||||
:data:`sys._xoptions` dictionary.
|
||||
|
@ -593,6 +599,9 @@ Miscellaneous options
|
|||
.. versionadded:: 3.12
|
||||
The ``-X perf`` option.
|
||||
|
||||
.. versionadded:: 3.13
|
||||
The ``-X cpu_count`` option.
|
||||
|
||||
|
||||
Options you shouldn't use
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
@ -1063,6 +1072,15 @@ conflict.
|
|||
|
||||
.. versionadded:: 3.12
|
||||
|
||||
.. envvar:: PYTHON_CPU_COUNT
|
||||
|
||||
If this variable is set to a positive integer, it overrides the return
|
||||
values of :func:`os.cpu_count` and :func:`os.process_cpu_count`.
|
||||
|
||||
See also the :option:`-X cpu_count <-X>` command-line option.
|
||||
|
||||
.. versionadded:: 3.13
|
||||
|
||||
|
||||
Debug-mode variables
|
||||
~~~~~~~~~~~~~~~~~~~~
|
||||
|
|
|
@ -188,6 +188,12 @@ os
|
|||
:const:`os.TFD_TIMER_ABSTIME`, and :const:`os.TFD_TIMER_CANCEL_ON_SET`
|
||||
(Contributed by Masaru Tsuchiyama in :gh:`108277`.)
|
||||
|
||||
* :func:`os.cpu_count` and :func:`os.process_cpu_count` can be overridden through
|
||||
the new environment variable :envvar:`PYTHON_CPU_COUNT` or the new command-line option
|
||||
:option:`-X cpu_count <-X>`. This option is useful for users who need to limit
|
||||
CPU resources of a container system without having to modify the container (application code).
|
||||
(Contributed by Donghee Na in :gh:`109595`)
|
||||
|
||||
pathlib
|
||||
-------
|
||||
|
||||
|
|
|
@ -180,6 +180,8 @@ typedef struct PyConfig {
|
|||
int safe_path;
|
||||
int int_max_str_digits;
|
||||
|
||||
int cpu_count;
|
||||
|
||||
/* --- Path configuration inputs ------------ */
|
||||
int pathconfig_warnings;
|
||||
wchar_t *program_name;
|
||||
|
|
|
@ -1138,7 +1138,7 @@ if name == 'nt':
|
|||
)
|
||||
|
||||
|
||||
if _exists('sched_getaffinity'):
|
||||
if _exists('sched_getaffinity') and sys._get_cpu_count_config() < 0:
|
||||
def process_cpu_count():
|
||||
"""
|
||||
Get the number of CPUs of the current process.
|
||||
|
|
|
@ -878,11 +878,8 @@ class CmdLineTest(unittest.TestCase):
|
|||
assert_python_failure('-c', code, PYTHONINTMAXSTRDIGITS='foo')
|
||||
assert_python_failure('-c', code, PYTHONINTMAXSTRDIGITS='100')
|
||||
|
||||
def res2int(res):
|
||||
out = res.out.strip().decode("utf-8")
|
||||
return tuple(int(i) for i in out.split())
|
||||
|
||||
res = assert_python_ok('-c', code)
|
||||
res2int = self.res2int
|
||||
current_max = sys.get_int_max_str_digits()
|
||||
self.assertEqual(res2int(res), (current_max, current_max))
|
||||
res = assert_python_ok('-X', 'int_max_str_digits=0', '-c', code)
|
||||
|
@ -902,6 +899,26 @@ class CmdLineTest(unittest.TestCase):
|
|||
)
|
||||
self.assertEqual(res2int(res), (6000, 6000))
|
||||
|
||||
def test_cpu_count(self):
|
||||
code = "import os; print(os.cpu_count(), os.process_cpu_count())"
|
||||
res = assert_python_ok('-X', 'cpu_count=4321', '-c', code)
|
||||
self.assertEqual(self.res2int(res), (4321, 4321))
|
||||
res = assert_python_ok('-c', code, PYTHON_CPU_COUNT='1234')
|
||||
self.assertEqual(self.res2int(res), (1234, 1234))
|
||||
|
||||
def test_cpu_count_default(self):
|
||||
code = "import os; print(os.cpu_count(), os.process_cpu_count())"
|
||||
res = assert_python_ok('-X', 'cpu_count=default', '-c', code)
|
||||
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
|
||||
res = assert_python_ok('-X', 'cpu_count=default', '-c', code, PYTHON_CPU_COUNT='1234')
|
||||
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
|
||||
es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='default')
|
||||
self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count()))
|
||||
|
||||
def res2int(self, res):
|
||||
out = res.out.strip().decode("utf-8")
|
||||
return tuple(int(i) for i in out.split())
|
||||
|
||||
|
||||
@unittest.skipIf(interpreter_requires_environment(),
|
||||
'Cannot run -I tests when PYTHON env vars are required.')
|
||||
|
|
|
@ -445,6 +445,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
|
|||
'use_hash_seed': 0,
|
||||
'hash_seed': 0,
|
||||
'int_max_str_digits': sys.int_info.default_max_str_digits,
|
||||
'cpu_count': -1,
|
||||
'faulthandler': 0,
|
||||
'tracemalloc': 0,
|
||||
'perf_profiling': 0,
|
||||
|
@ -893,6 +894,7 @@ class InitConfigTests(EmbeddingTestsMixin, unittest.TestCase):
|
|||
'module_search_paths': self.IGNORE_CONFIG,
|
||||
'safe_path': 1,
|
||||
'int_max_str_digits': 31337,
|
||||
'cpu_count': 4321,
|
||||
|
||||
'check_hash_pycs_mode': 'always',
|
||||
'pathconfig_warnings': 0,
|
||||
|
|
|
@ -0,0 +1,5 @@
|
|||
Add :option:`-X cpu_count <-X>` command line option to override return results of
|
||||
:func:`os.cpu_count` and :func:`os.process_cpu_count`.
|
||||
This option is useful for users who need to limit CPU resources of a container system
|
||||
without having to modify the container (application code).
|
||||
Patch by Donghee Na.
|
|
@ -14592,7 +14592,6 @@ os_get_terminal_size_impl(PyObject *module, int fd)
|
|||
}
|
||||
#endif /* defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL) */
|
||||
|
||||
|
||||
/*[clinic input]
|
||||
os.cpu_count
|
||||
|
||||
|
@ -14605,7 +14604,12 @@ static PyObject *
|
|||
os_cpu_count_impl(PyObject *module)
|
||||
/*[clinic end generated code: output=5fc29463c3936a9c input=ba2f6f8980a0e2eb]*/
|
||||
{
|
||||
int ncpu;
|
||||
const PyConfig *config = _Py_GetConfig();
|
||||
if (config->cpu_count > 0) {
|
||||
return PyLong_FromLong(config->cpu_count);
|
||||
}
|
||||
|
||||
int ncpu = 0;
|
||||
#ifdef MS_WINDOWS
|
||||
# ifdef MS_WINDOWS_DESKTOP
|
||||
ncpu = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
|
||||
|
|
|
@ -715,6 +715,7 @@ static int test_init_from_config(void)
|
|||
|
||||
putenv("PYTHONINTMAXSTRDIGITS=6666");
|
||||
config.int_max_str_digits = 31337;
|
||||
config.cpu_count = 4321;
|
||||
|
||||
init_from_config_clear(&config);
|
||||
|
||||
|
|
|
@ -1380,6 +1380,34 @@ exit:
|
|||
return return_value;
|
||||
}
|
||||
|
||||
PyDoc_STRVAR(sys__get_cpu_count_config__doc__,
|
||||
"_get_cpu_count_config($module, /)\n"
|
||||
"--\n"
|
||||
"\n"
|
||||
"Private function for getting PyConfig.cpu_count");
|
||||
|
||||
#define SYS__GET_CPU_COUNT_CONFIG_METHODDEF \
|
||||
{"_get_cpu_count_config", (PyCFunction)sys__get_cpu_count_config, METH_NOARGS, sys__get_cpu_count_config__doc__},
|
||||
|
||||
static int
|
||||
sys__get_cpu_count_config_impl(PyObject *module);
|
||||
|
||||
static PyObject *
|
||||
sys__get_cpu_count_config(PyObject *module, PyObject *Py_UNUSED(ignored))
|
||||
{
|
||||
PyObject *return_value = NULL;
|
||||
int _return_value;
|
||||
|
||||
_return_value = sys__get_cpu_count_config_impl(module);
|
||||
if ((_return_value == -1) && PyErr_Occurred()) {
|
||||
goto exit;
|
||||
}
|
||||
return_value = PyLong_FromLong((long)_return_value);
|
||||
|
||||
exit:
|
||||
return return_value;
|
||||
}
|
||||
|
||||
#ifndef SYS_GETWINDOWSVERSION_METHODDEF
|
||||
#define SYS_GETWINDOWSVERSION_METHODDEF
|
||||
#endif /* !defined(SYS_GETWINDOWSVERSION_METHODDEF) */
|
||||
|
@ -1423,4 +1451,4 @@ exit:
|
|||
#ifndef SYS_GETANDROIDAPILEVEL_METHODDEF
|
||||
#define SYS_GETANDROIDAPILEVEL_METHODDEF
|
||||
#endif /* !defined(SYS_GETANDROIDAPILEVEL_METHODDEF) */
|
||||
/*[clinic end generated code: output=549bb1f92a15f916 input=a9049054013a1b77]*/
|
||||
/*[clinic end generated code: output=3a7d3fbbcb281c22 input=a9049054013a1b77]*/
|
||||
|
|
|
@ -92,6 +92,7 @@ static const PyConfigSpec PYCONFIG_SPEC[] = {
|
|||
SPEC(use_frozen_modules, UINT),
|
||||
SPEC(safe_path, UINT),
|
||||
SPEC(int_max_str_digits, INT),
|
||||
SPEC(cpu_count, INT),
|
||||
SPEC(pathconfig_warnings, UINT),
|
||||
SPEC(program_name, WSTR),
|
||||
SPEC(pythonpath_env, WSTR_OPT),
|
||||
|
@ -229,7 +230,11 @@ The following implementation-specific options are available:\n\
|
|||
\n\
|
||||
-X int_max_str_digits=number: limit the size of int<->str conversions.\n\
|
||||
This helps avoid denial of service attacks when parsing untrusted data.\n\
|
||||
The default is sys.int_info.default_max_str_digits. 0 disables."
|
||||
The default is sys.int_info.default_max_str_digits. 0 disables.\n\
|
||||
\n\
|
||||
-X cpu_count=[n|default]: Override the return value of os.cpu_count(),\n\
|
||||
os.process_cpu_count(), and multiprocessing.cpu_count(). This can help users who need\n\
|
||||
to limit resources in a container."
|
||||
|
||||
#ifdef Py_STATS
|
||||
"\n\
|
||||
|
@ -267,6 +272,8 @@ static const char usage_envvars[] =
|
|||
" locale coercion and locale compatibility warnings on stderr.\n"
|
||||
"PYTHONBREAKPOINT: if this variable is set to 0, it disables the default\n"
|
||||
" debugger. It can be set to the callable of your debugger of choice.\n"
|
||||
"PYTHON_CPU_COUNT: Overrides the return value of os.process_cpu_count(),\n"
|
||||
" os.cpu_count(), and multiprocessing.cpu_count() if set to a positive integer.\n"
|
||||
"PYTHONDEVMODE: enable the development mode.\n"
|
||||
"PYTHONPYCACHEPREFIX: root directory for bytecode cache (pyc) files.\n"
|
||||
"PYTHONWARNDEFAULTENCODING: enable opt-in EncodingWarning for 'encoding=None'.\n"
|
||||
|
@ -732,6 +739,8 @@ config_check_consistency(const PyConfig *config)
|
|||
assert(config->_is_python_build >= 0);
|
||||
assert(config->safe_path >= 0);
|
||||
assert(config->int_max_str_digits >= 0);
|
||||
// cpu_count can be -1 if the user doesn't override it.
|
||||
assert(config->cpu_count != 0);
|
||||
// config->use_frozen_modules is initialized later
|
||||
// by _PyConfig_InitImportConfig().
|
||||
#ifdef Py_STATS
|
||||
|
@ -832,6 +841,7 @@ _PyConfig_InitCompatConfig(PyConfig *config)
|
|||
config->int_max_str_digits = -1;
|
||||
config->_is_python_build = 0;
|
||||
config->code_debug_ranges = 1;
|
||||
config->cpu_count = -1;
|
||||
}
|
||||
|
||||
|
||||
|
@ -1617,6 +1627,45 @@ config_read_env_vars(PyConfig *config)
|
|||
return _PyStatus_OK();
|
||||
}
|
||||
|
||||
static PyStatus
|
||||
config_init_cpu_count(PyConfig *config)
|
||||
{
|
||||
const char *env = config_get_env(config, "PYTHON_CPU_COUNT");
|
||||
if (env) {
|
||||
int cpu_count = -1;
|
||||
if (strcmp(env, "default") == 0) {
|
||||
cpu_count = -1;
|
||||
}
|
||||
else if (_Py_str_to_int(env, &cpu_count) < 0 || cpu_count < 1) {
|
||||
goto error;
|
||||
}
|
||||
config->cpu_count = cpu_count;
|
||||
}
|
||||
|
||||
const wchar_t *xoption = config_get_xoption(config, L"cpu_count");
|
||||
if (xoption) {
|
||||
int cpu_count = -1;
|
||||
const wchar_t *sep = wcschr(xoption, L'=');
|
||||
if (sep) {
|
||||
if (wcscmp(sep + 1, L"default") == 0) {
|
||||
cpu_count = -1;
|
||||
}
|
||||
else if (config_wstr_to_int(sep + 1, &cpu_count) < 0 || cpu_count < 1) {
|
||||
goto error;
|
||||
}
|
||||
}
|
||||
else {
|
||||
goto error;
|
||||
}
|
||||
config->cpu_count = cpu_count;
|
||||
}
|
||||
return _PyStatus_OK();
|
||||
|
||||
error:
|
||||
return _PyStatus_ERR("-X cpu_count=n option: n is missing or an invalid number, "
|
||||
"n must be greater than 0");
|
||||
}
|
||||
|
||||
static PyStatus
|
||||
config_init_perf_profiling(PyConfig *config)
|
||||
{
|
||||
|
@ -1799,6 +1848,13 @@ config_read_complex_options(PyConfig *config)
|
|||
}
|
||||
}
|
||||
|
||||
if (config->cpu_count < 0) {
|
||||
status = config_init_cpu_count(config);
|
||||
if (_PyStatus_EXCEPTION(status)) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
|
||||
if (config->pycache_prefix == NULL) {
|
||||
status = config_init_pycache_prefix(config);
|
||||
if (_PyStatus_EXCEPTION(status)) {
|
||||
|
|
|
@ -2306,6 +2306,20 @@ sys__getframemodulename_impl(PyObject *module, int depth)
|
|||
return Py_NewRef(r);
|
||||
}
|
||||
|
||||
/*[clinic input]
|
||||
sys._get_cpu_count_config -> int
|
||||
|
||||
Private function for getting PyConfig.cpu_count
|
||||
[clinic start generated code]*/
|
||||
|
||||
static int
|
||||
sys__get_cpu_count_config_impl(PyObject *module)
|
||||
/*[clinic end generated code: output=36611bb5efad16dc input=523e1ade2204084e]*/
|
||||
{
|
||||
const PyConfig *config = _Py_GetConfig();
|
||||
return config->cpu_count;
|
||||
}
|
||||
|
||||
static PerfMapState perf_map_state;
|
||||
|
||||
PyAPI_FUNC(int) PyUnstable_PerfMapState_Init(void) {
|
||||
|
@ -2440,6 +2454,7 @@ static PyMethodDef sys_methods[] = {
|
|||
SYS__STATS_CLEAR_METHODDEF
|
||||
SYS__STATS_DUMP_METHODDEF
|
||||
#endif
|
||||
SYS__GET_CPU_COUNT_CONFIG_METHODDEF
|
||||
{NULL, NULL} // sentinel
|
||||
};
|
||||
|
||||
|
|
Loading…
Reference in New Issue