feat: SLS-494 run users' worker initialization code

runpod/serverless/__init__.py

@@ -145,6 +145,12 @@ def start(config: Dict[str, Any]):
     config["handler"] (Callable): The handler function to run.
 
     config["rp_args"] (Dict[str, Any]): Arguments for the worker, populated by runtime arguments.
+
+    config["initializer"] (Callable, optional): Startup work before the job loop begins,
+        e.g. loading a model or starting an inference engine.
+
+    config["init_timeout"] (int, optional): Seconds to allow the initializer before
+        treating it as a failure. Omit for no timeout.
     """
     print(f"--- Starting Serverless Worker |  Version {runpod_version} ---")
 

runpod/serverless/modules/rp_capture.py

@@ -0,0 +1,94 @@
+"""
+runpod | serverless | rp_capture.py
+
+Captures stdout/stderr, to be reported upon handler or initializer failure.
+Swaps `sys.stdout`/`sys.stderr` for a tee proxy that writes to both the
+real stream and a buffer in a contextvar.
+"""
+
+import contextlib
+import contextvars
+import sys
+from collections.abc import Generator
+
+MAX_CAPTURED_CHARS = 16 * 1024
+
+# Capture buffer for the current context
+_current: "contextvars.ContextVar[_RingBuffer | None]" = contextvars.ContextVar(
+    "rp_stdio_capture", default=None
+)
+
+_installed = False
+
+
+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 flush(self) -> None:
+        self._real.flush()
+
+    def __getattr__(self, name):
+        # Delegate everything else to the real stream
+        return getattr(self._real, name)
+
+
+def install() -> None:
+    """Install the tee proxy on stdout/stderr. Idempotent."""
+    global _installed
+    if _installed:
+        return
+    sys.stdout = _TeeProxy(sys.stdout)
+    sys.stderr = _TeeProxy(sys.stderr)
+    _installed = True
+
+
+@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:]}"

runpod/serverless/modules/rp_http.py

@@ -1,7 +1,8 @@
 """
