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>

diff --git a/Doc/c-api/init_config.rst b/Doc/c-api/init_config.rst
index 56b4ee03c1a8fd17c8bb9a1719cef94048249e0c..0240e25b6f160735a9bbbf9246f20a8c981d0035 100644 (file)
--- 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 31710f6ff118d199141cd33535737e6551a93d49..3c87bbe8e598094c5d463de318d7f3ec024c2eae 100644 (file)
--- 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 a1595dfbc060f301a7526d09a5a1c8358da48972..4d1881f3c010c703cead71a602dc0c85e19115b8 100644 (file)
--- 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 f68a2251f06d4add945b3366253d64820dacf8bc..2767b0cb15451cfa56b33bb775720c1984c3b2b8 100644 (file)
--- 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 2fe1494a3d770ff308d7037d2aebf02ca91dcf08..9a24c1fabf05d5a384c78bb19567f48d074f0324 100644 (file)
--- 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 5d7b4e2d929e5b595cdd6e0542e02066e324547d..808c1056498b4990de360f309af98b35cd2a29dd 100644 (file)
--- 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 35842cedf14fc79ce12d9864320b2018533fdab1..a17946750ea7e7405e302444b46128c0771607fa 100644 (file)
--- a/Lib/os.py
+++ b/Lib/os.py
@@ -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.

diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
index f4754dbf735a1d9fbf2cceac817d3ef3169a8b5e..eaf19aa160e860b1a92a69ae95e7f923d4c925a7 100644 (file)
--- a/Lib/test/test_cmd_line.py
+++ b/Lib/test/test_cmd_line.py
@@ -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.')

diff --git a/Lib/test/test_embed.py b/Lib/test/test_embed.py
index 06f2d8b9a3621fcfa949efd87b58eed35c62bedb..5a8690a4836dd6daf57de6a7be5bd8b87eda4a2e 100644 (file)
--- 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 @@ 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,

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 (file)
index 0000000..f182f96
--- /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 0975ef71d44be5ec801f901fa3f8850a5122cdf9..650ae4bbd68656fe8b97534c87af3f1546907fb9 100644 (file)
--- 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 e66c51818227c46b70d5f8ac4f48b62769f25dfb..1f9aa4b3d449a109d601366738f682db49b63471 100644 (file)
--- 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 30691c3d08ae67604ec83e0f68d892a3a7700059..06105e221c1f89a27943d457c49ae71bb293572b 100644 (file)
--- 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 6b76b4dc681b74fd8561d37ea3f591ed44d21c06..f7eb8535e98a6a3e141bed29f85b777604c60a8f 100644 (file)
--- 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 a7ce07d28ae7df5dd3642d830ab95ce24adeb4ab..3debe7f7c139c610f6f46cc574cfc13f25410e36 100644 (file)
--- 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
 };