Fix aiaccel/hpo/apps/optimize.py .

aiaccel/hpo/apps/optimize.py

@@ -147,19 +147,16 @@ 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.get_finished_job_count() < config.n_trials:
+        n_max_jobs = min(jobs.available_slots(), config.n_trials - jobs.get_submitted_job_count())
         for _ in range(n_max_jobs):
             trial = study.ask()
 
             hparams = params.suggest_hparams(trial)
 
             jobs.job_name = str(jobs.job_filename) + f"_{trial.number}"
 
-            job = jobs.submit(
+            job = jobs.submit_wrapper(
                 args=[result_filename_template] + sum([[f"--{k}", f"{v:.5f}"] for k, v in hparams.items()], []),
                 tag=trial,
             )
@@ -172,8 +169,6 @@ def main() -> None:
 
             study.tell(trial, y)
 
-            finished_job_count += 1
-
 
 if __name__ == "__main__":
     main()

aiaccel/hpo/job_executors/base_job_executor.py

@@ -41,6 +41,9 @@ def __init__(
 
         self.job_list: list[Any] = []
 
+        self.submitted_job_count = 0
+        self.finished_job_count = 0
+
     @abstractmethod
     def submit(
         self,
@@ -61,6 +64,15 @@ def submit(
         """
         pass
 
+    def submit_wrapper(
+        self,
+        args: list[str],
+        tag: Any = None,
+        sleep_time: float = 5.0,
+    ) -> Any:
+        self.submitted_job_count += 1
+        self.submit(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
 
@@ -98,3 +111,9 @@ def get_running_jobs(self) -> list[Any]:
             list[JobFuture]: A list of running jobs.
         """
         return [job for job in self.job_list if job.status == JobStatus.RUNNING]
+
+    def get_submitted_job_count(self) -> int:
+        return self.submitted_job_count
+
+    def get_finished_job_count(self) -> int:
+        return self.finished_job_count