From bbbb7591ce53733ace452bec7aa7deadefa4c90e Mon Sep 17 00:00:00 2001
From: Olympia Dartsi <olympiadartsi@gmail.com>
Date: Tue, 2 Jul 2024 16:38:11 -0700
Subject: [PATCH] adding termination event

---
 alab_management/scripts/launch_lab.py | 37 +++++++++++++++++----------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/alab_management/scripts/launch_lab.py b/alab_management/scripts/launch_lab.py
index 8e66243c..a73284f7 100644
--- a/alab_management/scripts/launch_lab.py
+++ b/alab_management/scripts/launch_lab.py
@@ -9,27 +9,36 @@
 import multiprocessing
 with contextlib.suppress(RuntimeError):
     multiprocessing.set_start_method("spawn")
-termination_event = multiprocessing.Event()
 
 class RestartableProcess:
     """A class for creating processes that can be automatically restarted after failures."""
-    def __init__(self, target, args=(), live_time=None):
-        self.target_func = target
+
+    def __init__(self, target, args=(), live_time=None, termination_event=None):
+        self.target = target
         self.live_time = live_time
         self.process = None
-        self.termination_event = termination_event
+        self.termination_event = termination_event or multiprocessing.Event()
 
     def run(self):
-        while True:
-            self.process = multiprocessing.Process(target=self.target_func)
-            self.process.start()
-            self.process.join()  # Wait for process to finish
-
-            # Check exit code and restart if needed
-            if self.process.exitcode == 0:
-                print(f"Process {self.process.name} exited normally. Restarting...")
-            else:
-                print(f"Process {self.process.name} exited with code {self.process.exitcode}.")
+        start = time.time()
+        while not self.termination_event.is_set() and (self.live_time is None or time.time() - start < self.live_time):
+            try:
+                process = multiprocessing.Process(target=self.target, args=self.args)
+                process.start()
+                process.join()  # Wait for process to finish
+
+                # Check exit code, handle errors, and restart if needed
+                if process.exitcode == 0:
+                    print(f"Process {process.name} exited normally. Restarting...")
+                else:
+                    print(f"Process {process.name} exited with code {process.exitcode}.")
+            except Exception as e:
+                print(f"Error occurred while running process: {e}")
+
+            # Check for termination before restarting
+            if self.termination_event.is_set():
+                break
+
             time.sleep(self.live_time or 0)  # Restart after live_time or immediately if None
 
 def launch_dashboard(host: str, port: int, debug: bool = False):