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>
This commit is contained in:
Duane Griffin 2025-09-07 19:49:54 +12:00 committed by GitHub
parent 8ed5a2b56c
commit 0912b3a6db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 33 additions and 4 deletions

View file

@ -145,12 +145,13 @@ class ForkServer(object):
cmd = ('from multiprocessing.forkserver import main; ' +
'main(%d, %d, %r, **%r)')
main_kws = {}
if self._preload_modules:
desired_keys = {'main_path', 'sys_path'}
data = spawn.get_preparation_data('ignore')
main_kws = {x: y for x, y in data.items() if x in desired_keys}
else:
main_kws = {}
if 'sys_path' in data:
main_kws['sys_path'] = data['sys_path']
if 'init_main_from_path' in data:
main_kws['main_path'] = data['init_main_from_path']
with socket.socket(socket.AF_UNIX) as listener:
address = connection.arbitrary_address('AF_UNIX')

View file

@ -7028,6 +7028,18 @@ class MiscTestCase(unittest.TestCase):
self.assertEqual(q.get_nowait(), "done")
close_queue(q)
def test_preload_main(self):
# gh-126631: Check that __main__ can be pre-loaded
if multiprocessing.get_start_method() != "forkserver":
self.skipTest("forkserver specific test")
name = os.path.join(os.path.dirname(__file__), 'mp_preload_main.py')
_, out, err = test.support.script_helper.assert_python_ok(name)
self.assertEqual(err, b'')
# The trailing empty string comes from split() on output ending with \n
out = out.decode().split("\n")
self.assertEqual(out, ['__main__', '__mp_main__', 'f', 'f', ''])
#
# Mixins

View file

@ -0,0 +1,14 @@
import multiprocessing
print(f"{__name__}")
def f():
print("f")
if __name__ == "__main__":
ctx = multiprocessing.get_context("forkserver")
ctx.set_forkserver_preload(['__main__'])
for _ in range(2):
p = ctx.Process(target=f)
p.start()
p.join()

View file

@ -0,0 +1,2 @@
Fix :mod:`multiprocessing` ``forkserver`` bug which prevented ``__main__``
from being preloaded.