diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 05245207..6f268b5f 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_capture.py b/runpod/serverless/modules/rp_capture.py new file mode 100644 index 00000000..a9d6ac46 --- /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_initializer.py b/runpod/serverless/modules/rp_initializer.py new file mode 100644 index 00000000..c8f4e2af --- /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_job.py b/runpod/serverless/modules/rp_job.py index a45cebc6..a531f738 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/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 4cbf94ff..4cb9ea4a 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -4,12 +4,15 @@ """ import asyncio +import json import signal import sys import traceback from typing import Any, Dict, Set from ...http_client import AsyncClientSession, ClientSession, TooManyRequests +from .rp_capture import capture +from .rp_http import send_result from .rp_job import _job_stop_url, get_job, get_stop_signals, handle_job from .rp_logger import RunPodLogger, _reset_batch_id, _set_batch_id from .worker_state import JobsProgress, IS_LOCAL_TEST @@ -44,6 +47,9 @@ class JobScaler: def __init__(self, config: Dict[str, Any]): self._shutdown_event = asyncio.Event() + self._init_ready = asyncio.Event() + self._init_error: dict[str, Any] | None = None + self._claimed_request = False self.current_concurrency = 1 self.config = config self.job_progress = JobsProgress() # Cache the singleton instance @@ -60,6 +66,10 @@ def __init__(self, config: Dict[str, Any]): self.concurrency_modifier = _default_concurrency_modifier self.jobs_fetcher = get_job self.jobs_fetcher_timeout = 90 + # Bound on the single claim made after a failed init. Short on purpose: a + # queued request comes back at once, and an empty queue must not hold a dying + # worker open for the full long-poll. + self.init_claim_timeout = 10 self.jobs_handler = handle_job if concurrency_modifier := config.get("concurrency_modifier"): @@ -139,14 +149,26 @@ async def run(self): # Create an async session that will be closed when the worker is killed. async with AsyncClientSession() as session: # Create the worker's concurrent loops. + init_task = asyncio.create_task(self._run_init()) jobtake_task = asyncio.create_task(self.get_jobs(session)) jobrun_task = asyncio.create_task(self.run_jobs(session)) jobstop_task = asyncio.create_task(self.monitor_stop_signals(session)) - tasks = [jobtake_task, jobrun_task, jobstop_task] + # The initializer is not a loop: an initializer with no init_timeout can + # block forever, so only the request loops decide when the worker stops. + await asyncio.gather(jobtake_task, jobrun_task, jobstop_task) - # Run the worker's concurrent loops until shutdown. - await asyncio.gather(*tasks) + # Shutting down: abandon a still-running initializer. Its work sits on a + # daemon thread, so dropping it here lets the process exit. + init_task.cancel() + try: + await init_task + except asyncio.CancelledError: + pass + + # If the initializer fails, let the platform respawn the worker. + if self._init_error is not None: + sys.exit(1) def is_alive(self): """ @@ -170,6 +192,28 @@ def current_occupancy(self) -> int: ) return current_progress_count + current_queue_count + async def _claim_one_job_to_fail(self, session: ClientSession) -> None: + """Initialization failed before this worker ever held a request. Claim one + queued request and fail it with the reason. + + Without this, an instant failure - bad config, a missing file, an async + initializer that raises before its first await - never reaches a caller: the + worker exits, the platform respawns it into the same failure, and the request + that triggered the scale-up waits out its queue TTL with no explanation. + """ + try: + jobs = await asyncio.wait_for( + self.jobs_fetcher(session, 1), timeout=self.init_claim_timeout + ) + except asyncio.CancelledError: + raise + except Exception as error: # noqa: BLE001 - reporting must not mask the failure + log.debug(f"JobScaler.get_jobs | No request claimed to fail: {error}") + return + + for job in jobs or []: + await self._fail_job(session, job, self._init_error) + async def get_jobs(self, session: ClientSession): """ Retrieve multiple jobs from the server in batches using blocking requests. @@ -179,6 +223,10 @@ async def get_jobs(self, session: ClientSession): Adds jobs to the JobsQueue """ while self.is_alive(): + if self._init_error is not None: + # Initialization is terminal for this worker. Draining what we already + # hold is owned by run_jobs. + break await self.set_scale() jobs_needed = self.current_concurrency - self.current_occupancy() @@ -200,6 +248,14 @@ async def get_jobs(self, session: ClientSession): log.debug("JobScaler.get_jobs | No jobs acquired.") continue + self._claimed_request = True + + if self._init_error is not None: + # If initialization fails, fail all in-flight requests. + for job in acquired_jobs: + await self._fail_job(session, job, self._init_error) + return + for job in acquired_jobs: await self.jobs_queue.put(job) self.job_progress.add(job) @@ -227,6 +283,12 @@ async def get_jobs(self, session: ClientSession): # Yield control back to the event loop await asyncio.sleep(0) + if self._init_error is not None and not self._claimed_request: + # An init failure sets shutdown, so this loop can end before it ever + # claimed a request. Claim one on the way out, or the failure dies with + # the worker and the request that spawned it waits out its queue TTL. + await self._claim_one_job_to_fail(session) + async def run_jobs(self, session: ClientSession): """ Retrieve jobs from the jobs queue and process them concurrently. @@ -346,6 +408,21 @@ async def handle_job(self, session: ClientSession, job: dict): try: log.debug("Handling Job", job["id"]) + # Hold the handler until initialization finishes. + if self.config.get("initializer") is not None: + if not await self._wait_for_init(): + log.warn( + "Shutting down before initialization finished; leaving this " + "request for another worker.", + job["id"], + ) + return + + if self._init_error is not None: + # If initialization fails, fail the current request and don't run the handler. + await self._fail_job(session, job, self._init_error) + return + await self.jobs_handler(session, self.config, job) if self.config.get("refresh_worker", False): @@ -369,3 +446,57 @@ async def handle_job(self, session: ClientSession, job: dict): log.debug("Finished Job", job["id"]) _reset_batch_id(batch_id_token) + + async def _wait_for_init(self) -> bool: + """Wait for initialization to finish, or for the worker to start shutting down. + Returns whether initialization actually finished.""" + ready = asyncio.create_task(self._init_ready.wait()) + stopping = asyncio.create_task(self._shutdown_event.wait()) + try: + await asyncio.wait({ready, stopping}, return_when=asyncio.FIRST_COMPLETED) + finally: + ready.cancel() + stopping.cancel() + return self._init_ready.is_set() + + async def _fail_job( + self, session: ClientSession, job: dict, payload: Dict[str, Any] + ): + """Fail a request with a structured error (reason + logs).""" + log.error(f"Failing job due to init failure. | {job['id']}") + await send_result(session, {"error": json.dumps(payload)}, job, is_stream=False) + + async def _run_init(self): + """Run initializer concurrently with the loop. Upon completion, opens the gate to run + job handlers. On failure, records the reason, then drains and shuts the worker.""" + initializer = self.config.get("initializer") + if initializer is None: + self._init_ready.set() + return + + from .rp_initializer import ( + InitializerError, + InitializerTimeout, + build_init_failed_payload, + run_initializer_async, + ) + + try: + with capture() as cap: + try: + await run_initializer_async( + initializer, self.config.get("init_timeout") + ) + except (InitializerError, InitializerTimeout) as exc: + self._init_error = build_init_failed_payload(exc, cap.getvalue()) + if self._init_error is not None: + log.error(f"init_failed | {json.dumps(self._init_error)}") + finally: + # Always release held handlers. + self._init_ready.set() + + if self._init_error is not None: + # Stop long-running loops before waiting for acquired requests to drain. + self.kill_worker() + while self.current_occupancy() > 0: + await asyncio.sleep(0.1) diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 90053ec7..8c1bdcbe 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 00000000..8264c300 --- /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() diff --git a/tests/test_serverless/test_initializer.py b/tests/test_serverless/test_initializer.py new file mode 100644 index 00000000..f90e0b51 --- /dev/null +++ b/tests/test_serverless/test_initializer.py @@ -0,0 +1,460 @@ +"""Tests for the concurrent initializer: runs alongside the job loop, holds the handler +until ready, and fails the in-hand request (with captured stdout/stderr) on failure.""" + +# pylint: disable=protected-access + +import asyncio +import functools +import io +import json +import sys +import threading +import unittest +from unittest.mock import AsyncMock, MagicMock, patch + +from runpod.serverless.modules import rp_capture, rp_scale +from runpod.serverless.modules.rp_initializer import ( + InitializerError, + InitializerTimeout, + build_init_failed_payload, + run_initializer_async, +) +from runpod.serverless.modules.rp_scale import JobScaler + + +def _run(coro): + return asyncio.run(coro) + + +class TestRunInitializerAsync(unittest.TestCase): + """Runs the initializer to completion, or raises a clear error when it fails or times out.""" + + def test_sync_success_offloaded(self): + calls = [] + _run(run_initializer_async(lambda: calls.append("ran"))) + assert calls == ["ran"] + + def test_sync_callable_returning_awaitable_is_awaited(self): + state = {} + + async def load(): + await asyncio.sleep(0) + state["ready"] = True + + def initialize(): + return load() + + _run(run_initializer_async(initialize)) + assert state == {"ready": True} + + + def test_sync_failure_wraps_in_initializer_error(self): + def failing_initializer(): + raise ValueError("max_model_len must be positive, got 0") + + with self.assertRaises(InitializerError) as ctx: + _run(run_initializer_async(failing_initializer)) + assert isinstance(ctx.exception.original, ValueError) + assert "max_model_len" in str(ctx.exception) + + def test_async_success(self): + state = {} + + async def load(): + await asyncio.sleep(0) + state["ready"] = True + + _run(run_initializer_async(load)) + assert state == {"ready": True} + + def test_async_failure_wraps_in_initializer_error(self): + async def failing_initializer(): + raise RuntimeError("CUDA OOM") + + with self.assertRaises(InitializerError) as ctx: + _run(run_initializer_async(failing_initializer)) + assert isinstance(ctx.exception.original, RuntimeError) + + def test_async_timeout_raises_initializer_timeout(self): + async def slow(): + await asyncio.sleep(3) + + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(slow, timeout=1)) + + def test_zero_timeout_raises_initializer_timeout(self): + async def load(): + await asyncio.sleep(0) + + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(load, timeout=0)) + + + def test_sync_hang_times_out(self): + """A blocking sync load that never returns is cut off by init_timeout, not left stuck.""" + release = threading.Event() + + def hang(): + release.wait(10) + + try: + with self.assertRaises(InitializerTimeout): + _run(run_initializer_async(hang, timeout=1)) + finally: + release.set() # let the offloaded thread exit promptly + + def test_async_partial_is_awaited(self): + """functools.partial wrapping an async initializer is detected and awaited.""" + state = {} + + async def load(key): + await asyncio.sleep(0) + state[key] = True + + _run(run_initializer_async(functools.partial(load, "ready"))) + assert state == {"ready": True} + + def test_async_callable_object_is_awaited(self): + """An object whose __call__ is async is detected and awaited.""" + state = {} + + class Loader: + async def __call__(self): + await asyncio.sleep(0) + state["ready"] = True + + _run(run_initializer_async(Loader())) + assert state == {"ready": True} + + def test_base_exception_propagates_unwrapped(self): + """KeyboardInterrupt is process control, not an init failure: propagate.""" + + def interrupted(): + raise KeyboardInterrupt + + with self.assertRaises(KeyboardInterrupt): + _run(run_initializer_async(interrupted)) + + def test_system_exit_wraps_as_initializer_error(self): + """A load script calling sys.exit() is an init failure to surface with a reason, + not a clean exit - otherwise held in-hand jobs die without one.""" + + def bails(): + raise SystemExit(1) + + with self.assertRaises(InitializerError): + _run(run_initializer_async(bails)) + + +class TestInitFailedSignal(unittest.TestCase): + """Builds the structured init_failed payload, including captured logs.""" + + def test_payload_shape_from_sync_error(self): + try: + raise ValueError("bad config") + except ValueError as exc: + payload = build_init_failed_payload( + InitializerError(exc), logs="stderr tail" + ) + assert payload["event"] == "init_failed" + assert payload["error_type"] == "ValueError" + assert payload["error_message"] == "bad config" + assert "ValueError" in payload["error_traceback"] + assert payload["logs"] == "stderr tail" + assert "worker_id" in payload and "runpod_version" in payload + + def test_payload_omits_empty_logs(self): + payload = build_init_failed_payload(InitializerError(ValueError("x"))) + assert "logs" not in payload + + def test_payload_bounds_message_traceback_and_logs(self): + huge = "x" * (rp_capture.MAX_CAPTURED_CHARS * 3) + payload = build_init_failed_payload( + InitializerError(ValueError(huge)), + logs="y" * (rp_capture.MAX_CAPTURED_CHARS * 3), + ) + assert len(payload["error_message"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + assert len(payload["error_traceback"]) <= rp_capture.MAX_CAPTURED_CHARS + 100 + assert len(payload["logs"]) <= rp_capture.MAX_CAPTURED_CHARS + + +def _scaler(initializer=None): + config = {"handler": lambda j: j, "rp_args": {}} + if initializer is not None: + config["initializer"] = initializer + scaler = JobScaler(config) + scaler.job_progress = MagicMock() # avoid the process-wide singleton in unit tests + scaler.job_progress.get_job_count.return_value = 0 + return scaler + + +class TestRunInit(unittest.TestCase): + """The concurrent init task: opens the gate, and on failure records the reason and shuts down.""" + + def test_no_initializer_opens_gate_immediately(self): + scaler = _scaler(initializer=None) + _run(scaler._run_init()) + assert scaler._init_ready.is_set() + assert scaler._init_error is None + + def test_success_opens_gate_no_error(self): + ran = [] + scaler = _scaler(initializer=lambda: ran.append(1)) + _run(scaler._run_init()) + assert ran == [1] + assert scaler._init_ready.is_set() + assert scaler._init_error is None + assert not scaler._shutdown_event.is_set() + + def test_failure_records_reason_with_logs_and_shuts_down(self): + def failing_initializer(): + print("downloading model") + raise RuntimeError("CUDA OOM: model too big") + + scaler = _scaler(initializer=failing_initializer) + real = io.StringIO() + with ( + patch.object(sys, "stdout", rp_capture._TeeProxy(real)), + patch.object(rp_scale, "log") as mock_log, + ): + _run(scaler._run_init()) + + assert scaler._init_ready.is_set() # held handlers are released to fail fast + assert scaler._init_error is not None + assert scaler._init_error["error_message"] == "CUDA OOM: model too big" + assert "downloading model" in scaler._init_error["logs"] + assert any( + call.args[0].startswith("init_failed | ") + for call in mock_log.error.call_args_list + ) + assert ( + scaler._shutdown_event.is_set() + ) # broken worker shuts down (occupancy was 0) + + def test_failure_starts_shutdown_before_drain(self): + scaler = _scaler( + initializer=lambda: (_ for _ in ()).throw(RuntimeError("boom")) + ) + shutdown_states = [] + + def occupancy(): + shutdown_states.append(scaler._shutdown_event.is_set()) + return 0 + + scaler.current_occupancy = occupancy + _run(scaler._run_init()) + + assert shutdown_states == [True] + + +class TestHandleJobGate(unittest.TestCase): + """The handler is held until init is ready; init failure fails the in-hand request.""" + + def _prime(self, scaler, job): + # Balance the queue/progress bookkeeping handle_job's finally expects. + scaler.jobs_queue = asyncio.Queue(maxsize=4) + scaler.jobs_queue.put_nowait(job) + + def test_runs_handler_once_init_is_ready(self): + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._init_ready.set() # init already succeeded + job = {"id": "j1"} + + async def go(): + self._prime(scaler, job) + await scaler.handle_job(None, job) + + _run(go()) + scaler.jobs_handler.assert_awaited_once() + + def test_runs_handler_without_an_initializer(self): + scaler = _scaler(initializer=None) + scaler.jobs_handler = AsyncMock() + job = {"id": "j2"} + + async def go(): + self._prime(scaler, job) + await scaler.handle_job(None, job) + + _run(go()) + scaler.jobs_handler.assert_awaited_once() # no gate to wait on + + def test_init_failure_fails_request_without_running_handler(self): + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler._init_ready.set() + job = {"id": "j3"} + + async def go(): + self._prime(scaler, job) + with patch.object(rp_scale, "send_result", new=AsyncMock()) as mock_sr: + await scaler.handle_job(None, job) + mock_sr.assert_awaited_once() + sent = mock_sr.await_args[0][1] + assert json.loads(sent["error"])["error_message"] == "CUDA OOM" + + _run(go()) + scaler.jobs_handler.assert_not_awaited() # broken worker never runs the handler + + def test_init_failure_before_any_take_claims_a_request_to_fail(self): + """An instant init failure leaves no request in hand, and it sets shutdown, so + job-take ends immediately. It must still claim one request and fail it, or the + caller waits out the queue TTL for nothing.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() # _run_init sets shutdown on failure + scaler._fail_job = AsyncMock() + job = {"id": "orphan-1"} + scaler.jobs_fetcher = AsyncMock(return_value=[job]) + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + + scaler.jobs_fetcher.assert_awaited_once() + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + assert scaler.jobs_queue.qsize() == 0 # never queued into a broken worker + + def test_init_failure_with_empty_queue_exits_without_hanging(self): + """Nothing left to fail: the claim attempt returns empty and job-take stops.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() + scaler._fail_job = AsyncMock() + scaler.jobs_fetcher = AsyncMock(return_value=[]) + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + + scaler.jobs_fetcher.assert_awaited_once() + scaler._fail_job.assert_not_awaited() + + def test_init_failure_after_a_take_does_not_claim_another_request(self): + """Once this worker has claimed a request, the failure is reported against it. + Claiming a second request would fail work a healthy worker could serve.""" + scaler = _scaler(initializer=lambda: None) + scaler._fail_job = AsyncMock() + failure = {"error_message": "CUDA OOM", "event": "init_failed"} + calls = [] + + async def fetcher(_session, _needed): + calls.append(1) + return [{"id": f"job-{len(calls)}"}] + + scaler.jobs_fetcher = fetcher + + async def go(): + task = asyncio.create_task(scaler.get_jobs(AsyncMock())) + await asyncio.sleep(0) + scaler._init_error = failure # lands after the first take succeeded + scaler.kill_worker() + await asyncio.wait_for(task, timeout=0.5) + + _run(go()) + + assert len(calls) == 1 # no extra claim after the queued request + + def test_instant_async_init_failure_still_fails_the_queued_request(self): + """An async initializer that raises before its first await finishes before + job-take ever fetches, so nothing is in hand and shutdown is already set. + Driving the real `_run_init` failure path, the worker must still claim the + queued request and fail it rather than exit silently.""" + + async def bad_init(): + raise RuntimeError("bad config") + + scaler = _scaler(initializer=bad_init) + scaler._fail_job = AsyncMock() + job = {"id": "queued-1"} + scaler.jobs_fetcher = AsyncMock(return_value=[job]) + scaler.jobs_handler = AsyncMock() + + async def go(): + await scaler._run_init() # records the reason and sets shutdown + assert not scaler.is_alive() + await asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=1) + + _run(go()) + + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + reason = scaler._fail_job.await_args[0][2] + assert reason["event"] == "init_failed" + assert reason["error_message"] == "bad config" + scaler.jobs_handler.assert_not_awaited() # handler never runs on a broken worker + + def test_claim_attempt_is_bounded(self): + """A silent job-take must not hold a dying worker open. The claim gives up on + its own bound and the worker exits to be respawned.""" + scaler = _scaler(initializer=lambda: None) + scaler._init_error = {"error_message": "CUDA OOM", "event": "init_failed"} + scaler.kill_worker() + scaler._fail_job = AsyncMock() + scaler.init_claim_timeout = 0.05 + + async def never_returns(_session, _needed): + await asyncio.Event().wait() + + scaler.jobs_fetcher = never_returns + + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=1)) + + scaler._fail_job.assert_not_awaited() + + def test_jobs_acquired_after_init_failure_are_failed_not_queued(self): + """If init fails while a long-poll is in flight, fail returned jobs and stop + job-take without relying on another task to end the loop.""" + scaler = _scaler(initializer=lambda: None) + scaler._fail_job = AsyncMock() + job = {"id": "late-1"} + failure = {"error_message": "CUDA OOM", "event": "init_failed"} + + async def fetcher(_session, _needed): + # Init failure lands while this long-poll is in flight. + scaler._init_error = failure + return [job] + + scaler.jobs_fetcher = fetcher + _run(asyncio.wait_for(scaler.get_jobs(AsyncMock()), timeout=0.5)) + scaler._fail_job.assert_awaited_once() + assert scaler._fail_job.await_args[0][1] is job + assert scaler.jobs_queue.qsize() == 0 + scaler.job_progress.add.assert_not_called() + + def test_shutdown_before_init_ready_leaves_request_alone(self): + """SIGTERM while the initializer is still running must release a held request + without running the handler against an uninitialized worker.""" + scaler = _scaler(initializer=lambda: None) + scaler.jobs_handler = AsyncMock() + scaler.kill_worker() # shutdown lands while init is still in flight + job = {"id": "j4"} + + async def go(): + self._prime(scaler, job) + with patch.object(rp_scale, "send_result", new=AsyncMock()) as mock_sr: + await asyncio.wait_for(scaler.handle_job(None, job), timeout=0.5) + mock_sr.assert_not_awaited() # the platform retries it elsewhere + + _run(go()) + scaler.jobs_handler.assert_not_awaited() + + +class TestShutdownWithHangingInit(unittest.TestCase): + """A hung initializer with no init_timeout must not outlive shutdown: the daemon + thread is abandoned so the worker can exit.""" + + def test_run_returns_while_initializer_still_blocked(self): + release = threading.Event() + self.addCleanup(release.set) # let the daemon thread finish after the test + + scaler = _scaler(initializer=release.wait) # blocks with no init_timeout + scaler.kill_worker() # every loop exits immediately; only init is left + scaler.stop_signals_fetcher = AsyncMock(return_value=[]) + + _run(asyncio.wait_for(scaler.run(), timeout=2)) + + assert not release.is_set() # still blocked, and the worker left anyway + + +if __name__ == "__main__": + unittest.main()