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

docs/serverless/worker.md

@@ -18,14 +18,126 @@ 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:
+
+```python
+import runpod
+
+model = None
+
+@runpod.serverless.register_prestart_hook
+def load_model():
+    global model
+    model = load_weights()
+
+@runpod.serverless.register_prestart_hook
+async def warm_cache():
+    await cache.prime()
+
+def handler(job):
+    return model(job["input"])
+
+runpod.serverless.start({
+    "handler": handler,
+    "prestart_timeout": 600,
+})
+```
+
+Hooks run once per worker process, sequentially in registration order. Sync
+hooks run on a daemon thread; async hooks run on the active worker event loop.
+In production, the worker can acquire jobs during this phase, but it does not
+call the handler until every hook succeeds.
+
+External engines can publish a process handle for the handler to reuse:
+
+```python
+import subprocess
+
+engine = None
+
+@runpod.serverless.register_prestart_hook
+def start_engine():
+    global engine
+    engine = subprocess.Popen(["vllm", "serve", "my-model"])
+    wait_until_healthy(engine)
+```
+
+`prestart_timeout` covers the complete sequence, not each hook separately. If
+a hook raises or the phase times out, later hooks do not run. In production,
+the SDK reports the failing hook, exception, traceback, and captured output
+against jobs held by that worker, then exits so the platform can replace it.
+
+Support depends on the runtime mode:
+
+| Mode | Support | Behavior |
+|------|---------|----------|
+| Production queue | Full | Queue intake may start during prestart. Handlers wait behind a gate. Failure is attached to held or newly claimed work before the worker exits. |
+| Local test input | Supported | Hooks run before the synthetic request. Failure logs `prestart_failed`, skips the handler, and exits nonzero. |
+| 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. Failure logs `prestart_failed` and prevents API startup. |
+| 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. |
+
+The SDK rejects hosted API concurrency above one and any realtime configuration
+when hooks are registered. It never silently skips registered startup work.
+
+## Failure logs
+
+When a prestart hook or a handler fails, the SDK can attach the last 16 KB of
+whatever the worker wrote to `stdout`/`stderr` to the failure it reports, under a
+`logs` key:
+
+```json
+{
+  "error_type": "<class 'RuntimeError'>",
+  "error_message": "CUDA out of memory",
+  "error_traceback": "Traceback (most recent call last): ...",
+  "logs": "Loading weights from /models/my-model ...\n"
+}
+```
+
+That payload is returned to whoever called the request, so capture is **not on by
+default**. `RUNPOD_LOG_CAPTURE` controls it:
+
+| Value | Behavior |
+|-------|----------|
+| `auto` (default) | Capture only when prestart hooks are registered. A worker that has not adopted prestart hooks is untouched: streams are not wrapped and no `logs` key is added. |
+| `all` | Always capture, including handler failures on workers with no prestart hooks. |
+| `off` | Never capture, even with prestart hooks registered. |
+
+Only the `logs` field is bounded. `error_message` and `error_traceback` are
+reported in full, exactly as they were before this feature existed.
+
+### What capture can and cannot see
+
+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.** A subprocess inherits file descriptors 1 and 2 directly,
+  below Python. If your worker runs an engine as a subprocess and wants its output
+  captured, pipe it (`subprocess.Popen(..., stdout=PIPE)`) and re-emit the lines.
+- **Log handlers created before startup.** `logging.StreamHandler` resolves
+  `sys.stderr` once, when it is constructed. A `logging.basicConfig()` at module
+  import binds the real stream before the SDK starts, so those records bypass
+  capture. Call `logging.basicConfig()` from inside a prestart hook if you want
+  its output attached.
+- **Plain threads.** Output from a thread your code starts directly is not
+  captured; `asyncio.to_thread` is.
+
+In all of these cases output still reaches your worker logs as usual — only the
+`logs` field is affected.
+
 ## 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,140 @@
+"""
+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 only when prestart hooks are registered, so a worker
+                    that has not adopted prestart behaves exactly as before
+    all             always capture, including handler failures on workers with
+                    no prestart hooks
+    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)
+
+
+# Registered rather than subclassed: inheriting io.TextIOBase would shadow
+# encoding/errors/newlines with None instead of delegating to the real stream.
+# Registration is what keeps `isinstance(sys.stdout, io.TextIOBase)` true for
+# libraries that type-check the stream.
+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 install() -> None:
+    """Install the tee proxy on stdout/stderr, if this worker wants capture.
+
+    Idempotent. Without the proxy nothing ever reaches a capture buffer, so no
+    failure payload carries worker output and the streams are left untouched.
+    """
+    mode = capture_mode()
+    if mode == CAPTURE_OFF:
+        return
+
+    if mode == CAPTURE_AUTO:
+        # Local import: rp_prestart imports this module.
+        from .rp_prestart import has_prestart_hooks
+
+        if not has_prestart_hooks():
+            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() -> 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.
+
+    Yields the buffer; call `.getvalue()` for the captured text."""
+    buffer = _RingBuffer()
+    token = _current.set(buffer)
+    try:
+        yield buffer
+    finally:
+        # Suppress an abandoned async generator to avoid polluting stderr
+        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:]}"