feat: Add pre-fork hooks for multi-concurrent mode

Author: vip-amzn

RELEASE.CHANGELOG.md

@@ -1,3 +1,7 @@
+### September 24, 2026
+`4.1.0`
+- Add pre-fork hooks for multi-concurrent (Lambda Managed Instances) mode. A function can register callables with the `@register_pre_fork` decorator from `awslambdaric.lambda_concurrency_hooks`; they run once in the parent process, after the handler is imported and before worker processes are started. Workers re-import the handler in their own process, so hooks are for external side effects (starting a subprocess, warming a local service, writing to `/tmp`) and share no in-memory state with workers. A hook that raises is reported to the Runtime API as an INIT error with the type `Runtime.PreForkError` and no worker is started. No impact on the standard on-demand path.
+
 ### September 15, 2026
 `4.0.4`
 - Use the `level` key (instead of `log_level`) for the log level field in JSON-formatted uncaught error logs, aligning it with the key used by other structured log events ([#221](https://github.com/aws/aws-lambda-python-runtime-interface-client/pull/221))

awslambdaric/__init__.py

@@ -2,4 +2,4 @@
 Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 """
 
-__version__ = "4.0.4"
+__version__ = "4.1.0"

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,

awslambdaric/lambda_concurrency_hooks.py

@@ -0,0 +1,31 @@
+# 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

awslambdaric/lambda_multi_concurrent_utils.py

@@ -10,6 +10,7 @@
 
 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"
 
@@ -37,23 +38,76 @@ def run_single(
 
     @classmethod
     def _emit_worker_pool_event(cls, max_concurrency: int):
-        """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()
+        """Emit the worker pool DEBUG event. The sink is owned by _before_fork."""
         logging.getLogger().debug(
             {
                 "event": WORKER_POOL_INITIALIZING_EVENT,
                 "workerCount": max_concurrency,
                 "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
@@ -65,7 +119,7 @@ def run_concurrent(
         socket_path: str,
         max_concurrency: int,
     ):
-        cls._emit_worker_pool_event(max_concurrency)
+        cls._before_fork(handler, api_addr, max_concurrency)
 
         processes = []
         for _ in range(max_concurrency):

awslambdaric/lambda_runtime_exception.py

@@ -13,6 +13,7 @@ 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"
 

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()

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._emit_worker_pool_event"
+            "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._before_fork"
         ), patch(
             "awslambdaric.lambda_multi_concurrent_utils.bootstrap.run",
             side_effect=fake_bootstrap_run,