only ignore CancelledError inside stop

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):

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):