From d65dcdba0b6ca3713fbd48da11e437543cc8f812 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Sun, 16 Aug 2026 23:05:11 -0700 Subject: [PATCH 1/3] feat(serverless): report what a failing handler printed A handler that dies during a model load or a CUDA fault usually explains itself on stdout/stderr, but only the exception reached the platform, so the useful part was lost. Tee both streams into a per-context ring buffer and attach the tail to the error the worker reports. Every reported field is clipped so a huge message or log cannot push the job-done body past its limit. --- runpod/serverless/modules/rp_capture.py | 91 ++++++++++++++ runpod/serverless/modules/rp_job.py | 117 +++++++++--------- runpod/serverless/worker.py | 6 +- tests/test_serverless/test_capture.py | 157 ++++++++++++++++++++++++ 4 files changed, 315 insertions(+), 56 deletions(-) create mode 100644 runpod/serverless/modules/rp_capture.py create mode 100644 tests/test_serverless/test_capture.py diff --git a/runpod/serverless/modules/rp_capture.py b/runpod/serverless/modules/rp_capture.py new file mode 100644 index 000000000..a9d6ac468 --- /dev/null +++ b/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:]}" diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index a45cebc68..a531f738a 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/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"]) diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 90053ec72..8c1bdcbe1 100644 --- a/runpod/serverless/worker.py +++ b/runpod/serverless/worker.py @@ -7,7 +7,7 @@ import os from typing import Any, Dict -from runpod.serverless.modules import rp_logger, rp_local, rp_ping, rp_scale +from runpod.serverless.modules import rp_capture, rp_local, rp_logger, rp_ping, rp_scale from runpod.serverless.modules.rp_fitness import run_fitness_checks log = rp_logger.RunPodLogger() @@ -42,12 +42,16 @@ def run_worker(config: Dict[str, Any]) -> None: # One per-worker mirror: the job tracker writes it, the ping process reads # it. Attaching to JobsProgress means every add/remove syncs automatically. from runpod.serverless.modules.worker_state import JobsProgress, PingJobMirror + mirror = PingJobMirror() JobsProgress().set_mirror(mirror) # Start pinging Runpod to show that the worker is alive. heartbeat.start_ping(mirror) + # Capture stdout/stderr so handler and initializer failures report their logs. + rp_capture.install() + # Create a JobScaler responsible for adjusting the concurrency job_scaler = rp_scale.JobScaler(config) job_scaler.start() diff --git a/tests/test_serverless/test_capture.py b/tests/test_serverless/test_capture.py new file mode 100644 index 000000000..8264c300a --- /dev/null +++ b/tests/test_serverless/test_capture.py @@ -0,0 +1,157 @@ +"""Tests for stdout/stderr capture: what a failing handler printed is attached to the +error it reports back, and every reported field stays bounded.""" + +# pylint: disable=protected-access + +import asyncio +import io +import json +import sys +import unittest +from unittest.mock import patch + +from runpod.serverless.modules import rp_capture +from runpod.serverless.modules.rp_job import run_job, run_job_generator + + +def _run(coro): + return asyncio.run(coro) + + +def _run_gen(agen): + async def _drain(): + return [item async for item in agen] + + return asyncio.run(_drain()) + + +class TestStdioCapture(unittest.TestCase): + """Per-context stdout/stderr capture: records into the active buffer, passes through.""" + + def test_records_and_passes_through(self): + real = io.StringIO() + proxy = rp_capture._TeeProxy(real) + with patch.object(sys, "stdout", proxy), rp_capture.capture() as cap: + print("hello from handler") + assert "hello from handler" in cap.getvalue() + assert "hello from handler" in real.getvalue() # still reached the real stream + + def test_no_capture_outside_block(self): + real = io.StringIO() + proxy = rp_capture._TeeProxy(real) + with patch.object(sys, "stdout", proxy): + with rp_capture.capture() as cap: + pass + print("after the block") + assert "after the block" not in cap.getvalue() + + def test_ring_buffer_keeps_tail(self): + buf = rp_capture._RingBuffer(limit=10) + buf.write("0123456789ABCDEF") + assert buf.getvalue() == "6789ABCDEF" + + def test_install_is_idempotent(self): + stdout = io.StringIO() + stderr = io.StringIO() + with patch.object(sys, "stdout", stdout), patch.object(sys, "stderr", stderr): + rp_capture.install() + installed_stdout = sys.stdout + installed_stderr = sys.stderr + + rp_capture.install() + + assert sys.stdout is installed_stdout + assert sys.stderr is installed_stderr + + +class TestRunJobCapturesLogs(unittest.TestCase): + """A failing handler's stdout/stderr is attached to the job error output.""" + + def test_handler_error_attaches_logs(self): + def handler(_job): + print("loading weights") + print("boom trace", file=sys.stderr) + raise RuntimeError("kernel panic") + + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(sys, "stderr", rp_capture._TeeProxy(real)), + ): + result = _run(run_job(handler, {"id": "j1"})) + + error = json.loads(result["error"]) + assert error["error_message"] == "kernel panic" + assert "loading weights" in error["logs"] + assert "boom trace" in error["logs"] + + def test_generator_error_attaches_logs_in_error(self): + def handler(_job): + print("loading weights") + print("boom trace", file=sys.stderr) + raise RuntimeError("kernel panic") + yield # pragma: no cover - makes handler a generator + + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(sys, "stderr", rp_capture._TeeProxy(real)), + ): + result = _run_gen(run_job_generator(handler, {"id": "g1"})) + + assert len(result) == 1 + assert set(result[0]) == {"error"} + assert "kernel panic" in result[0]["error"] + assert "loading weights" in result[0]["error"] + assert "boom trace" in result[0]["error"] + + def test_handler_success_has_no_error(self): + result = _run(run_job(lambda _job: {"ok": True}, {"id": "j2"})) + assert result == {"output": {"ok": True}} + + +class TestErrorFieldBounding(unittest.TestCase): + """Error strings shipped back to the platform are bounded so a huge message/log can't + blow past the job-done body limit.""" + + def test_clip_keeps_head_and_tail(self): + text = "A" * 100 + "B" * 100 + clipped = rp_capture.clip(text, limit=40) + assert clipped.startswith("A" * 20) + assert clipped.endswith("B" * 20) + assert "truncated" in clipped + assert len(clipped) < len(text) + + def test_clip_passthrough_when_small(self): + assert rp_capture.clip("short", limit=100) == "short" + + def test_run_job_bounds_error_message(self): + huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) + + def handler(_job): + raise ValueError(huge) + + result = _run(run_job(handler, {"id": "big"})) + error = json.loads(result["error"]) + assert len(error["error_message"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + + def test_run_job_generator_bounds_combined_error(self): + huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) + + def handler(_job): + print("y" * (rp_capture.MAX_CAPTURED_CHARS * 3)) + raise ValueError(huge) + yield # pragma: no cover - makes handler a generator + + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(sys, "stderr", rp_capture._TeeProxy(real)), + ): + result = _run_gen(run_job_generator(handler, {"id": "big-generator"})) + + assert len(result[0]["error"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + + +if __name__ == "__main__": + unittest.main() From 7fc76814afce087278322013a1f229300730c885 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Sun, 16 Aug 2026 23:05:11 -0700 Subject: [PATCH 2/3] 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. --- runpod/serverless/__init__.py | 6 + runpod/serverless/modules/rp_initializer.py | 133 ++++++ runpod/serverless/modules/rp_scale.py | 137 +++++- tests/test_serverless/test_initializer.py | 460 ++++++++++++++++++++ 4 files changed, 733 insertions(+), 3 deletions(-) create mode 100644 runpod/serverless/modules/rp_initializer.py create mode 100644 tests/test_serverless/test_initializer.py diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 052452073..6f268b5f1 100644 --- a/runpod/serverless/__init__.py +++ b/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} ---") diff --git a/runpod/serverless/modules/rp_initializer.py b/runpod/serverless/modules/rp_initializer.py new file mode 100644 index 000000000..c8f4e2af8 --- /dev/null +++ b/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") diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 4cbf94ffb..4cb9ea4a9 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/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,9 @@ 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._claimed_request = False self.current_concurrency = 1 self.config = config self.job_progress = JobsProgress() # Cache the singleton instance @@ -60,6 +66,10 @@ def __init__(self, config: Dict[str, Any]): self.concurrency_modifier = _default_concurrency_modifier self.jobs_fetcher = get_job self.jobs_fetcher_timeout = 90 + # Bound on the single claim made after a failed init. Short on purpose: a + # queued request comes back at once, and an empty queue must not hold a dying + # worker open for the full long-poll. + self.init_claim_timeout = 10 self.jobs_handler = handle_job if concurrency_modifier := config.get("concurrency_modifier"): @@ -139,14 +149,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): """ @@ -170,6 +192,28 @@ def current_occupancy(self) -> int: ) return current_progress_count + current_queue_count + async def _claim_one_job_to_fail(self, session: ClientSession) -> None: + """Initialization failed before this worker ever held a request. Claim one + queued request and fail it with the reason. + + Without this, an instant failure - bad config, a missing file, an async + initializer that raises before its first await - never reaches a caller: the + worker exits, the platform respawns it into the same failure, and the request + that triggered the scale-up waits out its queue TTL with no explanation. + """ + try: + jobs = await asyncio.wait_for( + self.jobs_fetcher(session, 1), timeout=self.init_claim_timeout + ) + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 - reporting must not mask the failure + log.debug(f"JobScaler.get_jobs | No request claimed to fail: {error}") + return + + for job in jobs or []: + await self._fail_job(session, job, self._init_error) + async def get_jobs(self, session: ClientSession): """ Retrieve multiple jobs from the server in batches using blocking requests. @@ -179,6 +223,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. Draining what we already + # hold is owned by run_jobs. + break await self.set_scale() jobs_needed = self.current_concurrency - self.current_occupancy() @@ -200,6 +248,14 @@ async def get_jobs(self, session: ClientSession): log.debug("JobScaler.get_jobs | No jobs acquired.") continue + self._claimed_request = True + + 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) @@ -227,6 +283,12 @@ async def get_jobs(self, session: ClientSession): # Yield control back to the event loop await asyncio.sleep(0) + if self._init_error is not None and not self._claimed_request: + # An init failure sets shutdown, so this loop can end before it ever + # claimed a request. Claim one on the way out, or the failure dies with + # the worker and the request that spawned it waits out its queue TTL. + await self._claim_one_job_to_fail(session) + async def run_jobs(self, session: ClientSession): """ Retrieve jobs from the jobs queue and process them concurrently. @@ -346,6 +408,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 +446,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) diff --git a/tests/test_serverless/test_initializer.py b/tests/test_serverless/test_initializer.py new file mode 100644 index 000000000..f90e0b51d --- /dev/null +++ b/tests/test_serverless/test_initializer.py @@ -0,0 +1,460 @@ +"""Tests for the concurrent initializer: runs alongside the job loop, holds the handler +until ready, and fails the in-hand request (with captured stdout/stderr) on failure.""" + +# pylint: disable=protected-access + +import asyncio +import functools +import io +import json +import sys +import threading +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from runpod.serverless.modules import rp_capture, rp_scale +from runpod.serverless.modules.rp_initializer import ( + InitializerError, + InitializerTimeout, + build_init_failed_payload, + run_initializer_async, +) +from runpod.serverless.modules.rp_scale import JobScaler + + +def _run(coro): + return asyncio.run(coro) + + +class TestRunInitializerAsync(unittest.TestCase): + """Runs the initializer to completion, or raises a clear error when it fails or times out.""" + + def test_sync_success_offloaded(self): + calls = [] + _run(run_initializer_async(lambda: calls.append("ran"))) + assert calls == ["ran"] + + def test_sync_callable_returning_awaitable_is_awaited(self): + state = {} + + async def load(): + await asyncio.sleep(0) + state["ready"] = True + + def initialize(): + return load() + + _run(run_initializer_async(initialize)) + assert state == {"ready": True} + + + def test_sync_failure_wraps_in_initializer_error(self): + def failing_initializer(): + raise ValueError("max_model_len must be positive, got 0") + + with self.assertRaises(InitializerError) as ctx: + _run(run_initializer_async(failing_initializer)) + assert isinstance(ctx.exception.original, ValueError) + assert "max_model_len" in str(ctx.exception) + + def test_async_success(self): + state = {} + + async def load(): + await asyncio.sleep(0) + state["ready"] = True + + _run(run_initializer_async(load)) + assert state == {"ready": True} + + def test_async_failure_wraps_in_initializer_error(self): + async def failing_initializer(): + raise RuntimeError("CUDA OOM") + + with self.assertRaises(InitializerError) as ctx: + _run(run_initializer_async(failing_initializer)) + assert isinstance(ctx.exception.original, RuntimeError) + + def test_async_timeout_raises_initializer_timeout(self): + async def slow(): + await asyncio.sleep(3) + + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(slow, timeout=1)) + + def test_zero_timeout_raises_initializer_timeout(self): + async def load(): + await asyncio.sleep(0) + + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(load, timeout=0)) + + + def test_sync_hang_times_out(self): + """A blocking sync load that never returns is cut off by init_timeout, not left stuck.""" + release = threading.Event() + + def hang(): + release.wait(10) + + try: + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(hang, timeout=1)) + finally: + release.set() # let the offloaded thread exit promptly + + def test_async_partial_is_awaited(self): + """functools.partial wrapping an async initializer is detected and awaited.""" + state = {} + + async def load(key): + await asyncio.sleep(0) + state[key] = True + + _run(run_initializer_async(functools.partial(load, "ready"))) + assert state == {"ready": True} + + def test_async_callable_object_is_awaited(self): + """An object whose __call__ is async is detected and awaited.""" + state = {} + + class Loader: + async def __call__(self): + await asyncio.sleep(0) + state["ready"] = True + + _run(run_initializer_async(Loader())) + assert state == {"ready": True} + + def test_base_exception_propagates_unwrapped(self): + """KeyboardInterrupt is process control, not an init failure: propagate.""" + + def interrupted(): + raise KeyboardInterrupt + + with self.assertRaises(KeyboardInterrupt): + _run(run_initializer_async(interrupted)) + + def test_system_exit_wraps_as_initializer_error(self): + """A load script calling sys.exit() is an init failure to surface with a reason, + not a clean exit - otherwise held in-hand jobs die without one.""" + + def bails(): + raise SystemExit(1) + + with self.assertRaises(InitializerError): + _run(run_initializer_async(bails)) + + +class TestInitFailedSignal(unittest.TestCase): + """Builds the structured init_failed payload, including captured logs.""" + + def test_payload_shape_from_sync_error(self): + try: + raise ValueError("bad config") + except ValueError as exc: + payload = build_init_failed_payload( + InitializerError(exc), logs="stderr tail" + ) + assert payload["event"] == "init_failed" + assert payload["error_type"] == "ValueError" + assert payload["error_message"] == "bad config" + assert "ValueError" in payload["error_traceback"] + assert payload["logs"] == "stderr tail" + assert "worker_id" in payload and "runpod_version" in payload + + def test_payload_omits_empty_logs(self): + payload = build_init_failed_payload(InitializerError(ValueError("x"))) + assert "logs" not in payload + + def test_payload_bounds_message_traceback_and_logs(self): + huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) + payload = build_init_failed_payload( + InitializerError(ValueError(huge)), + logs="y" * (rp_capture.MAX_CAPTURED_CHARS * 3), + ) + assert len(payload["error_message"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + assert len(payload["error_traceback"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + assert len(payload["logs"]) <= rp_capture.MAX_CAPTURED_CHARS + + +def _scaler(initializer=None): + config = {"handler": lambda j: j, "rp_args": {}} + if initializer is not None: + config["initializer"] = initializer + scaler = JobScaler(config) + scaler.job_progress = MagicMock() # avoid the process-wide singleton in unit tests + scaler.job_progress.get_job_count.return_value = 0 + return scaler + + +class TestRunInit(unittest.TestCase): + """The concurrent init task: opens the gate, and on failure records the reason and shuts down.""" + + def test_no_initializer_opens_gate_immediately(self): + scaler = _scaler(initializer=None) + _run(scaler._run_init()) + assert scaler._init_ready.is_set() + assert scaler._init_error is None + + def test_success_opens_gate_no_error(self): + ran = [] + scaler = _scaler(initializer=lambda: ran.append(1)) + _run(scaler._run_init()) + assert ran == [1] + assert scaler._init_ready.is_set() + assert scaler._init_error is None + assert not scaler._shutdown_event.is_set() + + def test_failure_records_reason_with_logs_and_shuts_down(self): + def failing_initializer(): + print("downloading model") + raise RuntimeError("CUDA OOM: model too big") + + scaler = _scaler(initializer=failing_initializer) + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(rp_scale, "log") as mock_log, + ): + _run(scaler._run_init()) + + assert scaler._init_ready.is_set() # held handlers are released to fail fast + assert scaler._init_error is not None + assert scaler._init_error["error_message"] == "CUDA OOM: model too big" + assert "downloading model" in scaler._init_error["logs"] + assert any( + call.args[0].startswith("init_failed | ") + for call in mock_log.error.call_args_list + ) + assert ( + scaler._shutdown_event.is_set() + ) # broken worker shuts down (occupancy was 0) + + def test_failure_starts_shutdown_before_drain(self): + scaler = _scaler( + initializer=lambda: (_ for _ in ()).throw(RuntimeError("boom")) + ) + shutdown_states = [] + + def occupancy(): + shutdown_states.append(scaler._shutdown_event.is_set()) + return 0 + + scaler.current_occupancy = occupancy + _run(scaler._run_init()) + + assert shutdown_states == [True] + + +class TestHandleJobGate(unittest.TestCase): + """The handler is held until init is ready; init failure fails the in-hand request.""" + + def _prime(self, scaler, job): + # Balance the queue/progress bookkeeping handle_job's finally expects. + scaler.jobs_queue = asyncio.Queue(maxsize=4) + scaler.jobs_queue.put_nowait(job) + + def test_runs_handler_once_init_is_ready(self): + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._init_ready.set() # init already succeeded + job = {"id": "j1"} + + async def go(): + self._prime(scaler, job) + await scaler.handle_job(None, job) + + _run(go()) + scaler.jobs_handler.assert_awaited_once() + + def test_runs_handler_without_an_initializer(self): + scaler = _scaler(initializer=None) + scaler.jobs_handler = AsyncMock() + job = {"id": "j2"} + + async def go(): + self._prime(scaler, job) + await scaler.handle_job(None, job) + + _run(go()) + scaler.jobs_handler.assert_awaited_once() # no gate to wait on + + def test_init_failure_fails_request_without_running_handler(self): + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._init_ready.set() + job = {"id": "j3"} + + async def go(): + self._prime(scaler, job) + with patch.object(rp_scale, "send_result", new=AsyncMock()) as mock_sr: + await scaler.handle_job(None, job) + mock_sr.assert_awaited_once() + sent = mock_sr.await_args[0][1] + assert json.loads(sent["error"])["error_message"] == "CUDA OOM" + + _run(go()) + scaler.jobs_handler.assert_not_awaited() # broken worker never runs the handler + + def test_init_failure_before_any_take_claims_a_request_to_fail(self): + """An instant init failure leaves no request in hand, and it sets shutdown, so + job-take ends immediately. It must still claim one request and fail it, or the + caller waits out the queue TTL for nothing.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() # _run_init sets shutdown on failure + scaler._fail_job = AsyncMock() + job = {"id": "orphan-1"} + scaler.jobs_fetcher = AsyncMock(return_value=[job]) + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + + scaler.jobs_fetcher.assert_awaited_once() + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + assert scaler.jobs_queue.qsize() == 0 # never queued into a broken worker + + def test_init_failure_with_empty_queue_exits_without_hanging(self): + """Nothing left to fail: the claim attempt returns empty and job-take stops.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() + scaler._fail_job = AsyncMock() + scaler.jobs_fetcher = AsyncMock(return_value=[]) + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + + scaler.jobs_fetcher.assert_awaited_once() + scaler._fail_job.assert_not_awaited() + + def test_init_failure_after_a_take_does_not_claim_another_request(self): + """Once this worker has claimed a request, the failure is reported against it. + Claiming a second request would fail work a healthy worker could serve.""" + scaler = _scaler(initializer=lambda: None) + scaler._fail_job = AsyncMock() + failure = {"error_message": "CUDA OOM", "event": "init_failed"} + calls = [] + + async def fetcher(_session, _needed): + calls.append(1) + return [{"id": f"job-{len(calls)}"}] + + scaler.jobs_fetcher = fetcher + + async def go(): + task = asyncio.create_task(scaler.get_jobs(AsyncMock())) + await asyncio.sleep(0) + scaler._init_error = failure # lands after the first take succeeded + scaler.kill_worker() + await asyncio.wait_for(task, timeout=0.5) + + _run(go()) + + assert len(calls) == 1 # no extra claim after the queued request + + def test_instant_async_init_failure_still_fails_the_queued_request(self): + """An async initializer that raises before its first await finishes before + job-take ever fetches, so nothing is in hand and shutdown is already set. + Driving the real `_run_init` failure path, the worker must still claim the + queued request and fail it rather than exit silently.""" + + async def bad_init(): + raise RuntimeError("bad config") + + scaler = _scaler(initializer=bad_init) + scaler._fail_job = AsyncMock() + job = {"id": "queued-1"} + scaler.jobs_fetcher = AsyncMock(return_value=[job]) + scaler.jobs_handler = AsyncMock() + + async def go(): + await scaler._run_init() # records the reason and sets shutdown + assert not scaler.is_alive() + await asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=1) + + _run(go()) + + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + reason = scaler._fail_job.await_args[0][2] + assert reason["event"] == "init_failed" + assert reason["error_message"] == "bad config" + scaler.jobs_handler.assert_not_awaited() # handler never runs on a broken worker + + def test_claim_attempt_is_bounded(self): + """A silent job-take must not hold a dying worker open. The claim gives up on + its own bound and the worker exits to be respawned.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() + scaler._fail_job = AsyncMock() + scaler.init_claim_timeout = 0.05 + + async def never_returns(_session, _needed): + await asyncio.Event().wait() + + scaler.jobs_fetcher = never_returns + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=1)) + + scaler._fail_job.assert_not_awaited() + + def test_jobs_acquired_after_init_failure_are_failed_not_queued(self): + """If init fails while a long-poll is in flight, fail returned jobs and stop + job-take without relying on another task to end the loop.""" + scaler = _scaler(initializer=lambda: None) + scaler._fail_job = AsyncMock() + job = {"id": "late-1"} + failure = {"error_message": "CUDA OOM", "event": "init_failed"} + + async def fetcher(_session, _needed): + # Init failure lands while this long-poll is in flight. + scaler._init_error = failure + return [job] + + scaler.jobs_fetcher = fetcher + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + assert scaler.jobs_queue.qsize() == 0 + scaler.job_progress.add.assert_not_called() + + def test_shutdown_before_init_ready_leaves_request_alone(self): + """SIGTERM while the initializer is still running must release a held request + without running the handler against an uninitialized worker.""" + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler.kill_worker() # shutdown lands while init is still in flight + job = {"id": "j4"} + + async def go(): + self._prime(scaler, job) + with patch.object(rp_scale, "send_result", new=AsyncMock()) as mock_sr: + await asyncio.wait_for(scaler.handle_job(None, job), timeout=0.5) + mock_sr.assert_not_awaited() # the platform retries it elsewhere + + _run(go()) + scaler.jobs_handler.assert_not_awaited() + + +class TestShutdownWithHangingInit(unittest.TestCase): + """A hung initializer with no init_timeout must not outlive shutdown: the daemon + thread is abandoned so the worker can exit.""" + + def test_run_returns_while_initializer_still_blocked(self): + release = threading.Event() + self.addCleanup(release.set) # let the daemon thread finish after the test + + scaler = _scaler(initializer=release.wait) # blocks with no init_timeout + scaler.kill_worker() # every loop exits immediately; only init is left + scaler.stop_signals_fetcher = AsyncMock(return_value=[]) + + _run(asyncio.wait_for(scaler.run(), timeout=2)) + + assert not release.is_set() # still blocked, and the worker left anyway + + +if __name__ == "__main__": + unittest.main() From 9f851452b0fdcc19ab150742546a9f13e48536fb Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Mon, 17 Aug 2026 09:54:10 -0700 Subject: [PATCH 3/3] chore: retrigger code scanning