Home | History | Annotate | Download | only in multiprocessing
      1 #
      2 # Code used to start processes when using the spawn or forkserver
      3 # start methods.
      4 #
      5 # multiprocessing/spawn.py
      6 #
      7 # Copyright (c) 2006-2008, R Oudkerk
      8 # Licensed to PSF under a Contributor Agreement.
      9 #
     10 
     11 import os
     12 import sys
     13 import runpy
     14 import types
     15 
     16 from . import get_start_method, set_start_method
     17 from . import process
     18 from .context import reduction
     19 from . import util
     20 
     21 __all__ = ['_main', 'freeze_support', 'set_executable', 'get_executable',
     22            'get_preparation_data', 'get_command_line', 'import_main_path']
     23 
     24 #
     25 # _python_exe is the assumed path to the python executable.
     26 # People embedding Python want to modify it.
     27 #
     28 
     29 if sys.platform != 'win32':
     30     WINEXE = False
     31     WINSERVICE = False
     32 else:
     33     WINEXE = getattr(sys, 'frozen', False)
     34     WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
     35 
     36 if WINSERVICE:
     37     _python_exe = os.path.join(sys.exec_prefix, 'python.exe')
     38 else:
     39     _python_exe = sys.executable
     40 
     41 def set_executable(exe):
     42     global _python_exe
     43     _python_exe = exe
     44 
     45 def get_executable():
     46     return _python_exe
     47 
     48 #
     49 #
     50 #
     51 
     52 def is_forking(argv):
     53     '''
     54     Return whether commandline indicates we are forking
     55     '''
     56     if len(argv) >= 2 and argv[1] == '--multiprocessing-fork':
     57         return True
     58     else:
     59         return False
     60 
     61 
     62 def freeze_support():
     63     '''
     64     Run code for process object if this in not the main process
     65     '''
     66     if is_forking(sys.argv):
     67         kwds = {}
     68         for arg in sys.argv[2:]:
     69             name, value = arg.split('=')
     70             if value == 'None':
     71                 kwds[name] = None
     72             else:
     73                 kwds[name] = int(value)
     74         spawn_main(**kwds)
     75         sys.exit()
     76 
     77 
     78 def get_command_line(**kwds):
     79     '''
     80     Returns prefix of command line used for spawning a child process
     81     '''
     82     if getattr(sys, 'frozen', False):
     83         return ([sys.executable, '--multiprocessing-fork'] +
     84                 ['%s=%r' % item for item in kwds.items()])
     85     else:
     86         prog = 'from multiprocessing.spawn import spawn_main; spawn_main(%s)'
     87         prog %= ', '.join('%s=%r' % item for item in kwds.items())
     88         opts = util._args_from_interpreter_flags()
     89         return [_python_exe] + opts + ['-c', prog, '--multiprocessing-fork']
     90 
     91 
     92 def spawn_main(pipe_handle, parent_pid=None, tracker_fd=None):
     93     '''
     94     Run code specified by data received over pipe
     95     '''
     96     assert is_forking(sys.argv), "Not forking"
     97     if sys.platform == 'win32':
     98         import msvcrt
     99         new_handle = reduction.steal_handle(parent_pid, pipe_handle)
    100         fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY)
    101     else:
    102         from . import semaphore_tracker
    103         semaphore_tracker._semaphore_tracker._fd = tracker_fd
    104         fd = pipe_handle
    105     exitcode = _main(fd)
    106     sys.exit(exitcode)
    107 
    108 
    109 def _main(fd):
    110     with os.fdopen(fd, 'rb', closefd=True) as from_parent:
    111         process.current_process()._inheriting = True
    112         try:
    113             preparation_data = reduction.pickle.load(from_parent)
    114             prepare(preparation_data)
    115             self = reduction.pickle.load(from_parent)
    116         finally:
    117             del process.current_process()._inheriting
    118     return self._bootstrap()
    119 
    120 
    121 def _check_not_importing_main():
    122     if getattr(process.current_process(), '_inheriting', False):
    123         raise RuntimeError('''
    124         An attempt has been made to start a new process before the
    125         current process has finished its bootstrapping phase.
    126 
    127         This probably means that you are not using fork to start your
    128         child processes and you have forgotten to use the proper idiom
    129         in the main module:
    130 
    131             if __name__ == '__main__':
    132                 freeze_support()
    133                 ...
    134 
    135         The "freeze_support()" line can be omitted if the program
    136         is not going to be frozen to produce an executable.''')
    137 
    138 
    139 def get_preparation_data(name):
    140     '''
    141     Return info about parent needed by child to unpickle process object
    142     '''
    143     _check_not_importing_main()
    144     d = dict(
    145         log_to_stderr=util._log_to_stderr,
    146         authkey=process.current_process().authkey,
    147         )
    148 
    149     if util._logger is not None:
    150         d['log_level'] = util._logger.getEffectiveLevel()
    151 
    152     sys_path=sys.path.copy()
    153     try:
    154         i = sys_path.index('')
    155     except ValueError:
    156         pass
    157     else:
    158         sys_path[i] = process.ORIGINAL_DIR
    159 
    160     d.update(
    161         name=name,
    162         sys_path=sys_path,
    163         sys_argv=sys.argv,
    164         orig_dir=process.ORIGINAL_DIR,
    165         dir=os.getcwd(),
    166         start_method=get_start_method(),
    167         )
    168 
    169     # Figure out whether to initialise main in the subprocess as a module
    170     # or through direct execution (or to leave it alone entirely)
    171     main_module = sys.modules['__main__']
    172     main_mod_name = getattr(main_module.__spec__, "name", None)
    173     if main_mod_name is not None:
    174         d['init_main_from_name'] = main_mod_name
    175     elif sys.platform != 'win32' or (not WINEXE and not WINSERVICE):
    176         main_path = getattr(main_module, '__file__', None)
    177         if main_path is not None:
    178             if (not os.path.isabs(main_path) and
    179                         process.ORIGINAL_DIR is not None):
    180                 main_path = os.path.join(process.ORIGINAL_DIR, main_path)
    181             d['init_main_from_path'] = os.path.normpath(main_path)
    182 
    183     return d
    184 
    185 #
    186 # Prepare current process
    187 #
    188 
    189 old_main_modules = []
    190 
    191 def prepare(data):
    192     '''
    193     Try to get current process ready to unpickle process object
    194     '''
    195     if 'name' in data:
    196         process.current_process().name = data['name']
    197 
    198     if 'authkey' in data:
    199         process.current_process().authkey = data['authkey']
    200 
    201     if 'log_to_stderr' in data and data['log_to_stderr']:
    202         util.log_to_stderr()
    203 
    204     if 'log_level' in data:
    205         util.get_logger().setLevel(data['log_level'])
    206 
    207     if 'sys_path' in data:
    208         sys.path = data['sys_path']
    209 
    210     if 'sys_argv' in data:
    211         sys.argv = data['sys_argv']
    212 
    213     if 'dir' in data:
    214         os.chdir(data['dir'])
    215 
    216     if 'orig_dir' in data:
    217         process.ORIGINAL_DIR = data['orig_dir']
    218 
    219     if 'start_method' in data:
    220         set_start_method(data['start_method'], force=True)
    221 
    222     if 'init_main_from_name' in data:
    223         _fixup_main_from_name(data['init_main_from_name'])
    224     elif 'init_main_from_path' in data:
    225         _fixup_main_from_path(data['init_main_from_path'])
    226 
    227 # Multiprocessing module helpers to fix up the main module in
    228 # spawned subprocesses
    229 def _fixup_main_from_name(mod_name):
    230     # __main__.py files for packages, directories, zip archives, etc, run
    231     # their "main only" code unconditionally, so we don't even try to
    232     # populate anything in __main__, nor do we make any changes to
    233     # __main__ attributes
    234     current_main = sys.modules['__main__']
    235     if mod_name == "__main__" or mod_name.endswith(".__main__"):
    236         return
    237 
    238     # If this process was forked, __main__ may already be populated
    239     if getattr(current_main.__spec__, "name", None) == mod_name:
    240         return
    241 
    242     # Otherwise, __main__ may contain some non-main code where we need to
    243     # support unpickling it properly. We rerun it as __mp_main__ and make
    244     # the normal __main__ an alias to that
    245     old_main_modules.append(current_main)
    246     main_module = types.ModuleType("__mp_main__")
    247     main_content = runpy.run_module(mod_name,
    248                                     run_name="__mp_main__",
    249                                     alter_sys=True)
    250     main_module.__dict__.update(main_content)
    251     sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
    252 
    253 
    254 def _fixup_main_from_path(main_path):
    255     # If this process was forked, __main__ may already be populated
    256     current_main = sys.modules['__main__']
    257 
    258     # Unfortunately, the main ipython launch script historically had no
    259     # "if __name__ == '__main__'" guard, so we work around that
    260     # by treating it like a __main__.py file
    261     # See https://github.com/ipython/ipython/issues/4698
    262     main_name = os.path.splitext(os.path.basename(main_path))[0]
    263     if main_name == 'ipython':
    264         return
    265 
    266     # Otherwise, if __file__ already has the setting we expect,
    267     # there's nothing more to do
    268     if getattr(current_main, '__file__', None) == main_path:
    269         return
    270 
    271     # If the parent process has sent a path through rather than a module
    272     # name we assume it is an executable script that may contain
    273     # non-main code that needs to be executed
    274     old_main_modules.append(current_main)
    275     main_module = types.ModuleType("__mp_main__")
    276     main_content = runpy.run_path(main_path,
    277                                   run_name="__mp_main__")
    278     main_module.__dict__.update(main_content)
    279     sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
    280 
    281 
    282 def import_main_path(main_path):
    283     '''
    284     Set sys.modules['__main__'] to module at main_path
    285     '''
    286     _fixup_main_from_path(main_path)
    287