From 43308ccbdde7f1bf963584d051163477e0822100 Mon Sep 17 00:00:00 2001
From: Matteo Bettini <matbet@meta.com>
Date: Tue, 26 Sep 2023 11:12:15 +0100
Subject: [PATCH] [Feature] Add graceful experiment termination upon interrupt

Signed-off-by: Matteo Bettini <matbet@meta.com>
---
 benchmarl/experiment/experiment.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
index 2680a6c2..377d2513 100644
--- a/benchmarl/experiment/experiment.py
+++ b/benchmarl/experiment/experiment.py
@@ -322,7 +322,12 @@ def _setup_logger(self):
     def run(self):
         try:
             self._collection_loop()
+        except KeyboardInterrupt as interrupt:
+            print("\n\nExperiment was closed gracefully\n\n")
+            self.close()
+            raise interrupt
         except Exception as err:
+            print("\n\nExperiment failed and is closing gracefully\n\n")
             self.close()
             raise err