Skip to content

Commit 0912b3a

Browse files
duaneggpshead
andauthored
gh-126631: gh-137996: fix pre-loading of __main__ (GH-135295)
gh-126631: gh-137996: fix pre-loading of `__main__` The `main_path` parameter was renamed `init_main_from_name`, update the forkserver code accordingly. This was leading to slower startup times when people were trying to preload the main module. --------- Co-authored-by: Gregory P. Smith <greg@krypto.org>
1 parent 8ed5a2b commit 0912b3a

File tree

4 files changed

+33
-4
lines changed

4 files changed

+33
-4
lines changed

Lib/multiprocessing/forkserver.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -145,12 +145,13 @@ def ensure_running(self):
145145
cmd = ('from multiprocessing.forkserver import main; ' +
146146
'main(%d, %d, %r, **%r)')
147147

148+
main_kws = {}
148149
if self._preload_modules:
149-
desired_keys = {'main_path', 'sys_path'}
150150
data = spawn.get_preparation_data('ignore')
151-
main_kws = {x: y for x, y in data.items() if x in desired_keys}
152-
else:
153-
main_kws = {}
151+
if 'sys_path' in data:
152+
main_kws['sys_path'] = data['sys_path']
153+
if 'init_main_from_path' in data:
154+
main_kws['main_path'] = data['init_main_from_path']
154155

155156
with socket.socket(socket.AF_UNIX) as listener:
156157
address = connection.arbitrary_address('AF_UNIX')

Lib/test/_test_multiprocessing.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7028,6 +7028,18 @@ def child():
70287028
self.assertEqual(q.get_nowait(), "done")
70297029
close_queue(q)
70307030

7031+
def test_preload_main(self):
7032+
# gh-126631: Check that __main__ can be pre-loaded
7033+
if multiprocessing.get_start_method() != "forkserver":
7034+
self.skipTest("forkserver specific test")
7035+
7036+
name = os.path.join(os.path.dirname(__file__), 'mp_preload_main.py')
7037+
_, out, err = test.support.script_helper.assert_python_ok(name)
7038+
self.assertEqual(err, b'')
7039+
7040+
# The trailing empty string comes from split() on output ending with \n
7041+
out = out.decode().split("\n")
7042+
self.assertEqual(out, ['__main__', '__mp_main__', 'f', 'f', ''])
70317043

70327044
#
70337045
# Mixins

Lib/test/mp_preload_main.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import multiprocessing
2+
3+
print(f"{__name__}")
4+
5+
def f():
6+
print("f")
7+
8+
if __name__ == "__main__":
9+
ctx = multiprocessing.get_context("forkserver")
10+
ctx.set_forkserver_preload(['__main__'])
11+
for _ in range(2):
12+
p = ctx.Process(target=f)
13+
p.start()
14+
p.join()
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix :mod:`multiprocessing` ``forkserver`` bug which prevented ``__main__``
2+
from being preloaded.

0 commit comments

Comments
 (0)