diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index 0e42f8ae2b..45b1046677 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -629,6 +629,8 @@ def _resolve_operation_strict_schema(operation_env: str) -> bool: ENV_WORKER_MAX_RETRIES = "HINDSIGHT_API_WORKER_MAX_RETRIES" ENV_WORKER_TASK_RETRY_BACKOFF_SECONDS = "HINDSIGHT_API_WORKER_TASK_RETRY_BACKOFF_SECONDS" ENV_WORKER_HTTP_PORT = "HINDSIGHT_API_WORKER_HTTP_PORT" +ENV_WORKER_LIVENESS_PORT = "HINDSIGHT_API_WORKER_LIVENESS_PORT" +ENV_WORKER_LIVENESS_THRESHOLD_SECONDS = "HINDSIGHT_API_WORKER_LIVENESS_THRESHOLD_SECONDS" ENV_WORKER_MAX_SLOTS = "HINDSIGHT_API_WORKER_MAX_SLOTS" ENV_OPERATION_RETENTION_DAYS = "HINDSIGHT_API_OPERATION_RETENTION_DAYS" ENV_OPERATION_CLEANUP_BATCH_SIZE = "HINDSIGHT_API_OPERATION_CLEANUP_BATCH_SIZE" @@ -1107,7 +1109,14 @@ def _parse_strategy_boosts(raw: str | None) -> dict[str, str]: DEFAULT_WORKER_POLL_INTERVAL_MS = 500 # Poll database every 500ms DEFAULT_WORKER_MAX_RETRIES = 3 # Max retries before marking task failed DEFAULT_WORKER_TASK_RETRY_BACKOFF_SECONDS = 60 # Seconds between retries on transient task failure -DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health +DEFAULT_WORKER_HTTP_PORT = 8889 # HTTP port for worker metrics/health (async readiness) +DEFAULT_WORKER_LIVENESS_PORT = 8890 # HTTP port for thread-based liveness probe +# Liveness heartbeat staleness budget. The main event loop bumps a heartbeat +# ~1x/sec; the loop can be blocked for several seconds by litellm's synchronous +# botocore credential resolution + SigV4 signing during Bedrock inference. The +# threshold must exceed that worst-case block so the livenessProbe does not kill +# a busy-but-alive process, while still catching a genuinely wedged loop. +DEFAULT_WORKER_LIVENESS_THRESHOLD_SECONDS = 30 DEFAULT_WORKER_MAX_SLOTS = 10 # Total concurrent tasks per worker # Terminal rows keep their payload and metadata for one coherent debug/retry TTL. # Zero retention days disables automatic pruning entirely, and is the default: @@ -2027,6 +2036,8 @@ class HindsightConfig: worker_max_retries: int worker_task_retry_backoff_seconds: int worker_http_port: int + worker_liveness_port: int + worker_liveness_threshold_seconds: int worker_max_slots: int worker_slot_reservations: dict[str, int] worker_consolidation_bank_priority: dict[str, int] @@ -3088,6 +3099,13 @@ def from_env(cls) -> "HindsightConfig": ) ), worker_http_port=int(os.getenv(ENV_WORKER_HTTP_PORT, str(DEFAULT_WORKER_HTTP_PORT))), + worker_liveness_port=int(os.getenv(ENV_WORKER_LIVENESS_PORT, str(DEFAULT_WORKER_LIVENESS_PORT))), + worker_liveness_threshold_seconds=int( + os.getenv( + ENV_WORKER_LIVENESS_THRESHOLD_SECONDS, + str(DEFAULT_WORKER_LIVENESS_THRESHOLD_SECONDS), + ) + ), worker_max_slots=int(os.getenv(ENV_WORKER_MAX_SLOTS, str(DEFAULT_WORKER_MAX_SLOTS))), worker_slot_reservations={ op_type: int(os.getenv(env_var, str(default))) diff --git a/hindsight-api-slim/hindsight_api/worker/liveness.py b/hindsight-api-slim/hindsight_api/worker/liveness.py new file mode 100644 index 0000000000..71cbfe5ad9 --- /dev/null +++ b/hindsight-api-slim/hindsight_api/worker/liveness.py @@ -0,0 +1,93 @@ +"""Thread-based liveness server for the Hindsight worker. + +The worker runs its poller and its async ``/health`` (readiness) server on a +single asyncio event loop. During Bedrock inference, litellm performs +*synchronous* botocore credential resolution and SigV4 signing on the loop +thread, which can block the loop for several seconds. A Kubernetes +livenessProbe pointed at the async ``/health`` endpoint then times out and the +pod is SIGKILLed (exitCode 137) even though the process is busy, not dead. + +This module provides a liveness signal that survives a busy loop: a heartbeat +timestamp bumped ~once per second by a task on the main loop, served by a tiny +stdlib HTTP server on a *separate daemon thread*. While the loop is blocked in +sync signing the heartbeat goes momentarily stale, but the loop resumes and the +heartbeat catches up well within the threshold, so the probe keeps returning +200. A genuinely wedged loop never catches up and the probe fails, letting +Kubernetes restart the pod. + +The liveness server deliberately does NOT touch the asyncpg pool: the pool is +bound to the main event loop and cannot be awaited from another thread. The +DB check stays on the async ``/health`` endpoint as the *readiness* signal. +""" + +from __future__ import annotations + +import http.server +import logging +import threading +import time + +logger = logging.getLogger(__name__) + + +class Heartbeat: + """A monotonic timestamp the main loop bumps to prove it can make progress.""" + + def __init__(self) -> None: + self._last = time.monotonic() + + def beat(self) -> None: + """Record that the main loop just ran.""" + self._last = time.monotonic() + + def age(self) -> float: + """Seconds elapsed since the last :meth:`beat`.""" + return time.monotonic() - self._last + + +def is_alive(age_seconds: float, threshold_seconds: float) -> bool: + """Liveness predicate. + + Alive while the heartbeat is younger than ``threshold_seconds``. A busy loop + (e.g. blocked in sync SigV4 signing) stays alive as long as it resumes + within the threshold; a wedged loop that never resumes fails. + """ + return age_seconds < threshold_seconds + + +def start_liveness_server( + host: str, + port: int, + heartbeat: Heartbeat, + threshold_seconds: float, +) -> threading.Thread: + """Start a daemon-thread HTTP liveness server and return its thread. + + Any GET returns 200 while the heartbeat is fresh and 503 once it goes stale. + The server runs on its own thread with its own socket loop, so a blocked + main event loop cannot starve it. The thread is a daemon and needs no + explicit shutdown; it dies with the process. + """ + + class _LivenessHandler(http.server.BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 - name mandated by BaseHTTPRequestHandler + alive = is_alive(heartbeat.age(), threshold_seconds) + body = b"ok\n" if alive else b"stale\n" + self.send_response(200 if alive else 503) + self.send_header("Content-Type", "text/plain") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args, **kwargs): # silence per-request stderr logging + pass + + server = http.server.ThreadingHTTPServer((host, port), _LivenessHandler) + thread = threading.Thread( + target=server.serve_forever, + name="worker-liveness", + daemon=True, + ) + thread.start() + logger.info(f"Liveness server started on http://{host}:{port}/ (staleness threshold={threshold_seconds}s)") + return thread diff --git a/hindsight-api-slim/hindsight_api/worker/main.py b/hindsight-api-slim/hindsight_api/worker/main.py index 7799d32ea8..7930e98447 100644 --- a/hindsight-api-slim/hindsight_api/worker/main.py +++ b/hindsight-api-slim/hindsight_api/worker/main.py @@ -20,6 +20,7 @@ from ..config import get_config from ..engine.task_backend import WorkerTaskBackend +from .liveness import Heartbeat, start_liveness_server from .poller import WorkerPoller # Filter deprecation warnings from third-party libraries @@ -164,6 +165,15 @@ def main(): default="0.0.0.0", help="HTTP host to bind (default: 0.0.0.0)", ) + parser.add_argument( + "--liveness-port", + type=int, + default=config.worker_liveness_port, + help=( + f"HTTP port for the thread-based liveness probe " + f"(default: {config.worker_liveness_port}, env: HINDSIGHT_API_WORKER_LIVENESS_PORT)" + ), + ) # Logging options parser.add_argument( @@ -203,6 +213,7 @@ def main(): print(f" Operation retention: {config.operation_retention_days} days (terminal rows, payloads, and metadata)") print(f" Operation cleanup batch: {config.operation_cleanup_batch_size} rows/schema/cycle") print(f" HTTP server: {args.http_host}:{args.http_port}") + print(f" Liveness server: {args.http_host}:{args.liveness_port}") print() # Global references for cleanup @@ -320,11 +331,37 @@ def signal_handler(): ) server = uvicorn.Server(uvicorn_config) + # Liveness: a task on the main loop bumps a heartbeat ~1x/sec, and a + # separate daemon-thread HTTP server (below) serves it. Bedrock inference + # blocks the loop with sync botocore signing for several seconds; the + # thread server keeps answering the K8s livenessProbe during that block + # so a busy-but-alive process is not SIGKILLed. The async /health + # endpoint above stays the readiness signal (it does the DB check, which + # cannot run from the thread because the asyncpg pool is loop-bound). + heartbeat = Heartbeat() + + async def beat_heartbeat(): + while not shutdown_requested.is_set(): + heartbeat.beat() + try: + await asyncio.sleep(1.0) + except asyncio.CancelledError: + break + + heartbeat_task = asyncio.create_task(beat_heartbeat()) + start_liveness_server( + host=args.http_host, + port=args.liveness_port, + heartbeat=heartbeat, + threshold_seconds=config.worker_liveness_threshold_seconds, + ) + # Run the poller and HTTP server concurrently poller_task = asyncio.create_task(poller.run()) http_task = asyncio.create_task(server.serve()) print(f"Worker started. Metrics available at http://{args.http_host}:{args.http_port}/metrics") + print(f"Liveness available at http://{args.http_host}:{args.liveness_port}/") # Wait for shutdown signal try: @@ -344,6 +381,14 @@ def signal_handler(): except asyncio.CancelledError: pass + # Stop the liveness heartbeat (the liveness server is a daemon thread + # and exits with the process; no explicit shutdown needed). + heartbeat_task.cancel() + try: + await heartbeat_task + except asyncio.CancelledError: + pass + # Wait for HTTP server to finish try: await asyncio.wait_for(http_task, timeout=5.0) diff --git a/hindsight-api-slim/tests/test_worker_liveness.py b/hindsight-api-slim/tests/test_worker_liveness.py new file mode 100644 index 0000000000..8bcec589b3 --- /dev/null +++ b/hindsight-api-slim/tests/test_worker_liveness.py @@ -0,0 +1,34 @@ +"""Tests for hindsight_api.worker.liveness heartbeat-staleness logic.""" + +from hindsight_api.worker.liveness import Heartbeat, is_alive + + +def test_is_alive_fresh_heartbeat(): + """A heartbeat younger than the threshold is alive.""" + assert is_alive(age_seconds=0.0, threshold_seconds=30.0) is True + assert is_alive(age_seconds=29.999, threshold_seconds=30.0) is True + + +def test_is_alive_stale_heartbeat(): + """A heartbeat at or beyond the threshold is not alive. + + This is the case a genuinely wedged event loop hits: the heartbeat never + catches up, so the livenessProbe eventually fails and K8s restarts the pod. + """ + assert is_alive(age_seconds=30.0, threshold_seconds=30.0) is False + assert is_alive(age_seconds=45.0, threshold_seconds=30.0) is False + + +def test_is_alive_tolerates_multi_second_block(): + """A several-second loop block (sync SigV4 signing during Bedrock + inference) stays alive under the default 30s threshold — the whole point + of the thread-based liveness server.""" + assert is_alive(age_seconds=8.0, threshold_seconds=30.0) is True + + +def test_heartbeat_beat_resets_age(): + """beat() resets the age toward zero; a fresh Heartbeat is nearly zero age.""" + hb = Heartbeat() + assert hb.age() < 1.0 + hb.beat() + assert hb.age() < 1.0