-    This module is used to handle HTTP requests.
+This module is used to handle HTTP requests.
 """
 
+import asyncio
 import json
 import os
 
@@ -23,6 +24,12 @@
 )
 JOB_STREAM_URL = JOB_STREAM_URL_TEMPLATE.replace("$RUNPOD_POD_ID", WORKER_ID)
 
+HANDLER_STARTED_URL = (
+    JOB_DONE_URL.replace("/job-done/", "/handler-started/")
+    if "/job-done/" in JOB_DONE_URL
+    else None
+)
+
 log = RunPodLogger()
 
 
@@ -96,3 +103,31 @@ async def stream_result(session, job_data, job):
     await _handle_result(
         session, job_data, job, JOB_STREAM_URL, "Intermediate results sent."
     )
+
+
+async def report_handler_started(session, job):
+    """When a job is taken but still waiting for worker initialization, don't bill.
+    When the handler actually starts running, report to the platform so billing can start.
+    """
+    if not HANDLER_STARTED_URL:
+        return
+    url = HANDLER_STARTED_URL.replace("$ID", job["id"])
+
+    async def _post():
+        retry_client = RetryClient(
+            client_session=session, retry_options=FibonacciRetry(attempts=2)
+        )
+        async with retry_client.post(
+            url,
+            data="{}",
+            headers={
+                "charset": "utf-8",
+                "Content-Type": "application/x-www-form-urlencoded",
+            },
+        ) as client_response:
+            await client_response.text()
+
+    try:
+        await asyncio.wait_for(_post(), timeout=2)
+    except Exception as err:  # noqa: BLE001 - best-effort compute-start signal
+        log.debug(f"handler-started signal failed (non-fatal): {err}", job["id"])

runpod/serverless/modules/rp_initializer.py

@@ -0,0 +1,151 @@
+"""
+runpod | serverless | initializer
+
+Runs the user's startup initialization code concurrently with the job loop.
+The loop may take a request right away, but the handler is not called until the initializer
+finishes. On failure or timeout, the error + stdout/stderr are attached to
+the current request.
+
+A sync/blocking initializer is offloaded to a worker thread so it does not starve the
+loop; an async one is awaited directly.
+"""
+
+import asyncio
+import contextlib
+import contextvars
+import inspect
+import json
+import os
+import threading
+import traceback
+from collections.abc import Callable
+from typing import Any
+
+import requests
+
+from runpod.http_client import get_auth_header
+from runpod.serverless.modules.rp_capture import MAX_CAPTURED_CHARS, clip
+from runpod.serverless.modules.rp_logger import RunPodLogger
+from runpod.serverless.modules.worker_state import WORKER_ID
+from runpod.version import __version__ as runpod_version
+
+log = RunPodLogger()
+
+INIT_FAILED_EVENT = "init_failed"
+
+
+class InitializerTimeout(Exception):
+    """Raised when the initializer exceeds `init_timeout`."""
+
+
+class InitializerError(Exception):
+    """Wraps any exception raised by the user's initializer."""
+
+    def __init__(self, original: BaseException):
+        self.original = original
+        super().__init__(str(original))
+
+
+def build_init_failed_payload(exc: BaseException, logs: str = "") -> dict[str, Any]:
+    """Failure reason as a structured dict, using the same core fields as a handler error
+    (type, message, traceback). `logs` contains stdout/stderr."""
+    original = getattr(exc, "original", exc)
+    payload = {
+        "event": INIT_FAILED_EVENT,
+        "phase": "init",
+        "error_type": type(original).__name__,
+        "error_message": clip(str(original)),
+        "error_traceback": clip(
+            "".join(
+                traceback.format_exception(
+                    type(original), original, original.__traceback__
+                )
+            )
+        ),
+        "worker_id": WORKER_ID,
+        "runpod_version": runpod_version,
+    }
+    if logs:
+        payload["logs"] = logs[-MAX_CAPTURED_CHARS:]
+    return payload
+
+
+def emit_init_failed(payload: dict[str, Any]) -> None:
+    """Log the structured `init_failed` event and report it to the platform."""
+    log.error(f"{INIT_FAILED_EVENT} | {json.dumps(payload)}")
+    _report_to_platform(payload)
+
+
+def _report_to_platform(payload: dict[str, Any]) -> None:
+    """Best-effort POST of the failure reason back to the platform."""
+    ping_url = os.environ.get("RUNPOD_WEBHOOK_PING", "")
+    if not ping_url or "/ping/" not in ping_url:
+        return
+    report_url = ping_url.replace("$RUNPOD_POD_ID", WORKER_ID).replace(
+        "/ping/", "/init-failed/"
+    )
+    try:
+        requests.post(report_url, json=payload, headers=get_auth_header(), timeout=5)
+        log.info("Initializer | reported init_failed to platform")
+    except Exception as exc:  # noqa: BLE001 - best-effort; don't block
+        log.warn(f"Initializer | init_failed report failed (non-fatal): {exc}")
+
+
+async def _run_sync_in_daemon(fn: Callable) -> None:
+    """Run a blocking callable on a daemon thread instead of `asyncio.to_thread` so if stuck,
+    it can be abandoned an die without blocking the executor shutdown and process exit."""
+    loop = asyncio.get_running_loop()
+    done = loop.create_future()
+    box: dict[str, BaseException] = {}
+    ctx = contextvars.copy_context()
+
+    def worker():
+        try:
+            ctx.run(fn)
+        except BaseException as exc:  # noqa: BLE001 - re-raised on the loop below
+            box["exc"] = exc
+        finally:
+
+            def _signal():
+                if not done.done():
+                    done.set_result(None)
+
+            # loop already closed
+            with contextlib.suppress(RuntimeError):
+                loop.call_soon_threadsafe(_signal)
+
+    threading.Thread(target=worker, name="rp-initializer", daemon=True).start()
+    await done
+    if "exc" in box:
+        raise box["exc"]
+
+
+async def run_initializer_async(
+    initializer: Callable, timeout: int | None = None
+) -> None:
+    """Run the initializer to completion inside the running event loop, raising
+    `InitializerTimeout` on timeout or `InitializerError` for any other failure."""
+    log.info("Initializer | init started")
+    try:
+        if inspect.iscoroutinefunction(initializer) or inspect.iscoroutinefunction(
+            initializer.__call__
+        ):
+            awaitable = initializer()
+        else:
+            awaitable = _run_sync_in_daemon(initializer)
+
+        if timeout:
+            await asyncio.wait_for(asyncio.ensure_future(awaitable), timeout=timeout)
+        else:
+            await awaitable
+    except asyncio.TimeoutError as exc:
+        raise InitializerTimeout(
+            f"initializer exceeded init_timeout of {timeout}s"
+        ) from exc
+    except (InitializerError, InitializerTimeout):
+        raise
+    except SystemExit as exc:
+        raise InitializerError(exc) from exc
+    except Exception as exc:
+        raise InitializerError(exc) from exc
+    log.info("Initializer | ready")