From 0362cbf908aff2b87298f8a9422e7b368f890071 Mon Sep 17 00:00:00 2001 From: Donghee Na Date: Tue, 10 Oct 2023 19:00:09 +0900 Subject: [PATCH] gh-109595: Add -Xcpu_count= cmdline for container users (#109667) --------- Co-authored-by: Victor Stinner Co-authored-by: Gregory P. Smith [Google LLC] --- Doc/c-api/init_config.rst | 13 +++++ Doc/library/multiprocessing.rst | 11 +++- Doc/library/os.rst | 7 +++ Doc/using/cmdline.rst | 18 ++++++ Doc/whatsnew/3.13.rst | 6 ++ Include/cpython/initconfig.h | 2 + Lib/os.py | 2 +- Lib/test/test_cmd_line.py | 25 ++++++-- Lib/test/test_embed.py | 2 + ...-09-22-01-44-53.gh-issue-109595.fVINgD.rst | 5 ++ Modules/posixmodule.c | 8 ++- Programs/_testembed.c | 1 + Python/clinic/sysmodule.c.h | 30 +++++++++- Python/initconfig.c | 58 ++++++++++++++++++- Python/sysmodule.c | 15 +++++ 15 files changed, 192 insertions(+), 11 deletions(-) create mode 100644 Misc/NEWS.d/next/Core and Builtins/2023-09-22-01-44-53.gh-issue-109595.fVINgD.rst diff --git a/Doc/c-api/init_config.rst b/Doc/c-api/init_config.rst index 56b4ee03c1a..0240e25b6f1 100644 --- a/Doc/c-api/init_config.rst +++ b/Doc/c-api/init_config.rst @@ -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: diff --git a/Doc/library/multiprocessing.rst b/Doc/library/multiprocessing.rst index 31710f6ff11..3c87bbe8e59 100644 --- a/Doc/library/multiprocessing.rst +++ b/Doc/library/multiprocessing.rst @@ -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() diff --git a/Doc/library/os.rst b/Doc/library/os.rst index a1595dfbc06..4d1881f3c01 100644 --- a/Doc/library/os.rst +++ b/Doc/library/os.rst @@ -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 diff --git a/Doc/using/cmdline.rst b/Doc/using/cmdline.rst index f68a2251f06..2767b0cb154 100644 --- a/Doc/using/cmdline.rst +++ b/Doc/using/cmdline.rst @@ -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 ~~~~~~~~~~~~~~~~~~~~ diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst index 2fe1494a3d7..9a24c1fabf0 100644 --- a/Doc/whatsnew/3.13.rst +++ b/Doc/whatsnew/3.13.rst @@ -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 ------- diff --git a/Include/cpython/initconfig.h b/Include/cpython/initconfig.h index 5d7b4e2d929..808c1056498 100644 --- a/Include/cpython/initconfig.h +++ b/Include/cpython/initconfig.h @@ -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; diff --git a/Lib/os.py b/Lib/os.py index 35842cedf14..a17946750ea 100644 --- a/Lib/os.py +++ b/Lib/os.py @@ -1138,7 +1138,7 @@ def add_dll_directory(path): ) -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. diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py index f4754dbf735..eaf19aa160e 100644 --- a/Lib/test/test_cmd_line.py +++ b/Lib/test/test_cmd_line.py @@ -878,11 +878,8 @@ def test_int_max_str_digits(self): 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 @@ def res2int(res): ) 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.') diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py index 06f2d8b9a36..5a8690a4836 100644 --- a/Lib/test/test_embed.py +++ b/Lib/test/test_embed.py @@ -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 @@ def test_init_from_config(self): '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, diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-09-22-01-44-53.gh-issue-109595.fVINgD.rst b/Misc/NEWS.d/next/Core and Builtins/2023-09-22-01-44-53.gh-issue-109595.fVINgD.rst new file mode 100644 index 00000000000..f182f965834 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2023-09-22-01-44-53.gh-issue-109595.fVINgD.rst @@ -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. diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c index 0975ef71d44..650ae4bbd68 100644 --- a/Modules/posixmodule.c +++ b/Modules/posixmodule.c @@ -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); diff --git a/Programs/_testembed.c b/Programs/_testembed.c index e66c5181822..1f9aa4b3d44 100644 --- a/Programs/_testembed.c +++ b/Programs/_testembed.c @@ -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); diff --git a/Python/clinic/sysmodule.c.h b/Python/clinic/sysmodule.c.h index 30691c3d08a..06105e221c1 100644 --- a/Python/clinic/sysmodule.c.h +++ b/Python/clinic/sysmodule.c.h @@ -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]*/ diff --git a/Python/initconfig.c b/Python/initconfig.c index 6b76b4dc681..f7eb8535e98 100644 --- a/Python/initconfig.c +++ b/Python/initconfig.c @@ -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)) { diff --git a/Python/sysmodule.c b/Python/sysmodule.c index a7ce07d28ae..3debe7f7c13 100644 --- a/Python/sysmodule.c +++ b/Python/sysmodule.c @@ -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 };