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
11import os
12import sys
13import runpy
14import types
15
16from . import get_start_method, set_start_method
17from . import process
18from .context import reduction
19from . 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
29if sys.platform != 'win32':
30    WINEXE = False
31    WINSERVICE = False
32else:
33    WINEXE = getattr(sys, 'frozen', False)
34    WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
35
36if WINSERVICE:
37    _python_exe = os.path.join(sys.exec_prefix, 'python.exe')
38else:
39    _python_exe = sys.executable
40
41def set_executable(exe):
42    global _python_exe
43    _python_exe = exe
44
45def get_executable():
46    return _python_exe
47
48#
49#
50#
51
52def 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
62def 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
78def 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
92def 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        import _winapi
100
101        if parent_pid is not None:
102            source_process = _winapi.OpenProcess(
103                _winapi.SYNCHRONIZE | _winapi.PROCESS_DUP_HANDLE,
104                False, parent_pid)
105        else:
106            source_process = None
107        new_handle = reduction.duplicate(pipe_handle,
108                                         source_process=source_process)
109        fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY)
110        parent_sentinel = source_process
111    else:
112        from . import resource_tracker
113        resource_tracker._resource_tracker._fd = tracker_fd
114        fd = pipe_handle
115        parent_sentinel = os.dup(pipe_handle)
116    exitcode = _main(fd, parent_sentinel)
117    sys.exit(exitcode)
118
119
120def _main(fd, parent_sentinel):
121    with os.fdopen(fd, 'rb', closefd=True) as from_parent:
122        process.current_process()._inheriting = True
123        try:
124            preparation_data = reduction.pickle.load(from_parent)
125            prepare(preparation_data)
126            self = reduction.pickle.load(from_parent)
127        finally:
128            del process.current_process()._inheriting
129    return self._bootstrap(parent_sentinel)
130
131
132def _check_not_importing_main():
133    if getattr(process.current_process(), '_inheriting', False):
134        raise RuntimeError('''
135        An attempt has been made to start a new process before the
136        current process has finished its bootstrapping phase.
137
138        This probably means that you are not using fork to start your
139        child processes and you have forgotten to use the proper idiom
140        in the main module:
141
142            if __name__ == '__main__':
143                freeze_support()
144                ...
145
146        The "freeze_support()" line can be omitted if the program
147        is not going to be frozen to produce an executable.''')
148
149
150def get_preparation_data(name):
151    '''
152    Return info about parent needed by child to unpickle process object
153    '''
154    _check_not_importing_main()
155    d = dict(
156        log_to_stderr=util._log_to_stderr,
157        authkey=process.current_process().authkey,
158        )
159
160    if util._logger is not None:
161        d['log_level'] = util._logger.getEffectiveLevel()
162
163    sys_path=sys.path.copy()
164    try:
165        i = sys_path.index('')
166    except ValueError:
167        pass
168    else:
169        sys_path[i] = process.ORIGINAL_DIR
170
171    d.update(
172        name=name,
173        sys_path=sys_path,
174        sys_argv=sys.argv,
175        orig_dir=process.ORIGINAL_DIR,
176        dir=os.getcwd(),
177        start_method=get_start_method(),
178        )
179
180    # Figure out whether to initialise main in the subprocess as a module
181    # or through direct execution (or to leave it alone entirely)
182    main_module = sys.modules['__main__']
183    main_mod_name = getattr(main_module.__spec__, "name", None)
184    if main_mod_name is not None:
185        d['init_main_from_name'] = main_mod_name
186    elif sys.platform != 'win32' or (not WINEXE and not WINSERVICE):
187        main_path = getattr(main_module, '__file__', None)
188        if main_path is not None:
189            if (not os.path.isabs(main_path) and
190                        process.ORIGINAL_DIR is not None):
191                main_path = os.path.join(process.ORIGINAL_DIR, main_path)
192            d['init_main_from_path'] = os.path.normpath(main_path)
193
194    return d
195
196#
197# Prepare current process
198#
199
200old_main_modules = []
201
202def prepare(data):
203    '''
204    Try to get current process ready to unpickle process object
205    '''
206    if 'name' in data:
207        process.current_process().name = data['name']
208
209    if 'authkey' in data:
210        process.current_process().authkey = data['authkey']
211
212    if 'log_to_stderr' in data and data['log_to_stderr']:
213        util.log_to_stderr()
214
215    if 'log_level' in data:
216        util.get_logger().setLevel(data['log_level'])
217
218    if 'sys_path' in data:
219        sys.path = data['sys_path']
220
221    if 'sys_argv' in data:
222        sys.argv = data['sys_argv']
223
224    if 'dir' in data:
225        os.chdir(data['dir'])
226
227    if 'orig_dir' in data:
228        process.ORIGINAL_DIR = data['orig_dir']
229
230    if 'start_method' in data:
231        set_start_method(data['start_method'], force=True)
232
233    if 'init_main_from_name' in data:
234        _fixup_main_from_name(data['init_main_from_name'])
235    elif 'init_main_from_path' in data:
236        _fixup_main_from_path(data['init_main_from_path'])
237
238# Multiprocessing module helpers to fix up the main module in
239# spawned subprocesses
240def _fixup_main_from_name(mod_name):
241    # __main__.py files for packages, directories, zip archives, etc, run
242    # their "main only" code unconditionally, so we don't even try to
243    # populate anything in __main__, nor do we make any changes to
244    # __main__ attributes
245    current_main = sys.modules['__main__']
246    if mod_name == "__main__" or mod_name.endswith(".__main__"):
247        return
248
249    # If this process was forked, __main__ may already be populated
250    if getattr(current_main.__spec__, "name", None) == mod_name:
251        return
252
253    # Otherwise, __main__ may contain some non-main code where we need to
254    # support unpickling it properly. We rerun it as __mp_main__ and make
255    # the normal __main__ an alias to that
256    old_main_modules.append(current_main)
257    main_module = types.ModuleType("__mp_main__")
258    main_content = runpy.run_module(mod_name,
259                                    run_name="__mp_main__",
260                                    alter_sys=True)
261    main_module.__dict__.update(main_content)
262    sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
263
264
265def _fixup_main_from_path(main_path):
266    # If this process was forked, __main__ may already be populated
267    current_main = sys.modules['__main__']
268
269    # Unfortunately, the main ipython launch script historically had no
270    # "if __name__ == '__main__'" guard, so we work around that
271    # by treating it like a __main__.py file
272    # See https://github.com/ipython/ipython/issues/4698
273    main_name = os.path.splitext(os.path.basename(main_path))[0]
274    if main_name == 'ipython':
275        return
276
277    # Otherwise, if __file__ already has the setting we expect,
278    # there's nothing more to do
279    if getattr(current_main, '__file__', None) == main_path:
280        return
281
282    # If the parent process has sent a path through rather than a module
283    # name we assume it is an executable script that may contain
284    # non-main code that needs to be executed
285    old_main_modules.append(current_main)
286    main_module = types.ModuleType("__mp_main__")
287    main_content = runpy.run_path(main_path,
288                                  run_name="__mp_main__")
289    main_module.__dict__.update(main_content)
290    sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
291
292
293def import_main_path(main_path):
294    '''
295    Set sys.modules['__main__'] to module at main_path
296    '''
297    _fixup_main_from_path(main_path)
298