pythongh-109595: Add -Xcpu_count=<n> cmdline for container users (python#109667)

---------

Co-authored-by: Victor Stinner <[email protected]>
Co-authored-by: Gregory P. Smith [Google LLC] <[email protected]>

Author: 3 people

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:

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

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

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
 ~~~~~~~~~~~~~~~~~~~~

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

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

@@ -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.

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

@@ -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,

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.