diff --git a/.dockerignore b/.dockerignore index 50849a4..5d6aae9 100644 --- a/.dockerignore +++ b/.dockerignore @@ -16,3 +16,4 @@ mkdocs.yaml **/.mypy_cache **/__pycache__ **/tracing_database.json +**/tracing_database.json.lock diff --git a/.gitignore b/.gitignore index 837ce85..300754b 100644 --- a/.gitignore +++ b/.gitignore @@ -9,4 +9,5 @@ __pycache__ *.egg-info build tracing_database.json +tracing_database.json.lock .tox diff --git a/great_ai/persistence/parallel_tinydb_driver.py b/great_ai/persistence/parallel_tinydb_driver.py index 6e278a1..8a2896f 100644 --- a/great_ai/persistence/parallel_tinydb_driver.py +++ b/great_ai/persistence/parallel_tinydb_driver.py @@ -1,16 +1,15 @@ from datetime import datetime -from multiprocessing import Lock from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import pandas as pd +from filelock import FileLock from tinydb import TinyDB from ..views import Filter, SortBy, Trace from .tracing_database_driver import TracingDatabaseDriver DEFAULT_TRACING_DB_FILENAME = "tracing_database.json" -lock = Lock() operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"} @@ -22,7 +21,7 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): Saves the database as a JSON into a single file. Highly inefficient on inserting, not advised for production use. - A multiprocessing lock protects the database file to avoid parallelisation issues. + A file lock protects the database across threads and independently started processes. """ is_production_ready = False @@ -109,12 +108,10 @@ class ParallelTinyDbDriver(TracingDatabaseDriver): self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id)) def delete_batch(self, ids: List[str]) -> None: - with lock: - with TinyDB(self.path_to_db) as db: - for id in ids: - db.remove(lambda d: d["trace_id"] == id) + selected = set(ids) + self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] in selected)) def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any: - with lock: + with FileLock(str(self.path_to_db.resolve()) + ".lock"): with TinyDB(self.path_to_db) as db: return func(db) diff --git a/tests/test_database_regressions.py b/tests/test_database_regressions.py new file mode 100644 index 0000000..1e1167f --- /dev/null +++ b/tests/test_database_regressions.py @@ -0,0 +1,56 @@ +import multiprocessing as mp +import time +from pathlib import Path +from typing import Any +from unittest.mock import patch + +from tinydb.table import Table + +from great_ai import ParallelTinyDbDriver, Trace + + +def write_traces(path: str, index: int, barrier: Any) -> None: + ParallelTinyDbDriver.path_to_db = Path(path) + read_table = Table._read_table + + def slow_read(self: Table) -> Any: + data = read_table(self) + time.sleep(0.02) + return data + + barrier.wait(timeout=10) + with patch.object(Table, "_read_table", slow_read): + for i in range(3): + ParallelTinyDbDriver().save( + Trace( + trace_id=f"{index}-{i}", + created="2026-01-01", + original_execution_time_ms=0, + models=[], + tags=[], + logged_values={}, + ) + ) + + +def test_spawned_writers_share_database_lock(tmp_path: Path) -> None: + ctx = mp.get_context("spawn") + path = tmp_path / "traces.json" + barrier = ctx.Barrier(2) + processes = [ + ctx.Process(target=write_traces, args=(str(path), i, barrier)) for i in range(2) + ] + try: + for process in processes: + process.start() + for process in processes: + process.join(timeout=15) + assert process.exitcode == 0 + finally: + for process in processes: + if process.is_alive(): + process.kill() + process.join() + process.close() + with patch.object(ParallelTinyDbDriver, "path_to_db", path): + assert len(ParallelTinyDbDriver().query()[0]) == 6