[Python] How to implement timeout and retry for long-running Hugging Face model inference in...

I'm working with a locally hosted Hugging Face transformers model (mistral-7b, llama2-13b, etc.), using the pipeline interface on a GPU server (A100).

Sometimes inference takes much longer than expected depending on the prompt and system load. I want to implement a timeout so that if an inference call takes longer than 60 seconds, it cancels and optionally retries the request a few times.

Here is a simplified version of my code:


# Load model
generator = pipeline('text-generation', model="TheBloke/Mistral-7B-Instruct-v0.1-GGUF", device=0)

# Function that generates text
def generate_text(prompt: str):
output = generator(prompt, max_new_tokens=100)
return output[0]['generated_text']

I tried wrapping the call using asyncio and wait_for, like this:

import asyncio

async def generate_with_timeout(prompt: str, timeout_sec: int = 60):
loop = asyncio.get_event_loop()
try:
result = await asyncio.wait_for(
loop.run_in_executor(None, generate_text, prompt),
timeout=timeout_sec
)
return result
except asyncio.TimeoutError:
print(f"Timeout after {timeout_sec} seconds.")
return None


The problem is that when a timeout happens, the underlying thread (and GPU memory) is still busy because the generate() function can't really be cancelled midway. Over time, this can exhaust resources or even crash the server. Question:

What is the correct way to implement timeout and retry logic for Hugging Face model inference, so that it cancels cleanly and does not leak GPU memory or leave stuck threads?

Continue reading...