spawn.py revision 7d2d43c0b15b8062c9b5d672a78e653abe2e1d91 
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 pickle
13import sys
14
15from . import process
16from . import util
17from . import popen
18
19__all__ = ['_main', 'freeze_support', 'set_executable', 'get_executable',
20           'get_preparation_data', 'get_command_line', 'import_main_path']
21
22#
23# _python_exe is the assumed path to the python executable.
24# People embedding Python want to modify it.
25#
26
27if sys.platform != 'win32':
28    WINEXE = False
29    WINSERVICE = False
30else:
31    WINEXE = (sys.platform == 'win32' and getattr(sys, 'frozen', False))
32    WINSERVICE = sys.executable.lower().endswith("pythonservice.exe")
33
34if WINSERVICE:
35    _python_exe = os.path.join(sys.exec_prefix, 'python.exe')
36else:
37    _python_exe = sys.executable
38
39def set_executable(exe):
40    global _python_exe
41    _python_exe = exe
42
43def get_executable():
44    return _python_exe
45
46#
47#
48#
49
50def is_forking(argv):
51    '''
52    Return whether commandline indicates we are forking
53    '''
54    if len(argv) >= 2 and argv[1] == '--multiprocessing-fork':
55        return True
56    else:
57        return False
58
59
60def freeze_support():
61    '''
62    Run code for process object if this in not the main process
63    '''
64    if is_forking(sys.argv):
65        main()
66        sys.exit()
67
68
69def get_command_line(**kwds):
70    '''
71    Returns prefix of command line used for spawning a child process
72    '''
73    if getattr(sys, 'frozen', False):
74        return [sys.executable, '--multiprocessing-fork']
75    else:
76        prog = 'from multiprocessing.spawn import spawn_main; spawn_main(%s)'
77        prog %= ', '.join('%s=%r' % item for item in kwds.items())
78        opts = util._args_from_interpreter_flags()
79        return [_python_exe] + opts + ['-c', prog, '--multiprocessing-fork']
80
81
82def spawn_main(pipe_handle, parent_pid=None, tracker_fd=None):
83    '''
84    Run code specifed by data received over pipe
85    '''
86    assert is_forking(sys.argv)
87    if sys.platform == 'win32':
88        import msvcrt
89        from .reduction import steal_handle
90        new_handle = steal_handle(parent_pid, pipe_handle)
91        fd = msvcrt.open_osfhandle(new_handle, os.O_RDONLY)
92    else:
93        from . import semaphore_tracker
94        semaphore_tracker._semaphore_tracker_fd = tracker_fd
95        fd = pipe_handle
96    exitcode = _main(fd)
97    sys.exit(exitcode)
98
99
100def _main(fd):
101    with os.fdopen(fd, 'rb', closefd=True) as from_parent:
102        process.current_process()._inheriting = True
103        try:
104            preparation_data = pickle.load(from_parent)
105            prepare(preparation_data)
106            self = pickle.load(from_parent)
107        finally:
108            del process.current_process()._inheriting
109    return self._bootstrap()
110
111
112def _check_not_importing_main():
113    if getattr(process.current_process(), '_inheriting', False):
114        raise RuntimeError('''
115        An attempt has been made to start a new process before the
116        current process has finished its bootstrapping phase.
117
118        This probably means that you are not using fork to start your
119        child processes and you have forgotten to use the proper idiom
120        in the main module:
121
122            if __name__ == '__main__':
123                freeze_support()
124                ...
125
126        The "freeze_support()" line can be omitted if the program
127        is not going to be frozen to produce an executable.''')
128
129
130def get_preparation_data(name):
131    '''
132    Return info about parent needed by child to unpickle process object
133    '''
134    _check_not_importing_main()
135    d = dict(
136        log_to_stderr=util._log_to_stderr,
137        authkey=process.current_process().authkey,
138        )
139
140    if util._logger is not None:
141        d['log_level'] = util._logger.getEffectiveLevel()
142
143    sys_path=sys.path.copy()
144    try:
145        i = sys_path.index('')
146    except ValueError:
147        pass
148    else:
149        sys_path[i] = process.ORIGINAL_DIR
150
151    d.update(
152        name=name,
153        sys_path=sys_path,
154        sys_argv=sys.argv,
155        orig_dir=process.ORIGINAL_DIR,
156        dir=os.getcwd(),
157        start_method=popen.get_start_method(),
158        )
159
160    if sys.platform != 'win32' or (not WINEXE and not WINSERVICE):
161        main_path = getattr(sys.modules['__main__'], '__file__', None)
162        if not main_path and sys.argv[0] not in ('', '-c'):
163            main_path = sys.argv[0]
164        if main_path is not None:
165            if (not os.path.isabs(main_path) and
166                        process.ORIGINAL_DIR is not None):
167                main_path = os.path.join(process.ORIGINAL_DIR, main_path)
168            d['main_path'] = os.path.normpath(main_path)
169
170    return d
171
172#
173# Prepare current process
174#
175
176old_main_modules = []
177
178def prepare(data):
179    '''
180    Try to get current process ready to unpickle process object
181    '''
182    if 'name' in data:
183        process.current_process().name = data['name']
184
185    if 'authkey' in data:
186        process.current_process().authkey = data['authkey']
187
188    if 'log_to_stderr' in data and data['log_to_stderr']:
189        util.log_to_stderr()
190
191    if 'log_level' in data:
192        util.get_logger().setLevel(data['log_level'])
193
194    if 'sys_path' in data:
195        sys.path = data['sys_path']
196
197    if 'sys_argv' in data:
198        sys.argv = data['sys_argv']
199
200    if 'dir' in data:
201        os.chdir(data['dir'])
202
203    if 'orig_dir' in data:
204        process.ORIGINAL_DIR = data['orig_dir']
205
206    if 'start_method' in data:
207        popen.set_start_method(data['start_method'], start_helpers=False)
208
209    if 'main_path' in data:
210        import_main_path(data['main_path'])
211
212
213def import_main_path(main_path):
214    '''
215    Set sys.modules['__main__'] to module at main_path
216    '''
217    # XXX (ncoghlan): The following code makes several bogus
218    # assumptions regarding the relationship between __file__
219    # and a module's real name. See PEP 302 and issue #10845
220    if getattr(sys.modules['__main__'], '__file__', None) == main_path:
221        return
222
223    main_name = os.path.splitext(os.path.basename(main_path))[0]
224    if main_name == '__init__':
225        main_name = os.path.basename(os.path.dirname(main_path))
226
227    if main_name == '__main__':
228        main_module = sys.modules['__main__']
229        main_module.__file__ = main_path
230    elif main_name != 'ipython':
231        # Main modules not actually called __main__.py may
232        # contain additional code that should still be executed
233        import importlib
234        import types
235
236        if main_path is None:
237            dirs = None
238        elif os.path.basename(main_path).startswith('__init__.py'):
239            dirs = [os.path.dirname(os.path.dirname(main_path))]
240        else:
241            dirs = [os.path.dirname(main_path)]
242
243        assert main_name not in sys.modules, main_name
244        sys.modules.pop('__mp_main__', None)
245        # We should not try to load __main__
246        # since that would execute 'if __name__ == "__main__"'
247        # clauses, potentially causing a psuedo fork bomb.
248        loader = importlib.find_loader(main_name, path=dirs)
249        main_module = types.ModuleType(main_name)
250        try:
251            loader.init_module_attrs(main_module)
252        except AttributeError:  # init_module_attrs is optional
253            pass
254        main_module.__name__ = '__mp_main__'
255        code = loader.get_code(main_name)
256        exec(code, main_module.__dict__)
257
258        old_main_modules.append(sys.modules['__main__'])
259        sys.modules['__main__'] = sys.modules['__mp_main__'] = main_module
260