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