From d65dcdba0b6ca3713fbd48da11e437543cc8f812 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Sun, 16 Aug 2026 23:05:11 -0700 Subject: [PATCH 1/7] 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 15caa9c4625fc5c7fa73dfbd0cd2ac1235748500 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Sun, 16 Aug 2026 23:05:11 -0700 Subject: [PATCH 2/7] 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 | 139 ++++- tests/test_serverless/test_initializer.py | 534 ++++++++++++++++++++ 4 files changed, 809 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..a3da893bd 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,28 @@ 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: + # expected: we cancelled it on the line above + pass + + if self._init_error is not None: + from .rp_fitness import _terminate_unhealthy + + _terminate_unhealthy(1) def is_alive(self): """ @@ -170,6 +194,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 +225,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 +250,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 +285,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 +410,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 +448,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..519fb30dc --- /dev/null +++ b/tests/test_serverless/test_initializer.py @@ -0,0 +1,534 @@ +"""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 pathlib +import subprocess +import sys +import tempfile +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 + + +# Needs a real process: in-process the interpreter never joins threads, so the hang +# this guards against cannot be observed. +_HARD_EXIT_SCRIPT = """ +import asyncio, sys, threading, time + +MODE = sys.argv[1] +sys.argv = ["worker"] + +from runpod.serverless.modules.rp_scale import JobScaler + + +def sync_engine(): + # A sync initializer runs on a daemon thread, so its children inherit daemon + # status. An engine that sets daemon=False itself does not. + threading.Thread(target=lambda: time.sleep(60), daemon=False).start() + raise RuntimeError("engine start failed") + + +async def async_engine(): + # An async initializer is awaited on the main thread, so anything it spawns is + # non-daemon by inheritance. + threading.Thread(target=lambda: time.sleep(60)).start() + raise RuntimeError("engine start failed") + + +async def no_jobs(*args, **kwargs): + return [] + + +scaler = JobScaler( + { + "handler": lambda job: job, + "rp_args": {}, + "initializer": {"sync": sync_engine, "async": async_engine}[MODE], + } +) +scaler.jobs_fetcher = no_jobs +scaler.stop_signals_fetcher = no_jobs +scaler.init_claim_timeout = 1 + +asyncio.run(scaler.run()) +print("run() returned without exiting") +""" + + +class TestInitFailureExitsProcess(unittest.TestCase): + """Init failure must hard-exit, even with a non-daemon thread left running.""" + + def _run_worker(self, mode: str) -> subprocess.CompletedProcess: + with tempfile.TemporaryDirectory() as tmp: + script = pathlib.Path(tmp) / "worker.py" + script.write_text(_HARD_EXIT_SCRIPT) + return subprocess.run( + [sys.executable, str(script), mode], + capture_output=True, + text=True, + timeout=30, # generous: the fix exits in well under a second + check=False, + ) + + def test_sync_initializer_leaving_a_non_daemon_thread(self): + result = self._run_worker("sync") + assert result.returncode == 1 + assert "run() returned without exiting" not in result.stdout + + def test_async_initializer_leaving_a_non_daemon_thread(self): + result = self._run_worker("async") + assert result.returncode == 1 + assert "init_failed" in result.stdout + result.stderr # reported before exiting + + +if __name__ == "__main__": + unittest.main() From d3e6b01b5de3f6d0af03c30320a3879192abe106 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 19 Aug 2026 13:45:08 -0700 Subject: [PATCH 3/7] feat(serverless): add explicit prestart hooks --- docs/serverless/worker.md | 72 +- runpod/serverless/__init__.py | 56 +- runpod/serverless/modules/rp_capture.py | 3 +- runpod/serverless/modules/rp_fastapi.py | 86 ++- runpod/serverless/modules/rp_initializer.py | 133 ---- runpod/serverless/modules/rp_local.py | 26 +- runpod/serverless/modules/rp_prestart.py | 188 +++++ runpod/serverless/modules/rp_scale.py | 286 +++++--- runpod/serverless/worker.py | 2 +- tests/test_serverless/test_init.py | 4 +- tests/test_serverless/test_initializer.py | 534 -------------- .../test_modules/test_local.py | 67 +- tests/test_serverless/test_prestart.py | 267 +++++++ .../test_prestart_lifecycle.py | 691 ++++++++++++++++++ 14 files changed, 1598 insertions(+), 817 deletions(-) delete mode 100644 runpod/serverless/modules/rp_initializer.py create mode 100644 runpod/serverless/modules/rp_prestart.py delete mode 100644 tests/test_serverless/test_initializer.py create mode 100644 tests/test_serverless/test_prestart.py create mode 100644 tests/test_serverless/test_prestart_lifecycle.py diff --git a/docs/serverless/worker.md b/docs/serverless/worker.md index e1355260d..4f209d38b 100644 --- a/docs/serverless/worker.md +++ b/docs/serverless/worker.md @@ -18,14 +18,80 @@ 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 `init_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 `init_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. + ## 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: diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 6f268b5f1..d5132eb31 100644 --- a/runpod/serverless/__init__.py +++ b/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,24 +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. - config["initializer"] (Callable, optional): Startup work that runs alongside job intake. - The worker holds handler execution until the initializer finishes. + 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["init_timeout"] (int, optional): Seconds to allow the initializer before - treating it as a failure. Omit for no timeout. + config["prestart_timeout"] (int, optional): Seconds allowed for the complete + prestart phase. Omit for no timeout. """ print(f"--- Starting Serverless Worker | Version {runpod_version} ---") @@ -162,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( @@ -177,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( diff --git a/runpod/serverless/modules/rp_capture.py b/runpod/serverless/modules/rp_capture.py index a9d6ac468..4cb51d1bc 100644 --- a/runpod/serverless/modules/rp_capture.py +++ b/runpod/serverless/modules/rp_capture.py @@ -1,7 +1,7 @@ """ runpod | serverless | rp_capture.py -Captures stdout/stderr, to be reported upon handler or initializer failure. +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. """ @@ -19,7 +19,6 @@ ) - class _RingBuffer: """Keeps only the last `limit` characters since the tail is usually where the failure reason is.""" diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 5451ae40e..d77d08536 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -1,10 +1,13 @@ -""" Used to launch the FastAPI web server when worker is running in API mode. """ +"""Used to launch the FastAPI web server when worker is running in API mode.""" +import json import os import threading import uuid +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager from dataclasses import dataclass -from typing import Any, Dict, Optional, Union +from typing import Any import requests import uvicorn @@ -14,12 +17,24 @@ from ...http_client import SyncClientSession from ...version import __version__ as runpod_version +from . import rp_capture +from .rp_fitness import _terminate_unhealthy from .rp_handler import is_generator from .rp_job import run_job, run_job_generator +from .rp_logger import RunPodLogger from .rp_ping import Heartbeat +from .rp_prestart import ( + PrestartError, + PrestartTimeout, + build_init_failed_payload, + get_prestart_hooks, + run_prestart_hooks_async, +) from .worker_state import JobsProgress, PingJobMirror RUNPOD_ENDPOINT_ID = os.environ.get("RUNPOD_ENDPOINT_ID", None) +log = RunPodLogger() + TITLE = "Runpod | Development Worker API" @@ -29,17 +44,17 @@ Use this API for comprehensive testing of request submissions and result retrieval, mimicking the behavior of Runpod's operational environment. --- *Note: This API serves as a local testing tool and will not be utilized once your worker is operational on the Runpod platform.* -""" - -# Add CLI tool suggestion if RUNPOD_PROJECT_ID is not set. -if os.environ.get("RUNPOD_PROJECT_ID", None) is None: - DESCRIPTION += """ +""" + ( + """ ℹ️ | Consider developing with our CLI tool to streamline your worker development process. >_ wget -qO- cli.runpod.net | sudo bash >_ runpodctl project create """ + if os.environ.get("RUNPOD_PROJECT_ID") is None + else "" +) RUN_DESCRIPTION = """ Initiates processing jobs, returning a unique job ID. @@ -106,7 +121,7 @@ class Job: """Represents a job.""" id: str - input: Union[dict, list, str, int, float, bool] + input: dict[str, Any] | list[Any] | str | int | float | bool @dataclass @@ -115,17 +130,17 @@ class TestJob: input can be any type of data. """ - id: Optional[str] = None - input: Optional[Union[dict, list, str, int, float, bool]] = None - webhook: Optional[str] = None + id: str | None = None + input: dict[str, Any] | list[Any] | str | int | float | bool | None = None + webhook: str | None = None @dataclass class DefaultRequest: """Represents a test input.""" - input: Dict[str, Any] - webhook: Optional[str] = None + input: dict[str, Any] + webhook: str | None = None # ------------------------------ Output Objects ------------------------------ # @@ -135,8 +150,8 @@ class JobOutput: id: str status: str - output: Optional[Union[dict, list, str, int, float, bool]] = None - error: Optional[str] = None + output: dict[str, Any] | list[Any] | str | int | float | bool | None = None + error: str | None = None @dataclass @@ -145,18 +160,18 @@ class StreamOutput: id: str status: str = "IN_PROGRESS" - stream: Optional[Union[dict, list, str, int, float, bool]] = None - error: Optional[str] = None + stream: dict[str, Any] | list[Any] | str | int | float | bool | None = None + error: str | None = None # ------------------------------ Webhook Sender ------------------------------ # -def _send_webhook(url: str, payload: Dict[str, Any]) -> bool: +def _send_webhook(url: str, payload: dict[str, Any]) -> bool: """ Sends a webhook to the provided URL. Args: url (str): The URL to send the webhook to. - payload (Dict[str, Any]): The JSON payload to send. + payload (dict[str, Any]): The JSON payload to send. Returns: bool: True if the request was successful, False otherwise. @@ -177,7 +192,24 @@ def _send_webhook(url: str, payload: Dict[str, Any]) -> bool: class WorkerAPI: """Used to launch the FastAPI web server when the worker is running in API mode.""" - def __init__(self, config: Dict[str, Any]): + @asynccontextmanager + async def _lifespan(self, _app: FastAPI) -> AsyncGenerator[None, None]: + """Finish prestart before the development API accepts requests.""" + hooks = get_prestart_hooks() + if hooks: + with rp_capture.capture() as captured: + try: + await run_prestart_hooks_async( + hooks, self.config.get("prestart_timeout") + ) + except (PrestartError, PrestartTimeout) as exc: + failure = build_init_failed_payload(exc, captured.getvalue()) + log.error(f"init_failed | {json.dumps(failure)}") + _terminate_unhealthy(1) + raise + yield + + def __init__(self, config: dict[str, Any]): """ Initializes the WorkerAPI class. 1. Starts the heartbeat thread. @@ -194,6 +226,7 @@ def __init__(self, config: Dict[str, Any]): heartbeat.start_ping(mirror) self.config = config + rp_capture.install() tags_metadata = [ { @@ -217,6 +250,7 @@ def __init__(self, config: Dict[str, Any]): version=runpod_version, docs_url="/", openapi_tags=tags_metadata, + lifespan=self._lifespan, ) # Create an APIRouter and add the route for processing jobs. @@ -310,11 +344,13 @@ async def _realtime(self, job: Job): async def _sim_run(self, job_request: DefaultRequest) -> JobOutput: """Development endpoint to simulate run behavior.""" assigned_job_id = f"test-{uuid.uuid4()}" - job_list.add({ - "id": assigned_job_id, - "input": job_request.input, - "webhook": job_request.webhook - }) + job_list.add( + { + "id": assigned_job_id, + "input": job_request.input, + "webhook": job_request.webhook, + } + ) return jsonable_encoder({"id": assigned_job_id, "status": "IN_PROGRESS"}) # ---------------------------------- runsync --------------------------------- # diff --git a/runpod/serverless/modules/rp_initializer.py b/runpod/serverless/modules/rp_initializer.py deleted file mode 100644 index c8f4e2af8..000000000 --- a/runpod/serverless/modules/rp_initializer.py +++ /dev/null @@ -1,133 +0,0 @@ -""" -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_local.py b/runpod/serverless/modules/rp_local.py index 971e696e1..515b3f09b 100644 --- a/runpod/serverless/modules/rp_local.py +++ b/runpod/serverless/modules/rp_local.py @@ -6,16 +6,25 @@ import json import os import sys -from typing import Any, Dict +from typing import Any +from runpod.serverless.modules import rp_capture from runpod.serverless.modules.rp_logger import RunPodLogger +from .rp_fitness import _terminate_unhealthy from .rp_job import run_job +from .rp_prestart import ( + PrestartError, + PrestartTimeout, + build_init_failed_payload, + get_prestart_hooks, + run_prestart_hooks_async, +) log = RunPodLogger() -async def run_local(config: Dict[str, Any]) -> None: +async def run_local(config: dict[str, Any]) -> None: """ Runs the worker locally. """ @@ -29,7 +38,7 @@ async def run_local(config: Dict[str, Any]) -> None: sys.exit(1) log.info("Using test_input.json as job input.") - with open("test_input.json", "r", encoding="UTF-8") as file: + with open("test_input.json", encoding="UTF-8") as file: local_job = json.loads(file.read()) if local_job.get("input", None) is None: @@ -39,6 +48,17 @@ async def run_local(config: Dict[str, Any]) -> None: # Set the job ID local_job["id"] = local_job.get("id", "local_test") log.debug(f"Retrieved local job: {local_job}") + rp_capture.install() + with rp_capture.capture() as captured: + try: + await run_prestart_hooks_async( + get_prestart_hooks(), config.get("prestart_timeout") + ) + except (PrestartError, PrestartTimeout) as exc: + failure = build_init_failed_payload(exc, captured.getvalue()) + log.error(f"init_failed | {json.dumps(failure)}") + _terminate_unhealthy(1) + sys.exit(1) job_result = await run_job(config["handler"], local_job) diff --git a/runpod/serverless/modules/rp_prestart.py b/runpod/serverless/modules/rp_prestart.py new file mode 100644 index 000000000..14273d487 --- /dev/null +++ b/runpod/serverless/modules/rp_prestart.py @@ -0,0 +1,188 @@ +""" +Prestart hook registration and execution for supported Serverless modes. + +Hooks run sequentially before handler execution. Queue workers may acquire +requests concurrently while the handler gate stays closed. Local input and +single-worker hosted API modes finish prestart before invoking their handler +or serving HTTP. Sync hooks run on an abandonable daemon thread; async hooks +run on the active mode's event loop. +""" + +import asyncio +import contextlib +import contextvars +import inspect +import threading +import traceback +from collections.abc import Callable, Sequence +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" + +_prestart_hooks: list[Callable[[], Any]] = [] + + +def _hook_name(hook: Callable[[], Any] | None) -> str: + if hook is None: + return "unknown" + return getattr(hook, "__name__", type(hook).__name__) + + +def register_prestart_hook(hook: Callable[[], Any]) -> Callable[[], Any]: + """Register a sync or async hook to run before handler execution.""" + _prestart_hooks.append(hook) + log.debug(f"Registered prestart hook: {_hook_name(hook)}") + return hook + + +def get_prestart_hooks() -> tuple[Callable[[], Any], ...]: + """Return an immutable snapshot of hooks in registration order.""" + return tuple(_prestart_hooks) + + +def has_prestart_hooks() -> bool: + """Return whether any prestart hooks are registered.""" + return bool(_prestart_hooks) + + +def clear_prestart_hooks() -> None: + """Clear registered hooks. Intended for test isolation.""" + _prestart_hooks.clear() + + +class PrestartTimeout(Exception): + """Raised when the complete prestart phase exceeds `prestart_timeout`.""" + + def __init__(self, hook: str, timeout: float): + self.hook = hook + super().__init__( + f"prestart hook '{hook}' exceeded prestart_timeout of {timeout}s" + ) + + +class PrestartError(Exception): + """Wrap an exception raised by a prestart hook.""" + + def __init__(self, original: BaseException, hook: str): + self.original = original + self.hook = hook + super().__init__(str(original)) + + +def build_init_failed_payload(exc: BaseException, logs: str = "") -> dict[str, Any]: + """Build the structured `init_failed` reason sent through `/job-done`.""" + 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 hook := getattr(exc, "hook", None): + payload["hook"] = hook + if logs: + payload["logs"] = logs[-MAX_CAPTURED_CHARS:] + return payload + + +async def _run_sync_in_daemon(hook: Callable[[], Any]) -> Any: + """Run a blocking hook without creating an executor thread that blocks 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(hook)) + except BaseException as exc: # noqa: BLE001 - transferred to the event loop + errors.append(exc) + finally: + with contextlib.suppress(RuntimeError): + loop.call_soon_threadsafe(done.set) + + threading.Thread(target=worker, name="rp-prestart", daemon=True).start() + await done.wait() + if errors: + raise errors[0] + return results[0] if results else None + + +async def _invoke_hook(hook: Callable[[], Any]) -> None: + if inspect.iscoroutinefunction(hook) or inspect.iscoroutinefunction(hook.__call__): + result = hook() + else: + result = await _run_sync_in_daemon(hook) + + if inspect.isawaitable(result): + await result + + +async def _invoke_hook_preserving_external_cancellation( + hook: Callable[[], Any], +) -> None: + """Treat hook-raised cancellation as failure, but pass through cancellation + from a timeout or worker shutdown.""" + + async def invoke() -> None: + try: + await _invoke_hook(hook) + except asyncio.CancelledError: + raise + except SystemExit as exc: + raise PrestartError(exc, _hook_name(hook)) from exc + except Exception as exc: + raise PrestartError(exc, _hook_name(hook)) from exc + + hook_task = asyncio.create_task(invoke()) + try: + await asyncio.shield(hook_task) + except asyncio.CancelledError as exc: + if hook_task.cancelled(): + raise PrestartError(exc, _hook_name(hook)) from exc + + # The caller cancelled the prestart task. + # Now cancel and drain the child task before propagating shutdown. + hook_task.cancel() + with contextlib.suppress(BaseException): + await hook_task + raise + + +async def run_prestart_hooks_async( + hooks: Sequence[Callable[[], Any]], timeout: float | None = None +) -> None: + """Run hooks in order under one optional deadline.""" + current_hook: Callable[[], Any] | None = None + + async def run_all() -> None: + nonlocal current_hook + for current_hook in hooks: + log.info(f"Prestart | running hook: {_hook_name(current_hook)}") + await _invoke_hook_preserving_external_cancellation(current_hook) + + log.info("Prestart | phase started") + if timeout is None: + await run_all() + else: + try: + await asyncio.wait_for(run_all(), timeout=timeout) + except asyncio.TimeoutError as exc: + raise PrestartTimeout(_hook_name(current_hook), timeout) from exc + log.info("Prestart | ready") diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index a3da893bd..288d72f6c 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -8,14 +8,21 @@ import signal import sys import traceback -from typing import Any, Dict, Set +from typing import Any 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 +from .rp_prestart import ( + PrestartError, + PrestartTimeout, + build_init_failed_payload, + get_prestart_hooks, + run_prestart_hooks_async, +) +from .worker_state import IS_LOCAL_TEST, JobsProgress log = RunPodLogger() @@ -45,18 +52,21 @@ class JobScaler: Job Scaler. This class is responsible for scaling the number of concurrent requests. """ - def __init__(self, config: Dict[str, Any]): + 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._prestart_ready = asyncio.Event() + self._prestart_error: dict[str, Any] | None = None + # Whether the one best-effort job-take for reporting a prestart error + # has returned a job or used its bounded grace period. + self._failure_take_done = False self.current_concurrency = 1 self.config = config + self.prestart_hooks = get_prestart_hooks() self.job_progress = JobsProgress() # Cache the singleton instance # maps in-progress job ids to their running tasks so individual jobs # can be stopped without killing the whole worker - self.jobs_tasks: Dict[str, asyncio.Task] = {} + self.jobs_tasks: dict[str, asyncio.Task[Any]] = {} self.stop_signals_fetcher = get_stop_signals self.stop_signals_fetcher_timeout = 90 @@ -66,10 +76,7 @@ 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.prestart_claim_timeout = 10 self.jobs_handler = handle_job if concurrency_modifier := config.get("concurrency_modifier"): @@ -91,7 +98,9 @@ def __init__(self, config: Dict[str, Any]): if stop_signals_fetcher := self.config.get("stop_signals_fetcher"): self.stop_signals_fetcher = stop_signals_fetcher - if stop_signals_fetcher_timeout := self.config.get("stop_signals_fetcher_timeout"): + if stop_signals_fetcher_timeout := self.config.get( + "stop_signals_fetcher_timeout" + ): self.stop_signals_fetcher_timeout = stop_signals_fetcher_timeout async def set_scale(self): @@ -146,28 +155,49 @@ def handle_shutdown(self, signum, frame): self.kill_worker() async def run(self): - # Create an async session that will be closed when the worker is killed. + """Run prestart and the three persistent request loops concurrently.""" 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)) - - # 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) - - # 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() + # Keep prestart outside the request-loop gather. A hook may have no + # timeout, so request-loop shutdown must be able to cancel it. + prestart_task = asyncio.create_task(self._run_prestart()) + request_loop_tasks = ( + asyncio.create_task(self.get_jobs(session)), + asyncio.create_task(self.run_jobs(session)), + asyncio.create_task(self.monitor_stop_signals(session)), + ) + request_loops_future = asyncio.gather(*request_loop_tasks) + try: - await init_task - except asyncio.CancelledError: - # expected: we cancelled it on the line above - pass + # Wait for either lifecycle to end without cancelling the other. + done, _ = await asyncio.wait( + {prestart_task, request_loops_future}, + return_when=asyncio.FIRST_COMPLETED, + ) + + if prestart_task in done: + try: + await prestart_task + except BaseException: + # Normal prestart failures are handled inside _run_prestart. + self.kill_worker() + raise - if self._init_error is not None: + # Prestart completion does not end the worker; request loops do. + await request_loops_future + finally: + # Clean up every task before closing their shared HTTP session. + for task in request_loop_tasks: + if not task.done(): + task.cancel() + if not prestart_task.done(): + prestart_task.cancel() + await asyncio.gather( + *request_loop_tasks, prestart_task, return_exceptions=True + ) + await asyncio.gather(request_loops_future, return_exceptions=True) + + # Normal prestart failures are reported and drained before forced respawn. + if self._prestart_error is not None: from .rp_fitness import _terminate_unhealthy _terminate_unhealthy(1) @@ -194,18 +224,52 @@ def current_occupancy(self) -> int: ) return current_progress_count + current_queue_count + async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: int): + """After prestart failure, briefly keep an in-flight job-take alive so + its request can receive the error.""" + fetch_task = asyncio.create_task(self.jobs_fetcher(session, jobs_needed)) + shutdown_task = asyncio.create_task(self._shutdown_event.wait()) + try: + done, _ = await asyncio.wait( + {fetch_task, shutdown_task}, + timeout=self.jobs_fetcher_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if fetch_task in done: + return await fetch_task + + if shutdown_task in done: + if self._prestart_error is None: + return None + + self._failure_take_done = True + try: + return await asyncio.wait_for( + fetch_task, timeout=self.prestart_claim_timeout + ) + except asyncio.TimeoutError: + return None + + raise asyncio.TimeoutError + finally: + fetch_task.cancel() + shutdown_task.cancel() + await asyncio.gather(fetch_task, shutdown_task, return_exceptions=True) + 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. + """Prestart failed before this worker 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. + Without this, an instant failure 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. """ + payload = self._prestart_error + if payload is None: + return try: jobs = await asyncio.wait_for( - self.jobs_fetcher(session, 1), timeout=self.init_claim_timeout + self.jobs_fetcher(session, 1), timeout=self.prestart_claim_timeout ) except asyncio.CancelledError: raise @@ -214,7 +278,7 @@ async def _claim_one_job_to_fail(self, session: ClientSession) -> None: return for job in jobs or []: - await self._fail_job(session, job, self._init_error) + await self._fail_job(session, job, payload) async def get_jobs(self, session: ClientSession): """ @@ -225,9 +289,9 @@ 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. + if self._prestart_error is not None: + # Prestart failure is terminal for this worker. Draining requests + # already held by the worker is owned by run_jobs. break await self.set_scale() @@ -241,21 +305,20 @@ async def get_jobs(self, session: ClientSession): log.debug("JobScaler.get_jobs | Starting job acquisition.") # Keep the connection to the blocking call with timeout - acquired_jobs = await asyncio.wait_for( - self.jobs_fetcher(session, jobs_needed), - timeout=self.jobs_fetcher_timeout, + acquired_jobs = await self._fetch_jobs_until_stopped( + session, jobs_needed ) if not acquired_jobs: log.debug("JobScaler.get_jobs | No jobs acquired.") continue - self._claimed_request = True + self._failure_take_done = True - if self._init_error is not None: - # If initialization fails, fail all in-flight requests. + if self._prestart_error is not None: + # Fail every request acquired while prestart was running. for job in acquired_jobs: - await self._fail_job(session, job, self._init_error) + await self._fail_job(session, job, self._prestart_error) return for job in acquired_jobs: @@ -285,10 +348,10 @@ 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. + if self._prestart_error is not None and not self._failure_take_done: + # No job-take was in flight when prestart failed. Make one bounded + # take so the request that scaled this worker receives the startup + # error instead of waiting for its queue TTL. await self._claim_one_job_to_fail(session) async def run_jobs(self, session: ClientSession): @@ -297,7 +360,7 @@ async def run_jobs(self, session: ClientSession): Runs the block in an infinite loop while the worker is alive or jobs queue is not empty. """ - tasks: Set[asyncio.Task] = set() + tasks: set[asyncio.Task[Any]] = set() last_task_count = 0 while self.is_alive() or not self.jobs_queue.empty(): @@ -327,14 +390,45 @@ async def run_jobs(self, session: ClientSession): # don't busy wait await asyncio.sleep(0.1) - # Ensure all remaining tasks finish before stopping. Stopped jobs raise # CancelledError during this drain, which is expected, but a genuine # handler error must not be silently discarded. results = await asyncio.gather(*tasks, return_exceptions=True) for result in results: - if isinstance(result, Exception) and not isinstance(result, asyncio.CancelledError): - log.error(f"JobScaler.run_jobs | Task failed during shutdown drain: {result}") + if isinstance(result, Exception) and not isinstance( + result, asyncio.CancelledError + ): + log.error( + f"JobScaler.run_jobs | Task failed during shutdown drain: {result}" + ) + + async def _fetch_stop_signals_until_stopped( + self, session: ClientSession + ) -> list[str] | None: + """Run one blocking job-stop poll without letting it delay shutdown. + + Job-stop may long-poll for 90 seconds. Once shutdown begins, this worker + no longer needs cancellation messages for individual jobs, so cancel + the poll immediately. Without this, a fully drained failed worker + could remain alive, just waiting for the poll timeout. + """ + fetch_task = asyncio.create_task(self.stop_signals_fetcher(session)) + shutdown_task = asyncio.create_task(self._shutdown_event.wait()) + try: + done, _ = await asyncio.wait( + {fetch_task, shutdown_task}, + timeout=self.stop_signals_fetcher_timeout, + return_when=asyncio.FIRST_COMPLETED, + ) + if fetch_task in done: + return await fetch_task + if shutdown_task in done: + return None + raise asyncio.TimeoutError + finally: + fetch_task.cancel() + shutdown_task.cancel() + await asyncio.gather(fetch_task, shutdown_task, return_exceptions=True) async def monitor_stop_signals(self, session: ClientSession): """ @@ -354,12 +448,9 @@ async def monitor_stop_signals(self, session: ClientSession): while self.is_alive(): try: - # Bound the long-poll so shutdown is not blocked by the shared - # session's much longer default timeout. - job_ids = await asyncio.wait_for( - self.stop_signals_fetcher(session), - timeout=self.stop_signals_fetcher_timeout, - ) + job_ids = await self._fetch_stop_signals_until_stopped(session) + if job_ids is None: + return for job_id in job_ids: await self.stop_job(job_id) @@ -373,7 +464,9 @@ async def monitor_stop_signals(self, session: ClientSession): log.debug("JobScaler.monitor_stop_signals | Request was cancelled.") raise # CancelledError is a BaseException except asyncio.TimeoutError: - log.debug("JobScaler.monitor_stop_signals | Stop poll timed out. Retrying.") + log.debug( + "JobScaler.monitor_stop_signals | Stop poll timed out. Retrying." + ) except Exception as error: log.error( f"JobScaler.monitor_stop_signals | Error Type: {type(error).__name__} | Error Message: {str(error)}" @@ -402,7 +495,7 @@ async def stop_job(self, job_id: str) -> bool: task.cancel() return True - async def handle_job(self, session: ClientSession, job: dict): + async def handle_job(self, session: ClientSession, job: dict[str, Any]): """ Process an individual job. This function is run concurrently for multiple jobs. """ @@ -410,19 +503,18 @@ 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(): + # Hold the handler until every registered prestart hook finishes. + if self.prestart_hooks: + if not await self._wait_for_prestart(): log.warn( - "Shutting down before initialization finished; leaving this " + "Shutting down before prestart 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) + if self._prestart_error is not None: + await self._fail_job(session, job, self._prestart_error) return await self.jobs_handler(session, self.config, job) @@ -449,55 +541,47 @@ 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()) + async def _wait_for_prestart(self) -> bool: + """Wait for prestart to finish or for worker shutdown to begin.""" + ready = asyncio.create_task(self._prestart_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() + return self._prestart_ready.is_set() async def _fail_job( - self, session: ClientSession, job: dict, payload: Dict[str, Any] + self, session: ClientSession, job: dict[str, Any], payload: dict[str, Any] ): - """Fail a request with a structured error (reason + logs).""" - log.error(f"Failing job due to init failure. | {job['id']}") + """Fail a request with a structured startup error.""" + log.error(f"Failing job due to prestart 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() + async def _run_prestart(self): + """Run hooks beside queue intake, then open the handler gate.""" + if not self.prestart_hooks: + self._prestart_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") + await run_prestart_hooks_async( + self.prestart_hooks, self.config.get("prestart_timeout") + ) + except (PrestartError, PrestartTimeout) as exc: + self._prestart_error = build_init_failed_payload( + exc, cap.getvalue() ) - 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)}") + if self._prestart_error is not None: + log.error(f"init_failed | {json.dumps(self._prestart_error)}") finally: # Always release held handlers. - self._init_ready.set() + self._prestart_ready.set() - if self._init_error is not None: + if self._prestart_error is not None: # Stop long-running loops before waiting for acquired requests to drain. self.kill_worker() while self.current_occupancy() > 0: diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 8c1bdcbe1..a1cec532e 100644 --- a/runpod/serverless/worker.py +++ b/runpod/serverless/worker.py @@ -49,7 +49,7 @@ def run_worker(config: Dict[str, Any]) -> None: # 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. + # Capture stdout/stderr so handler and prestart failures report their logs. rp_capture.install() # Create a JobScaler responsible for adjusting the concurrency diff --git a/tests/test_serverless/test_init.py b/tests/test_serverless/test_init.py index a41c05a25..ac70d197d 100644 --- a/tests/test_serverless/test_init.py +++ b/tests/test_serverless/test_init.py @@ -1,6 +1,7 @@ """Tests for runpod.serverless.__init__ module exports.""" import inspect + import runpod.serverless @@ -24,6 +25,7 @@ def test_expected_public_symbols(self): 'start', 'progress_update', 'register_fitness_check', + 'register_prestart_hook', 'runpod_version', 'VolumeCache' } @@ -99,5 +101,5 @@ def test_all_covers_public_api_only(self): assert all_symbols.issubset(public_attrs), f"__all__ contains non-public symbols: {all_symbols - public_attrs}" # Expected public API should be exactly what's in __all__ - expected_public_api = {'start', 'progress_update', 'register_fitness_check', 'runpod_version', 'VolumeCache'} + expected_public_api = {'start', 'progress_update', 'register_fitness_check', 'register_prestart_hook', 'runpod_version', 'VolumeCache'} assert all_symbols == expected_public_api, f"Expected {expected_public_api}, got {all_symbols}" diff --git a/tests/test_serverless/test_initializer.py b/tests/test_serverless/test_initializer.py deleted file mode 100644 index 519fb30dc..000000000 --- a/tests/test_serverless/test_initializer.py +++ /dev/null @@ -1,534 +0,0 @@ -"""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 pathlib -import subprocess -import sys -import tempfile -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 - - -# Needs a real process: in-process the interpreter never joins threads, so the hang -# this guards against cannot be observed. -_HARD_EXIT_SCRIPT = """ -import asyncio, sys, threading, time - -MODE = sys.argv[1] -sys.argv = ["worker"] - -from runpod.serverless.modules.rp_scale import JobScaler - - -def sync_engine(): - # A sync initializer runs on a daemon thread, so its children inherit daemon - # status. An engine that sets daemon=False itself does not. - threading.Thread(target=lambda: time.sleep(60), daemon=False).start() - raise RuntimeError("engine start failed") - - -async def async_engine(): - # An async initializer is awaited on the main thread, so anything it spawns is - # non-daemon by inheritance. - threading.Thread(target=lambda: time.sleep(60)).start() - raise RuntimeError("engine start failed") - - -async def no_jobs(*args, **kwargs): - return [] - - -scaler = JobScaler( - { - "handler": lambda job: job, - "rp_args": {}, - "initializer": {"sync": sync_engine, "async": async_engine}[MODE], - } -) -scaler.jobs_fetcher = no_jobs -scaler.stop_signals_fetcher = no_jobs -scaler.init_claim_timeout = 1 - -asyncio.run(scaler.run()) -print("run() returned without exiting") -""" - - -class TestInitFailureExitsProcess(unittest.TestCase): - """Init failure must hard-exit, even with a non-daemon thread left running.""" - - def _run_worker(self, mode: str) -> subprocess.CompletedProcess: - with tempfile.TemporaryDirectory() as tmp: - script = pathlib.Path(tmp) / "worker.py" - script.write_text(_HARD_EXIT_SCRIPT) - return subprocess.run( - [sys.executable, str(script), mode], - capture_output=True, - text=True, - timeout=30, # generous: the fix exits in well under a second - check=False, - ) - - def test_sync_initializer_leaving_a_non_daemon_thread(self): - result = self._run_worker("sync") - assert result.returncode == 1 - assert "run() returned without exiting" not in result.stdout - - def test_async_initializer_leaving_a_non_daemon_thread(self): - result = self._run_worker("async") - assert result.returncode == 1 - assert "init_failed" in result.stdout + result.stderr # reported before exiting - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index 523c3cfbf..fe78911f8 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -1,14 +1,77 @@ -""" Tests for rp_local.py """ +"""Tests for rp_local.py""" from unittest import IsolatedAsyncioTestCase -from unittest.mock import mock_open, patch +from unittest.mock import AsyncMock, mock_open, patch from runpod.serverless.modules import rp_local +from runpod.serverless.modules.rp_prestart import ( + clear_prestart_hooks, + register_prestart_hook, +) class TestRunLocal(IsolatedAsyncioTestCase): """Tests for run_local function""" + def setUp(self): + clear_prestart_hooks() + + def tearDown(self): + clear_prestart_hooks() + + async def test_prestart_runs_before_local_handler(self): + calls = [] + + @register_prestart_hook + async def load_model(): + calls.append("prestart") + + async def run_job(*_args): + calls.append("handler") + return {"result": "success"} + + config = { + "handler": "handler", + "rp_args": {"test_input": {"input": "test"}}, + } + with ( + patch( + "runpod.serverless.modules.rp_local.run_job", + new=AsyncMock(side_effect=run_job), + ), + self.assertRaises(SystemExit) as sys_exit, + ): + await rp_local.run_local(config) + + self.assertEqual(sys_exit.exception.code, 0) + self.assertEqual(calls, ["prestart", "handler"]) + + async def test_prestart_failure_skips_local_handler(self): + @register_prestart_hook + def load_model(): + raise RuntimeError("model unavailable") + + config = { + "handler": "handler", + "rp_args": {"test_input": {"input": "test"}}, + } + with ( + patch( + "runpod.serverless.modules.rp_local.run_job", new=AsyncMock() + ) as run_job, + patch("runpod.serverless.modules.rp_local.log") as logger, + patch("runpod.serverless.modules.rp_local._terminate_unhealthy"), + self.assertRaises(SystemExit) as sys_exit, + ): + await rp_local.run_local(config) + + self.assertEqual(sys_exit.exception.code, 1) + run_job.assert_not_awaited() + failure_log = logger.error.call_args.args[0] + self.assertIn("init_failed", failure_log) + self.assertIn("load_model", failure_log) + self.assertIn("model unavailable", failure_log) + @patch( "runpod.serverless.modules.rp_local.run_job", return_value={"result": "success"} ) diff --git a/tests/test_serverless/test_prestart.py b/tests/test_serverless/test_prestart.py new file mode 100644 index 000000000..8afad66b0 --- /dev/null +++ b/tests/test_serverless/test_prestart.py @@ -0,0 +1,267 @@ +"""Public prestart-hook contract for queue-based Serverless workers.""" + +import asyncio +import os +import unittest +from unittest.mock import patch + +import runpod.serverless +from runpod.serverless.modules import rp_fastapi +from runpod.serverless.modules.rp_prestart import ( + PrestartError, + PrestartTimeout, + build_init_failed_payload, + clear_prestart_hooks, + get_prestart_hooks, + run_prestart_hooks_async, +) + + +def _run(coro): + return asyncio.run(coro) + + +class TestPrestartRegistry(unittest.TestCase): + def setUp(self): + clear_prestart_hooks() + + def tearDown(self): + clear_prestart_hooks() + + def test_decorator_registers_hooks_in_order_and_returns_each_callable(self): + def first(): + return None + + async def second(): + return None + + assert runpod.serverless.register_prestart_hook(first) is first + assert runpod.serverless.register_prestart_hook(second) is second + assert get_prestart_hooks() == (first, second) + + def test_mixed_hooks_run_sequentially(self): + calls = [] + + @runpod.serverless.register_prestart_hook + def first(): + calls.append("first") + + @runpod.serverless.register_prestart_hook + async def second(): + calls.append("second") + await asyncio.sleep(0) + calls.append("second-ready") + + @runpod.serverless.register_prestart_hook + def third(): + calls.append("third") + + _run(run_prestart_hooks_async(get_prestart_hooks())) + + assert calls == ["first", "second", "second-ready", "third"] + + def test_failure_stops_later_hooks_and_names_the_failing_hook(self): + calls = [] + + def load_model(): + calls.append("load_model") + raise RuntimeError("CUDA OOM") + + def warm_cache(): + calls.append("warm_cache") + + with self.assertRaises(PrestartError) as ctx: + _run(run_prestart_hooks_async((load_model, warm_cache))) + + assert calls == ["load_model"] + assert ctx.exception.hook == "load_model" + payload = build_init_failed_payload(ctx.exception) + assert payload["event"] == "init_failed" + assert payload["hook"] == "load_model" + assert payload["error_message"] == "CUDA OOM" + + def test_timeout_bounds_the_whole_phase_and_names_current_hook(self): + async def first(): + await asyncio.sleep(0.03) + + async def second(): + await asyncio.sleep(0.03) + + with self.assertRaises(PrestartTimeout) as ctx: + _run(run_prestart_hooks_async((first, second), timeout=0.04)) + + assert ctx.exception.hook == "second" + + def test_hook_raised_cancelled_error_is_a_prestart_failure(self): + async def cancelled_hook(): + raise asyncio.CancelledError + + with self.assertRaises(PrestartError) as ctx: + _run(run_prestart_hooks_async((cancelled_hook,))) + + assert isinstance(ctx.exception.original, asyncio.CancelledError) + assert ctx.exception.hook == "cancelled_hook" + + def test_external_task_cancellation_remains_cancellation(self): + async def scenario(): + started = asyncio.Event() + + async def blocked_hook(): + started.set() + await asyncio.Event().wait() + + task = asyncio.create_task(run_prestart_hooks_async((blocked_hook,))) + await started.wait() + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + + _run(scenario()) + + +class TestHostedAPIPrestart(unittest.TestCase): + def setUp(self): + clear_prestart_hooks() + + def tearDown(self): + clear_prestart_hooks() + + def test_lifespan_runs_hooks_before_serving(self): + state = [] + + @runpod.serverless.register_prestart_hook + async def load_model(): + state.append("ready") + + async def scenario(): + with patch.object(rp_fastapi.heartbeat, "start_ping"): + api = rp_fastapi.WorkerAPI( + {"handler": lambda job: job, "prestart_timeout": 1} + ) + async with api.rp_app.router.lifespan_context(api.rp_app): + self.assertEqual(state, ["ready"]) + + _run(scenario()) + + def test_lifespan_failure_prevents_serving(self): + @runpod.serverless.register_prestart_hook + def load_model(): + raise RuntimeError("model unavailable") + + async def scenario(): + with patch.object(rp_fastapi.heartbeat, "start_ping"): + api = rp_fastapi.WorkerAPI({"handler": lambda job: job}) + with ( + patch.object(rp_fastapi.log, "error") as logger, + patch.object(rp_fastapi, "_terminate_unhealthy"), + self.assertRaisesRegex(PrestartError, "model unavailable"), + ): + async with api.rp_app.router.lifespan_context(api.rp_app): + self.fail("API served despite prestart failure") + self.assertIn("init_failed", logger.call_args.args[0]) + self.assertIn("load_model", logger.call_args.args[0]) + + _run(scenario()) + + +class TestPrestartModeGuard(unittest.TestCase): + def setUp(self): + clear_prestart_hooks() + + @runpod.serverless.register_prestart_hook + def load_model(): + return None + + def tearDown(self): + clear_prestart_hooks() + + @staticmethod + def _config(*, serve_api=False): + return { + "handler": lambda job: job, + "rp_args": { + "rp_log_level": None, + "rp_debugger": None, + "rp_serve_api": serve_api, + "rp_api_port": 8000, + "rp_api_concurrency": 1, + "rp_api_host": "localhost", + "test_input": None, + }, + } + + def test_hosted_api_accepts_registered_hooks(self): + config = self._config(serve_api=True) + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + patch("runpod.serverless.modules.rp_fastapi.WorkerAPI") as worker_api, + ): + runpod.serverless.start(config) + + worker_api.assert_called_once_with(config) + worker_api.return_value.start_uvicorn.assert_called_once() + + def test_hosted_api_precedes_realtime_environment(self): + config = self._config(serve_api=True) + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + patch.dict(os.environ, {"RUNPOD_REALTIME_PORT": "8000"}), + patch("runpod.serverless.modules.rp_fastapi.WorkerAPI") as worker_api, + ): + runpod.serverless.start(config) + + worker_api.assert_called_once_with(config) + + def test_hosted_api_rejects_multiple_workers(self): + config = self._config(serve_api=True) + config["rp_args"]["rp_api_concurrency"] = 2 + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + self.assertRaisesRegex(RuntimeError, "rp_api_concurrency=1"), + ): + runpod.serverless.start(config) + + def test_realtime_rejects_registered_hooks(self): + config = self._config() + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + patch.dict(os.environ, {"RUNPOD_REALTIME_PORT": "8000"}), + self.assertRaisesRegex(RuntimeError, "realtime mode"), + ): + runpod.serverless.start(config) + + def test_local_accepts_registered_hooks(self): + config = self._config() + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + patch("runpod.serverless.worker.main") as worker_main, + patch.dict(os.environ, {}, clear=True), + ): + runpod.serverless.start(config) + + worker_main.assert_called_once_with(config) + + def test_queue_worker_accepts_registered_hooks(self): + config = self._config() + with ( + patch("runpod.serverless._set_config_args", return_value=config), + patch("runpod.serverless.signal.signal"), + patch("runpod.serverless.worker.main") as worker_main, + patch.dict( + os.environ, + {"RUNPOD_WEBHOOK_GET_JOB": "https://api.runpod.ai/job-take"}, + clear=True, + ), + ): + runpod.serverless.start(config) + + worker_main.assert_called_once_with(config) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_serverless/test_prestart_lifecycle.py b/tests/test_serverless/test_prestart_lifecycle.py new file mode 100644 index 000000000..d068e705b --- /dev/null +++ b/tests/test_serverless/test_prestart_lifecycle.py @@ -0,0 +1,691 @@ +"""Queue prestart lifecycle coverage: handler gate, failure delivery, drain, and exit.""" + +# pylint: disable=protected-access + +import asyncio +import functools +import io +import json +import pathlib +import subprocess +import sys +import tempfile +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_prestart import ( + PrestartError, + PrestartTimeout, + build_init_failed_payload, + run_prestart_hooks_async, +) +from runpod.serverless.modules.rp_scale import JobScaler + + +def _run(coro): + return asyncio.run(coro) + + +class TestRunPrestartHooksAsync(unittest.TestCase): + """Runs ordered hooks to completion or raises a named startup failure.""" + + def test_sync_success_offloaded(self): + calls = [] + _run(run_prestart_hooks_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_prestart_hooks_async((initialize,))) + assert state == {"ready": True} + + def test_sync_failure_wraps_in_prestart_error(self): + def failing_hook(): + raise ValueError("max_model_len must be positive, got 0") + + with self.assertRaises(PrestartError) as ctx: + _run(run_prestart_hooks_async((failing_hook,))) + 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_prestart_hooks_async((load,))) + assert state == {"ready": True} + + def test_async_failure_wraps_in_prestart_error(self): + async def failing_hook(): + raise RuntimeError("CUDA OOM") + + with self.assertRaises(PrestartError) as ctx: + _run(run_prestart_hooks_async((failing_hook,))) + assert isinstance(ctx.exception.original, RuntimeError) + + def test_hook_timeout_error_is_not_phase_timeout(self): + async def failing_hook(): + raise asyncio.TimeoutError("backend request timed out") + + with self.assertRaises(PrestartError) as ctx: + _run(run_prestart_hooks_async((failing_hook,), timeout=1)) + assert isinstance(ctx.exception.original, asyncio.TimeoutError) + + def test_async_timeout_raises_prestart_timeout(self): + async def slow(): + await asyncio.sleep(3) + + with self.assertRaises(PrestartTimeout): + _run(run_prestart_hooks_async((slow,), timeout=1)) + + def test_zero_timeout_raises_prestart_timeout(self): + async def load(): + await asyncio.sleep(0) + + with self.assertRaises(PrestartTimeout): + _run(run_prestart_hooks_async((load,), timeout=0)) + + def test_sync_hang_times_out(self): + """A blocking sync hook is cut off by prestart_timeout, not left stuck.""" + release = threading.Event() + + def hang(): + release.wait(10) + + try: + with self.assertRaises(PrestartTimeout): + _run(run_prestart_hooks_async((hang,), timeout=1)) + finally: + release.set() # let the offloaded thread exit promptly + + def test_async_partial_is_awaited(self): + """A partial wrapping an async hook is detected and awaited.""" + state = {} + + async def load(key): + await asyncio.sleep(0) + state[key] = True + + _run(run_prestart_hooks_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_prestart_hooks_async((Loader(),))) + assert state == {"ready": True} + + def test_base_exception_propagates_unwrapped(self): + """KeyboardInterrupt is process control, not a prestart failure.""" + + def interrupted(): + raise KeyboardInterrupt + + with self.assertRaises(KeyboardInterrupt): + _run(run_prestart_hooks_async((interrupted,))) + + def test_system_exit_wraps_as_prestart_error(self): + """A hook calling sys.exit() is a failure to surface with a reason, + not a clean exit that abandons held jobs.""" + + def bails(): + raise SystemExit(1) + + with self.assertRaises(PrestartError): + _run(run_prestart_hooks_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( + PrestartError(exc, "load_model"), 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 payload["hook"] == "load_model" + assert "worker_id" in payload and "runpod_version" in payload + + def test_payload_omits_empty_logs(self): + payload = build_init_failed_payload(PrestartError(ValueError("x"), "hook")) + 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( + PrestartError(ValueError(huge), "hook"), + 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(hook=None): + scaler = JobScaler({"handler": lambda j: j, "rp_args": {}}) + scaler.prestart_hooks = () if hook is None else (hook,) + scaler.job_progress = MagicMock() # avoid the process-wide singleton in unit tests + scaler.job_progress.get_job_count.return_value = 0 + return scaler + + +class TestRunPrestart(unittest.TestCase): + """Prestart opens the handler gate or records the failure and shuts down.""" + + def test_no_hooks_opens_gate_immediately(self): + scaler = _scaler(hook=None) + _run(scaler._run_prestart()) + assert scaler._prestart_ready.is_set() + assert scaler._prestart_error is None + + def test_success_opens_gate_no_error(self): + ran = [] + scaler = _scaler(hook=lambda: ran.append(1)) + _run(scaler._run_prestart()) + assert ran == [1] + assert scaler._prestart_ready.is_set() + assert scaler._prestart_error is None + assert not scaler._shutdown_event.is_set() + + def test_failure_records_hook_reason_with_logs_and_shuts_down(self): + def failing_hook(): + print("downloading model") + raise RuntimeError("CUDA OOM: model too big") + + scaler = _scaler(hook=failing_hook) + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(rp_scale, "log") as mock_log, + ): + _run(scaler._run_prestart()) + + assert ( + scaler._prestart_ready.is_set() + ) # held handlers are released to fail fast + assert scaler._prestart_error is not None + assert scaler._prestart_error["error_message"] == "CUDA OOM: model too big" + assert scaler._prestart_error["hook"] == "failing_hook" + assert "downloading model" in scaler._prestart_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(hook=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_prestart()) + + assert shutdown_states == [True] + + +class TestHandleJobGate(unittest.TestCase): + """The handler waits for prestart; 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_prestart_is_ready(self): + scaler = _scaler(hook=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._prestart_ready.set() + 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_handler_stays_blocked_until_prestart_completes(self): + started = asyncio.Event() + release = asyncio.Event() + + async def blocked_hook(): + started.set() + await release.wait() + + scaler = _scaler(hook=blocked_hook) + scaler.jobs_handler = AsyncMock() + job = {"id": "blocked"} + + async def go(): + self._prime(scaler, job) + prestart_task = asyncio.create_task(scaler._run_prestart()) + await started.wait() + handler_task = asyncio.create_task(scaler.handle_job(None, job)) + await asyncio.sleep(0) + scaler.jobs_handler.assert_not_awaited() + + release.set() + await asyncio.gather(prestart_task, handler_task) + + _run(go()) + scaler.jobs_handler.assert_awaited_once() + + def test_runs_handler_without_hooks(self): + scaler = _scaler(hook=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_prestart_failure_fails_request_without_running_handler(self): + scaler = _scaler(hook=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._prestart_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_prestart_failure_before_any_take_claims_a_request_to_fail(self): + """An instant failure leaves no request in hand and starts shutdown. + The worker must claim one request and fail it, or the caller waits out + the queue TTL without receiving the startup reason.""" + scaler = _scaler(hook=lambda: None) + scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() # _run_prestart 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_prestart_failure_with_empty_queue_exits_without_hanging(self): + """An empty claim attempt lets job-take stop without hanging.""" + scaler = _scaler(hook=lambda: None) + scaler._prestart_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_prestart_failure_after_a_take_does_not_claim_another_request(self): + """Once this worker claims a request, report the failure against it. + Claiming another would fail work that a healthy worker could serve.""" + scaler = _scaler(hook=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._prestart_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_failure_still_fails_the_queued_request(self): + """An async hook can fail before job-take fetches anything. The real + prestart failure path must still claim and fail the queued request.""" + + async def bad_hook(): + raise RuntimeError("bad config") + + scaler = _scaler(hook=bad_hook) + 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_prestart() # 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(hook=lambda: None) + scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() + scaler._fail_job = AsyncMock() + scaler.prestart_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_inflight_take_uses_failure_bound_after_shutdown(self): + """Prestart failure wakes an existing long-poll instead of waiting for + the normal 90-second job-take timeout before trying to exit.""" + scaler = _scaler(hook=lambda: None) + scaler._fail_job = AsyncMock() + scaler.prestart_claim_timeout = 0.05 + started = asyncio.Event() + + async def never_returns(_session, _needed): + started.set() + await asyncio.Event().wait() + + scaler.jobs_fetcher = never_returns + + async def go(): + task = asyncio.create_task(scaler.get_jobs(AsyncMock())) + await started.wait() + scaler._prestart_error = { + "error_message": "CUDA OOM", + "event": "init_failed", + } + scaler.kill_worker() + await asyncio.wait_for(task, timeout=0.5) + + _run(go()) + scaler._fail_job.assert_not_awaited() + + def test_jobs_acquired_after_prestart_failure_are_failed_not_queued(self): + """If prestart fails during a long-poll, fail returned jobs and stop + job-take without relying on another task to end the loop.""" + scaler = _scaler(hook=lambda: None) + scaler._fail_job = AsyncMock() + job = {"id": "late-1"} + failure = {"error_message": "CUDA OOM", "event": "init_failed"} + + async def fetcher(_session, _needed): + # Prestart failure lands while this long-poll is in flight. + scaler._prestart_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_prestart_ready_leaves_request_alone(self): + """SIGTERM during prestart releases a held request without running the + handler against an uninitialized worker.""" + scaler = _scaler(hook=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler.kill_worker() + 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 TestShutdownWithHangingPrestart(unittest.TestCase): + """A hung hook without a timeout must not outlive worker shutdown.""" + + def test_run_returns_while_hook_is_still_blocked(self): + release = threading.Event() + self.addCleanup(release.set) + + scaler = _scaler(hook=release.wait) + scaler.kill_worker() + 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 + + def test_prestart_failure_wakes_blocking_stop_poll(self): + async def blocked_stop_poll(_session): + await asyncio.Event().wait() + + def fail(): + raise RuntimeError("model unavailable") + + scaler = _scaler(hook=fail) + scaler.jobs_fetcher = AsyncMock(return_value=[]) + scaler.stop_signals_fetcher = blocked_stop_poll + scaler.prestart_claim_timeout = 0.01 + + with patch( + "runpod.serverless.modules.rp_fitness._terminate_unhealthy" + ) as terminate: + _run(asyncio.wait_for(scaler.run(), timeout=0.5)) + + terminate.assert_called_once_with(1) + + def test_process_control_exception_stops_queue_loops_and_propagates(self): + class StopWorker(BaseException): + pass + + async def interrupted(): + raise StopWorker("stop") + + async def blocked(_session, *_args): + await asyncio.Event().wait() + + scaler = _scaler(hook=interrupted) + scaler.jobs_fetcher = blocked + scaler.stop_signals_fetcher = blocked + scaler.jobs_handler = AsyncMock() + + with self.assertRaisesRegex(StopWorker, "stop"): + _run(asyncio.wait_for(scaler.run(), timeout=0.5)) + + scaler.jobs_handler.assert_not_awaited() + + def test_request_loop_failure_cancels_blocked_sibling_loops(self): + """A failed request loop must not outlive its shared HTTP session.""" + take_started = asyncio.Event() + stop_started = asyncio.Event() + take_stopped = asyncio.Event() + stop_stopped = asyncio.Event() + + async def blocked_take(_session): + take_started.set() + try: + await asyncio.Event().wait() + finally: + take_stopped.set() + + async def blocked_stop(_session): + stop_started.set() + try: + await asyncio.Event().wait() + finally: + stop_stopped.set() + + async def failed_run_jobs(_session): + await asyncio.gather(take_started.wait(), stop_started.wait()) + raise RuntimeError("request loop failed") + + scaler = _scaler(hook=lambda: None) + scaler.get_jobs = blocked_take + scaler.run_jobs = failed_run_jobs + scaler.monitor_stop_signals = blocked_stop + + with self.assertRaisesRegex(RuntimeError, "request loop failed"): + _run(asyncio.wait_for(scaler.run(), timeout=0.5)) + + assert take_stopped.is_set() + assert stop_stopped.is_set() + + +# Needs a real process: in-process the interpreter never joins threads, so the hang +# this guards against cannot be observed. +_HARD_EXIT_SCRIPT = """ +import asyncio, sys, threading, time + +ADAPTER, HOOK_MODE = sys.argv[1:3] +sys.argv = ["worker"] + +from runpod.serverless.modules.rp_fastapi import WorkerAPI +from runpod.serverless.modules.rp_local import run_local +from runpod.serverless.modules.rp_prestart import register_prestart_hook +from runpod.serverless.modules.rp_scale import JobScaler + + +def sync_engine(): + # A sync hook runs on a daemon thread, so its children inherit daemon status. + # An engine that explicitly sets daemon=False still blocks cooperative exit. + threading.Thread(target=lambda: time.sleep(60), daemon=False).start() + raise RuntimeError("engine start failed") + + +async def async_engine(): + # An async hook runs on the main event loop; child threads inherit non-daemon. + threading.Thread(target=lambda: time.sleep(60)).start() + raise RuntimeError("engine start failed") + + +async def no_jobs(*args, **kwargs): + return [] + + +hook = {"sync": sync_engine, "async": async_engine}[HOOK_MODE] +config = { + "handler": lambda job: job, + "rp_args": {"test_input": {"input": "test"}}, +} + +if ADAPTER == "queue": + scaler = JobScaler(config) + scaler.prestart_hooks = (hook,) + scaler.jobs_fetcher = no_jobs + scaler.stop_signals_fetcher = no_jobs + scaler.prestart_claim_timeout = 1 + asyncio.run(scaler.run()) +else: + register_prestart_hook(hook) + if ADAPTER == "local": + asyncio.run(run_local(config)) + else: + worker = object.__new__(WorkerAPI) + worker.config = config + + async def run_hosted(): + async with worker._lifespan(None): + pass + + asyncio.run(run_hosted()) + +print("run() returned without exiting") +""" + + +class TestPrestartFailureExitsProcess(unittest.TestCase): + """Prestart failure hard-exits even when a non-daemon thread remains.""" + + def _run_worker( + self, adapter: str, hook_mode: str = "async" + ) -> subprocess.CompletedProcess: + with tempfile.TemporaryDirectory() as tmp: + script = pathlib.Path(tmp) / "worker.py" + script.write_text(_HARD_EXIT_SCRIPT) + return subprocess.run( + [sys.executable, str(script), adapter, hook_mode], + capture_output=True, + text=True, + timeout=30, # generous: the fix exits in well under a second + check=False, + ) + + def assert_hard_exit(self, adapter: str, hook_mode: str = "async") -> None: + result = self._run_worker(adapter, hook_mode) + assert result.returncode == 1 + assert "run() returned without exiting" not in result.stdout + assert "init_failed" in result.stdout + result.stderr + + def test_queue_sync_hook_leaving_a_non_daemon_thread(self): + self.assert_hard_exit("queue", "sync") + + def test_queue_async_hook_leaving_a_non_daemon_thread(self): + self.assert_hard_exit("queue") + + def test_local_hook_leaving_a_non_daemon_thread(self): + self.assert_hard_exit("local") + + def test_hosted_api_hook_leaving_a_non_daemon_thread(self): + self.assert_hard_exit("hosted") + + +if __name__ == "__main__": + unittest.main() From 5bbad5fc913c9c73f7b5fbf9037815f9016445ad Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 19 Aug 2026 15:28:19 -0700 Subject: [PATCH 4/7] refactor(serverless): rename init failure to prestart failure --- docs/serverless/worker.md | 4 +- runpod/serverless/modules/rp_fastapi.py | 6 +-- runpod/serverless/modules/rp_local.py | 6 +-- runpod/serverless/modules/rp_prestart.py | 8 ++-- runpod/serverless/modules/rp_scale.py | 6 +-- .../test_modules/test_local.py | 2 +- tests/test_serverless/test_prestart.py | 8 ++-- .../test_prestart_lifecycle.py | 46 ++++++++++++------- 8 files changed, 49 insertions(+), 37 deletions(-) diff --git a/docs/serverless/worker.md b/docs/serverless/worker.md index 4f209d38b..e06967147 100644 --- a/docs/serverless/worker.md +++ b/docs/serverless/worker.md @@ -84,8 +84,8 @@ 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 `init_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 `init_failed` and prevents API startup. | +| 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. | diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index d77d08536..32206803f 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -26,7 +26,7 @@ from .rp_prestart import ( PrestartError, PrestartTimeout, - build_init_failed_payload, + build_prestart_failed_payload, get_prestart_hooks, run_prestart_hooks_async, ) @@ -203,8 +203,8 @@ async def _lifespan(self, _app: FastAPI) -> AsyncGenerator[None, None]: hooks, self.config.get("prestart_timeout") ) except (PrestartError, PrestartTimeout) as exc: - failure = build_init_failed_payload(exc, captured.getvalue()) - log.error(f"init_failed | {json.dumps(failure)}") + failure = build_prestart_failed_payload(exc, captured.getvalue()) + log.error(f"prestart_failed | {json.dumps(failure)}") _terminate_unhealthy(1) raise yield diff --git a/runpod/serverless/modules/rp_local.py b/runpod/serverless/modules/rp_local.py index 515b3f09b..cf682f453 100644 --- a/runpod/serverless/modules/rp_local.py +++ b/runpod/serverless/modules/rp_local.py @@ -16,7 +16,7 @@ from .rp_prestart import ( PrestartError, PrestartTimeout, - build_init_failed_payload, + build_prestart_failed_payload, get_prestart_hooks, run_prestart_hooks_async, ) @@ -55,8 +55,8 @@ async def run_local(config: dict[str, Any]) -> None: get_prestart_hooks(), config.get("prestart_timeout") ) except (PrestartError, PrestartTimeout) as exc: - failure = build_init_failed_payload(exc, captured.getvalue()) - log.error(f"init_failed | {json.dumps(failure)}") + failure = build_prestart_failed_payload(exc, captured.getvalue()) + log.error(f"prestart_failed | {json.dumps(failure)}") _terminate_unhealthy(1) sys.exit(1) diff --git a/runpod/serverless/modules/rp_prestart.py b/runpod/serverless/modules/rp_prestart.py index 14273d487..12377ebe4 100644 --- a/runpod/serverless/modules/rp_prestart.py +++ b/runpod/serverless/modules/rp_prestart.py @@ -24,7 +24,7 @@ log = RunPodLogger() -INIT_FAILED_EVENT = "init_failed" +PRESTART_FAILED_EVENT = "prestart_failed" _prestart_hooks: list[Callable[[], Any]] = [] @@ -76,11 +76,11 @@ def __init__(self, original: BaseException, hook: str): super().__init__(str(original)) -def build_init_failed_payload(exc: BaseException, logs: str = "") -> dict[str, Any]: - """Build the structured `init_failed` reason sent through `/job-done`.""" +def build_prestart_failed_payload(exc: BaseException, logs: str = "") -> dict[str, Any]: + """Build the structured `prestart_failed` reason sent through `/job-done`.""" original = getattr(exc, "original", exc) payload = { - "event": INIT_FAILED_EVENT, + "event": PRESTART_FAILED_EVENT, "error_type": type(original).__name__, "error_message": clip(str(original)), "error_traceback": clip( diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 288d72f6c..b98324768 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -18,7 +18,7 @@ from .rp_prestart import ( PrestartError, PrestartTimeout, - build_init_failed_payload, + build_prestart_failed_payload, get_prestart_hooks, run_prestart_hooks_async, ) @@ -572,11 +572,11 @@ async def _run_prestart(self): self.prestart_hooks, self.config.get("prestart_timeout") ) except (PrestartError, PrestartTimeout) as exc: - self._prestart_error = build_init_failed_payload( + self._prestart_error = build_prestart_failed_payload( exc, cap.getvalue() ) if self._prestart_error is not None: - log.error(f"init_failed | {json.dumps(self._prestart_error)}") + log.error(f"prestart_failed | {json.dumps(self._prestart_error)}") finally: # Always release held handlers. self._prestart_ready.set() diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index fe78911f8..661b550f4 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -68,7 +68,7 @@ def load_model(): self.assertEqual(sys_exit.exception.code, 1) run_job.assert_not_awaited() failure_log = logger.error.call_args.args[0] - self.assertIn("init_failed", failure_log) + self.assertIn("prestart_failed", failure_log) self.assertIn("load_model", failure_log) self.assertIn("model unavailable", failure_log) diff --git a/tests/test_serverless/test_prestart.py b/tests/test_serverless/test_prestart.py index 8afad66b0..af7ac13e3 100644 --- a/tests/test_serverless/test_prestart.py +++ b/tests/test_serverless/test_prestart.py @@ -10,7 +10,7 @@ from runpod.serverless.modules.rp_prestart import ( PrestartError, PrestartTimeout, - build_init_failed_payload, + build_prestart_failed_payload, clear_prestart_hooks, get_prestart_hooks, run_prestart_hooks_async, @@ -75,8 +75,8 @@ def warm_cache(): assert calls == ["load_model"] assert ctx.exception.hook == "load_model" - payload = build_init_failed_payload(ctx.exception) - assert payload["event"] == "init_failed" + payload = build_prestart_failed_payload(ctx.exception) + assert payload["event"] == "prestart_failed" assert payload["hook"] == "load_model" assert payload["error_message"] == "CUDA OOM" @@ -158,7 +158,7 @@ async def scenario(): ): async with api.rp_app.router.lifespan_context(api.rp_app): self.fail("API served despite prestart failure") - self.assertIn("init_failed", logger.call_args.args[0]) + self.assertIn("prestart_failed", logger.call_args.args[0]) self.assertIn("load_model", logger.call_args.args[0]) _run(scenario()) diff --git a/tests/test_serverless/test_prestart_lifecycle.py b/tests/test_serverless/test_prestart_lifecycle.py index d068e705b..caf088150 100644 --- a/tests/test_serverless/test_prestart_lifecycle.py +++ b/tests/test_serverless/test_prestart_lifecycle.py @@ -18,7 +18,7 @@ from runpod.serverless.modules.rp_prestart import ( PrestartError, PrestartTimeout, - build_init_failed_payload, + build_prestart_failed_payload, run_prestart_hooks_async, ) from runpod.serverless.modules.rp_scale import JobScaler @@ -154,17 +154,17 @@ def bails(): _run(run_prestart_hooks_async((bails,))) -class TestInitFailedSignal(unittest.TestCase): - """Builds the structured init_failed payload, including captured logs.""" +class TestPrestartFailedSignal(unittest.TestCase): + """Builds the structured prestart_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( + payload = build_prestart_failed_payload( PrestartError(exc, "load_model"), logs="stderr tail" ) - assert payload["event"] == "init_failed" + assert payload["event"] == "prestart_failed" assert payload["error_type"] == "ValueError" assert payload["error_message"] == "bad config" assert "ValueError" in payload["error_traceback"] @@ -173,12 +173,12 @@ def test_payload_shape_from_sync_error(self): assert "worker_id" in payload and "runpod_version" in payload def test_payload_omits_empty_logs(self): - payload = build_init_failed_payload(PrestartError(ValueError("x"), "hook")) + payload = build_prestart_failed_payload(PrestartError(ValueError("x"), "hook")) 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( + payload = build_prestart_failed_payload( PrestartError(ValueError(huge), "hook"), logs="y" * (rp_capture.MAX_CAPTURED_CHARS * 3), ) @@ -234,7 +234,7 @@ def failing_hook(): assert scaler._prestart_error["hook"] == "failing_hook" assert "downloading model" in scaler._prestart_error["logs"] assert any( - call.args[0].startswith("init_failed | ") + call.args[0].startswith("prestart_failed | ") for call in mock_log.error.call_args_list ) assert ( @@ -317,7 +317,10 @@ async def go(): def test_prestart_failure_fails_request_without_running_handler(self): scaler = _scaler(hook=lambda: None) scaler.jobs_handler = AsyncMock() - scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._prestart_error = { + "error_message": "CUDA OOM", + "event": "prestart_failed", + } scaler._prestart_ready.set() job = {"id": "j3"} @@ -337,7 +340,10 @@ def test_prestart_failure_before_any_take_claims_a_request_to_fail(self): The worker must claim one request and fail it, or the caller waits out the queue TTL without receiving the startup reason.""" scaler = _scaler(hook=lambda: None) - scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._prestart_error = { + "error_message": "CUDA OOM", + "event": "prestart_failed", + } scaler.kill_worker() # _run_prestart sets shutdown on failure scaler._fail_job = AsyncMock() job = {"id": "orphan-1"} @@ -353,7 +359,10 @@ def test_prestart_failure_before_any_take_claims_a_request_to_fail(self): def test_prestart_failure_with_empty_queue_exits_without_hanging(self): """An empty claim attempt lets job-take stop without hanging.""" scaler = _scaler(hook=lambda: None) - scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._prestart_error = { + "error_message": "CUDA OOM", + "event": "prestart_failed", + } scaler.kill_worker() scaler._fail_job = AsyncMock() scaler.jobs_fetcher = AsyncMock(return_value=[]) @@ -368,7 +377,7 @@ def test_prestart_failure_after_a_take_does_not_claim_another_request(self): Claiming another would fail work that a healthy worker could serve.""" scaler = _scaler(hook=lambda: None) scaler._fail_job = AsyncMock() - failure = {"error_message": "CUDA OOM", "event": "init_failed"} + failure = {"error_message": "CUDA OOM", "event": "prestart_failed"} calls = [] async def fetcher(_session, _needed): @@ -411,7 +420,7 @@ async def 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["event"] == "prestart_failed" assert reason["error_message"] == "bad config" scaler.jobs_handler.assert_not_awaited() # handler never runs on a broken worker @@ -419,7 +428,10 @@ 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(hook=lambda: None) - scaler._prestart_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._prestart_error = { + "error_message": "CUDA OOM", + "event": "prestart_failed", + } scaler.kill_worker() scaler._fail_job = AsyncMock() scaler.prestart_claim_timeout = 0.05 @@ -452,7 +464,7 @@ async def go(): await started.wait() scaler._prestart_error = { "error_message": "CUDA OOM", - "event": "init_failed", + "event": "prestart_failed", } scaler.kill_worker() await asyncio.wait_for(task, timeout=0.5) @@ -466,7 +478,7 @@ def test_jobs_acquired_after_prestart_failure_are_failed_not_queued(self): scaler = _scaler(hook=lambda: None) scaler._fail_job = AsyncMock() job = {"id": "late-1"} - failure = {"error_message": "CUDA OOM", "event": "init_failed"} + failure = {"error_message": "CUDA OOM", "event": "prestart_failed"} async def fetcher(_session, _needed): # Prestart failure lands while this long-poll is in flight. @@ -672,7 +684,7 @@ def assert_hard_exit(self, adapter: str, hook_mode: str = "async") -> None: result = self._run_worker(adapter, hook_mode) assert result.returncode == 1 assert "run() returned without exiting" not in result.stdout - assert "init_failed" in result.stdout + result.stderr + assert "prestart_failed" in result.stdout + result.stderr def test_queue_sync_hook_leaving_a_non_daemon_thread(self): self.assert_hard_exit("queue", "sync") From 88ccf4c1830b2876f1e9e2201bb972b895882eaa Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 19 Aug 2026 16:37:45 -0700 Subject: [PATCH 5/7] feat(serverless): let workers opt out of failure log capture Captured stdout/stderr is attached to the failure a worker reports, and that payload is returned to whoever called the request. On a private endpoint the caller already has the worker logs, but on a public or shared endpoint the caller is a third party, so a worker that prints credentials now hands them out. Add RUNPOD_SKIP_LOG_CAPTURE=true to skip installing the tee proxy, which leaves every failure payload without a logs field while output still reaches the worker logs. Document the field and the trade-off, and omit logs entirely when nothing was captured, matching the prestart payload. --- docs/serverless/worker.md | 30 +++++++++++++++++++++++++ runpod/serverless/modules/rp_capture.py | 14 +++++++++++- runpod/serverless/modules/rp_job.py | 3 ++- tests/test_serverless/test_capture.py | 30 +++++++++++++++++++++++++ 4 files changed, 75 insertions(+), 2 deletions(-) diff --git a/docs/serverless/worker.md b/docs/serverless/worker.md index e06967147..32eef376b 100644 --- a/docs/serverless/worker.md +++ b/docs/serverless/worker.md @@ -92,6 +92,36 @@ Support depends on the runtime mode: 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 handler or a prestart hook fails, the SDK attaches the last 16 KB of +whatever the worker wrote to `stdout`/`stderr` to the failure it returns, under +a `logs` key: + +```json +{ + "error_type": "", + "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 anything the worker +prints can reach the caller. On a private endpoint the caller is you, and the +same output is already in your worker logs. On a public or shared endpoint the +caller is a third party. If your worker prints credentials, connection strings, +or other data you do not want returned, either stop printing them or turn +capture off: + +```bash +RUNPOD_SKIP_LOG_CAPTURE=true +``` + +With capture off, failures still report the exception type, message, and +traceback; only the `logs` key is dropped. Output always continues to reach your +worker logs either way. + ## 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: diff --git a/runpod/serverless/modules/rp_capture.py b/runpod/serverless/modules/rp_capture.py index 4cb51d1bc..908b65aa3 100644 --- a/runpod/serverless/modules/rp_capture.py +++ b/runpod/serverless/modules/rp_capture.py @@ -4,10 +4,15 @@ 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. Set `RUNPOD_SKIP_LOG_CAPTURE=true` to keep worker output +inside the worker. """ import contextlib import contextvars +import os import sys from collections.abc import Generator @@ -58,7 +63,14 @@ def __getattr__(self, name): def install() -> None: - """Install the tee proxy on stdout/stderr. Idempotent.""" + """Install the tee proxy on stdout/stderr. Idempotent. + + Skipped when `RUNPOD_SKIP_LOG_CAPTURE=true`. Without the proxy nothing ever + reaches a capture buffer, so no failure payload carries worker output. + """ + if os.environ.get("RUNPOD_SKIP_LOG_CAPTURE", "").lower() == "true": + return + if not isinstance(sys.stdout, _TeeProxy): sys.stdout = _TeeProxy(sys.stdout) if not isinstance(sys.stderr, _TeeProxy): diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index a531f738a..72c0e359d 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/runpod/serverless/modules/rp_job.py @@ -295,8 +295,9 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]: "hostname": os.environ.get("RUNPOD_POD_HOSTNAME", "unknown"), "worker_id": os.environ.get("RUNPOD_POD_ID", "unknown"), "runpod_version": runpod_version, - "logs": captured_logs, } + if captured_logs: + error_info["logs"] = captured_logs log.error("Captured Handler Exception", job["id"]) log.error(json.dumps(error_info, indent=4)) diff --git a/tests/test_serverless/test_capture.py b/tests/test_serverless/test_capture.py index 8264c300a..fd93305ac 100644 --- a/tests/test_serverless/test_capture.py +++ b/tests/test_serverless/test_capture.py @@ -6,6 +6,7 @@ import asyncio import io import json +import os import sys import unittest from unittest.mock import patch @@ -64,6 +65,35 @@ def test_install_is_idempotent(self): assert sys.stderr is installed_stderr +class TestCaptureOptOut(unittest.TestCase): + """`RUNPOD_SKIP_LOG_CAPTURE` keeps worker output out of failure payloads.""" + + def test_install_is_skipped(self): + real_stdout, real_stderr = sys.stdout, sys.stderr + with ( + patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}), + patch.object(sys, "stdout", real_stdout), + patch.object(sys, "stderr", real_stderr), + ): + rp_capture.install() + + assert not isinstance(sys.stdout, rp_capture._TeeProxy) + assert not isinstance(sys.stderr, rp_capture._TeeProxy) + + def test_error_omits_logs_without_the_proxy(self): + def handler(_job): + print("HF_TOKEN=hf_secret") + raise RuntimeError("kernel panic") + + with patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}): + rp_capture.install() + result = _run(run_job(handler, {"id": "j3"})) + + error = json.loads(result["error"]) + assert error["error_message"] == "kernel panic" + assert "logs" not in error + + class TestRunJobCapturesLogs(unittest.TestCase): """A failing handler's stdout/stderr is attached to the job error output.""" From 00b11fd2b3a867c792023e9378042955669efc9b Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 19 Aug 2026 16:37:56 -0700 Subject: [PATCH 6/7] refactor(serverless): centralize the prestart failure contract Every mode repeated the same sequence: capture output, run the hooks, build the prestart_failed payload, log it. Renaming the failure event had to touch all three copies, and the copies had already drifted - local mode ran the phase even with no hooks registered, and the two single-process modes trailed their hard exit with a raise or sys.exit that _terminate_unhealthy makes unreachable. Move the shared sequence into run_prestart_phase, which returns the failure payload or None. Each mode keeps only its own disposition, because those differ: a queue worker has to fail the requests it already holds before exiting, so _terminate_unhealthy deliberately stays with the callers. The two tests covering those hard exits patched _terminate_unhealthy and then asserted the unreachable statement, so they verified a path production never takes. They now assert _terminate_unhealthy(1) itself, with SystemExit standing in for os._exit, which proves the handler is unreachable rather than merely unreached. --- runpod/serverless/modules/rp_fastapi.py | 29 ++++--------------- runpod/serverless/modules/rp_local.py | 22 ++++---------- runpod/serverless/modules/rp_prestart.py | 28 +++++++++++++++++- runpod/serverless/modules/rp_scale.py | 27 +++-------------- .../test_modules/test_local.py | 10 +++++-- tests/test_serverless/test_prestart.py | 12 +++++--- .../test_prestart_lifecycle.py | 4 +-- 7 files changed, 60 insertions(+), 72 deletions(-) diff --git a/runpod/serverless/modules/rp_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 32206803f..918e6467c 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -1,6 +1,5 @@ """Used to launch the FastAPI web server when worker is running in API mode.""" -import json import os import threading import uuid @@ -21,20 +20,11 @@ from .rp_fitness import _terminate_unhealthy from .rp_handler import is_generator from .rp_job import run_job, run_job_generator -from .rp_logger import RunPodLogger from .rp_ping import Heartbeat -from .rp_prestart import ( - PrestartError, - PrestartTimeout, - build_prestart_failed_payload, - get_prestart_hooks, - run_prestart_hooks_async, -) +from .rp_prestart import get_prestart_hooks, run_prestart_phase from .worker_state import JobsProgress, PingJobMirror RUNPOD_ENDPOINT_ID = os.environ.get("RUNPOD_ENDPOINT_ID", None) -log = RunPodLogger() - TITLE = "Runpod | Development Worker API" @@ -195,18 +185,11 @@ class WorkerAPI: @asynccontextmanager async def _lifespan(self, _app: FastAPI) -> AsyncGenerator[None, None]: """Finish prestart before the development API accepts requests.""" - hooks = get_prestart_hooks() - if hooks: - with rp_capture.capture() as captured: - try: - await run_prestart_hooks_async( - hooks, self.config.get("prestart_timeout") - ) - except (PrestartError, PrestartTimeout) as exc: - failure = build_prestart_failed_payload(exc, captured.getvalue()) - log.error(f"prestart_failed | {json.dumps(failure)}") - _terminate_unhealthy(1) - raise + if await run_prestart_phase( + get_prestart_hooks(), self.config.get("prestart_timeout") + ): + # Never serve on a broken startup. _terminate_unhealthy does not return. + _terminate_unhealthy(1) yield def __init__(self, config: dict[str, Any]): diff --git a/runpod/serverless/modules/rp_local.py b/runpod/serverless/modules/rp_local.py index cf682f453..0c8876980 100644 --- a/runpod/serverless/modules/rp_local.py +++ b/runpod/serverless/modules/rp_local.py @@ -13,13 +13,7 @@ from .rp_fitness import _terminate_unhealthy from .rp_job import run_job -from .rp_prestart import ( - PrestartError, - PrestartTimeout, - build_prestart_failed_payload, - get_prestart_hooks, - run_prestart_hooks_async, -) +from .rp_prestart import get_prestart_hooks, run_prestart_phase log = RunPodLogger() @@ -49,16 +43,10 @@ async def run_local(config: dict[str, Any]) -> None: local_job["id"] = local_job.get("id", "local_test") log.debug(f"Retrieved local job: {local_job}") rp_capture.install() - with rp_capture.capture() as captured: - try: - await run_prestart_hooks_async( - get_prestart_hooks(), config.get("prestart_timeout") - ) - except (PrestartError, PrestartTimeout) as exc: - failure = build_prestart_failed_payload(exc, captured.getvalue()) - log.error(f"prestart_failed | {json.dumps(failure)}") - _terminate_unhealthy(1) - sys.exit(1) + if await run_prestart_phase(get_prestart_hooks(), config.get("prestart_timeout")): + # Broken startup: replace the worker rather than run the handler. + # _terminate_unhealthy does not return. + _terminate_unhealthy(1) job_result = await run_job(config["handler"], local_job) diff --git a/runpod/serverless/modules/rp_prestart.py b/runpod/serverless/modules/rp_prestart.py index 12377ebe4..feb5a8a03 100644 --- a/runpod/serverless/modules/rp_prestart.py +++ b/runpod/serverless/modules/rp_prestart.py @@ -12,12 +12,13 @@ import contextlib import contextvars import inspect +import json import threading import traceback from collections.abc import Callable, Sequence from typing import Any -from runpod.serverless.modules.rp_capture import MAX_CAPTURED_CHARS, clip +from runpod.serverless.modules.rp_capture import MAX_CAPTURED_CHARS, capture, 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 @@ -186,3 +187,28 @@ async def run_all() -> None: except asyncio.TimeoutError as exc: raise PrestartTimeout(_hook_name(current_hook), timeout) from exc log.info("Prestart | ready") + + +async def run_prestart_phase( + hooks: Sequence[Callable[[], Any]], timeout: float | None = None +) -> dict[str, Any] | None: + """Run the whole prestart phase and report a failure once. + + Returns the `prestart_failed` payload, or None when every hook succeeded or + no hooks are registered. Every mode shares this so the failure contract + lives in one place; each mode decides on its own what to do with the + payload, because their obligations differ (a queue worker must fail the + requests it already holds before exiting). + """ + if not hooks: + return None + + with capture() as captured: + try: + await run_prestart_hooks_async(hooks, timeout) + return None + except (PrestartError, PrestartTimeout) as exc: + payload = build_prestart_failed_payload(exc, captured.getvalue()) + + log.error(f"prestart_failed | {json.dumps(payload)}") + return payload diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index b98324768..17c70ccdf 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -11,17 +11,10 @@ from typing import Any 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 .rp_prestart import ( - PrestartError, - PrestartTimeout, - build_prestart_failed_payload, - get_prestart_hooks, - run_prestart_hooks_async, -) +from .rp_prestart import get_prestart_hooks, run_prestart_phase from .worker_state import IS_LOCAL_TEST, JobsProgress log = RunPodLogger() @@ -561,22 +554,10 @@ async def _fail_job( async def _run_prestart(self): """Run hooks beside queue intake, then open the handler gate.""" - if not self.prestart_hooks: - self._prestart_ready.set() - return - try: - with capture() as cap: - try: - await run_prestart_hooks_async( - self.prestart_hooks, self.config.get("prestart_timeout") - ) - except (PrestartError, PrestartTimeout) as exc: - self._prestart_error = build_prestart_failed_payload( - exc, cap.getvalue() - ) - if self._prestart_error is not None: - log.error(f"prestart_failed | {json.dumps(self._prestart_error)}") + self._prestart_error = await run_prestart_phase( + self.prestart_hooks, self.config.get("prestart_timeout") + ) finally: # Always release held handlers. self._prestart_ready.set() diff --git a/tests/test_serverless/test_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index 661b550f4..2e14af197 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -59,13 +59,19 @@ def load_model(): patch( "runpod.serverless.modules.rp_local.run_job", new=AsyncMock() ) as run_job, - patch("runpod.serverless.modules.rp_local.log") as logger, - patch("runpod.serverless.modules.rp_local._terminate_unhealthy"), + patch("runpod.serverless.modules.rp_prestart.log") as logger, + # The real helper calls os._exit; SystemExit stands in for that so the + # test proves the handler is unreachable rather than merely unreached. + patch( + "runpod.serverless.modules.rp_local._terminate_unhealthy", + side_effect=SystemExit(1), + ) as terminate, self.assertRaises(SystemExit) as sys_exit, ): await rp_local.run_local(config) self.assertEqual(sys_exit.exception.code, 1) + terminate.assert_called_once_with(1) run_job.assert_not_awaited() failure_log = logger.error.call_args.args[0] self.assertIn("prestart_failed", failure_log) diff --git a/tests/test_serverless/test_prestart.py b/tests/test_serverless/test_prestart.py index af7ac13e3..748fbf164 100644 --- a/tests/test_serverless/test_prestart.py +++ b/tests/test_serverless/test_prestart.py @@ -6,7 +6,7 @@ from unittest.mock import patch import runpod.serverless -from runpod.serverless.modules import rp_fastapi +from runpod.serverless.modules import rp_fastapi, rp_prestart from runpod.serverless.modules.rp_prestart import ( PrestartError, PrestartTimeout, @@ -152,12 +152,16 @@ async def scenario(): with patch.object(rp_fastapi.heartbeat, "start_ping"): api = rp_fastapi.WorkerAPI({"handler": lambda job: job}) with ( - patch.object(rp_fastapi.log, "error") as logger, - patch.object(rp_fastapi, "_terminate_unhealthy"), - self.assertRaisesRegex(PrestartError, "model unavailable"), + patch.object(rp_prestart.log, "error") as logger, + # The real helper calls os._exit; SystemExit stands in for that. + patch.object( + rp_fastapi, "_terminate_unhealthy", side_effect=SystemExit(1) + ) as terminate, + self.assertRaises(SystemExit), ): async with api.rp_app.router.lifespan_context(api.rp_app): self.fail("API served despite prestart failure") + terminate.assert_called_once_with(1) self.assertIn("prestart_failed", logger.call_args.args[0]) self.assertIn("load_model", logger.call_args.args[0]) diff --git a/tests/test_serverless/test_prestart_lifecycle.py b/tests/test_serverless/test_prestart_lifecycle.py index caf088150..d19fcdcc6 100644 --- a/tests/test_serverless/test_prestart_lifecycle.py +++ b/tests/test_serverless/test_prestart_lifecycle.py @@ -14,7 +14,7 @@ import unittest from unittest.mock import AsyncMock, MagicMock, patch -from runpod.serverless.modules import rp_capture, rp_scale +from runpod.serverless.modules import rp_capture, rp_prestart, rp_scale from runpod.serverless.modules.rp_prestart import ( PrestartError, PrestartTimeout, @@ -222,7 +222,7 @@ def failing_hook(): real = io.StringIO() with ( patch.object(sys, "stdout", rp_capture._TeeProxy(real)), - patch.object(rp_scale, "log") as mock_log, + patch.object(rp_prestart, "log") as mock_log, ): _run(scaler._run_prestart()) From 86f06ef59133e61a4d6570c00e778b16dd54bb15 Mon Sep 17 00:00:00 2001 From: Jason Wang Date: Wed, 19 Aug 2026 19:45:48 -0700 Subject: [PATCH 7/7] fix(serverless): keep log capture and shutdown behavior off the existing path An audit of the most popular worker templates showed this feature reaching workers that never asked for it. Two of them install the SDK unpinned, so they take whatever we publish, and both fill the buffer with operational chatter while returning errors rather than raising, so the capture was pure exposure with no benefit. Meanwhile the workers the feature was built for run their engine as a subprocess, whose output arrives on inherited file descriptors and was never capturable in the first place. Make capture opt-in. RUNPOD_LOG_CAPTURE=auto, the default, installs the proxy only when prestart hooks are registered, so a worker that has not adopted prestart is untouched: streams unwrapped, payload unchanged. `all` and `off` force it either way. Stop truncating what this feature did not add. A handler's error_message and error_traceback are reported in full again; 16KB is well under the size of a real framework traceback, and the job-done body limit was not this change's problem to solve. Only the logs field stays bounded. Restore the in-flight job-take on ordinary shutdown. The platform assigns a job before writing the take response and does not release it when the client disconnects, so aborting a take that already carries one strands that request until the heartbeat sweeper reclaims it, spending its single retry. Prestart failure already had a grace period for this; every other shutdown now gets a short one. Also make the proxy indistinguishable from the stream it wraps: register it on io.TextIOBase so isinstance checks still pass (subclassing would shadow encoding with None), and capture writelines, which bypassed write(). Document what capture cannot see: child processes, log handlers constructed before startup, and plain threads. --- docs/serverless/worker.md | 46 ++++++--- runpod/serverless/modules/rp_capture.py | 50 ++++++++-- runpod/serverless/modules/rp_job.py | 8 +- runpod/serverless/modules/rp_scale.py | 27 ++++-- tests/test_serverless/test_capture.py | 97 ++++++++++++++++--- .../test_prestart_lifecycle.py | 52 ++++++++++ 6 files changed, 231 insertions(+), 49 deletions(-) diff --git a/docs/serverless/worker.md b/docs/serverless/worker.md index 32eef376b..d7f11c928 100644 --- a/docs/serverless/worker.md +++ b/docs/serverless/worker.md @@ -94,9 +94,9 @@ when hooks are registered. It never silently skips registered startup work. ## Failure logs -When a handler or a prestart hook fails, the SDK attaches the last 16 KB of -whatever the worker wrote to `stdout`/`stderr` to the failure it returns, under -a `logs` key: +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 { @@ -107,20 +107,36 @@ a `logs` key: } ``` -That payload is returned to whoever called the request, so anything the worker -prints can reach the caller. On a private endpoint the caller is you, and the -same output is already in your worker logs. On a public or shared endpoint the -caller is a third party. If your worker prints credentials, connection strings, -or other data you do not want returned, either stop printing them or turn -capture off: +That payload is returned to whoever called the request, so capture is **not on by +default**. `RUNPOD_LOG_CAPTURE` controls it: -```bash -RUNPOD_SKIP_LOG_CAPTURE=true -``` +| 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. -With capture off, failures still report the exception type, message, and -traceback; only the `logs` key is dropped. Output always continues to reach your -worker logs either way. +In all of these cases output still reaches your worker logs as usual — only the +`logs` field is affected. ## Worker Refresh diff --git a/runpod/serverless/modules/rp_capture.py b/runpod/serverless/modules/rp_capture.py index 908b65aa3..a926eff9b 100644 --- a/runpod/serverless/modules/rp_capture.py +++ b/runpod/serverless/modules/rp_capture.py @@ -6,18 +6,30 @@ real stream and a buffer in a contextvar. Captured output is attached to failure payloads, which are returned to whoever -called the request. Set `RUNPOD_SKIP_LOG_CAPTURE=true` to keep worker output -inside the worker. +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 @@ -54,6 +66,11 @@ def write(self, text) -> int: 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() @@ -62,15 +79,36 @@ def __getattr__(self, name): 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. Idempotent. + """Install the tee proxy on stdout/stderr, if this worker wants capture. - Skipped when `RUNPOD_SKIP_LOG_CAPTURE=true`. Without the proxy nothing ever - reaches a capture buffer, so no failure payload carries worker output. + Idempotent. Without the proxy nothing ever reaches a capture buffer, so no + failure payload carries worker output and the streams are left untouched. """ - if os.environ.get("RUNPOD_SKIP_LOG_CAPTURE", "").lower() == "true": + 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): diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index 72c0e359d..1ca15d96b 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/runpod/serverless/modules/rp_job.py @@ -15,7 +15,7 @@ from ...version import __version__ as runpod_version from ..utils import rp_debugger -from .rp_capture import capture, clip +from .rp_capture import capture from .rp_handler import is_generator from .rp_http import send_result, stream_result from .rp_tips import check_return_size @@ -290,8 +290,8 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]: captured_logs = cap.getvalue() error_info = { "error_type": str(type(err)), - "error_message": clip(str(err)), - "error_traceback": clip(traceback.format_exc()), + "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, @@ -339,6 +339,6 @@ async def run_job_generator( error = f"handler: {str(err)} \ntraceback: {traceback.format_exc()}" if captured_logs: error += f"\nlogs:\n{captured_logs}" - yield {"error": clip(error)} + yield {"error": error} finally: log.info("Finished running generator.", job["id"]) diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 17c70ccdf..0872dff8d 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -70,6 +70,9 @@ def __init__(self, config: dict[str, Any]): self.jobs_fetcher = get_job self.jobs_fetcher_timeout = 90 self.prestart_claim_timeout = 10 + # Only needs to outlast a response already in transit, so it is far + # shorter than prestart_claim_timeout, which waits for a job to appear. + self.shutdown_take_grace = 2 self.jobs_handler = handle_job if concurrency_modifier := config.get("concurrency_modifier"): @@ -218,8 +221,15 @@ def current_occupancy(self) -> int: return current_progress_count + current_queue_count async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: int): - """After prestart failure, briefly keep an in-flight job-take alive so - its request can receive the error.""" + """Run one blocking job-take without letting it delay shutdown. + + Shutdown cannot simply abandon the take. The platform assigns a job + before it writes the response and does not release it when the client + disconnects, so aborting a take that already carries a job strands that + request until the worker-heartbeat sweeper reclaims it, which also + spends the job's one retry. Give the response a bounded window to land: + long enough to report a prestart failure against, short otherwise. + """ fetch_task = asyncio.create_task(self.jobs_fetcher(session, jobs_needed)) shutdown_task = asyncio.create_task(self._shutdown_event.wait()) try: @@ -232,14 +242,13 @@ async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: i return await fetch_task if shutdown_task in done: - if self._prestart_error is None: - return None - - self._failure_take_done = True + if self._prestart_error is not None: + self._failure_take_done = True + grace = self.prestart_claim_timeout + else: + grace = self.shutdown_take_grace try: - return await asyncio.wait_for( - fetch_task, timeout=self.prestart_claim_timeout - ) + return await asyncio.wait_for(fetch_task, timeout=grace) except asyncio.TimeoutError: return None diff --git a/tests/test_serverless/test_capture.py b/tests/test_serverless/test_capture.py index fd93305ac..07759178a 100644 --- a/tests/test_serverless/test_capture.py +++ b/tests/test_serverless/test_capture.py @@ -13,6 +13,10 @@ from runpod.serverless.modules import rp_capture from runpod.serverless.modules.rp_job import run_job, run_job_generator +from runpod.serverless.modules.rp_prestart import ( + clear_prestart_hooks, + register_prestart_hook, +) def _run(coro): @@ -65,27 +69,57 @@ def test_install_is_idempotent(self): assert sys.stderr is installed_stderr -class TestCaptureOptOut(unittest.TestCase): - """`RUNPOD_SKIP_LOG_CAPTURE` keeps worker output out of failure payloads.""" +class TestCaptureIsOptIn(unittest.TestCase): + """A worker that has not adopted prestart hooks must see no change at all: + streams untouched and no logs field in the failure it reports.""" - def test_install_is_skipped(self): - real_stdout, real_stderr = sys.stdout, sys.stderr + def setUp(self): + clear_prestart_hooks() + self.addCleanup(clear_prestart_hooks) + + def _install_with(self, env) -> bool: + """Install under `env` and report whether the proxy was applied.""" + real_stdout, real_stderr = sys.__stdout__, sys.__stderr__ with ( - patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}), + patch.dict(os.environ, env, clear=False), patch.object(sys, "stdout", real_stdout), patch.object(sys, "stderr", real_stderr), ): rp_capture.install() + return isinstance(sys.stdout, rp_capture._TeeProxy) and isinstance( + sys.stderr, rp_capture._TeeProxy + ) + + def test_auto_skips_install_without_hooks(self): + assert self._install_with({"RUNPOD_LOG_CAPTURE": "auto"}) is False + + def test_auto_is_the_default(self): + env = {k: v for k, v in os.environ.items() if k != "RUNPOD_LOG_CAPTURE"} + with patch.dict(os.environ, env, clear=True): + assert rp_capture.capture_mode() == rp_capture.CAPTURE_AUTO + assert self._install_with({}) is False + + def test_auto_installs_once_a_hook_is_registered(self): + register_prestart_hook(lambda: None) + assert self._install_with({"RUNPOD_LOG_CAPTURE": "auto"}) is True + + def test_all_installs_without_hooks(self): + assert self._install_with({"RUNPOD_LOG_CAPTURE": "all"}) is True + + def test_off_skips_install_even_with_hooks(self): + register_prestart_hook(lambda: None) + assert self._install_with({"RUNPOD_LOG_CAPTURE": "off"}) is False - assert not isinstance(sys.stdout, rp_capture._TeeProxy) - assert not isinstance(sys.stderr, rp_capture._TeeProxy) + def test_unknown_mode_falls_back_to_auto(self): + with patch.dict(os.environ, {"RUNPOD_LOG_CAPTURE": "yes-please"}): + assert rp_capture.capture_mode() == rp_capture.CAPTURE_AUTO - def test_error_omits_logs_without_the_proxy(self): + def test_error_payload_is_unchanged_without_capture(self): def handler(_job): print("HF_TOKEN=hf_secret") raise RuntimeError("kernel panic") - with patch.dict(os.environ, {"RUNPOD_SKIP_LOG_CAPTURE": "true"}): + with patch.dict(os.environ, {"RUNPOD_LOG_CAPTURE": "off"}): rp_capture.install() result = _run(run_job(handler, {"id": "j3"})) @@ -141,8 +175,8 @@ def test_handler_success_has_no_error(self): 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.""" + """Only the log tail this SDK adds is bounded. A handler's own error message and + traceback are reported exactly as before, so existing workers see no truncation.""" def test_clip_keeps_head_and_tail(self): text = "A" * 100 + "B" * 100 @@ -155,7 +189,7 @@ def test_clip_keeps_head_and_tail(self): def test_clip_passthrough_when_small(self): assert rp_capture.clip("short", limit=100) == "short" - def test_run_job_bounds_error_message(self): + def test_run_job_reports_a_huge_handler_error_untruncated(self): huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) def handler(_job): @@ -163,9 +197,26 @@ def handler(_job): result = _run(run_job(handler, {"id": "big"})) error = json.loads(result["error"]) - assert len(error["error_message"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + assert error["error_message"] == huge + assert "truncated" not in error["error_traceback"] - def test_run_job_generator_bounds_combined_error(self): + def test_run_job_bounds_only_the_captured_logs(self): + def handler(_job): + print("y" * (rp_capture.MAX_CAPTURED_CHARS * 3)) + raise ValueError("boom") + + 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": "big-logs"})) + + error = json.loads(result["error"]) + assert error["error_message"] == "boom" + assert len(error["logs"]) <= rp_capture.MAX_CAPTURED_CHARS + + def test_run_job_generator_bounds_only_the_captured_logs(self): huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) def handler(_job): @@ -180,7 +231,23 @@ def handler(_job): ): result = _run_gen(run_job_generator(handler, {"id": "big-generator"})) - assert len(result[0]["error"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + # The handler's own error survives whole; only the appended log tail is capped. + assert huge in result[0]["error"] + logs = result[0]["error"].split("\nlogs:\n", 1)[1] + assert len(logs) <= rp_capture.MAX_CAPTURED_CHARS + + def test_proxy_is_indistinguishable_from_a_text_stream(self): + proxy = rp_capture._TeeProxy(io.StringIO()) + assert isinstance(proxy, io.TextIOBase) + assert isinstance(proxy, io.IOBase) + + def test_writelines_is_captured(self): + real = io.StringIO() + proxy = rp_capture._TeeProxy(real) + with patch.object(sys, "stdout", proxy), rp_capture.capture() as buf: + sys.stdout.writelines(["a\n", "b\n"]) + assert buf.getvalue() == "a\nb\n" + assert real.getvalue() == "a\nb\n" if __name__ == "__main__": diff --git a/tests/test_serverless/test_prestart_lifecycle.py b/tests/test_serverless/test_prestart_lifecycle.py index d19fcdcc6..45a788822 100644 --- a/tests/test_serverless/test_prestart_lifecycle.py +++ b/tests/test_serverless/test_prestart_lifecycle.py @@ -472,6 +472,58 @@ async def go(): _run(go()) scaler._fail_job.assert_not_awaited() + def test_ordinary_shutdown_still_collects_an_in_transit_job(self): + """A plain shutdown must not abandon a take whose response already + carries a job: the platform will not release it, so dropping it strands + the request and spends its one retry.""" + scaler = _scaler(hook=lambda: None) + scaler.shutdown_take_grace = 0.3 + assigned = [] + started = asyncio.Event() + + async def take(_session, _needed): + started.set() + assigned.append("job-1") # the platform has committed the job + await asyncio.sleep(0.05) # response in transit + return [{"id": "job-1", "input": {}}] + + scaler.jobs_fetcher = take + + async def go(): + result = asyncio.create_task( + scaler._fetch_jobs_until_stopped(AsyncMock(), 1) + ) + await started.wait() + scaler.kill_worker() # no prestart error: an ordinary shutdown + return await asyncio.wait_for(result, timeout=1) + + acquired = _run(go()) + assert assigned == ["job-1"] + assert acquired == [{"id": "job-1", "input": {}}] + + def test_ordinary_shutdown_does_not_wait_out_a_silent_take(self): + """The grace period is a bound, not a delay: nothing in flight means + shutdown proceeds immediately.""" + scaler = _scaler(hook=lambda: None) + scaler.shutdown_take_grace = 0.2 + started = asyncio.Event() + + async def never_returns(_session, _needed): + started.set() + await asyncio.Event().wait() + + scaler.jobs_fetcher = never_returns + + async def go(): + result = asyncio.create_task( + scaler._fetch_jobs_until_stopped(AsyncMock(), 1) + ) + await started.wait() + scaler.kill_worker() + return await asyncio.wait_for(result, timeout=1) + + assert _run(go()) is None + def test_jobs_acquired_after_prestart_failure_are_failed_not_queued(self): """If prestart fails during a long-poll, fail returned jobs and stop job-take without relying on another task to end the loop."""