From 0eb1957999ccd9e54685f0e69c3a99df908c3e12 Mon Sep 17 00:00:00 2001
From: Tom Scogland <scogland1@llnl.gov>
Date: Thu, 21 Mar 2024 01:32:28 -0700
Subject: cmd/python: use runpy to allow multiprocessing in scripts (#41789)
Running a `spack-python` script like this:
```python
import spack
import multiprocessing
def echo(args):
print(args)
if __name__ == "__main__":
pool = multiprocessing.Pool(2)
pool.map(echo, range(10))
```
will fail in `develop` with an error like this:
```console
_pickle.PicklingError: Can't pickle <function echo at 0x104865820>: attribute lookup echo on __main__ failed
```
Python expects to be able to look up the method `echo` in `sys.path["__main__"]` in
subprocesses spawned by `multiprocessing`, but because we use `InteractiveConsole` to
run `spack python`, the executed file isn't considered to be the `__main__` module, and
lookups in subprocesses fail. We tried to fake this by setting `__name__` to `__main__`
in the `spack python` command, but that doesn't fix the fact that no `__main__` module
exists.
Another annoyance with `InteractiveConsole` is that `__file__` is not defined in the
main script scope, so you can't use it in your scripts.
We can use the [runpy.run_path()](https://docs.python.org/3/library/runpy.html#runpy.run_path) function,
which has been around since Python 3.2, to fix this.
- [x] Use `runpy` module to launch non-interactive `spack python` invocations
- [x] Only use `InteractiveConsole` for interactive `spack python`
---
lib/spack/spack/cmd/python.py | 57 ++++++++++++++++++++--------------------
lib/spack/spack/cmd/unit_test.py | 17 ++++++++++++
2 files changed, 45 insertions(+), 29 deletions(-)
(limited to 'lib')
diff --git a/lib/spack/spack/cmd/python.py b/lib/spack/spack/cmd/python.py
index 71ce88eed6..a4f177fa38 100644
--- a/lib/spack/spack/cmd/python.py
+++ b/lib/spack/spack/cmd/python.py
@@ -116,39 +116,38 @@ def ipython_interpreter(args):
def python_interpreter(args):
"""A python interpreter is the default interpreter"""
- # Fake a main python shell by setting __name__ to __main__.
- console = code.InteractiveConsole({"__name__": "__main__", "spack": spack})
- if "PYTHONSTARTUP" in os.environ:
- startup_file = os.environ["PYTHONSTARTUP"]
- if os.path.isfile(startup_file):
- with open(startup_file) as startup:
- console.runsource(startup.read(), startup_file, "exec")
- if args.python_command:
- propagate_exceptions_from(console)
- console.runsource(args.python_command)
- elif args.python_args:
- propagate_exceptions_from(console)
+ if args.python_args and not args.python_command:
sys.argv = args.python_args
- with open(args.python_args[0]) as file:
- console.runsource(file.read(), args.python_args[0], "exec")
+ runpy.run_path(args.python_args[0], run_name="__main__")
else:
- # Provides readline support, allowing user to use arrow keys
- console.push("import readline")
- # Provide tabcompletion
- console.push("from rlcompleter import Completer")
- console.push("readline.set_completer(Completer(locals()).complete)")
- console.push('readline.parse_and_bind("tab: complete")')
-
- console.interact(
- "Spack version %s\nPython %s, %s %s"
- % (
- spack.spack_version,
- platform.python_version(),
- platform.system(),
- platform.machine(),
+ # Fake a main python shell by setting __name__ to __main__.
+ console = code.InteractiveConsole({"__name__": "__main__", "spack": spack})
+ if "PYTHONSTARTUP" in os.environ:
+ startup_file = os.environ["PYTHONSTARTUP"]
+ if os.path.isfile(startup_file):
+ with open(startup_file) as startup:
+ console.runsource(startup.read(), startup_file, "exec")
+ if args.python_command:
+ propagate_exceptions_from(console)
+ console.runsource(args.python_command)
+ else:
+ # Provides readline support, allowing user to use arrow keys
+ console.push("import readline")
+ # Provide tabcompletion
+ console.push("from rlcompleter import Completer")
+ console.push("readline.set_completer(Completer(locals()).complete)")
+ console.push('readline.parse_and_bind("tab: complete")')
+
+ console.interact(
+ "Spack version %s\nPython %s, %s %s"
+ % (
+ spack.spack_version,
+ platform.python_version(),
+ platform.system(),
+ platform.machine(),
+ )
)
- )
def propagate_exceptions_from(console):
diff --git a/lib/spack/spack/cmd/unit_test.py b/lib/spack/spack/cmd/unit_test.py
index 2931be5e74..db0c7ff0e5 100644
--- a/lib/spack/spack/cmd/unit_test.py
+++ b/lib/spack/spack/cmd/unit_test.py
@@ -34,6 +34,13 @@ def setup_parser(subparser):
default=False,
help="show full pytest help, with advanced options",
)
+ subparser.add_argument(
+ "-n",
+ "--numprocesses",
+ type=int,
+ default=1,
+ help="run tests in parallel up to this wide, default 1 for sequential",
+ )
# extra spack arguments to list tests
list_group = subparser.add_argument_group("listing tests")
@@ -229,6 +236,16 @@ def unit_test(parser, args, unknown_args):
if args.extension:
pytest_root = spack.extensions.load_extension(args.extension)
+ if args.numprocesses is not None and args.numprocesses > 1:
+ pytest_args.extend(
+ [
+ "--dist",
+ "loadfile",
+ "--tx",
+ f"{args.numprocesses}*popen//python=spack-tmpconfig spack python",
+ ]
+ )
+
# pytest.ini lives in the root of the spack repository.
with llnl.util.filesystem.working_dir(pytest_root):
if args.list:
--
cgit v1.2.3-70-g09d2