FIX support for frozen executable on all platforms

loky/__init__.py

@@ -16,28 +16,34 @@
 
 from ._base import Future
 from .backend.context import cpu_count
+from .backend.spawn import freeze_support
 from .backend.reduction import set_loky_pickler
 from .reusable_executor import get_reusable_executor
 from .cloudpickle_wrapper import wrap_non_picklable_objects
 from .process_executor import BrokenProcessPool, ProcessPoolExecutor
 
 
 __all__ = [
-    "get_reusable_executor",
-    "cpu_count",
-    "wait",
-    "as_completed",
-    "Future",
+    # Constants
+    "ALL_COMPLETED",
+    "FIRST_COMPLETED",
+    "FIRST_EXCEPTION",
+    # Classes
     "Executor",
+    "Future",
     "ProcessPoolExecutor",
+    # Functions
+    "as_completed",
+    "cpu_count",
+    "freeze_support",
+    "get_reusable_executor",
+    "set_loky_pickler",
+    "wait",
+    "wrap_non_picklable_objects",
+    # Errors
     "BrokenProcessPool",
     "CancelledError",
     "TimeoutError",
-    "FIRST_COMPLETED",
-    "FIRST_EXCEPTION",
-    "ALL_COMPLETED",
-    "wrap_non_picklable_objects",
-    "set_loky_pickler",
 ]
 
 

loky/backend/fork_exec.py

@@ -38,6 +38,6 @@ def fork_exec(cmd, keep_fds, env=None):
     pid = os.fork()
     if pid == 0:  # pragma: no cover
         close_fds(keep_fds)
-        os.execve(sys.executable, cmd, child_env)
+        os.execve(cmd[0], cmd, child_env)
     else:
         return pid

loky/backend/popen_loky_posix.py

@@ -6,9 +6,8 @@
 import os
 import sys
 import signal
-import pickle
 from io import BytesIO
-from multiprocessing import util, process
+from multiprocessing import util
 from multiprocessing.connection import wait
 from multiprocessing.context import set_spawning_popen
 
@@ -90,7 +89,6 @@ def terminate(self):
                     raise
 
     def _launch(self, process_obj):
-
         tracker_fd = resource_tracker._resource_tracker.getfd()
 
         fp = BytesIO()
@@ -109,15 +107,13 @@ def _launch(self, process_obj):
         try:
             parent_r, child_w = os.pipe()
             child_r, parent_w = os.pipe()
-            # for fd in self._fds:
-            #     _mk_inheritable(fd)
 
-            cmd_python = [sys.executable]
-            cmd_python += ["-m", self.__module__]
-            cmd_python += ["--process-name", str(process_obj.name)]
-            cmd_python += ["--pipe", str(reduction._mk_inheritable(child_r))]
+            reduction._mk_inheritable(child_r)
             reduction._mk_inheritable(child_w)
             reduction._mk_inheritable(tracker_fd)
+            cmd_python = spawn.get_command_line(
+                fd=child_r, process_name=process_obj.name
+            )
             self._fds += [child_r, child_w, tracker_fd]
             if sys.version_info >= (3, 8) and os.name == "posix":
                 mp_tracker_fd = prep_data["mp_tracker_args"]["fd"]
@@ -147,47 +143,3 @@ def _launch(self, process_obj):
     @staticmethod
     def thread_is_spawning():
         return True
-
-
-if __name__ == "__main__":
-    import argparse
-
-    parser = argparse.ArgumentParser("Command line parser")
-    parser.add_argument(
-        "--pipe", type=int, required=True, help="File handle for the pipe"
-    )
-    parser.add_argument(
-        "--process-name",
-        type=str,
-        default=None,
-        help="Identifier for debugging purpose",
-    )
-
-    args = parser.parse_args()
-
-    info = {}
-    exitcode = 1
-    try:
-        with os.fdopen(args.pipe, "rb") as from_parent:
-            process.current_process()._inheriting = True
-            try:
-                prep_data = pickle.load(from_parent)
-                spawn.prepare(prep_data)
-                process_obj = pickle.load(from_parent)
-            finally:
-                del process.current_process()._inheriting
-
-        exitcode = process_obj._bootstrap()
-    except Exception:
-        print("\n\n" + "-" * 80)
-        print(f"{args.process_name} failed with traceback: ")
-        print("-" * 80)
-        import traceback
-
-        print(traceback.format_exc())
-        print("\n" + "-" * 80)
-    finally:
-        if from_parent is not None:
-            from_parent.close()
-
-        sys.exit(exitcode)

