Fix aiaccel/hpo/apps/optimize.py .

aiaccel/hpo/apps/optimize.py

@@ -147,11 +147,8 @@ def main() -> None:
 
     result_filename_template = "{job.cwd}/{job.job_name}_result.pkl"
 
-    finished_job_count = 0
-
-    while finished_job_count < config.n_trials:
-        n_running_jobs = len(jobs.get_running_jobs())
-        n_max_jobs = min(jobs.available_slots(), config.n_trials - finished_job_count - n_running_jobs)
+    while jobs.finished_job_count < config.n_trials:
+        n_max_jobs = min(jobs.available_slots(), config.n_trials - jobs.submitted_job_count)
         for _ in range(n_max_jobs):
             trial = study.ask()
 
@@ -172,8 +169,6 @@ def main() -> None:
 
             study.tell(trial, y)
 
-            finished_job_count += 1
-
 
 if __name__ == "__main__":
     main()

aiaccel/hpo/job_executors/abci_job_executor.py

@@ -49,7 +49,7 @@ def __init__(
         self.job_group = job_group
         self.job_list: list[AbciJob] = []
 
-    def submit(
+    def submit_impl(
         self,
         args: list[str],
         tag: Any = None,

aiaccel/hpo/job_executors/base_job_executor.py

@@ -41,8 +41,11 @@ def __init__(
 
         self.job_list: list[Any] = []
 
+        self.submitted_job_count = 0
+        self.finished_job_count = 0
+
     @abstractmethod
-    def submit(
+    def submit_impl(
         self,
         args: list[str],
         tag: Any = None,
@@ -61,6 +64,15 @@ def submit(
         """
         pass
 
+    def submit(
+        self,
+        args: list[str],
+        tag: Any = None,
+        sleep_time: float = 5.0,
+    ) -> Any:
+        self.submitted_job_count += 1
+        return self.submit_impl(args, tag, sleep_time)
+
     def update_status_batch(self) -> None:
         """
         Updates the status of a batch of jobs.
@@ -87,6 +99,7 @@ def collect_finished(self) -> list[Any]:
         finished_jobs = [job for job in self.job_list if job.status >= JobStatus.FINISHED]
         for job in finished_jobs:
             self.job_list.remove(job)
+            self.finished_job_count += 1
 
         return finished_jobs
 

aiaccel/hpo/job_executors/local_job_executor.py

@@ -40,7 +40,7 @@ def __init__(
 
         self.job_list: list[LocalJob] = []
 
-    def submit(
+    def submit_impl(
         self,
         args: list[str],
         tag: Any = None,