Story 9

src/health_check.py

@@ -0,0 +1,134 @@
+"""
+Worker health check HTTP endpoint.
+
+Exposes GET /health returning a JSON status document.
+Designed to be started in a background thread alongside the main worker.
+"""
+
+import json
+import logging
+import os
+import threading
+import time
+from datetime import datetime, timezone
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+_START_TIME = time.monotonic()
+
+
+class HealthStatus:
+    """Mutable singleton updated by the worker as it processes videos."""
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._data: Dict[str, Any] = {
+            "status": "starting",
+            "gpu_available": False,
+            "gpu_memory_used_gb": 0.0,
+            "queue_depth": 0,
+            "videos_processed_today": 0,
+            "last_error": None,
+        }
+
+    def update(self, **kwargs):
+        with self._lock:
+            self._data.update(kwargs)
+
+    def snapshot(self) -> Dict[str, Any]:
+        with self._lock:
+            data = dict(self._data)
+        data["uptime_seconds"] = round(time.monotonic() - _START_TIME, 1)
+        return data
+
+
+# Module-level singleton shared between the HTTP handler and the worker
+health = HealthStatus()
+
+
+class _HealthHandler(BaseHTTPRequestHandler):
+    """Minimal HTTP handler — only responds to GET /health."""
+
+    def do_GET(self):  # noqa: N802
+        if self.path.rstrip("/") == "/health":
+            body = json.dumps(health.snapshot(), default=str).encode()
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+    def log_message(self, fmt, *args):  # silence default access log
+        logger.debug("health_check: " + fmt, *args)
+
+
+def start_health_server(port: int = 8080) -> threading.Thread:
+    """
+    Start the health HTTP server in a daemon thread.
+
+    Returns the thread so callers can join if needed.
+    """
+    server = HTTPServer(("0.0.0.0", port), _HealthHandler)
+
+    def _serve():
+        logger.info("Health check server listening on port %d", port)
+        server.serve_forever()
+
+    t = threading.Thread(target=_serve, daemon=True, name="health-server")
+    t.start()
+    return t
+
+
+def refresh_health_from_system(db_connector=None, scratch_path: str = "/scratch"):
+    """
+    Update the health singleton with current GPU, queue, and disk state.
+
+    Call this periodically from the worker main loop.
+    """
+    updates: Dict[str, Any] = {"status": "healthy"}
+
+    # GPU state
+    try:
+        import pynvml
+        pynvml.nvmlInit()
+        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+        mem = pynvml.nvmlDeviceGetMemoryInfo(handle)
+        updates["gpu_available"] = True
+        updates["gpu_memory_used_gb"] = round(mem.used / 1024**3, 2)
+        pynvml.nvmlShutdown()
+    except Exception:
+        updates["gpu_available"] = _torch_gpu_available()
+
+    # Queue depth from DB
+    if db_connector is not None:
+        try:
+            rows = db_connector.fetchall(
+                "SELECT COUNT(*) AS cnt FROM videos WHERE status IN ('PENDING','PROCESSING')"
+            )
+            updates["queue_depth"] = int((rows or [{"cnt": 0}])[0]["cnt"])
+
+            today = datetime.now(timezone.utc).date().isoformat()
+            rows2 = db_connector.fetchall(
+                "SELECT COUNT(*) AS cnt FROM videos "
+                "WHERE status = 'COMPLETED' AND DATE(last_processed_time) = %s",
+                (today,),
+            )
+            updates["videos_processed_today"] = int((rows2 or [{"cnt": 0}])[0]["cnt"])
+        except Exception as exc:
+            logger.debug("Health DB query failed: %s", exc)
+
+    health.update(**updates)
+
+
+def _torch_gpu_available() -> bool:
+    try:
+        import torch
+        return torch.cuda.is_available()
+    except ImportError:
+        return False