diff --git a/graalpython/com.oracle.graal.python.shell/src/com/oracle/graal/python/shell/GraalPythonMain.java b/graalpython/com.oracle.graal.python.shell/src/com/oracle/graal/python/shell/GraalPythonMain.java index 067e673786..dcc72a50c3 100644 --- a/graalpython/com.oracle.graal.python.shell/src/com/oracle/graal/python/shell/GraalPythonMain.java +++ b/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 givenArguments; private List relaunchArgs; @@ -465,6 +468,10 @@ protected List preprocessArguments(List givenArgs, Map 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 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 directory (or :).\n" + " The default module search path uses /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" + diff --git a/graalpython/com.oracle.graal.python.test/src/tests/test_cmd_line.py b/graalpython/com.oracle.graal.python.test/src/tests/test_cmd_line.py index 642c1153cf..c5219a4b97 100644 --- a/graalpython/com.oracle.graal.python.test/src/tests/test_cmd_line.py +++ b/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) - diff --git a/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_cmd_line.txt b/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_cmd_line.txt index c7c2b78e47..fd5ce4f10a 100644 --- a/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_cmd_line.txt +++ b/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 diff --git a/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_os.txt b/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_os.txt index b0db8da62a..cfaf0bff30 100644 --- a/graalpython/com.oracle.graal.python.test/src/tests/unittest_tags/test_os.txt +++ b/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 diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/PosixModuleBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/PosixModuleBuiltins.java index 8c8eeb7e64..e8d00837c6 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/PosixModuleBuiltins.java +++ b/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(); } } diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/SysModuleBuiltins.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/SysModuleBuiltins.java index f0c5813458..21a1151de8 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/SysModuleBuiltins.java +++ b/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 diff --git a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonOptions.java b/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonOptions.java index ff99173823..0e706cb2b7 100644 --- a/graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonOptions.java +++ b/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 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 DontWriteBytecodeFlag = new OptionKey<>(true);