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

docs/serverless/worker.md

@@ -18,14 +18,68 @@ runpod.serverless.start({"handler": handler})
 
 The `config` parameter is a dictionary containing the following keys:
 
-| Key       | Type       | Description                                                  |
-|-----------|------------|--------------------------------------------------------------|
-| `handler` | `function` | The handler function that will be called with the job input. |
+| Key               | Type       | Description                                                        |
+|-------------------|------------|--------------------------------------------------------------------|
+| `handler`         | `function` | The handler function called with each job input.                   |
+| `prestart_timeout`| `number`   | Optional deadline in seconds for all registered prestart hooks.    |
 
 ### handler
 
 The handler function can either have a standard return or be a generator function. If the handler is a generator function, it will be called with the job input and the generator will be iterated over until it is exhausted.
 
+## Prestart hooks
+
+Queue-based workers can register startup work that must finish before the
+handler receives jobs. Hooks run once per worker process, sequentially in registration order.
+
+```python
+import runpod
+
+model = None
+
+@runpod.serverless.register_prestart_hook
+def load_model():
+    global model
+    model = load_weights()
+
+def handler(job):
+    return model(job["input"])
+
+runpod.serverless.start({
+    "handler": handler,
+    "prestart_timeout": 600,
+})
+```
+
+Support depends on the runtime mode:
+
+| Mode | Support | Behavior |
+|------|---------|----------|
+| Production queue | Full | Queue job intake runs concurrently with prestart. Handlers run after prestart finished. Prestart failure is attached to held requests. |
+| Local test input | Supported | Hooks run before the synthetic request. |
+| Hosted development API | Supported with `--rp_api_concurrency 1` | FastAPI lifespan runs hooks before serving. One Uvicorn worker keeps startup state and the handler in the same process. |
+| Realtime | Unsupported | Realtime has separate worker cardinality, readiness, and persistent-connection failure semantics that this hook contract does not define. |
+| Load-balanced endpoints | Not applicable | These images own their HTTP server lifecycle and do not start through this SDK worker entrypoint. |
+
+If shutdown begins while a claimed request waits for prestart, that request
+fails with `prestart_cancelled`.
+
+## Failure logs
+
+When a prestart hook fails, the SDK can attach its `stdout`/`stderr` to the
+reported failure as `logs`. Handler failure logs require explicit opt-in:
+
+| Value | Behavior |
+|-------|----------|
+| `auto` (default) | Capture prestart failures when hooks are registered. |
+| `all` | Capture prestart and handler failures. |
+| `off` | Never capture. |
+
+Capture replaces `sys.stdout`/`sys.stderr` at worker startup, so it sees direct
+stream writes made within an enabled hook or handler capture. It does not see
+child processes, log handlers created before startup, or threads your code
+starts directly (`asyncio.to_thread` is captured).
+
 ## Worker Refresh
 
 For more complex operations where you are downloading files or making changes to the worker, it can be beneficial to refresh the worker between jobs. This can be accomplished by enabling a `refresh_worker` worker flag in one of two ways:

runpod/serverless/__init__.py

@@ -10,21 +10,24 @@
 import signal
 import sys
 import time
-from typing import Any, Dict
+from typing import Any
 
 from ..version import __version__ as runpod_version
 from . import worker
+from .modules.rp_fitness import register_fitness_check
 from .modules.rp_logger import RunPodLogger
+from .modules.rp_prestart import has_prestart_hooks as _has_prestart_hooks
+from .modules.rp_prestart import register_prestart_hook
 from .modules.rp_progress import progress_update
-from .modules.rp_fitness import register_fitness_check
 from .utils.rp_volume_cache import VolumeCache
 
 __all__ = [
-    "start",
+    "VolumeCache",
     "progress_update",
     "register_fitness_check",
+    "register_prestart_hook",
     "runpod_version",
-    "VolumeCache",
+    "start",
 ]
 
 log = RunPodLogger()
@@ -84,7 +87,7 @@
 )
 
 
