diff --git a/server.py b/server.py index c172792..bf6747b 100644 --- a/server.py +++ b/server.py @@ -55,6 +55,8 @@ class TelemetryBroker: self.running = False self.samples = 0 self.last_sample_ms = 0.0 + # Set on shutdown so open SSE generators finish instead of holding the server up. + self.closing = False async def start(self) -> None: if self.running: @@ -64,6 +66,13 @@ class TelemetryBroker: async def stop(self) -> None: self.running = False + self.closing = True + # Wake every subscriber so their generator can return. Without this, uvicorn waits + # on the open SSE responses during graceful shutdown and systemd eventually + # SIGKILLs the unit -- which skips the in-process GPU restore hook entirely. + for q in list(self.subscribers): + with contextlib.suppress(asyncio.QueueFull): + q.put_nowait(None) if self.task: self.task.cancel() with contextlib.suppress(asyncio.CancelledError): @@ -250,11 +259,13 @@ async def sse_telemetry_stream(request: Request): try: snap = await broker.get() yield f"data: {json.dumps(snap)}\n\n" - while True: + while not broker.closing: if await request.is_disconnected(): break try: snap = await asyncio.wait_for(q.get(), timeout=15.0) + if snap is None: # shutdown sentinel + break yield f"data: {json.dumps(snap)}\n\n" except asyncio.TimeoutError: yield ": keepalive\n\n" @@ -497,4 +508,5 @@ async def root_index(): if __name__ == "__main__": import uvicorn - uvicorn.run("server:app", host="0.0.0.0", port=9090, reload=False, log_level="info") + uvicorn.run("server:app", host="0.0.0.0", port=9090, reload=False, log_level="info", + timeout_graceful_shutdown=10)