v1.1.7; Add go2 lab eval; Add server client connection checker

robogauge/scripts/server.py

@@ -21,11 +21,12 @@ os.environ["OMP_NUM_THREADS"] = "1"
 os.environ["MKL_NUM_THREADS"] = "1"
 
 import multiprocessing
+import threading
 import uvicorn
 import queue
 import time
 import uuid
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from typing import Dict, Optional
 import argparse
@@ -47,6 +48,8 @@ default_args_list = [
     '--headless',
 ]
 
+HEARTBEAT_TIMEOUT_SECONDS = 10.0
+
 @dataclass
 class EvalTaskData:
     model_path: str
@@ -67,15 +70,42 @@ class ResponseStatus:
     ERROR = "error"
     NOT_FOUND = "not_found"
 
-def run_api_server(input_queue: multiprocessing.Queue, result_dict: dict, port=9973):
+def update_main_heartbeat(health_dict: dict, stop_event: threading.Event, interval: float = 1.0):
+    while not stop_event.is_set():
+        health_dict["main_pid"] = os.getpid()
+        health_dict["last_heartbeat"] = time.time()
+        stop_event.wait(interval)
+
+def is_main_process_healthy(health_dict: dict) -> bool:
+    last_heartbeat = float(health_dict.get("last_heartbeat", 0.0))
+    return time.time() - last_heartbeat <= HEARTBEAT_TIMEOUT_SECONDS
+
+def run_api_server(input_queue: multiprocessing.Queue, result_dict: dict, health_dict: dict, port=9973):
     """
     Running in a separate subprocess.
     I/O Process: submit requests -> put into queue -> return ID.
     """
     app = FastAPI()
 
+    @app.get("/health")
+    def health():
+        last_heartbeat = float(health_dict.get("last_heartbeat", 0.0))
+        heartbeat_age = time.time() - last_heartbeat
+        if heartbeat_age > HEARTBEAT_TIMEOUT_SECONDS:
+            raise HTTPException(
+                status_code=503,
+                detail=f"RoboGauge main process heartbeat is stale ({heartbeat_age:.1f}s).",
+            )
+        return {
+            "status": "ok",
+            "main_pid": health_dict.get("main_pid"),
+            "heartbeat_age": heartbeat_age,
+        }
+
     @app.post("/submit_eval")
     def submit_eval(req: EvalRequest):
+        if not is_main_process_healthy(health_dict):
+            raise HTTPException(status_code=503, detail="RoboGauge main process is not healthy.")
         task_id = str(uuid.uuid4())
         task_data = EvalTaskData(
             model_path=req.model_path,
@@ -110,10 +140,18 @@ def main():
     manager = ctx.Manager()
     task_queue = manager.Queue()
     results_store = manager.dict()
+    health_store = manager.dict()
+    heartbeat_stop = threading.Event()
+    heartbeat_thread = threading.Thread(
+        target=update_main_heartbeat,
+        args=(health_store, heartbeat_stop),
+        daemon=True,
+    )
+    heartbeat_thread.start()
 
     api_p = ctx.Process(
         target=run_api_server, 
-        args=(task_queue, results_store, args_cli.port),
+        args=(task_queue, results_store, health_store, args_cli.port),
         daemon=True
     )
     api_p.start()
@@ -165,6 +203,8 @@ def main():
         print("\n🛑 Shutting down...")
         api_p.terminate()
         api_p.join()
+        heartbeat_stop.set()
+        heartbeat_thread.join(timeout=2.0)
 
 if __name__ == "__main__":
     main()