mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-04-19 13:52:07 +00:00
16 lines
482 B
Python
16 lines
482 B
Python
import time
|
|
|
|
from aiohttp import ClientConnectorError, ClientOSError, ServerDisconnectedError
|
|
from text_generation import AsyncClient
|
|
|
|
|
|
async def health_check(client: AsyncClient, timeout: int = 60):
|
|
assert timeout > 0
|
|
for _ in range(timeout):
|
|
try:
|
|
await client.generate("test")
|
|
return
|
|
except (ClientConnectorError, ClientOSError, ServerDisconnectedError) as e:
|
|
time.sleep(1)
|
|
raise RuntimeError("Health check failed")
|