FIX support for frozen executable on all platforms

loky/__init__.py

@@ -16,17 +16,38 @@
 
 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", "Executor", "ProcessPoolExecutor",
-           "BrokenProcessPool", "CancelledError", "TimeoutError",
-           "FIRST_COMPLETED", "FIRST_EXCEPTION", "ALL_COMPLETED",
-           "wrap_non_picklable_objects", "set_loky_pickler"]
+__all__ = [
+    # 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",
+]
 
 
 __version__ = '3.4.0.dev0'

loky/backend/fork_exec.py

@@ -37,6 +37,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
 
@@ -96,7 +95,8 @@ def _launch(self, process_obj):
         try:
             prep_data = spawn.get_preparation_data(
                 process_obj._name,
-                getattr(process_obj, "init_main_module", True))
+                getattr(process_obj, "init_main_module", True)
+            )
             reduction.dump(prep_data, fp)
             reduction.dump(process_obj, fp)
 
@@ -106,15 +106,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']
@@ -143,40 +141,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
@@ -47,11 +46,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}
 
@@ -76,7 +74,8 @@ def __init__(self, process_obj):
                     hp, ht, pid, _ = _winapi.CreateProcess(
                         python_exe, cmd,
                         None, None, inherit, 0,
-                        child_env, None, None)
+                        child_env, None, None
+                    )
                     _winapi.CloseHandle(ht)
                 except BaseException:
                     _winapi.CloseHandle(rhandle)
@@ -108,49 +107,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
 
@@ -130,13 +130,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
@@ -147,7 +147,7 @@ def ensure_running(self):
                     if _HAVE_SIGMASK:
                         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(signal.SIG_UNBLOCK,
@@ -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)
@@ -271,8 +276,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
@@ -348,23 +353,24 @@ def _unlink_resources(rtype_registry, rtype):
 # Start a program with only specified fds kept open
 #
 
-def spawnv_passfds(path, args, passfds):
+def spawnv_passfds(cmd, passfds):
     passfds = sorted(passfds)
     if sys.platform != "win32":
         errpipe_read, errpipe_write = os.pipe()
         try:
             from .reduction import _mk_inheritable
             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:
             pass

loky/backend/spawn.py

@@ -10,6 +10,8 @@
 import sys
 import runpy
 import types
+import pickle
+import importlib
 from multiprocessing import process, util
 
 
@@ -240,3 +242,68 @@ def _fixup_main_from_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):
+    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)