diff --git a/runpod/serverless/modules/rp_fitness.py b/runpod/serverless/modules/rp_fitness.py index a60f0b0f..369a84f0 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()). @@ -148,7 +175,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 +188,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 +234,10 @@ async def run_fitness_checks() -> None: ) log.debug(f"Traceback:\n{full_traceback}") - # Exit immediately with failure code + # 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..31255180 --- /dev/null +++ b/tests/test_serverless/test_modules/test_fitness/test_force_kill.py @@ -0,0 +1,94 @@ +""" +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. +""" + +import subprocess +import sys +from unittest.mock import patch + +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)