fix shutdown

fe0f1cdb1ee8cdf99e04a23d70be3344a494b856 · 2024-11-22 19:38:22 +0400 · Alex Cheema

Files touched

M exo/api/chatgpt_api.py
M exo/helpers.py
M exo/main.py

Diff

commit fe0f1cdb1ee8cdf99e04a23d70be3344a494b856
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Fri Nov 22 19:38:22 2024 +0400

    fix shutdown
---
 exo/api/chatgpt_api.py | 5 ++---
 exo/helpers.py         | 2 +-
 exo/main.py            | 4 ++--
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/exo/api/chatgpt_api.py b/exo/api/chatgpt_api.py
index 979e2d4b..dfa2e09b 100644
--- a/exo/api/chatgpt_api.py
+++ b/exo/api/chatgpt_api.py
@@ -8,8 +8,7 @@ from typing import List, Literal, Union, Dict
 from aiohttp import web
 import aiohttp_cors
 import traceback
-import os
-import sys
+import signal
 from exo import DEBUG, VERSION
 from exo.download.download_progress import RepoProgressEvent
 from exo.helpers import PrefixDict, shutdown
@@ -188,7 +187,7 @@ class ChatGPTAPI:
     response = web.json_response({"detail": "Quit signal received"}, status=200)
     await response.prepare(request)
     await response.write_eof()
-    await shutdown(signal.SIGINT, asyncio.get_event_loop())
+    await shutdown(signal.SIGINT, asyncio.get_event_loop(), self.node.server)
 
   async def timeout_middleware(self, app, handler):
     async def middleware(request):
diff --git a/exo/helpers.py b/exo/helpers.py
index 4defb98b..a808e993 100644
--- a/exo/helpers.py
+++ b/exo/helpers.py
@@ -237,7 +237,7 @@ def get_all_ip_addresses():
     return ["localhost"]
 
 
-async def shutdown(signal, loop):
+async def shutdown(signal, loop, server):
   """Gracefully shutdown the server and close the asyncio loop."""
   print(f"Received exit signal {signal.name}...")
   print("Thank you for using exo.")
diff --git a/exo/main.py b/exo/main.py
index ffd16eda..66fa20da 100644
--- a/exo/main.py
+++ b/exo/main.py
@@ -217,7 +217,7 @@ async def main():
 
   # Use a more direct approach to handle signals
   def handle_exit():
-    asyncio.ensure_future(shutdown(signal.SIGTERM, loop))
+    asyncio.ensure_future(shutdown(signal.SIGTERM, loop, node.server))
 
   if platform.system() != "Windows":
     for s in [signal.SIGINT, signal.SIGTERM]:
@@ -244,7 +244,7 @@ def run():
   except KeyboardInterrupt:
     print("Received keyboard interrupt. Shutting down...")
   finally:
-    loop.run_until_complete(shutdown(signal.SIGTERM, loop))
+    loop.run_until_complete(shutdown(signal.SIGTERM, loop, node.server))
     loop.close()

← d3505e03 run tests with --disable-tui for faster tests · back to Exo · fixiing required engines definition 39139c14 →