From 62f0a060d5aab2160be1c95ac98abb6472a988ac Mon Sep 17 00:00:00 2001
From: Jizong's Fox <24821325+jizongFox@users.noreply.github.com>
Date: Tue, 30 May 2023 18:54:13 -0400
Subject: [PATCH] Update main.py

---
 gpu_queue/main.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/gpu_queue/main.py b/gpu_queue/main.py
index 19c1c17..73d3778 100644
--- a/gpu_queue/main.py
+++ b/gpu_queue/main.py
@@ -64,7 +64,7 @@ def get_args():
 class JobSubmitter:
 
     def __init__(self, job_array: Job_Array_Type, available_gpus: List[str] = ["0"], save_dir="log",
-                 verbose=False) -> None:
+                 verbose=False, wait_second:int=3) -> None:
         super().__init__()
         os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
         self.job_array = job_array
@@ -79,6 +79,7 @@ def __init__(self, job_array: Job_Array_Type, available_gpus: List[str] = ["0"],
             self.gpu_queue.put(gpu)
         print("%d jobs has been loaded" % len(self.job_array))
         self.result_dict = {}
+        self.wait_second = wait_second
 
     def submit_jobs(self):
         while True:
@@ -87,7 +88,7 @@ def submit_jobs(self):
                 job = self.job_queue.get(timeout=1)  # if it is going te be empty, end the program
                 gpu = self.gpu_queue.get(timeout=None, block=True)  # this will wait forever
                 self._process_daemeon(job, gpu)
-                time.sleep(2)
+                time.sleep(self.wait_second)
 
             except TimeoutError:
                 pass