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