loky/backend/popen_loky_win32.py

@@ -2,8 +2,7 @@
 import sys
 import msvcrt
 import _winapi
-from pickle import load
-from multiprocessing import process, util
+from multiprocessing import util
 from multiprocessing.context import get_spawning_popen, set_spawning_popen
 from multiprocessing.popen_spawn_win32 import Popen as _Popen
 from multiprocessing.reduction import duplicate
@@ -50,11 +49,10 @@ def __init__(self, process_obj):
         rhandle = duplicate(msvcrt.get_osfhandle(rfd), inheritable=True)
         os.close(rfd)
 
-        cmd = get_command_line(parent_pid=os.getpid(), pipe_handle=rhandle)
+        cmd = spawn.get_command_line(fd=rhandle)
+        python_exe = cmd[0]
         cmd = " ".join(f'"{x}"' for x in cmd)
 
-        python_exe = spawn.get_executable()
-
         # copy the environment variables to set in the child process
         child_env = {**os.environ, **process_obj.env}
 
@@ -118,49 +116,3 @@ def __init__(self, process_obj):
     def duplicate_for_child(self, handle):
         assert self is get_spawning_popen()
         return duplicate(handle, self.sentinel)
-
-
-def get_command_line(pipe_handle, **kwds):
-    """Returns prefix of command line used for spawning a child process."""
-    if getattr(sys, "frozen", False):
-        return [sys.executable, "--multiprocessing-fork", pipe_handle]
-    else:
-        prog = "from loky.backend.popen_loky_win32 import main; main()"
-        opts = util._args_from_interpreter_flags()
-        return [
-            spawn.get_executable(),
-            *opts,
-            "-c",
-            prog,
-            "--multiprocessing-fork",
-            pipe_handle,
-        ]
-
-
-def is_forking(argv):
-    """Return whether commandline indicates we are forking."""
-    if len(argv) >= 2 and argv[1] == "--multiprocessing-fork":
-        assert len(argv) == 3
-        return True
-    else:
-        return False
-
-
-def main():
-    """Run code specified by data received over pipe."""
-    assert is_forking(sys.argv)
-
-    handle = int(sys.argv[-1])
-    fd = msvcrt.open_osfhandle(handle, os.O_RDONLY)
-    from_parent = os.fdopen(fd, "rb")
-
-    process.current_process()._inheriting = True
-    preparation_data = load(from_parent)
-    spawn.prepare(preparation_data)
-    self = load(from_parent)
-    process.current_process()._inheriting = False
-
-    from_parent.close()
-
-    exitcode = self._bootstrap()
-    sys.exit(exitcode)

loky/backend/resource_tracker.py

@@ -49,8 +49,8 @@
 import signal
 import warnings
 import threading
-from _multiprocessing import sem_unlink
 from multiprocessing import util
+from _multiprocessing import sem_unlink
 
 from . import spawn
 
@@ -128,13 +128,13 @@ def ensure_running(self):
                 os.close(r)
                 r = _r
 
-            cmd = f"from {main.__module__} import main; main({r}, {VERBOSE})"
             try:
                 fds_to_pass.append(r)
                 # process will out live us, so no need to wait on pid
-                exe = spawn.get_executable()
-                args = [exe, *util._args_from_interpreter_flags(), "-c", cmd]
-                util.debug(f"launching resource tracker: {args}")
+                cmd = spawn.get_command_line(
+                    main_prog=main, fd=r, verbose=int(VERBOSE)
+                )
+                util.debug(f"launching resource tracker: {cmd}")
                 # bpo-33613: Register a signal mask that will block the
                 # signals.  This signal mask will be inherited by the child
                 # that is going to be spawned and will protect the child from a
@@ -146,7 +146,7 @@ def ensure_running(self):
                         signal.pthread_sigmask(
                             signal.SIG_BLOCK, _IGNORED_SIGNALS
                         )
