improved schedule

Author: arcuri82

experiments/schedule.py

@@ -12,11 +12,12 @@
 import time
 import platform
 import pathlib
+import psutil
 
 # Note: here we for flush on ALL prints, otherwise we would end up with messed up logs
 
-if len(sys.argv) != 3:
-    print("Usage:\nschedule.py <N> <FOLDER>", flush=True)
+if len(sys.argv) != 4:
+    print("Usage:\nschedule.py <N> <FOLDER> <TIMEOUT_MINUTES>", flush=True)
     exit(1)
 
 # The number of jobs to run in parallel
@@ -29,10 +30,13 @@
 # Location of experiment folder
 FOLDER = sys.argv[2]
 
+TIMEOUT_MINUTES = int(sys.argv[3])
+
+if TIMEOUT_MINUTES < 1:
+    print("Invalid value for TIMEOUT_MINUTES: " + str(TIMEOUT_MINUTES), flush=True)
+
 SHELL = platform.system() == 'Windows'
 
-# Changed, as now we might have few subfolders with scripts, eg for generated tests
-# SCRIPTS_FOLDER = os.path.join(FOLDER, "scripts")
 SCRIPTS_FOLDER = pathlib.PurePath(FOLDER).as_posix()
 
 def checkDocker():
@@ -80,8 +84,33 @@ def runScript(s):
     handler = subprocess.Popen(command, shell=SHELL, cwd=SCRIPTS_FOLDER, start_new_session=True)
     buffer.append(handler)
 
+def killProcess(h):
+    print("Terminating process.", flush=True)
+    parent = psutil.Process(h.pid)
+    children = parent.children(recursive=True)
+
+    # Graceful terminate
+    for p in children:
+        p.terminate()
+    parent.terminate()
+
+    gone, alive = psutil.wait_procs(children + [parent], timeout=10)
+
+    # Force kill remaining
+    for p in alive:
+        print(f"Force killing PID {p.pid}")
+        p.kill()
+
+    h.wait()
+
+
+########################################################################################################################
+
+last_start = time.time()
+
 for s in scripts:
     if len(buffer) < N:
+       last_start = time.time()
        runScript(s)
     else:
         while len(buffer) == N:
@@ -93,20 +122,36 @@ def runScript(s):
             # keep the ones running... those have return code not set yet
             buffer = [h for h in buffer if h.returncode is None]
             if len(buffer) == N :
+                # all running in buffer... but has any timeout?
+                # TODO for simplicity we just check latest added... so timeout is not enforced for ALL jobs.
+                # however, note that internally the jobs have their own timeouts... these here are just extra checks
+                elapsed_time = time.time() - last_start
+                if elapsed_time > TIMEOUT_MINUTES * 60:
+                    killProcess(buffer[0])
+                # wait before checking again
                 time.sleep(5)
             else:
+                last_start = time.time()
                 runScript(s)
                 break
 
 print("Waiting for last scripts to end", flush=True)
 
+budget = TIMEOUT_MINUTES * 60
+
 for h in buffer:
-    h.wait()
-    if h.returncode != 0:
-        print("Process terminated with code: " + str(h.returncode), flush=True)
+    start = time.time()
+    try:
+        h.wait(budget)
+        if h.returncode != 0:
+            print("Process terminated with code: " + str(h.returncode), flush=True)
+    except subprocess.TimeoutExpired:
+        print("Timeout reached.", flush=True)
+        killProcess(h)
+    elapsed = time.time() - start
+    budget = max(0, budget - elapsed)
 
 print("All jobs are completed", flush=True)
 
-#TODO how to make sure no subprocess is left hanging?