[GR-75710] Implement -X cpu_count option

graalpython/com.oracle.graal.python.shell/src/com/oracle/graal/python/shell/GraalPythonMain.java

@@ -162,6 +162,9 @@ private static void appendDecodedChars(StringBuilder builder, CharBuffer out) {
     private boolean snaptshotStartup = false;
     private boolean warnDefaultEncoding = false;
     private int intMaxStrDigits = -1;
+    private int cpuCount = -1;
+    private boolean cpuCountXOptionSet = false;
+    private boolean cpuCountSet = false;
     private VersionAction versionAction = VersionAction.None;
     private List<String> givenArguments;
     private List<String> relaunchArgs;
@@ -465,6 +468,10 @@ protected List<String> preprocessArguments(List<String> givenArgs, Map<String, S
                                 if (eq > 0) {
                                     intMaxStrDigits = validateIntMaxStrDigits(xOption.substring(eq), "-X int_max_str_digits");
                                 }
+                            } else if (xOption.equals("cpu_count") || xOption.startsWith("cpu_count=")) {
+                                cpuCount = validateCpuCount(xOption, true);
+                                cpuCountXOptionSet = true;
+                                cpuCountSet = true;
                             } else if ("jit".equals(xOption) || xOption.startsWith("jit=")) {
                                 applyJitModePreset(polyglotOptions, xOption);
                                 jitModePresetSpecified = true;
@@ -792,6 +799,11 @@ protected void launch(Builder contextBuilder) {
             if (intMaxStrDigits < 0 && maxStrDigitsEnv != null) {
                 intMaxStrDigits = validateIntMaxStrDigits(maxStrDigitsEnv, "PYTHONINTMAXSTRDIGITS");
             }
+            String cpuCountEnv = getEnv("PYTHON_CPU_COUNT");
+            if (!cpuCountXOptionSet && cpuCountEnv != null) {
+                cpuCount = validateCpuCount(cpuCountEnv, false);
+                cpuCountSet = true;
+            }
 
             String hashSeed = getEnv("PYTHONHASHSEED");
             if (hashSeed != null) {
@@ -853,6 +865,9 @@ protected void launch(Builder contextBuilder) {
         if (intMaxStrDigits > 0) {
             contextBuilder.option("python.IntMaxStrDigits", Integer.toString(intMaxStrDigits));
         }
+        if (cpuCountSet) {
+            contextBuilder.option("python.CpuCount", Integer.toString(cpuCount));
+        }
         contextBuilder.option("python.DontWriteBytecodeFlag", Boolean.toString(dontWriteBytecode));
         contextBuilder.option("python.QuietFlag", Boolean.toString(quietFlag));
         contextBuilder.option("python.NoUserSiteFlag", Boolean.toString(noUserSite));
@@ -979,6 +994,29 @@ private int validateIntMaxStrDigits(String input, String name) {
         throw abort(String.format("%s: invalid limit; must be >= %d  or 0 for unlimited.", name, INT_MAX_STR_DIGITS_THRESHOLD), 1);
     }
 
+    private int validateCpuCount(String input, boolean xOption) {
+        String value = input;
+        if (xOption) {
+            int eq = input.indexOf('=');
+            if (eq < 0) {
+                throw abort("-X cpu_count=n option: n is missing or an invalid number, n must be greater than 0", 1);
+            }
+            value = input.substring(eq + 1);
+        }
+        if (value.equals("default")) {
+            return -1;
+        }
+        try {
+            int count = Integer.parseInt(value.stripLeading());
+            if (count > 0) {
+                return count;
+            }
+        } catch (NumberFormatException e) {
+            // fallthrough
+        }
+        throw abort("-X cpu_count=n option: n is missing or an invalid number, n must be greater than 0", 1);
+    }
+
     private void applyJitModePreset(Map<String, String> polyglotOptions, String xOption) {
         boolean fallbackRuntime = usesFallbackRuntime();
         switch (xOption) {
@@ -1233,7 +1271,7 @@ protected void printHelp(OptionCategory maxCategory) {
                         "-V     : print the Python version number and exit (also --version)\n" +
                         "         when given twice, print more information about the build\n" +
                         "-X opt : set implementation-specific option\n" +
-                        "         CPython-compatible options supported by GraalPy: warn_default_encoding, int_max_str_digits\n" +
+                        "         CPython-compatible options supported by GraalPy: warn_default_encoding, int_max_str_digits, cpu_count\n" +
                         "         GraalPy implementation-specific options: jit=0|1|2 (default: jit=1)\n" +
                         "-W arg : warning control; arg is action:message:category:module:lineno\n" +
                         "         also PYTHONWARNINGS=arg\n" +
@@ -1247,6 +1285,7 @@ protected void printHelp(OptionCategory maxCategory) {
                         "               default module search path.  The result is sys.path.\n" +
                         "PYTHONHOME   : alternate <prefix> directory (or <prefix>:<exec_prefix>).\n" +
                         "               The default module search path uses <prefix>/pythonX.X.\n" +
+                        "PYTHON_CPU_COUNT: override the return value of os.cpu_count() (-X cpu_count).\n" +
                         "PYTHONCASEOK : ignore case in 'import' statements (Windows).\n" +
                         "PYTHONIOENCODING: Encoding[:errors] used for stdin/stdout/stderr.\n" +
                         "PYTHONHASHSEED: if this variable is set to 'random', the effect is the same\n" +

graalpython/com.oracle.graal.python.test/src/tests/test_cmd_line.py

@@ -37,6 +37,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import os
 import subprocess
 import sys
 import unittest
@@ -46,6 +47,57 @@
 
 class CmdLineTest(unittest.TestCase):
 
+    CPU_COUNT_CODE = (
+        "import os, sys; "
+        "print(sys._get_cpu_count_config(), os.cpu_count(), os.process_cpu_count())"
+    )
+
+    def run_cpu_count(self, *args, cpu_count_env=None):
+        env = os.environ.copy()
+        env.pop('PYTHON_CPU_COUNT', None)
+        if cpu_count_env is not None:
+            env['PYTHON_CPU_COUNT'] = cpu_count_env
+        return subprocess.run(
+            [sys.executable, *args, '-c', self.CPU_COUNT_CODE],
+            capture_output=True,
+            text=True,
+            env=env,
+        )
+
+    def test_cpu_count_overrides(self):
+        result = self.run_cpu_count('-X', 'cpu_count=4321')
+        self.assertEqual(0, result.returncode, result)
+        self.assertEqual('4321 4321 4321', result.stdout.strip())
+
+        result = self.run_cpu_count(cpu_count_env='1234')
+        self.assertEqual(0, result.returncode, result)
+        self.assertEqual('1234 1234 1234', result.stdout.strip())
+
+        result = self.run_cpu_count('-X', 'cpu_count=5678', cpu_count_env='1234')
+        self.assertEqual(0, result.returncode, result)
+        self.assertEqual('5678 5678 5678', result.stdout.strip())
+
+    def test_cpu_count_default_and_ignore_environment(self):
+        for args in (('-X', 'cpu_count=default'), ('-E',)):
+            with self.subTest(args=args):
+                result = self.run_cpu_count(*args, cpu_count_env='1234')
+                self.assertEqual(0, result.returncode, result)
+                config, cpu_count, process_cpu_count = map(int, result.stdout.split())
+                self.assertEqual(-1, config)
+                self.assertGreater(cpu_count, 0)
+                self.assertEqual(cpu_count, process_cpu_count)
+
+    def test_cpu_count_invalid(self):
+        for value in ('cpu_count', 'cpu_count=', 'cpu_count=0', 'cpu_count=-1', 'cpu_count=invalid'):
+            with self.subTest(xoption=value):
+                result = self.run_cpu_count('-X', value)
+                self.assertNotEqual(0, result.returncode)
+
+        for value in ('0', '-1', 'invalid', 'cpu_count=1'):
+            with self.subTest(environment=value):
+                result = self.run_cpu_count(cpu_count_env=value)
+                self.assertNotEqual(0, result.returncode)
+
     def test_stdin_script_exit_code(self):
         code = "import sys\nsys.exit(42)\n"
         result = subprocess.run([sys.executable], input=code, text=True)
@@ -70,4 +122,3 @@ def test_jit_mode_invalid_value(self):
         )
         self.assertNotEqual(0, result.returncode)
         self.assertIn('expected jit=0, jit=1, or jit=2', result.stderr)
-

graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_cmd_line.txt

@@ -2,6 +2,8 @@ test.test_cmd_line.CmdLineTest.test_argv0_normalization @ win32-AMD64,win32-AMD6
 test.test_cmd_line.CmdLineTest.test_builtin_input @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github
 test.test_cmd_line.CmdLineTest.test_closed_stdout @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
 test.test_cmd_line.CmdLineTest.test_coding @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github
+test.test_cmd_line.CmdLineTest.test_cpu_count @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
+test.test_cmd_line.CmdLineTest.test_cpu_count_default @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
 test.test_cmd_line.CmdLineTest.test_del___main__ @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github
 test.test_cmd_line.CmdLineTest.test_directories @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
 test.test_cmd_line.CmdLineTest.test_empty_PYTHONPATH_issue16309 @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github

graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_os.txt

@@ -21,6 +21,7 @@ test.test_os.BytesWalkTests.test_walk_prune @ darwin-arm64,linux-aarch64,linux-a
 test.test_os.BytesWalkTests.test_walk_symlink @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
 test.test_os.BytesWalkTests.test_walk_topdown @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github
 test.test_os.CPUCountTests.test_cpu_count @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
+test.test_os.CPUCountTests.test_process_cpu_count @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github,win32-AMD64,win32-AMD64-github
 test.test_os.ChownFileTests.test_chown_gid @ darwin-arm64,linux-aarch64-github,linux-x86_64-github
 test.test_os.ChownFileTests.test_chown_uid_gid_arguments_must_be_index @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github
 test.test_os.ChownFileTests.test_chown_without_permission @ darwin-arm64,linux-aarch64,linux-aarch64-github,linux-x86_64,linux-x86_64-github

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/PosixModuleBuiltins.java

@@ -2970,7 +2970,11 @@ protected ArgumentClinicProvider getArgumentClinic() {
     abstract static class CpuCountNode extends PythonBuiltinNode {
         @TruffleBoundary
         @Specialization
-        static int getCpuCount() {
+        static int getCpuCount(@Bind PythonContext context) {
+            int configuredCpuCount = context.getOption(PythonOptions.CpuCount);
+            if (configuredCpuCount > 0) {
+                return configuredCpuCount;
+            }
             return Runtime.getRuntime().availableProcessors();
         }
     }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/SysModuleBuiltins.java

@@ -2341,6 +2341,15 @@ int get() {
         }
     }
 
+    @Builtin(name = "_get_cpu_count_config")
+    @GenerateNodeFactory
+    abstract static class GetCpuCountConfig extends PythonBuiltinNode {
+        @Specialization
+        int get() {
+            return getContext().getOption(PythonOptions.CpuCount);
+        }
+    }
+
     @Builtin(name = "set_int_max_str_digits", minNumOfPositionalArgs = 1, parameterNames = {"maxdigits"})
     @ArgumentClinic(name = "maxdigits", conversion = ClinicConversion.Int)
     @GenerateNodeFactory

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonOptions.java

@@ -166,6 +166,20 @@ private PythonOptions() {
                         throw new IllegalArgumentException(String.format("IntMaxStrDigits: invalid limit; must be >= %d or 0 for unlimited.", INT_MAX_STR_DIGITS_THRESHOLD));
                     }));
 
+    @Option(category = OptionCategory.USER, help = "Override the return value of os.cpu_count(). Equivalent to the Python -X cpu_count option.", stability = OptionStability.STABLE) //
+    public static final OptionKey<Integer> CpuCount = new OptionKey<>(-1,
+                    new OptionType<>("CpuCount", (input) -> {
+                        try {
+                            int value = Integer.parseInt(input.stripLeading());
+                            if (value == -1 || value > 0) {
+                                return value;
+                            }
+                        } catch (NumberFormatException e) {
+                            // fallthrough
+                        }
+                        throw new IllegalArgumentException("CpuCount: invalid count; must be greater than 0, or -1 for the default.");
+                    }));
+
     @Option(category = OptionCategory.USER, help = "Equivalent to the Python -B flag. Don't write bytecode files.", usageSyntax = "true|false", stability = OptionStability.STABLE) //
     public static final OptionKey<Boolean> DontWriteBytecodeFlag = new OptionKey<>(true);