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.

1. 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

• Docker
• The CLI directly

Tasks

• 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.