Performance: Replace inefficient polling in server loop with event-driven shutdown and optimize Health check

## Summary The current implementation of the server loop in `text_generation_server/server.py` uses an inefficient polling mechanism to handle shutdown signals. Additionally, the Health check endpoint performs redundant GPU allocations, leading to unnecessary overhead and memory fragmentation. --- ## Root Cause 1. Polling Latency: The server loop relies on a blocking-style polling mechanism: ``` while signal_handler.KEEP_PROCESSING: await asyncio.sleep(0.5) ``` Even after the flag is set to False, the loop must wait for the current sleep(0.5) cycle to complete before exiting. 2. Health Check Overhead: The Health check unnecessarily touches device memory on every probe: ``` if self.model.device.type == "cuda": torch.zeros((2, 2)).cuda() ``` This causes redundant CUDA kernel launches and potential memory churn. --- ## Suggested Fix 1. Event-Driven Shutdown: Replace the boolean flag with asyncio.Event in the SignalHandler and use `await signal_handler.exit_event.wait()` in the main loop. 2. Efficient Health Check: Verify device availability without allocation: async def Health(self, request, context): ``` if self.model.device.type == "cuda": if not torch.cuda.is_available(): raise RuntimeError("GPU device not available") return generate_pb2.HealthResponse() ``` --- ## Environment - OS: Linux (Containerized) - Python: 3.12+ - Backend: Text Generation Inference (TGI) --- ### Information - [x] Docker - [x] The CLI directly ### Tasks - [x] An officially supported command - [ ] My own modifications ## Reproduction Script ``` import asyncio import time async def reproduce_tgi_latency(): keep_processing = True start_time = time.time() # Simulate OS Signal (SIGINT/SIGTERM) arriving at 0.1s async def trigger_signal(): await asyncio.sleep(0.1) nonlocal keep_processing print(f"\n[TEST] Signal Received at {time.time() - start_time:.4f}s") keep_processing = False asyncio.create_task(trigger_signal()) # --- EXACT TGI LOGIC FROM server.py --- print("[*] TGI Server Loop Started (Polling every 0.5s)") while keep_processing: await asyncio.sleep(0.5) end_time = time.time() print(f"[*] Server finally stopped at {end_time - start_time:.4f}s") await reproduce_tgi_latency() ``` --- ## Actual Behavior ``` [TEST] Signal Received at 0.1010s [*] Server finally stopped at 0.5017s [!] Latency: ~400ms delay (Loop blocked by asyncio.sleep) ``` --- ## Expected Behavior The server should respond to signals immediately (< 1ms) using an event-driven approach instead of polling with a sleep timer. --- ## Conclusion This is a performance and correctness issue. Implementing an event-driven shutdown ensures microsecond responsiveness to OS signals and prevents redundant GPU resource usage.