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, and only when the execution environment uses more than one worker. 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,30 @@
+# 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.
+
+    Only runs when the execution environment uses more than one worker.
+
+        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,79 @@ 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()
+
+        # One worker behaves like non-concurrent mode, where module-level
+        # initialization already runs once in that process.
+        if max_concurrency > 1:
+            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 +122,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,