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`

2 files changed, 45 insertions, 29 deletions

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: