feat: SLS-497 Add explicit prestart hooks to the runpod-python SDK

[jasonwang-runpod]

Problem

Workers often load models or start inference engines before calling runpod.serverless.start(). The SDK cannot observe that prestart work, so failures leave requests in IN_QUEUE status while useful info remains in worker logs.

Solution

Add explicit prestart hooks that the SDK supervises before handler execution.

• Workers can register sync or async hooks with @runpod.serverless.register_prestart_hook and provide an optional timeout for the prestart phase.
• Queue-based workers take jobs concurrently with prestart, but don't run the handler function until prestart succeeds.
• A failure or timeout is attached to a request as prestart_failed. Then the worker drains and exits.
• Local input and single-worker hosted API mode run prestart before invoking or serving the handler. Realtime and multi-process hosted API modes are unsupported.
• Failure payloads can carry the last 16 KB of the worker's stdout and stderr. RUNPOD_LOG_CAPTURE gates this: auto (default) captures only when hooks are registered, all always captures, off never does. Note: log capture doesn't see child processes.
• A worker that doesn't register any hooks is unaffected.

Testing

uv run pytest -q: 710 passed, 94.7% coverage.

An sglang endpoint was tested w/ different scenarios:

1. no hooks: COMPLETED, output only, no logs field
2. hook success: COMPLETED, engine boot ran inside the prestart hook
3. failing hook: FAILED (as expected) with reason prestart_failed
4. prestart_timeout=5 with slow hook: FAILED (as expected) in 4.9s with PrestartTimeout; logs had the captured stdout
5. hook ok + handler raises, capture=auto: FAILED (as expected), with no logs field

Supersedes #567.

[Copilot]

Pull request overview

Adds supervised prestart hooks across supported Serverless worker modes, including startup-failure reporting and bounded stdout/stderr capture.

Changes:

• Adds ordered sync/async prestart hooks with optional phase timeout.
• Gates handlers during initialization and drains failed queue workers.
• Captures bounded logs for startup and handler failures.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
runpod/serverless/__init__.py	Exposes hooks and validates runtime modes.
runpod/serverless/worker.py	Installs output capture for workers.
runpod/serverless/modules/rp_prestart.py	Implements hook registration and execution.
runpod/serverless/modules/rp_capture.py	Implements bounded contextual output capture.
runpod/serverless/modules/rp_scale.py	Integrates queue gating, failure delivery, and shutdown.
runpod/serverless/modules/rp_local.py	Runs hooks before local handlers.
runpod/serverless/modules/rp_fastapi.py	Runs hooks through FastAPI lifespan.
runpod/serverless/modules/rp_job.py	Attaches captured logs to handler failures.
docs/serverless/worker.md	Documents prestart configuration and modes.
tests/test_serverless/test_prestart.py	Tests the public hook contract and mode guards.
tests/test_serverless/test_prestart_lifecycle.py	Tests queue lifecycle and process termination.
tests/test_serverless/test_capture.py	Tests capture and output bounds.
tests/test_serverless/test_modules/test_local.py	Tests local prestart behavior.
tests/test_serverless/test_init.py	Verifies the new public export.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated no new comments.

Suppressed comments (3)

docs/serverless/worker.md:77

• This overstates what the implementation captures. The active buffer is a ContextVar, and a normal threading.Thread starts without that context, so direct stdout/stderr writes from engine/background threads created by a hook or handler are omitted even while that hook or handler is running. Document the child-thread limitation (or propagate the context) so users do not rely on missing diagnostics.

Capture replaces `sys.stdout`/`sys.stderr` at worker startup, so it sees `print`
and direct stream writes made while a hook or handler runs. It does not see
child processes and log handlers created before startup.

runpod/serverless/modules/rp_capture.py:26

• MAX_CAPTURED_CHARS limits Unicode code points, not the encoded size sent to /job-done. A 16,384-character non-ASCII tail can occupy up to 64 KB in UTF-8 (and the inner json.dumps can expand characters further), so the advertised 16 KB bound and the payload-size protection are not enforced. Bound the encoded/serialized byte length while preserving valid character boundaries; the ring buffer, clip, and prestart payload slicing need to use the same byte-based contract.

MAX_CAPTURED_CHARS = 16 * 1024

tests/test_serverless/test_prestart.py:91

