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
**/__pycache__
**/tracing_database.json
**/tracing_database.json.lock

1
.gitignore vendored
View file

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

View file

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

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