feat(serverless): run the user initializer alongside job intake

Startup work placed before runpod.serverless.start() ran outside the SDK's view: a
model load that hung or crashed left requests sitting in IN_QUEUE until the TTL
expired, with the reason buried in worker logs.

Accept an optional initializer and init_timeout. The worker now runs the
initializer as a fourth concurrent task, keeps taking requests, and holds handler
execution until initialization finishes. If initialization fails, the worker
reports the reason and its captured logs against the request it is holding through
the existing job-done route, fails any request a long-poll returns afterwards, and
exits so the platform respawns it under existing backoff.

Author: jasonwang-runpod

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_initializer.py

@@ -0,0 +1,134 @@
+"""
+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,
+        "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
+
+
+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_scale.py

@@ -4,12 +4,15 @@
 """
 
 import asyncio
+import json
 import signal
 import sys
 import traceback
 from typing import Any, Dict, Set
 
 from ...http_client import AsyncClientSession, ClientSession, TooManyRequests
+from .rp_capture import capture
+from .rp_http import send_result
 from .rp_job import _job_stop_url, get_job, get_stop_signals, handle_job
 from .rp_logger import RunPodLogger, _reset_batch_id, _set_batch_id
 from .worker_state import JobsProgress, IS_LOCAL_TEST
@@ -44,6 +47,8 @@ class JobScaler:
 
     def __init__(self, config: Dict[str, Any]):
         self._shutdown_event = asyncio.Event()
+        self._init_ready = asyncio.Event()
+        self._init_error: dict[str, Any] | None = None
         self.current_concurrency = 1
         self.config = config
         self.job_progress = JobsProgress()  # Cache the singleton instance
@@ -139,14 +144,26 @@ async def run(self):
         # Create an async session that will be closed when the worker is killed.
         async with AsyncClientSession() as session:
             # Create the worker's concurrent loops.
+            init_task = asyncio.create_task(self._run_init())
             jobtake_task = asyncio.create_task(self.get_jobs(session))
             jobrun_task = asyncio.create_task(self.run_jobs(session))
             jobstop_task = asyncio.create_task(self.monitor_stop_signals(session))
 
-            tasks = [jobtake_task, jobrun_task, jobstop_task]
+            # The initializer is not a loop: an initializer with no init_timeout can
+            # block forever, so only the request loops decide when the worker stops.
+            await asyncio.gather(jobtake_task, jobrun_task, jobstop_task)
 
-            # Run the worker's concurrent loops until shutdown.
-            await asyncio.gather(*tasks)
+            # Shutting down: abandon a still-running initializer. Its work sits on a
+            # daemon thread, so dropping it here lets the process exit.
+            init_task.cancel()
+            try:
+                await init_task
+            except asyncio.CancelledError:
+                pass
+
+        # If the initializer fails, let the platform respawn the worker.
+        if self._init_error is not None:
+            sys.exit(1)
 
     def is_alive(self):
         """
@@ -179,6 +196,10 @@ async def get_jobs(self, session: ClientSession):
         Adds jobs to the JobsQueue
         """
         while self.is_alive():
+            if self._init_error is not None:
+                # Initialization is terminal for this worker. Request draining is
+                # owned by run_jobs; job-take has no more work to do.
+                return
             await self.set_scale()
 
             jobs_needed = self.current_concurrency - self.current_occupancy()
@@ -200,6 +221,12 @@ async def get_jobs(self, session: ClientSession):
                     log.debug("JobScaler.get_jobs | No jobs acquired.")
                     continue
 
+                if self._init_error is not None:
+                    # If initialization fails, fail all in-flight requests.
+                    for job in acquired_jobs:
+                        await self._fail_job(session, job, self._init_error)
+                    return
+
                 for job in acquired_jobs:
                     await self.jobs_queue.put(job)
                     self.job_progress.add(job)
@@ -346,6 +373,21 @@ async def handle_job(self, session: ClientSession, job: dict):
         try:
             log.debug("Handling Job", job["id"])
 
+            # Hold the handler until initialization finishes.
+            if self.config.get("initializer") is not None:
+                if not await self._wait_for_init():
+                    log.warn(
+                        "Shutting down before initialization finished; leaving this "
+                        "request for another worker.",
+                        job["id"],
+                    )
+                    return
+
+                if self._init_error is not None:
+                    # If initialization fails, fail the current request and don't run the handler.
+                    await self._fail_job(session, job, self._init_error)
+                    return
+
             await self.jobs_handler(session, self.config, job)
 
             if self.config.get("refresh_worker", False):
@@ -369,3 +411,57 @@ async def handle_job(self, session: ClientSession, job: dict):
 
             log.debug("Finished Job", job["id"])
             _reset_batch_id(batch_id_token)
+
+    async def _wait_for_init(self) -> bool:
+        """Wait for initialization to finish, or for the worker to start shutting down.
+        Returns whether initialization actually finished."""
+        ready = asyncio.create_task(self._init_ready.wait())
+        stopping = asyncio.create_task(self._shutdown_event.wait())
+        try:
+            await asyncio.wait({ready, stopping}, return_when=asyncio.FIRST_COMPLETED)
+        finally:
+            ready.cancel()
+            stopping.cancel()
+        return self._init_ready.is_set()
+
+    async def _fail_job(
+        self, session: ClientSession, job: dict, payload: Dict[str, Any]
+    ):
+        """Fail a request with a structured error (reason + logs)."""
+        log.error(f"Failing job due to init failure. | {job['id']}")
+        await send_result(session, {"error": json.dumps(payload)}, job, is_stream=False)
+
+    async def _run_init(self):
+        """Run initializer concurrently with the loop. Upon completion, opens the gate to run
+        job handlers. On failure, records the reason, then drains and shuts the worker."""
+        initializer = self.config.get("initializer")
+        if initializer is None:
+            self._init_ready.set()
+            return
+
+        from .rp_initializer import (
+            InitializerError,
+            InitializerTimeout,
+            build_init_failed_payload,
+            run_initializer_async,
+        )
+
+        try:
+            with capture() as cap:
+                try:
+                    await run_initializer_async(
+                        initializer, self.config.get("init_timeout")
+                    )
+                except (InitializerError, InitializerTimeout) as exc:
+                    self._init_error = build_init_failed_payload(exc, cap.getvalue())
+            if self._init_error is not None:
+                log.error(f"init_failed | {json.dumps(self._init_error)}")
+        finally:
+            # Always release held handlers.
+            self._init_ready.set()
+
+        if self._init_error is not None:
+            # Stop long-running loops before waiting for acquired requests to drain.
+            self.kill_worker()
+            while self.current_occupancy() > 0:
+                await asyncio.sleep(0.1)