Automatically trigger loop when possible

Author: basnijholt

adaptive_scheduler/_server_support/database_manager.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import pickle
 from dataclasses import asdict, dataclass, field
@@ -194,6 +195,7 @@ def __init__(
         self._pickling_time: float | None = None
         self._total_learner_size: int | None = None
         self._db: SimpleDatabase | None = None
+        self._trigger_event = asyncio.Event()
 
     def _setup(self) -> None:
         if self.db_fname.exists() and not self.overwrite_db:
@@ -214,7 +216,9 @@ def update(self, queue: dict[str, dict[str, str]] | None = None) -> None:
             queue = self.scheduler.queue(me_only=True)
         job_names_in_queue = [x["job_name"] for x in queue.values()]
         failed = self._db.get_all(
-            lambda e: e.job_name is not None and e.job_name not in job_names_in_queue,  # type: ignore[operator]
+            lambda e: not e.is_done
+            and e.job_name is not None
+            and e.job_name not in job_names_in_queue,  # type: ignore[operator]
         )
         self.failed.extend([asdict(entry) for _, entry in failed])
         indices = [index for index, _ in failed]
@@ -272,6 +276,12 @@ def _output_logs(self, job_id: str, job_name: str) -> list[Path]:
             for f in output_fnames
         ]
 
+    def _done_but_still_running(self) -> list[_DBEntry]:
+        if self._db is None:
+            return []
+        entries = self._db.get_all(lambda e: e.is_done and e.job_id is not None)
+        return [entry for _, entry in entries]
+
     def _choose_fname(self) -> tuple[int, str | list[str] | None]:
         assert self._db is not None
         entry = self._db.get(
@@ -338,16 +348,17 @@ def _start_request(
 
     def _stop_request(self, fname: str | list[str] | Path | list[Path]) -> None:
         fname_str = _ensure_str(fname)
-        reset = {"job_id": None, "is_done": True, "job_name": None, "is_pending": False}
+        reset = {"is_done": True, "is_pending": False}
         assert self._db is not None
         entry_indices = [index for index, _ in self._db.get_all(lambda e: e.fname == fname_str)]
         self._db.update(reset, entry_indices)
+        print(f"Done with {fname_str}")
 
     def _stop_requests(self, fnames: FnamesTypes) -> None:
         # Same as `_stop_request` but optimized for processing many `fnames` at once
         assert self._db is not None
         fnames_str = {str(fname) for fname in _ensure_str(fnames)}
-        reset = {"job_id": None, "is_done": True, "job_name": None, "is_pending": False}
+        reset = {"is_done": True, "is_pending": False}
         entry_indices = [
             index for index, _ in self._db.get_all(lambda e: str(e.fname) in fnames_str)
         ]
@@ -381,7 +392,8 @@ def _dispatch(
             if request_type == "stop":
                 fname = request_arg[0]  # workers send us the fname they were given
                 log.debug("got a stop request", fname=fname)
-                self._stop_request(fname)  # reset the job_id to None
+                self._stop_request(fname)  # set is_done
+                self.trigger()
                 return None
         except Exception as e:  # noqa: BLE001
             return e
@@ -426,3 +438,7 @@ async def _manage(self) -> None:
                     break
         finally:
             socket.close()
+
+    def trigger(self) -> None:
+        """External method to trigger the _manage loop to continue."""
+        self._trigger_event.set()

adaptive_scheduler/_server_support/job_manager.py

@@ -174,7 +174,6 @@ def __init__(
         # Other attributes
         self.n_started = 0
         self._request_times: dict[str, str] = {}
-        self._trigger_event = asyncio.Event()
 
         # Command line launcher options
         self.save_dataframe = save_dataframe
@@ -221,7 +220,9 @@ async def _update_database_and_get_not_queued(
         n_done = self.database_manager.n_done()
         if n_done == len(self.job_names):
             return None  # we are finished!
-        n_to_schedule = max(0, len(not_queued) - n_done)
+        n_done_but_running = len(self.database_manager._done_but_still_running())
+        n_to_schedule = max(0, n_done_but_running + len(not_queued) - n_done)
+        print(f"n_done={n_done}, {not_queued=}, {queued=}, {n_to_schedule=}")
         return queued, set(list(not_queued)[:n_to_schedule])
 
     async def _start_new_jobs(
@@ -233,6 +234,7 @@ async def _start_new_jobs(
             len(not_queued),
             self.max_simultaneous_jobs - len(queued),
         )
+        print(f"num_jobs_to_start={num_jobs_to_start}")
         for _ in range(num_jobs_to_start):
             index, fname = self.database_manager._choose_fname()
             if index == -1:
@@ -243,6 +245,7 @@ async def _start_new_jobs(
             log.debug(
                 f"Starting `job_name={job_name}` with `index={index}` and `fname={fname}`",
             )
+            print(f"Starting `job_name={job_name}` with `index={index}` and `fname={fname}`")
             await asyncio.to_thread(self.scheduler.start_job, job_name, index=index)
             self.database_manager._confirm_submitted(index, job_name)
 
@@ -263,7 +266,7 @@ async def _manage(self) -> None:
                 if await sleep_unless_task_is_done(
                     self.database_manager.task,  # type: ignore[arg-type]
                     self.interval,
-                    self._trigger_event,
+                    self.database_manager._trigger_event,
                 ):  # if true, we are done
                     return
             except asyncio.CancelledError:  # noqa: PERF203
@@ -283,10 +286,10 @@ async def _manage(self) -> None:
                 if await sleep_unless_task_is_done(
                     self.database_manager.task,  # type: ignore[arg-type]
                     5,
-                    self._trigger_event,
+                    self.database_manager._trigger_event,
                 ):  # if true, we are done
                     return
 
     def trigger(self) -> None:
         """External method to trigger the _manage loop to continue."""
-        self._trigger_event.set()
+        self.database_manager.trigger()