• This assertion has only 10 ms of scheduling slack. On a loaded CI runner the first 30 ms sleep can resume after the 40 ms phase deadline, making the reported hook first and intermittently failing the test. Make the first hook only yield once and have the second block indefinitely so the timeout deterministically occurs in second.

    def test_timeout_bounds_the_whole_phase_and_names_current_hook(self):
        async def first():
            await asyncio.sleep(0.03)

        async def second():
            await asyncio.sleep(0.03)

        with self.assertRaises(PrestartTimeout) as ctx:
            _run(run_prestart_hooks_async((first, second), timeout=0.04))

[deanq]

Reviewed the prestart-hooks + capture changes. The concurrency design (concurrent intake with a gated handler, failing the claimed request on startup failure) is sound and well-tested. Flagging a few correctness/security gaps on the failure and shutdown paths, plus one nit.

One cross-cutting note: the prestart-failure path routes termination through rp_fitness._terminate_unhealthy (the fitness force-kill helper) but defers its os._exit behind the graceful loop-drain -- see the inline note on run(). That partially undoes the hard-exit guarantee fitness relies on. Details inline.

[deanq]

Registering any prestart hook flips hooks_registered=True, which in the default auto mode enables the stdout/stderr tee for the whole worker. Since run_job already wraps every handler in capture(), a handler that prints sensitive data and then raises will have those lines captured into error_info['logs'] and returned to the caller in the job result.

Concretely: a worker adopts an unrelated prestart hook -> capture turns on -> a handler does print(f"HF_TOKEN={token}") then raises -> the token lands in the error payload polled by the client.

The coupling (prestart-hook presence silently enables handler-failure log capture) is non-obvious and has data-exfiltration impact. Suggest decoupling: gate handler-failure capture on an explicit opt-in rather than on hook registration, or redact/scrub captured logs before embedding them in the returned error.

[deanq]

If shutdown begins while prestart is still running, _wait_for_prestart() returns False, and handle_job logs "leaving this request for another worker" and returns -- without calling send_result or re-queuing the job.

But this worker already claimed the job. If the platform does not re-dispatch an already-assigned request, it's now neither completed nor failed: the caller waits out its queue TTL with no result -- a silent failure.

Suggest failing the claimed request explicitly on this path (send an error result) rather than relying on server-side re-dispatch that may never happen.

[deanq]

_run_prestart sets _prestart_ready in a finally block. If prestart is cancelled during shutdown (mid-hook), run_prestart_phase raises CancelledError, which still propagates through finally: self._prestart_ready.set() -- so _prestart_ready becomes set while _prestart_error remains None.

A handler concurrently in _wait_for_prestart then observes _prestart_ready.is_set() == True, returns True, sees _prestart_error is None, and runs the handler against a half-initialized environment (e.g. model still None) even though prestart never actually completed.

Suggest distinguishing "prestart finished" from "prestart cancelled" -- e.g. set an explicit error/cancelled state before setting the ready event, so _wait_for_prestart doesn't release the handler on cancellation.

[deanq]

On prestart failure the worker can't actually terminate until in-flight long-polls return. kill_worker() only sets an event the loops check at their top, but run() blocks here on await request_loops_future -- and get_jobs / monitor_stop_signals are parked mid-wait_for with 90s fetch timeouts. So _terminate_unhealthy(1) (line 196) is deferred up to ~90s after the failure was already reported, leaving a known-broken worker alive and schedulable.

This partially defeats the point of rp_fitness._terminate_unhealthy, whose whole rationale is os._exit precisely because "sys.exit can hang forever and the worker keeps serving jobs." Routing the prestart-failure exit through the graceful loop-drain reintroduces exactly that hang.

Suggest calling the hard-exit path (or cancelling the request loops) immediately on prestart failure, rather than awaiting request_loops_future first.

[deanq]

In the cancellation branch, after hook_task.cancel() this does with contextlib.suppress(BaseException): await hook_task. If an async hook catches CancelledError and keeps running (e.g. a retry loop that ignores cancellation), await hook_task never returns -- so the prestart-timeout / shutdown path that got us here never completes, and the worker deadlocks instead of terminating.

Suggest bounding this await (e.g. await asyncio.wait_for(hook_task, timeout=...) and proceeding regardless on timeout) so an uncooperative hook can't wedge the timeout/shutdown path.

[deanq]

self.prestart_claim_timeout = 10 is a bare inline magic number with no name explaining intent and no unit. This PR already defines named constants elsewhere (e.g. MAX_CAPTURED_CHARS in rp_capture.py), so this is inconsistent. Suggest a module-level named constant with the unit in the name, e.g. PRESTART_CLAIM_TIMEOUT_SECONDS = 10.