fix(serverless): force-kill worker on fitness check failure

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)")

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

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)