From 26601c811bdaf207e55195b20b78d3efff2f7eb8 Mon Sep 17 00:00:00 2001 From: Davide Melfi Date: Tue, 13 Oct 2026 09:53:18 +0000 Subject: [PATCH] fix: revert pre-fork hooks for multi-concurrent mode This reverts commit fa006b674ee3fbaf9dcb7ce74532811c4babd0d6 (#226). The pre-fork hooks change introduced in 4.1.0 is breaking, so it is removed: the `awslambdaric.lambda_concurrency_hooks` module, the `@register_pre_fork` decorator and the `Runtime.PreForkError` error type are gone, restoring the 4.0.4 multi-concurrent startup flow. The on-demand path is unaffected. Bumps the version to 4.2.1 and adds a release note. --- RELEASE.CHANGELOG.md | 4 + awslambdaric/__init__.py | 2 +- awslambdaric/bootstrap.py | 4 +- awslambdaric/lambda_concurrency_hooks.py | 31 -- awslambdaric/lambda_multi_concurrent_utils.py | 74 +---- awslambdaric/lambda_runtime_exception.py | 1 - tests/test_bootstrap.py | 16 +- tests/test_concurrency.py | 2 +- tests/test_concurrency_hooks.py | 309 ------------------ tests/test_multi_concurrent_runner.py | 30 +- 10 files changed, 43 insertions(+), 430 deletions(-) delete mode 100644 awslambdaric/lambda_concurrency_hooks.py delete mode 100644 tests/test_concurrency_hooks.py diff --git a/RELEASE.CHANGELOG.md b/RELEASE.CHANGELOG.md index 361b652..65e5939 100644 --- a/RELEASE.CHANGELOG.md +++ b/RELEASE.CHANGELOG.md @@ -1,3 +1,7 @@ +### October 13, 2026 +`4.2.1` +- Revert pre-fork hooks for multi-concurrent (Lambda Managed Instances) mode introduced in `4.1.0` ([#226](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/226)), as the change was breaking. The `awslambdaric.lambda_concurrency_hooks` module, the `@register_pre_fork` decorator and the `Runtime.PreForkError` error type are removed, restoring the `4.0.4` multi-concurrent startup behavior. No impact on the standard on-demand path. + ### October 7, 2026 `4.2.0` - Add a `w3c()` helper to `LambdaContext` that returns the W3C trace-context fields (`traceparent`, `tracestate`, `baggage`) carried on `clientContext.w3c` at invoke time ([#228](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/228)) diff --git a/awslambdaric/__init__.py b/awslambdaric/__init__.py index 5ab16d3..54fc0d5 100644 --- a/awslambdaric/__init__.py +++ b/awslambdaric/__init__.py @@ -2,4 +2,4 @@ Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved. """ -__version__ = "4.2.0" +__version__ = "4.2.1" diff --git a/awslambdaric/bootstrap.py b/awslambdaric/bootstrap.py index c83eccf..9554b39 100644 --- a/awslambdaric/bootstrap.py +++ b/awslambdaric/bootstrap.py @@ -35,7 +35,7 @@ INIT_TYPE_SNAP_START = "snap-start" -def get_handler(handler): +def _get_handler(handler): try: modname, fname = handler.rsplit(".", 1) except ValueError as e: @@ -516,7 +516,7 @@ def run(handler, lambda_runtime_client): _log_preview_runtime_warning() - request_handler = get_handler(handler) + request_handler = _get_handler(handler) except FaultException as e: error_result = make_error( e.msg, diff --git a/awslambdaric/lambda_concurrency_hooks.py b/awslambdaric/lambda_concurrency_hooks.py deleted file mode 100644 index b2194ab..0000000 --- a/awslambdaric/lambda_concurrency_hooks.py +++ /dev/null @@ -1,31 +0,0 @@ -# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. -# SPDX-License-Identifier: Apache-2.0 - -from typing import Any, Callable - -# The customer-facing name lives only here: the runner consumes get_pre_fork(), -# so a later rename is one line plus a backwards-compatible alias. -__all__ = ["register_pre_fork"] - -_pre_fork_registry: list[tuple[Callable[..., Any], tuple, dict]] = [] - - -def register_pre_fork(func: Callable[..., Any]) -> Callable[..., Any]: - """ - Register a function to run once in the parent, before workers are forked. - - Runs for any worker count, so a hook's side effects are in place whatever - concurrency the execution environment is configured with. - - from awslambdaric.lambda_concurrency_hooks import register_pre_fork - - @register_pre_fork - def start_inference_server(): - subprocess.Popen(["python", "serve.py", "--port", "8000"]) - """ - _pre_fork_registry.append((func, (), {})) - return func - - -def get_pre_fork() -> list[tuple[Callable[..., Any], tuple, dict]]: - return _pre_fork_registry diff --git a/awslambdaric/lambda_multi_concurrent_utils.py b/awslambdaric/lambda_multi_concurrent_utils.py index 4065b75..053e624 100644 --- a/awslambdaric/lambda_multi_concurrent_utils.py +++ b/awslambdaric/lambda_multi_concurrent_utils.py @@ -10,7 +10,6 @@ from . import bootstrap from .lambda_runtime_client import LambdaMultiConcurrentRuntimeClient -from .lambda_concurrency_hooks import get_pre_fork WORKER_POOL_INITIALIZING_EVENT = "runtime_worker_pool_initializing" @@ -38,7 +37,13 @@ def run_single( @classmethod def _emit_worker_pool_event(cls, max_concurrency: int): - """Emit the worker pool DEBUG event. The sink is owned by _before_fork.""" + """Emit worker pool DEBUG event once from the parent before forking. + + No output redirection here: RAPID wires the runtime main process's + stdout/stderr to the log egress at spawn. The FD provider socket is + only for the forked workers, which redirect in run_single. + """ + log_sink = bootstrap.init_logging() logging.getLogger().debug( { "event": WORKER_POOL_INITIALIZING_EVENT, @@ -46,68 +51,9 @@ def _emit_worker_pool_event(cls, max_concurrency: int): "executionEnvironmentMaxConcurrency": max_concurrency, } ) - - @classmethod - def _init_handler(cls, handler: str, client, log_sink): - """Import the handler, mirroring the guard in bootstrap.run: report an - init error to RAPID and exit if it fails.""" - try: - return bootstrap.get_handler(handler) - except bootstrap.FaultException as e: - error_result = bootstrap.make_error(e.msg, e.exception_type, e.trace) - except Exception: - error_result = bootstrap.build_fault_result(sys.exc_info(), None) - - bootstrap.log_error(error_result, log_sink) - client.post_init_error(error_result) - sys.exit(1) - - @classmethod - def _run_pre_fork_hooks(cls, handler: str, api_addr: str, log_sink): - """Run @register_pre_fork hooks once in the parent, in registration order. - - Importing the handler here is what runs its module-level - @register_pre_fork decorators. Workers re-import it in their own - process, so hooks are for external side effects (a subprocess, a warmed - service, a file in /tmp) and share no in-memory state with workers. - - A failing hook is reported as an INIT error and exits: no worker should - run against a precondition the hook failed to establish. - """ - client = LambdaMultiConcurrentRuntimeClient(api_addr, False) - cls._init_handler(handler, client, log_sink) - - try: - for func, args, kwargs in get_pre_fork(): - func(*args, **kwargs) - except Exception: - error_result = bootstrap.build_fault_result(sys.exc_info(), None) - bootstrap.log_error(error_result, log_sink) - client.post_init_error( - error_result, bootstrap.FaultException.PRE_FORK_ERROR - ) - sys.exit(1) - - @classmethod - def _before_fork(cls, handler: str, api_addr: str, max_concurrency: int): - """Run the parent's work that must happen before forking workers. - - One sink covers both steps, released before returning: forked workers - inherit the parent's handler (fork is the POSIX default before 3.14) - and would log every line twice. Not released on the failure path, where - the process is exiting anyway. - - No redirection here: RAPID wires the parent's stdout/stderr to the log - egress at spawn; the FD provider socket is for workers (run_single). - """ - log_sink = bootstrap.init_logging() - - cls._run_pre_fork_hooks(handler, api_addr, log_sink) - - # After the hooks, so it is never emitted for a pool that fails to start. - cls._emit_worker_pool_event(max_concurrency) - logging.getLogger().handlers.clear() + # Close the sink deterministically now that its handler is gone + # (no-op for StandardLogSink; releases the fd for framed sinks). log_sink.__exit__(None, None, None) @classmethod @@ -119,7 +65,7 @@ def run_concurrent( socket_path: str, max_concurrency: int, ): - cls._before_fork(handler, api_addr, max_concurrency) + cls._emit_worker_pool_event(max_concurrency) processes = [] for _ in range(max_concurrency): diff --git a/awslambdaric/lambda_runtime_exception.py b/awslambdaric/lambda_runtime_exception.py index 6840ae8..3ea5b29 100644 --- a/awslambdaric/lambda_runtime_exception.py +++ b/awslambdaric/lambda_runtime_exception.py @@ -13,7 +13,6 @@ class FaultException(Exception): MALFORMED_HANDLER_NAME = "Runtime.MalformedHandlerName" BEFORE_SNAPSHOT_ERROR = "Runtime.BeforeSnapshotError" AFTER_RESTORE_ERROR = "Runtime.AfterRestoreError" - PRE_FORK_ERROR = "Runtime.PreForkError" LAMBDA_CONTEXT_UNMARSHAL_ERROR = "Runtime.LambdaContextUnmarshalError" LAMBDA_RUNTIME_CLIENT_ERROR = "Runtime.LambdaRuntimeClientError" diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 1558941..4a6c54f 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -740,7 +740,7 @@ def __eq__(self, other): def test_get_event_handler_bad_handler(self): handler_name = "bad_handler" with self.assertRaises(FaultException) as cm: - response_handler = bootstrap.get_handler(handler_name) + response_handler = bootstrap._get_handler(handler_name) returned_exception = cm.exception self.assertEqual( self.FaultExceptionMatcher( @@ -753,7 +753,7 @@ def test_get_event_handler_bad_handler(self): def test_get_event_handler_import_error(self): handler_name = "no_module.handler" with self.assertRaises(FaultException) as cm: - response_handler = bootstrap.get_handler(handler_name) + response_handler = bootstrap._get_handler(handler_name) returned_exception = cm.exception self.assertEqual( self.FaultExceptionMatcher( @@ -778,7 +778,7 @@ def test_get_event_handler_syntax_error(self): handler_name = "{}.syntax_error".format(filename) with self.assertRaises(FaultException) as cm: - response_handler = bootstrap.get_handler(handler_name) + response_handler = bootstrap._get_handler(handler_name) returned_exception = cm.exception self.assertEqual( self.FaultExceptionMatcher( @@ -801,7 +801,7 @@ def test_get_event_handler_missing_error(self): filename, _ = os.path.splitext(filename_w_ext) handler_name = "{}.my_handler".format(filename) with self.assertRaises(FaultException) as cm: - response_handler = bootstrap.get_handler(handler_name) + response_handler = bootstrap._get_handler(handler_name) returned_exception = cm.exception self.assertEqual( self.FaultExceptionMatcher( @@ -814,12 +814,12 @@ def test_get_event_handler_missing_error(self): def test_get_event_handler_slash(self): importlib.invalidate_caches() handler_name = "tests/test_handler_with_slash/test_handler.my_handler" - response_handler = bootstrap.get_handler(handler_name) + response_handler = bootstrap._get_handler(handler_name) response_handler() def test_get_event_handler_build_in_conflict(self): with self.assertRaises(FaultException) as cm: - response_handler = bootstrap.get_handler("sys.hello") + response_handler = bootstrap._get_handler("sys.hello") returned_exception = cm.exception self.assertEqual( self.FaultExceptionMatcher( @@ -830,13 +830,13 @@ def test_get_event_handler_build_in_conflict(self): ) def test_get_event_handler_doesnt_throw_build_in_module_name_slash(self): - response_handler = bootstrap.get_handler( + response_handler = bootstrap._get_handler( "tests/test_built_in_module_name/sys.my_handler" ) response_handler() def test_get_event_handler_doent_throw_build_in_module_name(self): - response_handler = bootstrap.get_handler( + response_handler = bootstrap._get_handler( "tests.test_built_in_module_name.sys.my_handler" ) response_handler() diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index f3cd507..61ceeca 100644 --- a/tests/test_concurrency.py +++ b/tests/test_concurrency.py @@ -39,7 +39,7 @@ def fake_bootstrap_run(handler, lambda_runtime_client): with patch( "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._redirect_output" ), patch( - "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._before_fork" + "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._emit_worker_pool_event" ), patch( "awslambdaric.lambda_multi_concurrent_utils.bootstrap.run", side_effect=fake_bootstrap_run, diff --git a/tests/test_concurrency_hooks.py b/tests/test_concurrency_hooks.py deleted file mode 100644 index 3649018..0000000 --- a/tests/test_concurrency_hooks.py +++ /dev/null @@ -1,309 +0,0 @@ -""" -Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved. -""" - -import unittest -from unittest.mock import patch, MagicMock - -from awslambdaric import lambda_concurrency_hooks as hooks -from awslambdaric.lambda_concurrency_hooks import register_pre_fork -from awslambdaric.lambda_multi_concurrent_utils import MultiConcurrentRunner -from awslambdaric.lambda_runtime_exception import FaultException - - -class PreForkRegistryTest(unittest.TestCase): - def setUp(self): - hooks._pre_fork_registry.clear() - - def tearDown(self): - hooks._pre_fork_registry.clear() - - def test_register_returns_the_function_unchanged(self): - def my_hook(): - pass - - # The decorator must hand the function back, or it disappears from the - # customer's module. - self.assertIs(register_pre_fork(my_hook), my_hook) - - def test_register_stores_func_with_empty_args_and_kwargs(self): - def my_hook(): - pass - - register_pre_fork(my_hook) - # The runner unpacks this 3-tuple shape. - self.assertEqual(hooks.get_pre_fork(), [(my_hook, (), {})]) - - -class RunPreForkHooksTest(unittest.TestCase): - def setUp(self): - hooks._pre_fork_registry.clear() - self.log_sink = MagicMock() - - def tearDown(self): - hooks._pre_fork_registry.clear() - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - def test_hooks_run_once_in_registration_order( - self, mock_client_cls, mock_init_handler - ): - order = [] - - @register_pre_fork - def first(): - order.append(1) - - @register_pre_fork - def second(): - order.append(2) - - MultiConcurrentRunner._run_pre_fork_hooks("app.handler", "addr", self.log_sink) - - self.assertEqual(order, [1, 2]) - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - def test_handler_is_imported_before_hooks_run( - self, mock_client_cls, mock_init_handler - ): - order = [] - mock_init_handler.side_effect = lambda *a, **k: order.append("import") - - @register_pre_fork - def hook(): - order.append("hook") - - MultiConcurrentRunner._run_pre_fork_hooks("app.handler", "addr", self.log_sink) - - # Importing the handler is what registers the hooks, so it must come first. - self.assertEqual(order, ["import", "hook"]) - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_failing_hook_is_reported_as_init_error_and_exits( - self, mock_bootstrap, mock_client_cls, mock_init_handler - ): - mock_bootstrap.FaultException = FaultException - mock_bootstrap.build_fault_result.return_value = {"errorType": "RuntimeError"} - - @register_pre_fork - def bad_hook(): - raise RuntimeError("boom") - - with self.assertRaises(SystemExit) as raised: - MultiConcurrentRunner._run_pre_fork_hooks( - "app.handler", "addr", self.log_sink - ) - - self.assertEqual(raised.exception.code, 1) - # Reported like a failed handler import, so the customer sees a real - # INIT error carrying the hook's traceback rather than a bare crash. - mock_bootstrap.log_error.assert_called_once_with( - {"errorType": "RuntimeError"}, self.log_sink - ) - # The hook phase is named explicitly: the bare Python exception type - # would be normalized to Runtime.Unknown by the platform. - mock_client_cls.return_value.post_init_error.assert_called_once_with( - {"errorType": "RuntimeError"}, "Runtime.PreForkError" - ) - - def test_pre_fork_error_type_survives_platform_normalization(self): - # Lambda normalizes anything outside to - # Runtime.Unknown, which would misreport a well-formed hook failure. - self.assertRegex( - FaultException.PRE_FORK_ERROR, r"^(Runtime|Function)\.[A-Z][a-zA-Z]+$" - ) - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_later_hooks_do_not_run_after_one_fails( - self, mock_bootstrap, mock_client_cls, mock_init_handler - ): - mock_bootstrap.FaultException = FaultException - ran = [] - - @register_pre_fork - def bad_hook(): - raise RuntimeError("boom") - - @register_pre_fork - def later_hook(): - ran.append("later") - - with self.assertRaises(SystemExit): - MultiConcurrentRunner._run_pre_fork_hooks( - "app.handler", "addr", self.log_sink - ) - - self.assertEqual(ran, []) - - -class BeforeForkTest(unittest.TestCase): - """_before_fork owns the parent's log sink for the pre-fork work.""" - - @patch.object(MultiConcurrentRunner, "_emit_worker_pool_event") - @patch.object(MultiConcurrentRunner, "_run_pre_fork_hooks") - @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_sets_up_sink_once_and_releases_it_before_returning( - self, mock_bootstrap, mock_logging, mock_run_hooks, mock_emit - ): - log_sink = mock_bootstrap.init_logging.return_value - - MultiConcurrentRunner._before_fork("app.handler", "addr", 3) - - mock_bootstrap.init_logging.assert_called_once_with() - # Both steps log through that one sink. - mock_run_hooks.assert_called_once_with("app.handler", "addr", log_sink) - mock_emit.assert_called_once_with(3) - # Released before returning, so no forked worker inherits the handler. - mock_logging.getLogger.return_value.handlers.clear.assert_called_once_with() - log_sink.__exit__.assert_called_once_with(None, None, None) - - @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_hooks_run_before_the_pool_event(self, mock_bootstrap, mock_logging): - order_tracker = MagicMock() - - with patch.object( - MultiConcurrentRunner, "_run_pre_fork_hooks" - ) as mock_run_hooks, patch.object( - MultiConcurrentRunner, "_emit_worker_pool_event" - ) as mock_emit: - order_tracker.attach_mock(mock_run_hooks, "run_hooks") - order_tracker.attach_mock(mock_emit, "emit") - MultiConcurrentRunner._before_fork("app.handler", "addr", 2) - - # The event is emitted after the hooks, so it is never emitted for a - # pool that then fails to start. - self.assertEqual( - [c[0] for c in order_tracker.mock_calls[:2]], ["run_hooks", "emit"] - ) - - -class RunConcurrentHookTest(unittest.TestCase): - """Drive the public entry point with a really registered hook.""" - - def setUp(self): - hooks._pre_fork_registry.clear() - - def tearDown(self): - hooks._pre_fork_registry.clear() - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - @patch.object(MultiConcurrentRunner, "_emit_worker_pool_event") - @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - @patch("multiprocessing.Process") - def test_registered_hook_runs_exactly_once_through_run_concurrent( - self, - mock_process, - mock_bootstrap, - mock_logging, - mock_emit, - mock_client_cls, - mock_init_handler, - ): - mock_process.return_value = MagicMock() - hook_fn = MagicMock() - register_pre_fork(hook_fn) - - MultiConcurrentRunner.run_concurrent( - "app.handler", "addr", False, "/sock", max_concurrency=4 - ) - - # One call in the parent, not one per worker. - hook_fn.assert_called_once_with() - self.assertEqual(mock_process.call_count, 4) - - @patch.object(MultiConcurrentRunner, "_init_handler") - @patch( - "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient" - ) - @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - @patch("multiprocessing.Process") - def test_failing_hook_starts_no_workers( - self, - mock_process, - mock_bootstrap, - mock_logging, - mock_client_cls, - mock_init_handler, - ): - @register_pre_fork - def bad_hook(): - raise RuntimeError("s3 download failed") - - with self.assertRaises(SystemExit): - MultiConcurrentRunner.run_concurrent( - "app.handler", "addr", False, "/sock", max_concurrency=4 - ) - - # Workers must never run against a precondition the hook failed to - # establish - one clear init failure beats N per-invocation failures. - mock_process.assert_not_called() - - -class InitHandlerTest(unittest.TestCase): - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_returns_handler_on_success(self, mock_bootstrap): - sentinel = object() - mock_bootstrap.get_handler.return_value = sentinel - client = MagicMock() - - result = MultiConcurrentRunner._init_handler("app.handler", client, MagicMock()) - - self.assertIs(result, sentinel) - client.post_init_error.assert_not_called() - - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_reports_fault_exception_and_exits(self, mock_bootstrap): - from awslambdaric.lambda_runtime_exception import FaultException - - mock_bootstrap.FaultException = FaultException - mock_bootstrap.get_handler.side_effect = FaultException( - FaultException.IMPORT_MODULE_ERROR, "no module", None - ) - mock_bootstrap.make_error.return_value = {"errorType": "x"} - client = MagicMock() - log_sink = MagicMock() - - with self.assertRaises(SystemExit): - MultiConcurrentRunner._init_handler("app.handler", client, log_sink) - - mock_bootstrap.log_error.assert_called_once_with({"errorType": "x"}, log_sink) - client.post_init_error.assert_called_once_with({"errorType": "x"}) - - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") - def test_reports_unexpected_exception_and_exits(self, mock_bootstrap): - from awslambdaric.lambda_runtime_exception import FaultException - - mock_bootstrap.FaultException = FaultException - mock_bootstrap.get_handler.side_effect = ValueError("unexpected") - mock_bootstrap.build_fault_result.return_value = {"errorType": "ValueError"} - client = MagicMock() - - with self.assertRaises(SystemExit): - MultiConcurrentRunner._init_handler("app.handler", client, MagicMock()) - - mock_bootstrap.build_fault_result.assert_called_once() - client.post_init_error.assert_called_once() - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_multi_concurrent_runner.py b/tests/test_multi_concurrent_runner.py index 829591f..5cc1371 100644 --- a/tests/test_multi_concurrent_runner.py +++ b/tests/test_multi_concurrent_runner.py @@ -58,9 +58,9 @@ def test_run_single_creates_client_and_calls_bootstrap( mock_client_cls.assert_called_once_with("addr", True) mock_bootstrap.run.assert_called_once_with("h.fn", mock_client) - @patch.object(MultiConcurrentRunner, "_before_fork") + @patch.object(MultiConcurrentRunner, "_emit_worker_pool_event") @patch("multiprocessing.Process") - def test_run_concurrent_spawns_and_joins(self, mock_process, mock_before_fork): + def test_run_concurrent_spawns_and_joins(self, mock_process, mock_emit): fake_proc = MagicMock() mock_process.return_value = fake_proc @@ -78,11 +78,9 @@ def test_run_concurrent_spawns_and_joins(self, mock_process, mock_before_fork): self.assertEqual(target, MultiConcurrentRunner.run_single) self.assertEqual(args, ("h", "a", False, "/sock")) - @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") @patch("multiprocessing.Process") def test_run_concurrent_emits_worker_pool_event_once_before_spawning( - self, mock_process, mock_bootstrap, mock_logging + self, mock_process ): mock_process.return_value = MagicMock() order_tracker = MagicMock() @@ -91,28 +89,34 @@ def test_run_concurrent_emits_worker_pool_event_once_before_spawning( with patch.object( MultiConcurrentRunner, "_emit_worker_pool_event" ) as mock_emit: - with patch.object(MultiConcurrentRunner, "_run_pre_fork_hooks"): - order_tracker.attach_mock(mock_emit, "emit") - MultiConcurrentRunner.run_concurrent( - "h", "a", False, "/sock", max_concurrency=3 - ) + order_tracker.attach_mock(mock_emit, "emit") + MultiConcurrentRunner.run_concurrent( + "h", "a", False, "/sock", max_concurrency=3 + ) mock_emit.assert_called_once_with(3) self.assertEqual(order_tracker.mock_calls[0], call.emit(3)) @patch("awslambdaric.lambda_multi_concurrent_utils.logging") - def test_emit_worker_pool_event_emits_pool_size(self, mock_logging): + @patch("awslambdaric.lambda_multi_concurrent_utils.bootstrap") + def test_emit_worker_pool_event_sets_up_parent_logging_and_emits( + self, mock_bootstrap, mock_logging + ): with patch.object(MultiConcurrentRunner, "_redirect_output") as mock_redirect: MultiConcurrentRunner._emit_worker_pool_event(16) # Parent never redirects: RAPID wires its stdout at spawn. mock_redirect.assert_not_called() + mock_bootstrap.init_logging.assert_called_once_with() mock_logging.getLogger.return_value.debug.assert_called_once() event = mock_logging.getLogger.return_value.debug.call_args[0][0] self.assertEqual(event["workerCount"], 16) self.assertEqual(event["executionEnvironmentMaxConcurrency"], 16) - # Setting up and releasing the sink is _before_fork's job now. - mock_logging.getLogger.return_value.handlers.clear.assert_not_called() + mock_logging.getLogger.return_value.handlers.clear.assert_called_once_with() + # Sink is closed deterministically after the handler is removed. + mock_bootstrap.init_logging.return_value.__exit__.assert_called_once_with( + None, None, None + ) @patch( "awslambdaric.lambda_multi_concurrent_utils.LambdaMultiConcurrentRuntimeClient"