diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index a60f0b0f..77df97e7 100644 --- a/runpod/serverless/modules/rp_fitness.py +++ b/runpod/serverless/modules/rp_fitness.py @@ -2,15 +2,17 @@ Fitness check system for worker startup validation. Fitness checks run before handler initialization on the actual RunPod serverless -platform to validate the worker environment. Any check failure causes immediate -exit with sys.exit(1), signaling unhealthy state to the container orchestrator. +platform to validate the worker environment. Any check failure force-kills the +worker via os._exit(1), signaling unhealthy state to the container orchestrator. Fitness checks do NOT run in local development mode or testing mode. """ from __future__ import annotations +import contextlib import inspect +import os import sys import time import traceback @@ -20,6 +22,31 @@ log = RunPodLogger() + +def _terminate_unhealthy(code: int = 1) -> None: + """ + Force-kill the worker after a fitness check failure. + + Uses os._exit rather than sys.exit because a fitness failure means the + environment is broken and the worker must die immediately so the + orchestrator can restart it. sys.exit only raises SystemExit, which + triggers cooperative interpreter shutdown and blocks joining non-daemon + threads. Workers routinely have such threads alive by the time checks run + (e.g. vLLM's AsyncLLMEngine, constructed at import before the checks), so + sys.exit can hang forever and the worker keeps serving jobs. os._exit + bypasses thread joins, atexit handlers, and asyncgen cleanup. + + Args: + code: Process exit code (default 1, signaling unhealthy). + """ + # Best-effort flush of buffered logs before the hard exit skips normal + # cleanup. A broken worker may have a closed/None stdio stream; never let a + # flush failure stop the exit, which is the whole point of this helper. + for stream in (sys.stdout, sys.stderr): + with contextlib.suppress(Exception): + stream.flush() + os._exit(code) + # Global registry for fitness check functions, preserves registration order _fitness_checks: list[Callable] = [] @@ -29,7 +56,7 @@ def register_fitness_check(func: Callable) -> Callable: Decorator to register a fitness check function. Fitness checks validate worker health at startup before handler initialization. - If any check fails, the worker exits with sys.exit(1). + If any check fails, the worker is force-killed with os._exit(1). Supports both sync and async functions (auto-detected via inspect.iscoroutinefunction()). @@ -83,6 +110,51 @@ def _reset_registration_state() -> None: _registration_state["system_checks"] = False +# Bound how long the best-effort unhealthy report may delay the exit. +_REPORT_TIMEOUT_SECONDS = 2 + + +def _report_unhealthy(check: str, reason: str) -> None: + """ + Best-effort report of a fitness-check failure to the host before exit. + + Sends a single GET to the ping URL (same URL/credentials the heartbeat + uses) with status=unhealthy plus the failing check name and reason, so the + host can emit a queryable worker.fitness_failed event. Any failure — no + ping URL, no API key, HTTP error, timeout — is swallowed, so this can never + prevent the os._exit that follows. It is synchronous, so it may delay that + exit by up to _REPORT_TIMEOUT_SECONDS (network phases only; it adds no + delay when there is no ping URL/API key to report to). + """ + ping_url = os.environ.get("RUNPOD_WEBHOOK_PING") + api_key = os.environ.get("RUNPOD_AI_API_KEY") + if not ping_url or ping_url == "PING_NOT_SET" or not api_key: + return + + try: + # Deferred imports: keep module import light and avoid import cycles. + from runpod.http_client import SyncClientSession + from runpod.serverless.modules.worker_state import WORKER_ID + from runpod.version import __version__ as runpod_version + + ping_url = ping_url.replace("$RUNPOD_POD_ID", WORKER_ID) + params = { + "status": "unhealthy", + "check": check, + "reason": reason[:256], + "runpod_version": runpod_version, + } + session = SyncClientSession() + try: + session.headers.update({"Authorization": api_key}) + session.get(ping_url, params=params, timeout=_REPORT_TIMEOUT_SECONDS) + finally: + session.close() + except Exception: + # Best-effort only; the exit is the guarantee, not this report. + pass + + def _ensure_gpu_check_registered() -> None: """ Ensure GPU fitness check is registered. @@ -148,7 +220,10 @@ async def run_fitness_checks() -> None: 5. On any exception: - Log detailed error with check name, exception type, and message - Log traceback at DEBUG level - - Call sys.exit(1) immediately (fail-fast) + - Force-kill the worker via os._exit(1) immediately (fail-fast). This is + a hard exit, not a cooperative sys.exit/SystemExit: it does not unwind + the stack or run cleanup, so callers cannot catch it and it cannot be + blocked by live non-daemon threads. 6. On successful completion of all checks: - Log completion message with total execution time @@ -158,8 +233,9 @@ async def run_fitness_checks() -> None: and handles checks with dependencies correctly. Timing uses high-precision perf_counter for accurate measurements. - Raises: - SystemExit: Calls sys.exit(1) if any check fails. + Note: + A failing check terminates the process via os._exit(1); this function + does not return in that case and does not raise SystemExit. """ # Defer GPU check auto-registration until fitness checks are about to run # This avoids circular import issues during module initialization @@ -203,9 +279,18 @@ async def run_fitness_checks() -> None: ) log.debug(f"Traceback:\n{full_traceback}") - # Exit immediately with failure code + # Best-effort report to the host so the failure is queryable. It is + # bounded (see _REPORT_TIMEOUT_SECONDS) and fully swallowed, so it + # can delay the force-exit below but can never prevent it. + try: + _report_unhealthy(check_name, f"{error_type}: {error_message}") + except Exception: # a report failure must never prevent the exit + pass + + # Force-kill immediately; see _terminate_unhealthy for why this is + # os._exit rather than sys.exit. log.error("Worker is unhealthy, exiting.") - sys.exit(1) + _terminate_unhealthy(1) total_elapsed_ms = (time.perf_counter() - total_start_time) * 1000 log.info(f"All fitness checks passed. ({total_elapsed_ms:.2f}ms)") diff --git a/tests/test_serverless/test_modules/test_fitness/conftest.py b/tests/test_serverless/test_modules/test_fitness/conftest.py index dfad1af9..f8df8614 100644 --- a/tests/test_serverless/test_modules/test_fitness/conftest.py +++ b/tests/test_serverless/test_modules/test_fitness/conftest.py @@ -2,6 +2,7 @@ import pytest +from runpod.serverless.modules import rp_fitness from runpod.serverless.modules.rp_fitness import ( clear_fitness_checks, _reset_registration_state, @@ -14,9 +15,20 @@ def cleanup_fitness_checks(monkeypatch): Disables auto-registration of system checks to avoid interference with fitness check framework tests. + + Also neutralizes the unhealthy force-kill: in production a failed check + calls os._exit, which would kill the pytest process itself. Redirect it + to raise SystemExit(1) so tests can assert exit behavior in-process. + Tests that need the real os._exit patch it themselves. """ monkeypatch.setenv("RUNPOD_SKIP_AUTO_SYSTEM_CHECKS", "true") monkeypatch.setenv("RUNPOD_SKIP_GPU_CHECK", "true") + + def _raise_system_exit(code=0): + raise SystemExit(code) + + monkeypatch.setattr(rp_fitness.os, "_exit", _raise_system_exit) + _reset_registration_state() clear_fitness_checks() yield diff --git a/tests/test_serverless/test_modules/test_fitness/test_force_kill.py b/tests/test_serverless/test_modules/test_fitness/test_force_kill.py new file mode 100644 index 00000000..9bb4f213 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_force_kill.py @@ -0,0 +1,174 @@ +""" +Regression tests for SLS-379. + +A failed fitness check must force-kill the worker even when non-daemon +background threads are alive (e.g. vLLM's AsyncLLMEngine, which is +constructed at import time before fitness checks run). ``sys.exit(1)`` only +raises ``SystemExit`` and then blocks in interpreter shutdown joining those +threads, so the worker logs "unhealthy, exiting." but never terminates. +The exit must go through ``os._exit`` to terminate unconditionally. + +Also covers SLS-380: the best-effort unhealthy report the worker sends to the +host before force-exiting. +""" + +import subprocess +import sys +from unittest.mock import MagicMock, patch + +import pytest + +from runpod.serverless.modules import rp_fitness + + +# Child program mirrors worker.py: a non-daemon thread is alive (like vLLM's +# engine loop) when a failing fitness check triggers the unhealthy exit. +_CHILD_PROGRAM = """ +import asyncio +import os +import threading +import time + +os.environ["RUNPOD_SKIP_AUTO_SYSTEM_CHECKS"] = "true" +os.environ["RUNPOD_SKIP_GPU_CHECK"] = "true" + +from runpod.serverless.modules.rp_fitness import ( + register_fitness_check, + run_fitness_checks, +) + + +def _never_returns(): + while True: + time.sleep(0.5) + + +@register_fitness_check +def _failing_check(): + raise RuntimeError("simulated fitness failure") + + +threading.Thread(target=_never_returns, daemon=False).start() +asyncio.run(run_fitness_checks()) +""" + + +def test_unhealthy_exit_terminates_with_live_non_daemon_thread(): + """Worker must die (exit 1) within the timeout, not hang, when a + non-daemon thread is alive at the time of the fitness failure.""" + try: + result = subprocess.run( + [sys.executable, "-c", _CHILD_PROGRAM], + capture_output=True, + text=True, + timeout=15, + ) + except subprocess.TimeoutExpired as exc: + raise AssertionError( + "Worker hung instead of exiting on fitness failure " + "(sys.exit could not join the live non-daemon thread)." + ) from exc + + assert result.returncode == 1, ( + f"expected exit code 1, got {result.returncode}. " + f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + assert "Worker is unhealthy, exiting." in (result.stdout + result.stderr) + + +def test_terminate_unhealthy_uses_os_exit(): + """The unhealthy exit must call os._exit (unconditional) rather than + sys.exit (cooperative).""" + with patch("runpod.serverless.modules.rp_fitness.os._exit") as mock_exit: + rp_fitness._terminate_unhealthy(1) + + mock_exit.assert_called_once_with(1) + + +def test_terminate_unhealthy_exits_even_if_flush_raises(): + """A closed/broken stdio stream must not stop the hard exit; a failing + flush is swallowed so os._exit always runs.""" + with patch("runpod.serverless.modules.rp_fitness.os._exit") as mock_exit, \ + patch("sys.stdout") as mock_stdout, \ + patch("sys.stderr") as mock_stderr: + mock_stdout.flush.side_effect = ValueError("I/O operation on closed file") + mock_stderr.flush.side_effect = ValueError("I/O operation on closed file") + + rp_fitness._terminate_unhealthy(1) + + mock_exit.assert_called_once_with(1) + + +def test_report_unhealthy_posts_check_and_reason(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping/$RUNPOD_POD_ID") + monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") + + fake_session = MagicMock() + with patch("runpod.http_client.SyncClientSession", return_value=fake_session), \ + patch("runpod.serverless.modules.worker_state.WORKER_ID", "podABC"): + rp_fitness._report_unhealthy("_cuda_init_check", "RuntimeError: boom") + + assert fake_session.get.call_count == 1 + call = fake_session.get.call_args + url = call.args[0] if call.args else call.kwargs["url"] + params = call.kwargs["params"] + assert "podABC" in url # $RUNPOD_POD_ID substituted + assert params["status"] == "unhealthy" + assert params["check"] == "_cuda_init_check" + assert params["reason"] == "RuntimeError: boom" + fake_session.headers.update.assert_called_once_with({"Authorization": "key-123"}) + + +def test_report_unhealthy_skipped_without_ping_url(monkeypatch): + monkeypatch.delenv("RUNPOD_WEBHOOK_PING", raising=False) + monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") + with patch("runpod.http_client.SyncClientSession") as session_cls: + rp_fitness._report_unhealthy("_memory_check", "RuntimeError: low") + session_cls.assert_not_called() + + +def test_report_unhealthy_truncates_long_reason(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping") + monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") + + fake_session = MagicMock() + with patch("runpod.http_client.SyncClientSession", return_value=fake_session): + rp_fitness._report_unhealthy("_disk_check", "x" * 300) + + params = fake_session.get.call_args.kwargs["params"] + assert len(params["reason"]) == 256 + + +def test_report_unhealthy_skipped_without_api_key(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping") + monkeypatch.delenv("RUNPOD_AI_API_KEY", raising=False) + with patch("runpod.http_client.SyncClientSession") as session_cls: + rp_fitness._report_unhealthy("_memory_check", "RuntimeError: low") + session_cls.assert_not_called() + + +def test_report_unhealthy_swallows_errors(monkeypatch): + monkeypatch.setenv("RUNPOD_WEBHOOK_PING", "https://api.test/ping") + monkeypatch.setenv("RUNPOD_AI_API_KEY", "key-123") + fake_session = MagicMock() + fake_session.get.side_effect = RuntimeError("network down") + with patch("runpod.http_client.SyncClientSession", return_value=fake_session): + rp_fitness._report_unhealthy("_disk_check", "RuntimeError: full") # must not raise + + +@pytest.mark.asyncio +async def test_failure_reports_before_exit(): + from runpod.serverless.modules.rp_fitness import register_fitness_check, run_fitness_checks + + @register_fitness_check + def _cuda_init_check(): + raise RuntimeError("device busy") + + with patch("runpod.serverless.modules.rp_fitness._report_unhealthy") as mock_report: + with pytest.raises(SystemExit): + await run_fitness_checks() + + mock_report.assert_called_once() + args = mock_report.call_args.args + assert args[0] == "_cuda_init_check" + assert "RuntimeError" in args[1] and "device busy" in args[1]