-def _set_config_args(config) -> dict:
+def _set_config_args(config: dict[str, Any]) -> dict[str, Any]:
     """
     Sets the config rp_args, removing any recognized arguments from sys.argv.
     Returns: config
@@ -133,18 +136,49 @@ def _signal_handler(sig, frame):
     sys.exit(0)
 
 
+def _validate_prestart_mode(config: dict[str, Any], realtime_port: int) -> None:
+    """Check whether registered hooks have a safe adapter for the selected mode.
+
+    Queue and local-input modes are single-process SDK lifecycles. Hosted API
+    mode is supported only with one Uvicorn worker so the hook runs exactly once
+    in the same process as the handler. Realtime is rejected because its worker
+    cardinality, readiness, and persistent-connection failure contract are not
+    defined for prestart hooks.
+    """
+    if not _has_prestart_hooks():
+        return
+
+    if config["rp_args"]["rp_serve_api"]:
+        if config["rp_args"]["rp_api_concurrency"] != 1:
+            raise RuntimeError(
+                "Prestart hooks require rp_api_concurrency=1 in hosted API mode."
+            )
+        return
+
+    if realtime_port:
+        raise RuntimeError("Prestart hooks are not supported in realtime mode.")
+
+
 # ---------------------------------------------------------------------------- #
 #                            Start Serverless Worker                           #
 # ---------------------------------------------------------------------------- #
-def start(config: Dict[str, Any]):
+def start(config: dict[str, Any]):
     """
     Starts the serverless worker.
 
-    config (Dict[str, Any]): Configuration parameters for the worker.
+    config (dict[str, Any]): Configuration parameters for the worker.
 
     config["handler"] (Callable): The handler function to run.
 
