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:
Donghee Na 2023-10-10 19:00:09 +09:00 committed by GitHub
parent 5aa62a8de1
commit 0362cbf908
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
15 changed files with 192 additions and 11 deletions

View file

@ -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)) {