python · corona10 · Oct 10, 2023 · Sep 21, 2023 · Sep 21, 2023 · Sep 21, 2023
diff --git a/Doc/library/os.rst b/Doc/library/os.rst
@@ -5164,6 +5164,10 @@ Miscellaneous System Information
 
    .. versionadded:: 3.4
 
+   .. versionchanged:: 3.13
+      if :samp:`-X cpu_count={n}` is given, :func:`cpu_count` will return the overrided
+      value *n*.
+
 
 .. function:: getloadavg()
 

diff --git a/Doc/using/cmdline.rst b/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}` will override the number of CPU count from system.
+     If this option is provided, :func:`os.cpu_count` will return the overrided value.
+     And *n* must be greater equal then 1.
-   * :samp:`-X cpu_count={n}` will override the number of CPU count from system.
-     If this option is provided, :func:`os.cpu_count` will return the overrided value.
-     And *n* must be greater equal then 1.
+   * :samp:`-X cpu_count={n}` overrides the number of CPU count from system:
+     :func:`os.cpu_count` returns *n*. The value *n* must be greater than
+     or equal to 1.
-   * :samp:`-X cpu_count={n}` will override the number of CPU count from system.
-     If this option is provided, :func:`os.cpu_count` will return the overrided value.
-     And *n* must be greater equal then 1.
+   * :samp:`-X cpu_count={n}` overrides the number of CPU count from system:
+     :func:`os.cpu_count` returns *n*. The value *n* must be greater than
+     or equal to 1.
+     This option will be useful for users who need to limit CPU resources of a container system.
+
 
    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
 ~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -180,6 +180,9 @@ typedef struct PyConfig {
     int safe_path;
     int int_max_str_digits;
 
+    // Python 3.13
+    int cpu_count;
+
     /* --- Path configuration inputs ------------ */
     int pathconfig_warnings;
     wchar_t *program_name;

diff --git a/Lib/test/test_cmd_line.py b/Lib/test/test_cmd_line.py
@@ -901,6 +901,11 @@ def res2int(res):
         )
         self.assertEqual(res2int(res), (6000, 6000))
 
+    def test_cpu_count(self):
+        code = "import os; print(os.cpu_count())"
+        res = assert_python_ok('-X', 'cpu_count=4321', '-c', code)
+        self.assertEqual(res.out.strip().decode("utf-8"), '4321')
+
 
 @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
@@ -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,
@@ -894,6 +895,7 @@ def test_init_from_config(self):
             'module_search_paths': self.IGNORE_CONFIG,
             'safe_path': 1,
             'int_max_str_digits': 31337,
+            'cpu_count': -1,
 
             'check_hash_pycs_mode': 'always',
             'pathconfig_warnings': 0,
@@ -966,6 +968,7 @@ def test_init_python_env(self):
             'module_search_paths': self.IGNORE_CONFIG,
             'safe_path': 1,
             'int_max_str_digits': 4567,
+            'cpu_count': -1,
         }
         if Py_STATS:
             config['_pystats'] = 1

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
@@ -0,0 +1,2 @@
+Add :samp:`-Xcpu_count={n}` command line option for overriding system CPU resources from
+Python program. Patch by Donghee Na.
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -14349,6 +14349,10 @@ static PyObject *
 os_cpu_count_impl(PyObject *module)
 /*[clinic end generated code: output=5fc29463c3936a9c input=e7c8f4ba6dbbadd3]*/
 {
+    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

diff --git 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 = -1;
 
     init_from_config_clear(&config);
 

@@ -130,7 +130,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: override CPU count of os.cpu_count().\n\
+    This helps for users who need to limit CPU resources of a container system."
 
 #ifdef Py_STATS
 "\n\
@@ -633,6 +636,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
@@ -732,6 +737,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;
 }
 
 
@@ -959,6 +965,7 @@ _PyConfig_Copy(PyConfig *config, const PyConfig *config2)
     COPY_WSTRLIST(orig_argv);
     COPY_ATTR(_is_python_build);
     COPY_ATTR(int_max_str_digits);
+    COPY_ATTR(cpu_count);
 #ifdef Py_STATS
     COPY_ATTR(_pystats);
 #endif
@@ -1069,6 +1076,7 @@ _PyConfig_AsDict(const PyConfig *config)
     SET_ITEM_INT(safe_path);
     SET_ITEM_INT(_is_python_build);
     SET_ITEM_INT(int_max_str_digits);
+    SET_ITEM_INT(cpu_count);
 #ifdef Py_STATS
     SET_ITEM_INT(_pystats);
 #endif
@@ -1379,6 +1387,7 @@ _PyConfig_FromDict(PyConfig *config, PyObject *dict)
     GET_UINT(safe_path);
     GET_UINT(_is_python_build);
     GET_INT(int_max_str_digits);
+    GET_INT(cpu_count);
 #ifdef Py_STATS
     GET_UINT(_pystats);
 #endif
@@ -1678,6 +1687,29 @@ config_read_env_vars(PyConfig *config)
     return _PyStatus_OK();
 }
 
+static PyStatus
+config_init_cpu_count(PyConfig *config)
+{
+    int cpu_count = -1;
+    const wchar_t *xoption = config_get_xoption(config, L"cpu_count");
+    if (xoption) {
+        const wchar_t *sep = wcschr(xoption, L'=');
+        if (sep) {
+            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");
+}
+
 static PyStatus
 config_init_perf_profiling(PyConfig *config)
 {
@@ -1860,6 +1892,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)) {
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Add :samp:`-Xcpu_count={n}` command line option for overriding system CPU resources from
		Python program. Patch by Donghee Na.
Comment thread corona10 marked this conversation as resolved. Outdated