fix: keep health and metrics responsive with idle connections

src/prime_rl/utils/metrics_server.py

@@ -7,7 +7,7 @@
 
 import threading
 import time
-from http.server import BaseHTTPRequestHandler, HTTPServer
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING
 
 from loguru import logger
@@ -26,7 +26,7 @@ class HealthServer:
     def __init__(self, port: int, host: str = "0.0.0.0"):
         self.port = port
         self.host = host
-        self._server: HTTPServer | None = None
+        self._server: ThreadingHTTPServer | None = None
         self._thread: threading.Thread | None = None
         self._started = False
 
@@ -55,7 +55,7 @@ def start(self) -> None:
             logger.warning(f"{self.__class__.__name__} already started")
             return
 
-        self._server = HTTPServer((self.host, self.port), self._make_handler())
+        self._server = ThreadingHTTPServer((self.host, self.port), self._make_handler())
         self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
         self._thread.start()
         self._started = True
@@ -141,7 +141,7 @@ def start(self) -> None:
             logger.warning("Metrics server already started")
             return
 
-        self._server = HTTPServer((self.host, self.port), self._make_handler())
+        self._server = ThreadingHTTPServer((self.host, self.port), self._make_handler())
         self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
         self._thread.start()
         self._started = True

tests/unit/utils/test_metrics_server.py

@@ -45,12 +45,14 @@ def test_server_start_stop():
 
     server.start()
     assert server._started
-    time.sleep(0.1)
-
-    response = urllib.request.urlopen(f"http://localhost:{port}/metrics", timeout=2)
-    assert response.status == 200
-
-    server.stop()
+    try:
+        with socket.create_connection(("localhost", port), timeout=2):
+            time.sleep(0.1)
+            with urllib.request.urlopen(f"http://localhost:{port}/metrics", timeout=2) as response:
+                assert response.status == 200
+            server.stop()
+    finally:
+        server.stop()
     assert not server._started
 
 
@@ -169,13 +171,15 @@ def test_health_server_start_stop():
 
     server.start()
     assert server._started
-    time.sleep(0.1)
-
-    response = urllib.request.urlopen(f"http://localhost:{port}/health", timeout=2)
-    assert response.status == 200
-    assert response.read() == b"ok\n"
-
-    server.stop()
+    try:
+        with socket.create_connection(("localhost", port), timeout=2):
+            time.sleep(0.1)
+            with urllib.request.urlopen(f"http://localhost:{port}/health", timeout=2) as response:
+                assert response.status == 200
+                assert response.read() == b"ok\n"
+            server.stop()
+    finally:
+        server.stop()
     assert not server._started