Repository: pex-tool/pex

Issue:
The `--venv` mode `pex` script does not have a `__name__ == '__main__'` guard breaking multiprocessing.
The multiprocessing module relies on being able to import the file run as `__main__` from spawned subprocesses without the side-effect of actually running `__main__` code. See "Safe importing of main module" here: https://docs.python.org/3/library/multiprocessing.html#the-spawn-and-forkserver-start-methods

To play well with multiprocessing, the `pex` script should get a `__main__` guard.

This was spotted over in #1235.


Discussion:

Resolving patch:
diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py
--- a/pex/tools/commands/venv.py
+++ b/pex/tools/commands/venv.py
@@ -140,108 +140,109 @@ def record_provenance(src_to_dst):
         """\
         #!{venv_python} -sE
 
-        import os
-        import sys
-
-        python = {venv_python!r}
-        if sys.executable != python:
-            sys.stderr.write("Re-execing from {{}}\\n".format(sys.executable))
-            os.execv(python, [python, "-sE"] + sys.argv)
-
-        os.environ["VIRTUAL_ENV"] = {venv_dir!r}
-        sys.path.extend(os.environ.get("PEX_EXTRA_SYS_PATH", "").split(os.pathsep))
-
-        bin_dir = {venv_bin_dir!r}
-        bin_path = os.environ.get("PEX_VENV_BIN_PATH", {bin_path!r})
-        if bin_path != "false":
-            PATH = os.environ.get("PATH", "").split(os.pathsep)
-            if bin_path == "prepend":
-                PATH.insert(0, bin_dir)
-            elif bin_path == "append":
-                PATH.append(bin_dir)
-            else:
+        if __name__ == "__main__":
+            import os
+            import sys
+
+            python = {venv_python!r}
+            if sys.executable != python:
+                sys.stderr.write("Re-execing from {{}}\\n".format(sys.executable))
+                os.execv(python, [python, "-sE"] + sys.argv)
+
+            os.environ["VIRTUAL_ENV"] = {venv_dir!r}
+            sys.path.extend(os.environ.get("PEX_EXTRA_SYS_PATH", "").split(os.pathsep))
+
+            bin_dir = {venv_bin_dir!r}
+            bin_path = os.environ.get("PEX_VENV_BIN_PATH", {bin_path!r})
+            if bin_path != "false":
+                PATH = os.environ.get("PATH", "").split(os.pathsep)
+                if bin_path == "prepend":
+                    PATH.insert(0, bin_dir)
+                elif bin_path == "append":
+                    PATH.append(bin_dir)
+                else:
+                    sys.stderr.write(
+                        "PEX_VENV_BIN_PATH must be one of 'false', 'prepend' or 'append', given: "
+                        "{{!r}}\\n".format(
+                            bin_path
+                        )
+                    )
+                    sys.exit(1)
+                os.environ["PATH"] = os.pathsep.join(PATH)
+
+            PEX_EXEC_OVERRIDE_KEYS = ("PEX_INTERPRETER", "PEX_SCRIPT", "PEX_MODULE")
+            pex_overrides = {{
+                key: os.environ.pop(key) for key in PEX_EXEC_OVERRIDE_KEYS if key in os.environ
+            }}
+            if len(pex_overrides) > 1:
                 sys.stderr.write(
-                    "PEX_VENV_BIN_PATH must be one of 'false', 'prepend' or 'append', given: "
-                    "{{!r}}\\n".format(
-                        bin_path
+                    "Can only specify one of {{overrides}}; found: {{found}}\\n".format(
+                        overrides=", ".join(PEX_EXEC_OVERRIDE_KEYS),
+                        found=" ".join("{{}}={{}}".format(k, v) for k, v in pex_overrides.items())
                     )
                 )
                 sys.exit(1)
-            os.environ["PATH"] = os.pathsep.join(PATH)
-
-        PEX_EXEC_OVERRIDE_KEYS = ("PEX_INTERPRETER", "PEX_SCRIPT", "PEX_MODULE")
-        pex_overrides = {{
-            key: os.environ.pop(key) for key in PEX_EXEC_OVERRIDE_KEYS if key in os.environ
-        }}
-        if len(pex_overrides) > 1:
-            sys.stderr.write(
-                "Can only specify one of {{overrides}}; found: {{found}}\\n".format(
-                    overrides=", ".join(PEX_EXEC_OVERRIDE_KEYS),
-                    found=" ".join("{{}}={{}}".format(k, v) for k, v in pex_overrides.items())
-                )
+
+            pex_script = pex_overrides.get("PEX_SCRIPT")
+            if pex_script:
+                script_path = os.path.join(bin_dir, pex_script)
+                os.execv(script_path, [script_path] + sys.argv[1:])
+
+            pex_interpreter = pex_overrides.get("PEX_INTERPRETER", "").lower() in ("1", "true")
+            PEX_INTERPRETER_ENTRYPOINT = "code:interact"
+            entry_point = (
+                PEX_INTERPRETER_ENTRYPOINT
+                if pex_interpreter
+                else pex_overrides.get("PEX_MODULE", {entry_point!r} or PEX_INTERPRETER_ENTRYPOINT)
             )
-            sys.exit(1)
-
-        pex_script = pex_overrides.get("PEX_SCRIPT")
-        if pex_script:
-            script_path = os.path.join(bin_dir, pex_script)
-            os.execv(script_path, [script_path] + sys.argv[1:])
-
-        pex_interpreter = pex_overrides.get("PEX_INTERPRETER", "").lower() in ("1", "true")
-        PEX_INTERPRETER_ENTRYPOINT = "code:interact"
-        entry_point = (
-            PEX_INTERPRETER_ENTRYPOINT
-            if pex_interpreter
-            else pex_overrides.get("PEX_MODULE", {entry_point!r} or PEX_INTERPRETER_ENTRYPOINT)
-        )
-        if entry_point == PEX_INTERPRETER_ENTRYPOINT and len(sys.argv) > 1:
-            args = sys.argv[1:]
-            arg = args[0]
-            if arg == "-m":
-                if len(args) < 2:
-                    sys.stderr.write("Argument expected for the -m option\\n")
-                    sys.exit(2)
-                entry_point = module = args[1]
-                sys.argv = args[1:]
-                # Fall through to entry_point handling below.
-            else:
-                filename = arg
-                sys.argv = args
-                if arg == "-c":
+            if entry_point == PEX_INTERPRETER_ENTRYPOINT and len(sys.argv) > 1:
+                args = sys.argv[1:]
+                arg = args[0]
+                if arg == "-m":
                     if len(args) < 2:
-                        sys.stderr.write("Argument expected for the -c option\\n")
+                        sys.stderr.write("Argument expected for the -m option\\n")
                         sys.exit(2)
-                    filename = "-c <cmd>"
-                    content = args[1]
-                    sys.argv = ["-c"] + args[2:]
-                elif arg == "-":
-                    content = sys.stdin.read()
+                    entry_point = module = args[1]
+                    sys.argv = args[1:]
+                    # Fall through to entry_point handling below.
                 else:
-                    with open(arg) as fp:
-                        content = fp.read()
-
-                ast = compile(content, filename, "exec", flags=0, dont_inherit=1)
-                globals_map = globals().copy()
-                globals_map["__name__"] = "__main__"
-                globals_map["__file__"] = filename
-                locals_map = globals_map
-                {exec_ast}
-                sys.exit(0)
-
-        module_name, _, function = entry_point.partition(":")
-        if not function:
-            import runpy
-            runpy.run_module(module_name, run_name="__main__")
-        else:
-            import importlib
-            module = importlib.import_module(module_name)
-            # N.B.: Functions may be hung off top-level objects in the module namespace,
-            # e.g.: Class.method; so we drill down through any attributes to the final function
-            # object.
-            namespace, func = module, None
-            for attr in function.split("."):
-                func = namespace = getattr(namespace, attr)
-            func()
+                    filename = arg
+                    sys.argv = args
+                    if arg == "-c":
+                        if len(args) < 2:
+                            sys.stderr.write("Argument expected for the -c option\\n")
+                            sys.exit(2)
+                        filename = "-c <cmd>"
+                        content = args[1]
+                        sys.argv = ["-c"] + args[2:]
+                    elif arg == "-":
+                        content = sys.stdin.read()
+                    else:
+                        with open(arg) as fp:
+                            content = fp.read()
+
+                    ast = compile(content, filename, "exec", flags=0, dont_inherit=1)
+                    globals_map = globals().copy()
+                    globals_map["__name__"] = "__main__"
+                    globals_map["__file__"] = filename
+                    locals_map = globals_map
+                    {exec_ast}
+                    sys.exit(0)
+
+            module_name, _, function = entry_point.partition(":")
+            if not function:
+                import runpy
+                runpy.run_module(module_name, run_name="__main__")
+            else:
+                import importlib
+                module = importlib.import_module(module_name)
+                # N.B.: Functions may be hung off top-level objects in the module namespace,
+                # e.g.: Class.method; so we drill down through any attributes to the final function
+                # object.
+                namespace, func = module, None
+                for attr in function.split("."):
+                    func = namespace = getattr(namespace, attr)
+                func()
         """.format(
             venv_python=venv_python,
             venv_bin_dir=venv_bin_dir,


Test patch:
diff --git a/tests/tools/commands/test_venv.py b/tests/tools/commands/test_venv.py
--- a/tests/tools/commands/test_venv.py
+++ b/tests/tools/commands/test_venv.py
@@ -3,6 +3,7 @@
 
 from __future__ import absolute_import
 
+import multiprocessing
 import os
 import subprocess
 import tempfile
@@ -11,7 +12,7 @@
 
 import pytest
 
-from pex.common import temporary_dir, touch
+from pex.common import safe_open, temporary_dir, touch
 from pex.executor import Executor
 from pex.testing import run_pex_command
 from pex.tools.commands.virtualenv import Virtualenv
@@ -289,3 +290,54 @@ def test_venv_pex_interpreter_special_modes(create_pex_venv):
         )
         assert 0 == returncode, stderr
         assert expected_fabric_file_path == stdout.strip()
