Terminate AsyncLLMEngine task gracefully (#10)

triton-inference-server · Oct 12, 2023 · ab5a811 · ab5a811
1 parent 27a07d3
commit ab5a811
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/model.py b/src/model.py
@@ -118,6 +118,10 @@ async def await_shutdown(self):
             )
             await asyncio.sleep(5)
 
+        for task in asyncio.all_tasks(loop=self._loop):
+            if task is not asyncio.current_task():
+                task.cancel()
+
         self.logger.log_info("Shutdown complete")
 
     def get_sampling_params_dict(self, params_json):