Destroy process group on Keyboard Interrupt

src/exo/worker/runner/vllm_inference/runner.py

@@ -297,15 +297,19 @@ class Runner:
         self.event_sender.send(TaskAcknowledged(task_id=task.task_id))
 
     def main(self):
-        with self.task_receiver:
-            for task in self.task_receiver:
-                if task.task_id in self.seen:
-                    logger.warning("repeat task - potential error")
-                    continue
-                self.seen.add(task.task_id)
-                self.handle_first_task(task)
-                if isinstance(self.current_status, RunnerShutdown):
-                    break
+        try:
+            with self.task_receiver:
+                for task in self.task_receiver:
+                    if task.task_id in self.seen:
+                        logger.warning("repeat task - potential error")
+                        continue
+                    self.seen.add(task.task_id)
+                    self.handle_first_task(task)
+                    if isinstance(self.current_status, RunnerShutdown):
+                        break
+        finally:
+            if torch.distributed.is_initialized():
+                torch.distributed.destroy_process_group()
 
     def handle_first_task(self, task: Task):
         self.send_task_status(task.task_id, TaskStatus.Running)