only ignore CancelledError inside stop

dd24e7db1e099d13d37a28d18cbb8e4b26d84ac5 · 2024-08-21 19:32:10 +0100 · Alex Cheema

Files touched

M exo/networking/grpc/grpc_server.py
M main.py

Diff

commit dd24e7db1e099d13d37a28d18cbb8e4b26d84ac5
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Wed Aug 21 19:32:10 2024 +0100

    only ignore CancelledError inside stop
---
 exo/networking/grpc/grpc_server.py | 8 ++++++--
 main.py                            | 5 +----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/exo/networking/grpc/grpc_server.py b/exo/networking/grpc/grpc_server.py
index fb56cf9f..ae571d11 100644
--- a/exo/networking/grpc/grpc_server.py
+++ b/exo/networking/grpc/grpc_server.py
@@ -1,6 +1,7 @@
 import grpc
 from concurrent import futures
 import numpy as np
+from asyncio import CancelledError
 
 from . import node_service_pb2
 from . import node_service_pb2_grpc
@@ -33,8 +34,11 @@ class GRPCServer(node_service_pb2_grpc.NodeServiceServicer):
 
   async def stop(self) -> None:
     if self.server:
-      await self.server.stop(grace=5)
-      await self.server.wait_for_termination()
+      try:
+        await self.server.stop(grace=5)
+        await self.server.wait_for_termination()
+      except CancelledError:
+        pass
       if DEBUG >= 1: print("Server stopped and all connections are closed")
 
   async def SendPrompt(self, request, context):
diff --git a/main.py b/main.py
index d6e0dee7..ad969a9a 100644
--- a/main.py
+++ b/main.py
@@ -120,10 +120,7 @@ async def shutdown(signal, loop):
     [task.cancel() for task in server_tasks]
     print(f"Cancelling {len(server_tasks)} outstanding tasks")
     await asyncio.gather(*server_tasks, return_exceptions=True)
-    try:
-        await server.stop()
-    except CancelledError:
-        pass
+    await server.stop()
     loop.stop()
 
 async def run_model_cli(node: Node, inference_engine: InferenceEngine, model_name: str, prompt: str):

← 2e123335 ignore CancelledError when stopping the server · back to Exo · increase cli generation timeout 647ffb94 →