print a warning if stream task ever times out

d8c40bb4f1df86e3a129f191fecf1d7f54531d28 · 2024-07-17 19:09:11 -0700 · Alex Cheema

Files touched

M exo/api/chatgpt_api.py

Diff

commit d8c40bb4f1df86e3a129f191fecf1d7f54531d28
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Jul 17 19:09:11 2024 -0700

    print a warning if stream task ever times out
---
 exo/api/chatgpt_api.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index f7b3d242..1480590b 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -200,7 +200,10 @@ class ChatGPTAPI:
                 _, tokens, _ = await callback.wait(on_result, timeout=self.response_timeout_secs)
                 if stream_task: # in case there is still a stream task running, wait for it to complete
                     if DEBUG >= 2: print(f"Pending stream task. Waiting for stream task to complete.")
-                    await stream_task
+                    try:
+                        await asyncio.wait_for(stream_task, timeout=30)
+                    except asyncio.TimeoutError:
+                        print("WARNING: Stream task timed out. This should not happen.")
                 await response.write_eof()
                 return response
             else:

← 1e1e11cd check if inference_engine has tokenizer before printing with · back to Exo · Squashed 'tinychat/' content from commit fa7e734b4 0870e6bf →