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