+
+
+@pytest.mark.parametrize(
+    "start_method", getattr(multiprocessing, "get_all_start_methods", lambda: [None])()
+)
+def test_venv_multiprocessing_issues_1236(
+    tmpdir,  # type: Any
+    start_method,  # type: Optional[str]
+):
+    # type: (...) -> None
+    src = os.path.join(str(tmpdir), "src")
+    with safe_open(os.path.join(src, "foo.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                def bar():
+                    print('hello')
+                """
+            )
+        )
+    with safe_open(os.path.join(src, "main.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                import multiprocessing
+                from foo import bar
+
+                if __name__ == '__main__':
+                    if {start_method!r}:
+                        multiprocessing.set_start_method({start_method!r})
+                    p = multiprocessing.Process(target=bar)
+                    p.start()
+                """.format(
+                    start_method=start_method
+                )
+            )
+        )
+
+    pex_file = os.path.join(str(tmpdir), "mp.pex")
+    result = run_pex_command(args=["-D", src, "-m", "main", "-o", pex_file, "--include-tools"])
+    result.assert_success()
+
+    # Confirm multiprocessing works via normal PEX file execution.
+    output = subprocess.check_output(args=[pex_file])
+    assert "hello" == output.decode("utf-8").strip()
+
+    # Confirm multiprocessing works via the `pex` venv script.
+    venv = os.path.join(str(tmpdir), "venv")
+    subprocess.check_call(args=[pex_file, "venv", venv], env=make_env(PEX_TOOLS=True))
+    output = subprocess.check_output(args=[os.path.join(venv, "pex")])
+    assert "hello" == output.decode("utf-8").strip()