-    config["rp_args"] (Dict[str, Any]): Arguments for the worker, populated by runtime arguments.
+    config["rp_args"] (dict[str, Any]): Arguments populated by runtime arguments.
+
+    Prestart hooks registered with `register_prestart_hook` run once before
+    handler execution in queue-based, local test, and hosted API modes.
+    Production queue intake continues while hooks run; local and hosted API
+    handlers do not accept work until every hook finishes.
+
+    config["prestart_timeout"] (int, optional): Seconds allowed for the complete
+        prestart phase. Omit for no timeout.
     """
     print(f"--- Starting Serverless Worker |  Version {runpod_version} ---")
 
@@ -156,9 +190,12 @@ def start(config: Dict[str, Any]):
     realtime_port = _get_realtime_port()
     realtime_concurrency = _get_realtime_concurrency()
 
+    _validate_prestart_mode(config, realtime_port)
+
     if config["rp_args"]["rp_serve_api"]:
         log.info("Starting API server.")
         from .modules import rp_fastapi
+
         api_server = rp_fastapi.WorkerAPI(config)
 
         api_server.start_uvicorn(
@@ -171,6 +208,7 @@ def start(config: Dict[str, Any]):
     if realtime_port:
         log.info(f"Starting API server for realtime on port {realtime_port}.")
         from .modules import rp_fastapi
+
         api_server = rp_fastapi.WorkerAPI(config)
 
         api_server.start_uvicorn(

runpod/serverless/modules/rp_capture.py

@@ -0,0 +1,149 @@
+"""
+runpod | serverless | rp_capture.py
+
+Captures stdout/stderr to report upon handler or prestart failure.
+Swaps `sys.stdout`/`sys.stderr` for a tee proxy that writes to both the
+real stream and a buffer in a contextvar.
+
+Captured output is attached to failure payloads, which are returned to whoever
+called the request, so capture is not on by default. `RUNPOD_LOG_CAPTURE`
+selects when it runs:
+
+    auto (default)  capture prestart failures when hooks are registered
+    all             capture prestart and handler failures
+    off             never capture
+"""
+
+import contextlib
+import contextvars
+import io
+import os
+import sys
+from collections.abc import Generator
+
+MAX_CAPTURED_CHARS = 16 * 1024
+
+CAPTURE_AUTO = "auto"
+CAPTURE_ALL = "all"
+CAPTURE_OFF = "off"
+_CAPTURE_MODES = (CAPTURE_AUTO, CAPTURE_ALL, CAPTURE_OFF)
+
+# Capture buffer for the current context
+_current: "contextvars.ContextVar[_RingBuffer | None]" = contextvars.ContextVar(
+    "rp_stdio_capture", default=None
+)
+
+
+class _RingBuffer:
+    """Keeps only the last `limit` characters since the tail is usually where the
+    failure reason is."""
+
+    def __init__(self, limit: int = MAX_CAPTURED_CHARS):
+        self.limit = limit
+        self._buf = ""
+
+    def write(self, text: str) -> int:
+        self._buf = (self._buf + text)[-self.limit :]
+        return len(text)
+
+    def getvalue(self) -> str:
+        return self._buf
+
+
+class _TeeProxy:
+    """Forwards to the real stream and mirrors into its buffer."""
+
+    def __init__(self, real):
+        self._real = real
+
+    def write(self, text) -> int:
+        n = self._real.write(text)
+        buffer = _current.get()
+        if buffer is not None:
+            with contextlib.suppress(Exception):
+                buffer.write(text)
+        return n
+
+    def writelines(self, lines) -> None:
+        # Not covered by write(); callers that use it would otherwise bypass capture.
+        for line in lines:
+            self.write(line)
+
+    def flush(self) -> None:
+        self._real.flush()
+
+    def __getattr__(self, name):
+        # Delegate everything else to the real stream
+        return getattr(self._real, name)
+
+
+io.TextIOBase.register(_TeeProxy)
+
+
+def capture_mode() -> str:
+    """Resolve `RUNPOD_LOG_CAPTURE`, falling back to `auto` on anything unknown."""
+    mode = os.environ.get("RUNPOD_LOG_CAPTURE", CAPTURE_AUTO).strip().lower()
+    return mode if mode in _CAPTURE_MODES else CAPTURE_AUTO
+
+
+def handler_capture_enabled() -> bool:
+    """Return whether handler failure logs were explicitly enabled."""
+    return capture_mode() == CAPTURE_ALL
+
+
+def install(*, hooks_registered: bool = False) -> None:
+    """Install the tee proxy on stdout/stderr, if this worker wants capture.
+
+    Idempotent. `hooks_registered` is what `auto` mode keys off, and the caller
+    passes it in so this module stays independent of the prestart registry.
+    """
+    mode = capture_mode()
+    if mode == CAPTURE_OFF:
+        return
+
+    if mode == CAPTURE_AUTO and not hooks_registered:
+        return
+
+    if not isinstance(sys.stdout, _TeeProxy):
+        sys.stdout = _TeeProxy(sys.stdout)
+    if not isinstance(sys.stderr, _TeeProxy):
+        sys.stderr = _TeeProxy(sys.stderr)
+
+
+@contextlib.contextmanager
+def capture(*, enabled: bool = True) -> Generator[_RingBuffer]:
+    """Capture stdout/stderr written within this context (and within threads it spawns via
+    `asyncio.to_thread`), while still passing everything through to the real streams."""
+    buffer = _RingBuffer()
+    if not enabled:
+        yield buffer
+        return
+
+    token = _current.set(buffer)
+    try:
+        yield buffer
+    finally:
+        # Suppress an abandoned async generator to avoid polluting stderr
+        with contextlib.suppress(ValueError):
+            _current.reset(token)
+
+
+@contextlib.contextmanager
+def paused() -> Generator[None]:
+    """Temporarily stop mirroring logs into the buffer. Currently only used to
+    omit diagnostic logs."""
+    token = _current.set(None)
+    try:
+        yield
+    finally:
+        with contextlib.suppress(ValueError):
+            _current.reset(token)
+
+
+def clip(text: str, limit: int = MAX_CAPTURED_CHARS) -> str:
+    """Truncate an error string, keeping the head and tail (the useful parts)."""
+    if not text or len(text) <= limit:
+        return text
+    keep = limit // 2
+    omitted = len(text) - 2 * keep
+    return f"{text[:keep]}\n...[{omitted} characters truncated]...\n{text[-keep:]}"