-                    pid = spawnv_passfds(exe, args, fds_to_pass)
+                    pid = spawnv_passfds(cmd, fds_to_pass)
                 finally:
                     if _HAVE_SIGMASK:
                         signal.pthread_sigmask(
@@ -208,6 +208,11 @@ def _send(self, cmd, name, rtype):
 
 def main(fd, verbose=0):
     """Run resource tracker."""
+    # Mak sure the arguments have the right type as theyr are
+    # passed as strings through the command line.
+    fd = int(fd)
+    verbose = int(verbose)
+
     # protect the process from ^C and "killall python" etc
     if verbose:
         util.log_to_stderr(level=util.DEBUG)
@@ -274,8 +279,8 @@ def main(fd, verbose=0):
                         del registry[rtype][name]
                         if verbose:
                             util.debug(
-                                f"[ResourceTracker] unregister {name} {rtype}: "
-                                f"registry({len(registry)})"
+                                f"[ResourceTracker] unregister {name} {rtype}:"
+                                f" registry({len(registry)})"
                             )
                     elif cmd == "MAYBE_UNLINK":
                         registry[rtype][name] -= 1
@@ -353,7 +358,7 @@ def _unlink_resources(rtype_registry, rtype):
 #
 
 
-def spawnv_passfds(path, args, passfds):
+def spawnv_passfds(cmd, passfds):
     passfds = sorted(passfds)
     if sys.platform != "win32":
         errpipe_read, errpipe_write = os.pipe()
@@ -362,15 +367,16 @@ def spawnv_passfds(path, args, passfds):
             from .fork_exec import fork_exec
 
             _pass = [_mk_inheritable(fd) for fd in passfds]
-            return fork_exec(args, _pass)
+            return fork_exec(cmd, _pass)
         finally:
             os.close(errpipe_read)
             os.close(errpipe_write)
     else:
-        cmd = " ".join(f'"{x}"' for x in args)
+        exe = cmd[0]
+        cmd = " ".join(f'"{x}"' for x in cmd)
         try:
             _, ht, pid, _ = _winapi.CreateProcess(
-                path, cmd, None, None, True, 0, None, None, None
+                exe, cmd, None, None, True, 0, None, None, None
             )
             _winapi.CloseHandle(ht)
         except BaseException:

loky/backend/spawn.py

@@ -11,7 +11,11 @@
 import runpy
 import textwrap
 import types
+import pickle
+import importlib
+
 from multiprocessing import process, util
+from multiprocessing import freeze_support as _freeze_support_mp
 
 
 if sys.platform != "win32":
@@ -250,3 +254,72 @@ def _fixup_main_from_path(main_path):
     main_content = runpy.run_path(main_path, run_name="__mp_main__")
     main_module.__dict__.update(main_content)
     sys.modules["__main__"] = sys.modules["__mp_main__"] = main_module
+
+
+def main(fd, process_name=None):
+    fd = int(fd)
+    if sys.platform == "win32":
+        fd = msvcrt.open_osfhandle(fd, os.O_RDONLY)
+
+    exitcode = 1
+    try:
+        with os.fdopen(fd, "rb") as from_parent:
+            process.current_process()._inheriting = True
+            try:
+                prep_data = pickle.load(from_parent)
+                prepare(prep_data)
+                process_obj = pickle.load(from_parent)
+            finally:
+                del process.current_process()._inheriting
+
+        exitcode = process_obj._bootstrap()
+    except Exception:
+        print("\n\n" + "-" * 80)
+        print(f"{process_name} failed with traceback: ")
+        print("-" * 80)
+        import traceback
+
+        print(traceback.format_exc())
+        print("\n" + "-" * 80)
+    finally:
+        if from_parent is not None:
+            from_parent.close()
+
+        sys.exit(exitcode)
+
+
+def get_command_line(main_prog=main, **kwargs):
+    """
+    Returns prefix of command line used for spawning a child process
+    """
+
+    if getattr(sys, "frozen", False):
+        list_kwargs = [f"{k}={v}" for k, v in kwargs.items()]
+        argv = [
+            sys.executable,
+            "--multiprocessing-fork",
+            main_prog.__module__,
+            *list_kwargs,
+        ]
+    else:
+        list_kwargs = [f'{k}="{v}"' for k, v in kwargs.items()]
+        prog = (
+            f"from {main_prog.__module__} import main; "
+            f'main({", ".join(list_kwargs)})'
+        )
+        opts = util._args_from_interpreter_flags()
+        argv = [get_executable(), *opts, "-c", prog]
+    return argv
+
+
+def freeze_support():
+    if len(sys.argv) >= 2 and sys.argv[1] == "--multiprocessing-fork":
+        module_main = sys.argv[2]
+        main = importlib.import_module(module_main).main
+        kwargs = {}
+        for p in sys.argv[3:]:
+            k, v = p.split("=")
+            kwargs[k] = v
+        exitcode = main(**kwargs)
+        sys.exit(exitcode)
+    _freeze_support_mp()