Skip to content

Commit 34df31c

Browse files
committed
fix: pin worker start method to fork in LMI mode
multiprocessing.Process used the interpreter default start method, which is fork on 3.13 and forkserver on 3.14+, so the worker process model differed by runtime. Pin "fork" via get_context so workers are direct children of the runtime process on every version and the parent default is left untouched for customer code.
1 parent c1ce090 commit 34df31c

3 files changed

Lines changed: 75 additions & 10 deletions

File tree

‎awslambdaric/lambda_multi_concurrent_utils.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,10 @@ def run_concurrent(
6767
):
6868
cls._emit_worker_pool_event(max_concurrency)
6969

70+
ctx = multiprocessing.get_context("fork")
7071
processes = []
7172
for _ in range(max_concurrency):
72-
p = multiprocessing.Process(
73+
p = ctx.Process(
7374
target=cls.run_single,
7475
args=(handler, api_addr, use_thread, socket_path),
7576
)

‎tests/test_concurrency.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,10 @@
33
"""
44

55
import multiprocessing
6+
import os
7+
import sys
8+
import threading
9+
import types
610
import unittest
711
from unittest.mock import patch, MagicMock
812

@@ -52,6 +56,50 @@ def fake_bootstrap_run(handler, lambda_runtime_client):
5256
self.assertEqual(success_counter.value, 6)
5357
self.assertEqual(fail_counter.value, 2)
5458

59+
def test_workers_are_forked_children_regardless_of_default_start_method(self):
60+
previous = multiprocessing.get_start_method(allow_none=True)
61+
multiprocessing.set_start_method("forkserver", force=True)
62+
try:
63+
parent_pid = os.getpid()
64+
parent_threads_at_fork = threading.active_count()
65+
marker = "_fork_pin_marker_module"
66+
sys.modules[marker] = types.ModuleType(marker)
67+
n = 3
68+
ppids = multiprocessing.Array("i", n)
69+
inherited = multiprocessing.Array("i", n)
70+
slot = multiprocessing.Value("i", 0)
71+
72+
def fake_bootstrap_run(handler, lambda_runtime_client):
73+
with slot.get_lock():
74+
i = slot.value
75+
slot.value += 1
76+
ppids[i] = os.getppid()
77+
inherited[i] = int(marker in sys.modules)
78+
79+
with patch(
80+
"awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._redirect_output"
81+
), patch(
82+
"awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._emit_worker_pool_event"
83+
), patch(
84+
"awslambdaric.lambda_multi_concurrent_utils.bootstrap.run",
85+
side_effect=fake_bootstrap_run,
86+
):
87+
MultiConcurrentRunner.run_concurrent(
88+
self.handler,
89+
self.addr,
90+
self.use_thread,
91+
self.socket,
92+
max_concurrency=n,
93+
)
94+
95+
self.assertEqual(multiprocessing.get_start_method(), "forkserver")
96+
self.assertEqual(parent_threads_at_fork, 1)
97+
self.assertEqual(list(ppids), [parent_pid] * n)
98+
self.assertEqual(list(inherited), [1] * n)
99+
finally:
100+
sys.modules.pop(marker, None)
101+
multiprocessing.set_start_method(previous, force=True)
102+
55103

56104
if __name__ == "__main__":
57105
unittest.main()

‎tests/test_multi_concurrent_runner.py‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -59,32 +59,48 @@ def test_run_single_creates_client_and_calls_bootstrap(
5959
mock_bootstrap.run.assert_called_once_with("h.fn", mock_client)
6060

6161
@patch.object(MultiConcurrentRunner, "_emit_worker_pool_event")
62-
@patch("multiprocessing.Process")
63-
def test_run_concurrent_spawns_and_joins(self, mock_process, mock_emit):
62+
@patch("multiprocessing.get_context")
63+
def test_run_concurrent_spawns_and_joins(self, mock_get_context, mock_emit):
6464
fake_proc = MagicMock()
65-
mock_process.return_value = fake_proc
65+
mock_ctx = MagicMock()
66+
mock_ctx.Process.return_value = fake_proc
67+
mock_get_context.return_value = mock_ctx
6668

6769
MultiConcurrentRunner.run_concurrent(
6870
"h", "a", False, "/sock", max_concurrency=3
6971
)
7072

71-
self.assertEqual(mock_process.call_count, 3)
73+
self.assertEqual(mock_ctx.Process.call_count, 3)
7274
self.assertEqual(fake_proc.start.call_count, 3)
7375
self.assertEqual(fake_proc.join.call_count, 3)
7476

75-
for call_args in mock_process.call_args_list:
77+
for call_args in mock_ctx.Process.call_args_list:
7678
target = call_args.kwargs.get("target") or call_args[1].get("target")
7779
args = call_args.kwargs.get("args") or call_args[1].get("args")
7880
self.assertEqual(target, MultiConcurrentRunner.run_single)
7981
self.assertEqual(args, ("h", "a", False, "/sock"))
8082

81-
@patch("multiprocessing.Process")
83+
@patch.object(MultiConcurrentRunner, "_emit_worker_pool_event")
84+
@patch("multiprocessing.get_context")
85+
def test_run_concurrent_pins_fork_start_method(self, mock_get_context, mock_emit):
86+
mock_ctx = MagicMock()
87+
mock_get_context.return_value = mock_ctx
88+
89+
MultiConcurrentRunner.run_concurrent(
90+
"h", "a", False, "/sock", max_concurrency=2
91+
)
92+
93+
mock_get_context.assert_called_once_with("fork")
94+
self.assertEqual(mock_ctx.Process.call_count, 2)
95+
96+
@patch("multiprocessing.get_context")
8297
def test_run_concurrent_emits_worker_pool_event_once_before_spawning(
83-
self, mock_process
98+
self, mock_get_context
8499
):
85-
mock_process.return_value = MagicMock()
100+
mock_ctx = MagicMock()
101+
mock_get_context.return_value = mock_ctx
86102
order_tracker = MagicMock()
87-
order_tracker.attach_mock(mock_process, "process")
103+
order_tracker.attach_mock(mock_ctx.Process, "process")
88104

89105
with patch.object(
90106
MultiConcurrentRunner, "_emit_worker_pool_event"

0 commit comments

Comments
 (0)