gh-109595: Add -Xcpu_count=<n> cmdline for container users

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 value of :func:`os.cpu_count` and :func:`os.process_cpu_count` functions
+      into *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:

Doc/library/os.rst

@@ -5195,6 +5195,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()
 
@@ -5214,6 +5218,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

Doc/using/cmdline.rst

@@ -546,6 +546,11 @@ 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` and :func:`os.process_cpu_count`.
+     *n* must be greater than or equal to 1.
+     This option is useful for users who need to limit CPU resources of a container system.
+     See also :envvar:`PYTHON_CPU_COUNT`.
+     If *n* is ``default``, :func:`os.cpu_count` and  :func:`os.process_cpu_count` are not overridden.
 
    It also allows passing arbitrary values and retrieving them through the
    :data:`sys._xoptions` dictionary.
@@ -593,6 +598,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 +1071,15 @@ conflict.
 
    .. versionadded:: 3.12
 
+.. envvar:: PYTHON_CPU_COUNT
+
+   If this variable is set to a positive integer, it overrides
+   :func:`os.cpu_count` and and :func:`os.process_cpu_count` return result.
+
+   See also the :option:`-X cpu_count <-X>` command-line option.
+
+   .. versionadded:: 3.13
+
 
 Debug-mode variables
 ~~~~~~~~~~~~~~~~~~~~

Doc/whatsnew/3.13.rst

@@ -170,6 +170,12 @@ os
   usable by the calling thread of the current process.
   (Contributed by Victor Stinner in :gh:`109649`.)
 
+* :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
 -------
 

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;

Lib/os.py

@@ -55,6 +55,7 @@ def _get_exports_list(module):
     from posix import *
     try:
         from posix import _exit
+        from posix import _get_cpu_count_config
         __all__.append('_exit')
     except ImportError:
         pass
@@ -1138,7 +1139,7 @@ def add_dll_directory(path):
         )
 
 
-if _exists('sched_getaffinity'):
+if _exists('sched_getaffinity') and _get_cpu_count_config() != 'overridden':
     def process_cpu_count():
         """
         Get the number of CPUs of the current process.

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.')

Lib/test/test_embed.py

@@ -448,6 +448,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,
@@ -895,6 +896,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,

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.

Modules/posixmodule.c

@@ -14335,6 +14335,22 @@ os_get_terminal_size_impl(PyObject *module, int fd)
 }
 #endif /* defined(TERMSIZE_USE_CONIO) || defined(TERMSIZE_USE_IOCTL) */
 
+/*[clinic input]
+os._get_cpu_count_config
+
+Private function for get PyConfig.cpu_count
+[clinic start generated code]*/
+
+static PyObject *
+os__get_cpu_count_config_impl(PyObject *module)
+/*[clinic end generated code: output=3b0ba7e445c9c6b2 input=c52e5bc863732c21]*/
+{
+    const PyConfig *config = _Py_GetConfig();
+    if (config->cpu_count > 0) {
+        return PyUnicode_FromString("overridden");
+    }
+    return PyUnicode_FromString("default");
+}
 
 /*[clinic input]
 os.cpu_count
@@ -14348,7 +14364,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);
@@ -15977,6 +15998,7 @@ static PyMethodDef posix_methods[] = {
     OS_LISTXATTR_METHODDEF
 
     OS_GET_TERMINAL_SIZE_METHODDEF
+    OS__GET_CPU_COUNT_CONFIG_METHODDEF
     OS_CPU_COUNT_METHODDEF
     OS_GET_INHERITABLE_METHODDEF
     OS_SET_INHERITABLE_METHODDEF

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);
 

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),
@@ -228,7 +229,10 @@ 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 CPU count of os.cpu_count() and os.process_cpu_count().\n\
+    This helps for users who need to limit CPU resources of a container system."
 
 #ifdef Py_STATS
 "\n\
@@ -731,6 +735,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
@@ -830,6 +836,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;
 }
 
 
@@ -1615,6 +1622,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)
 {
@@ -1797,6 +1843,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)) {