Serialize TinyDB access across independently spawned processes

This commit is contained in:
Andras Schmelczer 2026-10-27 19:51:10 +00:00
commit 9c02d2ecf7
4 changed files with 63 additions and 8 deletions

View file

@ -16,3 +16,4 @@ mkdocs.yaml
**/.mypy_cache **/.mypy_cache
**/__pycache__ **/__pycache__
**/tracing_database.json **/tracing_database.json
**/tracing_database.json.lock

1
.gitignore vendored
View file

@ -9,4 +9,5 @@ __pycache__
*.egg-info *.egg-info
build build
tracing_database.json tracing_database.json
tracing_database.json.lock
.tox .tox

View file

@ -1,16 +1,15 @@
from datetime import datetime from datetime import datetime
from multiprocessing import Lock
from pathlib import Path from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
import pandas as pd import pandas as pd
from filelock import FileLock
from tinydb import TinyDB from tinydb import TinyDB
from ..views import Filter, SortBy, Trace from ..views import Filter, SortBy, Trace
from .tracing_database_driver import TracingDatabaseDriver from .tracing_database_driver import TracingDatabaseDriver
DEFAULT_TRACING_DB_FILENAME = "tracing_database.json" DEFAULT_TRACING_DB_FILENAME = "tracing_database.json"
lock = Lock()
operator_mapping = {"=": "eq", "!=": "ne", "<": "lt", "<=": "le", ">": "gt", ">=": "ge"} 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, Saves the database as a JSON into a single file. Highly inefficient on inserting,
not advised for production use. 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 is_production_ready = False
@ -109,12 +108,10 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id)) self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] == id))
def delete_batch(self, ids: List[str]) -> None: def delete_batch(self, ids: List[str]) -> None:
with lock: selected = set(ids)
with TinyDB(self.path_to_db) as db: self._safe_execute(lambda db: db.remove(lambda d: d["trace_id"] in selected))
for id in ids:
db.remove(lambda d: d["trace_id"] == id)
def _safe_execute(self, func: Callable[[TinyDB], Any]) -> Any: 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: with TinyDB(self.path_to_db) as db:
return func(db) return func(db)

View file

@ -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