135 lines
4.0 KiB
Python
135 lines
4.0 KiB
Python
"""
|
|
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
|