feat: SLS-494 sdk supervises 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 that runs alongside job intake.
+        The worker holds handler execution until the initializer finishes.
+
+    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,91 @@
+"""
+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
+)
+
+
+
+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."""
+    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:]}"

runpod/serverless/modules/rp_initializer.py

@@ -0,0 +1,133 @@
+"""
+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 threading
+import traceback
+from collections.abc import Callable
+from typing import Any
+
+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,
+        "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
+
+
+async def _run_sync_in_daemon(fn: Callable) -> Any:
+    """Run a blocking callable on a daemon thread instead of `asyncio.to_thread` so if stuck,
+    it can be abandoned and die without blocking executor shutdown or process exit."""
+    loop = asyncio.get_running_loop()
+    done = asyncio.Event()
+    results: list[Any] = []
+    errors: list[BaseException] = []
+    ctx = contextvars.copy_context()
+
+    def worker():
+        try:
+            results.append(ctx.run(fn))
+        except Exception as exc:  # noqa: BLE001 - transferred to the event loop below
+            errors.append(exc)
+        except (
+            KeyboardInterrupt,
+            SystemExit,
+            GeneratorExit,
+            asyncio.CancelledError,
+        ) as exc:
+            errors.append(exc)
+        finally:
+            with contextlib.suppress(RuntimeError):
+                loop.call_soon_threadsafe(done.set)
+
+    threading.Thread(target=worker, name="rp-initializer", daemon=True).start()
+    await done.wait()
+    if errors:
+        raise errors[0]
+    return results[0] if results else None
+
+
+async def _invoke_initializer(initializer: Callable) -> None:
+    if inspect.iscoroutinefunction(initializer) or inspect.iscoroutinefunction(
+        initializer.__call__
+    ):
+        result = initializer()
+    else:
+        result = await _run_sync_in_daemon(initializer)
+
+    if inspect.isawaitable(result):
+        await result
+
+
+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:
+        awaitable = _invoke_initializer(initializer)
+        if timeout is not None:
+            await asyncio.wait_for(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")

runpod/serverless/modules/rp_job.py

@@ -15,6 +15,7 @@
 
 from ...version import __version__ as runpod_version
 from ..utils import rp_debugger
+from .rp_capture import capture, clip
 from .rp_handler import is_generator
 from .rp_http import send_result, stream_result
 from .rp_tips import check_return_size
@@ -253,54 +254,55 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]:
     log.info("Started.", job["id"])
     run_result = {}
 
-    try:
-        handler_return = handler(job)
-        job_output = (
-            await handler_return
-            if inspect.isawaitable(handler_return)
-            else handler_return
-        )
-
-        log.debug(f"Handler output: {job_output}", job["id"])
+    with capture() as cap:
+        try:
+            handler_return = handler(job)
+            job_output = (
+                await handler_return
+                if inspect.isawaitable(handler_return)
+                else handler_return
+            )
 
-        if isinstance(job_output, dict):
-            error_msg = job_output.pop("error", None)
-            refresh_worker = job_output.pop("refresh_worker", None)
-            run_result["output"] = job_output
+            log.debug(f"Handler output: {job_output}", job["id"])
 
-            if error_msg:
-                run_result["error"] = error_msg
-            if refresh_worker:
-                run_result["stopPod"] = True
+            if isinstance(job_output, dict):
+                error_msg = job_output.pop("error", None)
+                refresh_worker = job_output.pop("refresh_worker", None)
+                run_result["output"] = job_output
 
-        elif isinstance(job_output, bool):
-            run_result = {"output": job_output}
+                if error_msg:
+                    run_result["error"] = error_msg
+                if refresh_worker:
+                    run_result["stopPod"] = True
 
-        else:
-            run_result = {"output": job_output}
+            elif isinstance(job_output, bool):
+                run_result = {"output": job_output}
 
-        if run_result.get("output") == {}:
-            run_result.pop("output")
+            else:
+                run_result = {"output": job_output}
 
-        check_return_size(run_result)  # Checks the size of the return body.
+            if run_result.get("output") == {}:
+                run_result.pop("output")
 
-    except Exception as err:
-        error_info = {
-            "error_type": str(type(err)),
-            "error_message": str(err),
-            "error_traceback": traceback.format_exc(),
-            "hostname": os.environ.get("RUNPOD_POD_HOSTNAME", "unknown"),
-            "worker_id": os.environ.get("RUNPOD_POD_ID", "unknown"),
-            "runpod_version": runpod_version,
-        }
+            check_return_size(run_result)  # Checks the size of the return body.
 
-        log.error("Captured Handler Exception", job["id"])
-        log.error(json.dumps(error_info, indent=4))
-        run_result = {"error": json.dumps(error_info)}
+        except Exception as err:  # noqa: BLE001 - user handler may raise anything; surface it
+            captured_logs = cap.getvalue()
+            error_info = {
+                "error_type": str(type(err)),
+                "error_message": clip(str(err)),
+                "error_traceback": clip(traceback.format_exc()),
+                "hostname": os.environ.get("RUNPOD_POD_HOSTNAME", "unknown"),
+                "worker_id": os.environ.get("RUNPOD_POD_ID", "unknown"),
+                "runpod_version": runpod_version,
+                "logs": captured_logs,
+            }
 
-    finally:
-        log.debug(f"run_job return: {run_result}", job["id"])
+            log.error("Captured Handler Exception", job["id"])
+            log.error(json.dumps(error_info, indent=4))
+            run_result = {"error": json.dumps(error_info)}
 
+    log.debug(f"run_job return: {run_result}", job["id"])
     return run_result
 
 
@@ -317,20 +319,25 @@ async def run_job_generator(
         job["id"],
     )
 
-    try:
-        job_output = handler(job)
-
-        if is_async_gen:
-            async for output_partial in job_output:
-                log.debug(f"Async Generator output: {output_partial}", job["id"])
-                yield {"output": output_partial}
-        else:
-            for output_partial in job_output:
-                log.debug(f"Generator output: {output_partial}", job["id"])
-                yield {"output": output_partial}
-
-    except Exception as err:
-        log.error(err, job["id"])
-        yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"}
-    finally:
-        log.info("Finished running generator.", job["id"])
+    with capture() as cap:
+        try:
+            job_output = handler(job)
+
+            if is_async_gen:
+                async for output_partial in job_output:
+                    log.debug(f"Async Generator output: {output_partial}", job["id"])
+                    yield {"output": output_partial}
+            else:
+                for output_partial in job_output:
+                    log.debug(f"Generator output: {output_partial}", job["id"])
+                    yield {"output": output_partial}
+
+        except Exception as err:  # noqa: BLE001 - user handler may raise anything; surface it
+            captured_logs = cap.getvalue()
+            log.error(err, job["id"])
+            error = f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"
+            if captured_logs:
+                error += f"\nlogs:\n{captured_logs}"
+            yield {"error": clip(error)}
+        finally:
+            log.info("Finished running generator.", job["id"])