diff --git a/Doc/c-api/init_config.rst b/Doc/c-api/init_config.rst index 0240e25b6f1607..fcc87a13471aea 100644 --- a/Doc/c-api/init_config.rst +++ b/Doc/c-api/init_config.rst @@ -884,7 +884,7 @@ PyConfig 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 + Configured by the :samp:`-X cpu_count={n|default|process}` command line flag or the :envvar:`PYTHON_CPU_COUNT` environment variable. Default: ``-1``. diff --git a/Doc/using/cmdline.rst b/Doc/using/cmdline.rst index 2767b0cb15451c..4dcd99b17908ac 100644 --- a/Doc/using/cmdline.rst +++ b/Doc/using/cmdline.rst @@ -551,7 +551,8 @@ Miscellaneous options *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. + If *n* is ``default``, nothing is overridden. If *n* is ``process``, + :func:`os.cpu_count` becomes an alias of :func:`os.process_cpu_count` It also allows passing arbitrary values and retrieving them through the :data:`sys._xoptions` dictionary. @@ -1076,6 +1077,8 @@ conflict. 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`. + If this variable is set to ``process``, :func:`os.cpu_count` becomes + an alias of :func:`os.process_cpu_count`. See also the :option:`-X cpu_count <-X>` command-line option. diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst index 9a24c1fabf05d5..35627be28f1b39 100644 --- a/Doc/whatsnew/3.13.rst +++ b/Doc/whatsnew/3.13.rst @@ -192,7 +192,7 @@ os 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`) + (Contributed by Donghee Na in :gh:`109595` and :gh:`110649`) pathlib ------- diff --git a/Lib/os.py b/Lib/os.py index a17946750ea7e7..68eefbaa7e687f 100644 --- a/Lib/os.py +++ b/Lib/os.py @@ -1147,6 +1147,9 @@ def process_cpu_count(): current process. Return None if indeterminable. """ return len(sched_getaffinity(0)) + if sys._get_cpu_count_config() == -2: + # Just an alias to process_count() (same docstring) + cpu_count = process_cpu_count else: # Just an alias to cpu_count() (same docstring) process_cpu_count = cpu_count diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py index eaf19aa160e860..08896face20e71 100644 --- a/Lib/test/test_cmd_line.py +++ b/Lib/test/test_cmd_line.py @@ -915,6 +915,15 @@ def test_cpu_count_default(self): es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='default') self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count())) + def test_cpu_count_process(self): + code = "import os; print(os.cpu_count(), os.process_cpu_count())" + res = assert_python_ok('-X', 'cpu_count=process', '-c', code) + self.assertEqual(self.res2int(res), (os.process_cpu_count(), os.process_cpu_count())) + res = assert_python_ok('-X', 'cpu_count=default', '-c', code, PYTHON_CPU_COUNT='process') + self.assertEqual(self.res2int(res), (os.cpu_count(), os.process_cpu_count())) + es = assert_python_ok('-c', code, PYTHON_CPU_COUNT='process') + self.assertEqual(self.res2int(res), (os.process_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()) diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-10-10-23-37-42.gh-issue-110649.jPp6s9.rst b/Misc/NEWS.d/next/Core and Builtins/2023-10-10-23-37-42.gh-issue-110649.jPp6s9.rst new file mode 100644 index 00000000000000..2d29160b6d69cf --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2023-10-10-23-37-42.gh-issue-110649.jPp6s9.rst @@ -0,0 +1,3 @@ +Update :option:`-X cpu_count <-X>` command line option to support +``process`` mode that :func:`os.cpu_count` becomes an alias of +:func:`os.process_cpu_count`. Patch by Donghee Na. diff --git a/Python/initconfig.c b/Python/initconfig.c index f7eb8535e98a6a..d3d18ff7965f46 100644 --- a/Python/initconfig.c +++ b/Python/initconfig.c @@ -232,7 +232,7 @@ The following implementation-specific options are available:\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.\n\ \n\ --X cpu_count=[n|default]: Override the return value of os.cpu_count(),\n\ +-X cpu_count=[n|default|process]: 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." @@ -1636,6 +1636,9 @@ config_init_cpu_count(PyConfig *config) if (strcmp(env, "default") == 0) { cpu_count = -1; } + else if(strcmp(env, "process") == 0) { + cpu_count = -2; + } else if (_Py_str_to_int(env, &cpu_count) < 0 || cpu_count < 1) { goto error; } @@ -1650,6 +1653,9 @@ config_init_cpu_count(PyConfig *config) if (wcscmp(sep + 1, L"default") == 0) { cpu_count = -1; } + else if (wcscmp(sep + 1, L"process") == 0) { + cpu_count = -2; + } else if (config_wstr_to_int(sep + 1, &cpu_count) < 0 || cpu_count < 1) { goto error; }