diff --git a/docs/serverless/worker.md b/docs/serverless/worker.md index e1355260d..d7f11c928 100644 --- a/docs/serverless/worker.md +++ b/docs/serverless/worker.md @@ -18,14 +18,126 @@ 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 `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. | + +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 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 +{ + "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 capture is **not on by +default**. `RUNPOD_LOG_CAPTURE` controls it: + +| 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. + +In all of these cases output still reaches your worker logs as usual — only the +`logs` field is affected. + ## 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 052452073..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,18 +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. + + 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["prestart_timeout"] (int, optional): Seconds allowed for the complete + prestart phase. Omit for no timeout. """ print(f"--- Starting Serverless Worker | Version {runpod_version} ---") @@ -156,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( @@ -171,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 new file mode 100644 index 000000000..a926eff9b --- /dev/null +++ b/runpod/serverless/modules/rp_capture.py @@ -0,0 +1,140 @@ +""" +runpod | serverless | rp_capture.py + +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, 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 +) + + +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 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() + + def __getattr__(self, name): + # Delegate everything else to the real stream + 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, if this worker wants capture. + + Idempotent. Without the proxy nothing ever reaches a capture buffer, so no + failure payload carries worker output and the streams are left untouched. + """ + 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): + 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_fastapi.py b/runpod/serverless/modules/rp_fastapi.py index 5451ae40e..918e6467c 100644 --- a/runpod/serverless/modules/rp_fastapi.py +++ b/runpod/serverless/modules/rp_fastapi.py @@ -1,10 +1,12 @@ -""" 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 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,9 +16,12 @@ 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_ping import Heartbeat +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) @@ -29,17 +34,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 +111,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 +120,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 +140,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 +150,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 +182,17 @@ 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.""" + 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]): """ Initializes the WorkerAPI class. 1. Starts the heartbeat thread. @@ -194,6 +209,7 @@ def __init__(self, config: Dict[str, Any]): heartbeat.start_ping(mirror) self.config = config + rp_capture.install() tags_metadata = [ { @@ -217,6 +233,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 +327,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_job.py b/runpod/serverless/modules/rp_job.py index a45cebc68..1ca15d96b 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 from .rp_handler import is_generator from .rp_http import send_result, stream_result from .rp_tips import check_return_size @@ -253,54 +254,56 @@ 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": 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, + } + if captured_logs: + error_info["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 +320,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": error} + finally: + log.info("Finished running generator.", job["id"]) diff --git a/runpod/serverless/modules/rp_local.py b/runpod/serverless/modules/rp_local.py index 971e696e1..0c8876980 100644 --- a/runpod/serverless/modules/rp_local.py +++ b/runpod/serverless/modules/rp_local.py @@ -6,16 +6,19 @@ 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 get_prestart_hooks, run_prestart_phase 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 +32,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 +42,11 @@ 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() + 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 new file mode 100644 index 000000000..feb5a8a03 --- /dev/null +++ b/runpod/serverless/modules/rp_prestart.py @@ -0,0 +1,214 @@ +""" +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 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, 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 + +log = RunPodLogger() + +PRESTART_FAILED_EVENT = "prestart_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_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": PRESTART_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") + + +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 4cbf94ffb..0872dff8d 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -4,15 +4,18 @@ """ import asyncio +import json 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_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 get_prestart_hooks, run_prestart_phase +from .worker_state import IS_LOCAL_TEST, JobsProgress log = RunPodLogger() @@ -42,15 +45,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._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 @@ -60,6 +69,10 @@ def __init__(self, config: Dict[str, Any]): self.concurrency_modifier = _default_concurrency_modifier 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"): @@ -81,7 +94,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): @@ -136,17 +151,52 @@ 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. - 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)) + # 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: + # Wait for either lifecycle to end without cancelling the other. + done, _ = await asyncio.wait( + {prestart_task, request_loops_future}, + return_when=asyncio.FIRST_COMPLETED, + ) - tasks = [jobtake_task, jobrun_task, jobstop_task] + if prestart_task in done: + try: + await prestart_task + except BaseException: + # Normal prestart failures are handled inside _run_prestart. + self.kill_worker() + raise + + # 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) - # Run the worker's concurrent loops until shutdown. - await asyncio.gather(*tasks) + # 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) def is_alive(self): """ @@ -170,6 +220,68 @@ def current_occupancy(self) -> int: ) return current_progress_count + current_queue_count + async def _fetch_jobs_until_stopped(self, session: ClientSession, jobs_needed: int): + """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: + 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 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=grace) + 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: + """Prestart failed before this worker held a request. Claim one queued + request and fail it with the reason. + + 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.prestart_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, payload) + async def get_jobs(self, session: ClientSession): """ Retrieve multiple jobs from the server in batches using blocking requests. @@ -179,6 +291,10 @@ async def get_jobs(self, session: ClientSession): Adds jobs to the JobsQueue """ while self.is_alive(): + 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() jobs_needed = self.current_concurrency - self.current_occupancy() @@ -191,15 +307,22 @@ 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._failure_take_done = True + + 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._prestart_error) + return + for job in acquired_jobs: await self.jobs_queue.put(job) self.job_progress.add(job) @@ -227,13 +350,19 @@ async def get_jobs(self, session: ClientSession): # Yield control back to the event loop await asyncio.sleep(0) + 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): """ Retrieve jobs from the jobs queue and process them concurrently. 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(): @@ -263,14 +392,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): """ @@ -290,12 +450,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) @@ -309,7 +466,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)}" @@ -338,7 +497,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. """ @@ -346,6 +505,20 @@ async def handle_job(self, session: ClientSession, job: dict): try: log.debug("Handling Job", job["id"]) + # 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 prestart finished; leaving this " + "request for another worker.", + job["id"], + ) + return + + 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) if self.config.get("refresh_worker", False): @@ -369,3 +542,37 @@ 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_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._prestart_ready.is_set() + + async def _fail_job( + self, session: ClientSession, job: dict[str, Any], payload: dict[str, Any] + ): + """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_prestart(self): + """Run hooks beside queue intake, then open the handler gate.""" + try: + self._prestart_error = await run_prestart_phase( + self.prestart_hooks, self.config.get("prestart_timeout") + ) + finally: + # Always release held handlers. + self._prestart_ready.set() + + 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: + await asyncio.sleep(0.1) diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 90053ec72..a1cec532e 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 prestart 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..07759178a --- /dev/null +++ b/tests/test_serverless/test_capture.py @@ -0,0 +1,254 @@ +"""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 os +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 +from runpod.serverless.modules.rp_prestart import ( + clear_prestart_hooks, + register_prestart_hook, +) + + +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 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 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, 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 + + 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_payload_is_unchanged_without_capture(self): + def handler(_job): + print("HF_TOKEN=hf_secret") + raise RuntimeError("kernel panic") + + with patch.dict(os.environ, {"RUNPOD_LOG_CAPTURE": "off"}): + 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.""" + + 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): + """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 + 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_reports_a_huge_handler_error_untruncated(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 error["error_message"] == huge + assert "truncated" not in error["error_traceback"] + + 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): + 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"})) + + # 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__": + unittest.main() 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_modules/test_local.py b/tests/test_serverless/test_modules/test_local.py index 523c3cfbf..2e14af197 100644 --- a/tests/test_serverless/test_modules/test_local.py +++ b/tests/test_serverless/test_modules/test_local.py @@ -1,14 +1,83 @@ -""" 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_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) + 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..748fbf164 --- /dev/null +++ b/tests/test_serverless/test_prestart.py @@ -0,0 +1,271 @@ +"""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, rp_prestart +from runpod.serverless.modules.rp_prestart import ( + PrestartError, + PrestartTimeout, + build_prestart_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_prestart_failed_payload(ctx.exception) + assert payload["event"] == "prestart_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_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]) + + _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..45a788822 --- /dev/null +++ b/tests/test_serverless/test_prestart_lifecycle.py @@ -0,0 +1,755 @@ +"""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_prestart, rp_scale +from runpod.serverless.modules.rp_prestart import ( + PrestartError, + PrestartTimeout, + build_prestart_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 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_prestart_failed_payload( + PrestartError(exc, "load_model"), logs="stderr tail" + ) + assert payload["event"] == "prestart_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_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_prestart_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_prestart, "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("prestart_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": "prestart_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": "prestart_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": "prestart_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": "prestart_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"] == "prestart_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": "prestart_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": "prestart_failed", + } + scaler.kill_worker() + await asyncio.wait_for(task, timeout=0.5) + + _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.""" + scaler = _scaler(hook=lambda: None) + scaler._fail_job = AsyncMock() + job = {"id": "late-1"} + failure = {"error_message": "CUDA OOM", "event": "prestart_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 "prestart_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()