Implement SimpleDatabase and remove tinydb dependency (#161)

* Implement SimpleDatabase and remove tiny dependency

* add metadata to JSON

* add clear existing

* tests

* Fix types

* fix str

* remove unused classes

* Fix database

* Remove maybe_list

Author: basnijholt

adaptive_scheduler/_mock_scheduler.py

@@ -6,7 +6,7 @@
 import logging
 import os
 import subprocess
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Tuple, Union
 
 import structlog
 import zmq
@@ -25,6 +25,10 @@
 
 DEFAULT_URL = "tcp://127.0.0.1:60547"
 
+_RequestSubmitType = Tuple[str, str, Union[str, List[str]]]
+_RequestCancelType = Tuple[str, str]
+_RequestQueueType = Tuple[str]
+
 
 class MockScheduler:
     """Emulates a HPC-like scheduler.
@@ -156,18 +160,18 @@ async def _command_listener(self) -> Coroutine[None, None, None]:
 
     def _dispatch(
         self,
-        request: tuple[str, ...],
+        request: _RequestSubmitType | _RequestCancelType | _RequestQueueType,
     ) -> str | None | dict[str, dict[str, Any]] | Exception:
         log.debug("got a request", request=request)
         request_type, *request_arg = request
         try:
             if request_type == "submit":
                 job_name, fname = request_arg
                 log.debug("submitting a task", fname=fname, job_name=job_name)
-                job_id = self.submit(job_name, fname)
+                job_id = self.submit(job_name, fname)  # type: ignore[arg-type]
                 return job_id
             if request_type == "cancel":
-                job_id = request_arg[0]
+                job_id = request_arg[0]  # type: ignore[assignment]
                 log.debug("got a cancel request", job_id=job_id)
                 self.cancel(job_id)
                 return None

adaptive_scheduler/_server_support/database_manager.py

@@ -1,15 +1,16 @@
 """The DatabaseManager."""
 from __future__ import annotations
 
+import json
 import pickle
+from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, List, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Union
 
 import pandas as pd
 import zmq
 import zmq.asyncio
 import zmq.ssh
-from tinydb import Query, TinyDB
 
 from adaptive_scheduler.utils import (
     _deserialize,
@@ -26,8 +27,9 @@
 
     from adaptive_scheduler.scheduler import BaseScheduler
 
-ctx = zmq.asyncio.Context()
 
+ctx = zmq.asyncio.Context()
+FnameType = Union[str, Path, List[str], List[Path]]
 FnamesTypes = Union[List[str], List[Path], List[List[str]], List[List[Path]]]
 
 
@@ -56,6 +58,72 @@ def _ensure_str(
     raise ValueError(msg)
 
 
+@dataclass
+class _DBEntry:
+    fname: str | list[str]
+    job_id: str | None = None
+    is_done: bool = False
+    log_fname: str | None = None
+    job_name: str | None = None
+    output_logs: list[str] = field(default_factory=list)
+    start_time: float | None = None
+
+
+class SimpleDatabase:
+    def __init__(self, db_fname: str | Path, *, clear_existing: bool = False) -> None:
+        self.db_fname = Path(db_fname)
+        self._data: list[_DBEntry] = []
+        self._meta: dict[str, Any] = {}
+
+        if self.db_fname.exists():
+            if clear_existing:
+                self.db_fname.unlink()
+            else:
+                with self.db_fname.open() as f:
+                    raw_data = json.load(f)
+                    self._data = [_DBEntry(**entry) for entry in raw_data["data"]]
+
+    def all(self) -> list[_DBEntry]:  # noqa: A003
+        return self._data
+
+    def insert_multiple(self, entries: list[_DBEntry]) -> None:
+        self._data.extend(entries)
+        self._save()
+
+    def update(self, update_dict: dict, indices: list[int] | None = None) -> None:
+        for index, entry in enumerate(self._data):
+            if indices is None or index in indices:
+                for key, value in update_dict.items():
+                    assert hasattr(entry, key)
+                    setattr(entry, key, value)
+        self._save()
+
+    def count(self, condition: Callable[[_DBEntry], bool]) -> int:
+        return sum(1 for entry in self._data if condition(entry))
+
+    def get(self, condition: Callable[[_DBEntry], bool]) -> _DBEntry | None:
+        for entry in self._data:
+            if condition(entry):
+                return entry
+        return None
+
+    def get_all(
+        self,
+        condition: Callable[[_DBEntry], bool],
+    ) -> list[tuple[int, _DBEntry]]:
+        return [(i, entry) for i, entry in enumerate(self._data) if condition(entry)]
+
+    def contains(self, condition: Callable[[_DBEntry], bool]) -> bool:
+        return any(condition(entry) for entry in self._data)
+
+    def as_dicts(self) -> list[dict[str, Any]]:
+        return [asdict(entry) for entry in self._data]
+
+    def _save(self) -> None:
+        with self.db_fname.open("w") as f:
+            json.dump({"data": self.as_dicts(), "meta": self._meta}, f)
+
+
 class DatabaseManager(BaseManager):
     """Database manager.
 
@@ -100,20 +168,12 @@ def __init__(  # noqa: PLR0913
         self.fnames = fnames
         self.overwrite_db = overwrite_db
 
-        self.defaults: dict[str, Any] = {
-            "job_id": None,
-            "is_done": False,
-            "log_fname": None,
-            "job_name": None,
-            "output_logs": [],
-            "start_time": None,
-        }
-
-        self._last_reply: str | Exception | None = None
+        self._last_reply: str | list[str] | Exception | None = None
         self._last_request: tuple[str, ...] | None = None
         self.failed: list[dict[str, Any]] = []
         self._pickling_time: float | None = None
         self._total_learner_size: int | None = None
+        self._db: SimpleDatabase | None = None
 
     def _setup(self) -> None:
         if self.db_fname.exists() and not self.overwrite_db:
@@ -127,24 +187,21 @@ def _setup(self) -> None:
 
     def update(self, queue: dict[str, dict[str, str]] | None = None) -> None:
         """If the ``job_id`` isn't running anymore, replace it with None."""
+        assert self._db is not None
         if queue is None:
             queue = self.scheduler.queue(me_only=True)
-
-        with TinyDB(self.db_fname) as db:
-            failed = [
-                entry
-                for entry in db.all()
-                if (entry["job_id"] is not None) and (entry["job_id"] not in queue)
-            ]
-            self.failed.extend(failed)
-            doc_ids = [e.doc_id for e in failed]
-            db.update({"job_id": None, "job_name": None}, doc_ids=doc_ids)
+        failed = self._db.get_all(
+            lambda e: (e.job_id is not None) and (e.job_id not in queue),  # type: ignore[operator]
+        )
+        self.failed.extend([asdict(entry) for _, entry in failed])
+        indices = [index for index, _ in failed]
+        self._db.update({"job_id": None, "job_name": None}, indices)
 
     def n_done(self) -> int:
         """Return the number of jobs that are done."""
-        entry = Query()
-        with TinyDB(self.db_fname) as db:
-            return db.count(entry.is_done == True)  # noqa: E712
+        if self._db is None:
+            return 0
+        return self._db.count(lambda e: e.is_done)
 
     def is_done(self) -> bool:
         """Return True if all jobs are done."""
@@ -155,18 +212,18 @@ def create_empty_db(self) -> None:
 
         It keeps track of ``fname -> (job_id, is_done, log_fname, job_name)``.
         """
-        entries = [
-            dict(fname=_ensure_str(fname), **self.defaults) for fname in self.fnames
+        entries: list[_DBEntry] = [
+            _DBEntry(fname=fname) for fname in _ensure_str(self.fnames)
         ]
         if self.db_fname.exists():
             self.db_fname.unlink()
-        with TinyDB(self.db_fname) as db:
-            db.insert_multiple(entries)
+        self._db = SimpleDatabase(self.db_fname)
+        self._db.insert_multiple(entries)
 
     def as_dicts(self) -> list[dict[str, str]]:
         """Return the database as a list of dictionaries."""
-        with TinyDB(self.db_fname) as db:
-            return db.all()
+        assert self._db is not None
+        return self._db.as_dicts()
 
     def as_df(self) -> pd.DataFrame:
         """Return the database as a `pandas.DataFrame`."""
@@ -180,75 +237,86 @@ def _output_logs(self, job_id: str, job_name: str) -> list[Path]:
             for f in output_fnames
         ]
 
-    def _start_request(self, job_id: str, log_fname: str, job_name: str) -> str | None:
-        entry = Query()
-        with TinyDB(self.db_fname) as db:
-            if db.contains(entry.job_id == job_id):
-                entry = db.get(entry.job_id == job_id)
-                fname = entry["fname"]  # already running
-                msg = (
-                    f"The job_id {job_id} already exists in the database and "
-                    f"runs {fname}. You might have forgotten to use the "
-                    "`if __name__ == '__main__': ...` idom in your code. Read the "
-                    "warning in the [mpi4py](https://bit.ly/2HAk0GG) documentation.",
-                )
-                raise JobIDExistsInDbError(msg)
-            entry = db.get(
-                (entry.job_id == None) & (entry.is_done == False),  # noqa: E711,E712
-            )
-            log.debug("choose fname", entry=entry)
-            if entry is None:
-                return None
-            db.update(
-                {
-                    "job_id": job_id,
-                    "log_fname": log_fname,
-                    "job_name": job_name,
-                    "output_logs": _ensure_str(self._output_logs(job_id, job_name)),
-                    "start_time": _now(),
-                },
-                doc_ids=[entry.doc_id],
+    def _start_request(
+        self,
+        job_id: str,
+        log_fname: str,
+        job_name: str,
+    ) -> str | list[str] | None:
+        assert self._db is not None
+        if self._db.contains(lambda e: e.job_id == job_id):
+            entry = self._db.get(lambda e: e.job_id == job_id)
+            assert entry is not None
+            fname = entry.fname  # already running
+            msg = (
+                f"The job_id {job_id} already exists in the database and "
+                f"runs {fname}. You might have forgotten to use the "
+                "`if __name__ == '__main__': ...` idiom in your code. Read the "
+                "warning in the [mpi4py](https://bit.ly/2HAk0GG) documentation.",
             )
-        return entry["fname"]
+            raise JobIDExistsInDbError(msg)
+        entry = self._db.get(
+            lambda e: e.job_id is None and not e.is_done,
+        )
+        log.debug("choose fname", entry=entry)
+        if entry is None:
+            return None
+        index = self._db.all().index(entry)
+        self._db.update(
+            {
+                "job_id": job_id,
+                "log_fname": log_fname,
+                "job_name": job_name,
+                "output_logs": _ensure_str(self._output_logs(job_id, job_name)),
+                "start_time": _now(),
+            },
+            indices=[index],
+        )
+        return _ensure_str(entry.fname)  # type: ignore[return-value]
 
     def _stop_request(self, fname: str | list[str] | Path | list[Path]) -> None:
         fname_str = _ensure_str(fname)
-        entry = Query()
-        with TinyDB(self.db_fname) as db:
-            reset = {"job_id": None, "is_done": True, "job_name": None}
-            assert (
-                db.get(entry.fname == fname_str) is not None
-            )  # make sure the entry exists
-            db.update(reset, entry.fname == fname_str)
+        reset = {"job_id": None, "is_done": True, "job_name": None}
+        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)
 
     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)}
-        with TinyDB(self.db_fname) as db:
-            reset = {"job_id": None, "is_done": True, "job_name": None}
-            doc_ids = [e.doc_id for e in db.all() if str(e["fname"]) in fnames_str]
-            db.update(reset, doc_ids=doc_ids)
+        reset = {"job_id": None, "is_done": True, "job_name": None}
+        entry_indices = [
+            index for index, _ in self._db.get_all(lambda e: str(e.fname) in fnames_str)
+        ]
+        self._db.update(reset, entry_indices)
 
-    def _dispatch(self, request: tuple[str, ...]) -> str | Exception | None:
+    def _dispatch(
+        self,
+        request: tuple[str, str | list[str]] | tuple[str],
+    ) -> str | list[str] | Exception | None:
         request_type, *request_arg = request
         log.debug("got a request", request=request)
         try:
             if request_type == "start":
                 # workers send us their slurm ID for us to fill in
                 job_id, log_fname, job_name = request_arg
-                kwargs = {
-                    "job_id": job_id,
-                    "log_fname": log_fname,
-                    "job_name": job_name,
-                }
                 # give the worker a job and send back the fname to the worker
-                fname = self._start_request(**kwargs)
+                fname = self._start_request(job_id, log_fname, job_name)  # type: ignore[arg-type]
                 if fname is None:
                     # This should never happen because the _manage co-routine
                     # should have stopped the workers before this happens.
                     msg = "No more learners to run in the database."
                     raise RuntimeError(msg)  # noqa: TRY301
-                log.debug("choose a fname", fname=fname, **kwargs)
+                log.debug(
+                    "choose a fname",
+                    fname=fname,
+                    job_id=job_id,
+                    log_fname=log_fname,
+                    job_name=job_name,
+                )
                 return fname
             if request_type == "stop":
                 fname = request_arg[0]  # workers send us the fname they were given
@@ -291,7 +359,7 @@ async def _manage(self) -> None:
                         )
                 else:
                     assert self._last_request is not None  # for mypy
-                    self._last_reply = self._dispatch(self._last_request)
+                    self._last_reply = self._dispatch(self._last_request)  # type: ignore[arg-type]
                     await socket.send_serialized(self._last_reply, _serialize)
                 if self.is_done():
                     break

adaptive_scheduler/client_support.py

@@ -17,7 +17,6 @@
     _serialize,
     fname_to_learner,
     log_exception,
-    maybe_lst,
     sleep_unless_task_is_done,
 )
 
@@ -106,7 +105,7 @@ def get_learner(
         log.info("got fname and loaded learner")
 
     log.info("picked a learner")
-    return learner, maybe_lst(fname)
+    return learner, fname
 
 
 def tell_done(url: str, fname: str | list[str]) -> None: