Skip to content

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

Open
vip-amzn wants to merge 1 commit into
aws:mainfrom
vip-amzn:feat/pre-fork-hooks
Open

vip-amzn wants to merge 1 commit into
aws:mainfrom
vip-amzn:feat/pre-fork-hooks

Conversation

@vip-amzn

@vip-amzn vip-amzn commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

Add pre-fork hooks for multi-concurrent mode

Summary

Adds @register_pre_fork, a hook that runs once in the parent process before
worker processes are started, for functions running in multi-concurrent mode
(Lambda Managed Instances).

Motivation

In multi-concurrent mode the runtime interface client starts N worker processes,
and each worker imports the handler module independently. Module-level
initialization therefore runs N times, once per worker.

That is usually fine, but it is wasteful or outright incorrect for one-time setup
whose result is shared outside the process:

  • Downloading a large model or dataset to /tmp — N workers download the
    same object concurrently and race on the same path.
  • Starting a local inference server or sidecar subprocess — N copies start
    and compete for the same port.
  • Warming a local service, or writing a config file every worker reads —
    duplicated work, and the workers can observe a half-written file.

A pre-fork hook lets that work happen exactly once, before any worker exists.

Usage

import subprocess

from awslambdaric.lambda_concurrency_hooks import register_pre_fork


@register_pre_fork
def start_inference_server():
    subprocess.Popen(["python", "serve.py", "--port", "8000"])
    wait_for_ready("http://localhost:8000/health")


def handler(event, context):
    # Every worker talks to the single server started above.
    ...

Behaviour

  • Hooks run in the parent process, after the handler module is imported
    (importing it is what registers them) and before any worker is started.
  • Multiple hooks run in registration order.
  • Hooks run for any worker count, so a hook's side effects are in place
    whatever concurrency the execution environment is configured with.
  • No change to the standard on-demand path.

Hooks cannot share in-memory state with workers

Worker processes re-import the handler module independently, so a pre-fork hook
cannot hand in-memory objects to workers. Hooks are for external side
effects — a subprocess, a file in /tmp, a warmed local service. Anything a
worker needs in memory must still be created in the worker.

Failure handling

A hook that raises fails initialization:

  • the error is reported to the Runtime API as an INIT error with the type
    Runtime.PreForkError,
  • the underlying exception and traceback are logged so the failing hook is
    identifiable,
  • remaining hooks do not run, and no worker is started.

This is deliberate. If a hook could not establish the precondition its workers
depend on — the model was never downloaded, the sidecar never came up — starting
the pool anyway would turn one clear initialization failure into a stream of
confusing per-invocation failures.

Testing

  • Unit tests covering registration, registration order, running exactly once in
    the parent, and the failure path (reported error type, remaining hooks
    skipped, no workers started).
  • 100% line coverage on the new and changed modules.

@vip-amzn
vip-amzn requested review from maxday and trivenay September 24, 2026 15:04
Comment thread awslambdaric/lambda_multi_concurrent_utils.py Outdated

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants