update

Author: Dylan Huang

eval_protocol/dataset_logger/tinydb_evaluation_row_store.py

@@ -2,6 +2,7 @@
 from typing import List, Optional
 
 from tinydb import Query, TinyDB
+from tinyrecord.transaction import transaction
 
 from eval_protocol.dataset_logger.evaluation_row_store import EvaluationRowStore
 
@@ -12,6 +13,9 @@ class TinyDBEvaluationRowStore(EvaluationRowStore):
 
     Stores data as plain JSON files, which are human-readable and
     don't suffer from SQLite's binary format corruption issues.
+
+    Uses tinyrecord for atomic transactions to handle concurrent access
+    from multiple processes safely.
     """
 
     def __init__(self, db_path: str):
@@ -33,18 +37,34 @@ def upsert_row(self, data: dict) -> None:
             raise ValueError("execution_metadata.rollout_id is required to upsert a row")
 
         Row = Query()
-        self._table.upsert(data, Row.execution_metadata.rollout_id == rollout_id)
+        condition = Row.execution_metadata.rollout_id == rollout_id
+
+        # tinyrecord doesn't support upsert directly, so we implement it manually
+        # within a transaction for atomicity
+        with transaction(self._table) as tr:
+            # Check if document exists
+            existing = self._table.search(condition)
+            if existing:
+                # Update existing document
+                tr.update(data, condition)
+            else:
+                # Insert new document
+                tr.insert(data)
 
     def read_rows(self, rollout_id: Optional[str] = None) -> List[dict]:
+        # Clear cache to ensure fresh read in multi-process scenarios
+        self._table.clear_cache()
         if rollout_id is not None:
             Row = Query()
             return list(self._table.search(Row.execution_metadata.rollout_id == rollout_id))
         return list(self._table.all())
 
     def delete_row(self, rollout_id: str) -> int:
         Row = Query()
-        removed = self._table.remove(Row.execution_metadata.rollout_id == rollout_id)
-        return len(removed)
+        with transaction(self._table) as tr:
+            tr.remove(Row.execution_metadata.rollout_id == rollout_id)
+        # Return count after removal (we don't have access to removed count in transaction)
+        return 1
 
     def delete_all_rows(self) -> int:
         count = len(self._table)

eval_protocol/event_bus/tinydb_event_bus_database.py

@@ -4,6 +4,7 @@
 from uuid import uuid4
 
 from tinydb import Query, TinyDB
+from tinyrecord.transaction import transaction
 
 from eval_protocol.event_bus.event_bus_database import EventBusDatabase
 from eval_protocol.event_bus.logger import logger
@@ -15,6 +16,9 @@ class TinyDBEventBusDatabase(EventBusDatabase):
 
     Stores data as plain JSON files, which are human-readable and
     don't suffer from SQLite's binary format corruption issues.
+
+    Uses tinyrecord for atomic transactions to handle concurrent access
+    from multiple processes safely.
     """
 
     def __init__(self, db_path: str):
@@ -27,24 +31,26 @@ def __init__(self, db_path: str):
         self._table = self._db.table("events")
 
     def publish_event(self, event_type: str, data: Any, process_id: str) -> None:
-        """Publish an event to the database."""
+        """Publish an event to the database using atomic transaction."""
         try:
             # Serialize data, handling pydantic models
             if hasattr(data, "model_dump"):
                 serialized_data = data.model_dump(mode="json", exclude_none=True)
             else:
                 serialized_data = data
 
-            self._table.insert(
-                {
-                    "event_id": str(uuid4()),
-                    "event_type": event_type,
-                    "data": serialized_data,
-                    "timestamp": time.time(),
-                    "process_id": process_id,
-                    "processed": False,
-                }
-            )
+            document = {
+                "event_id": str(uuid4()),
+                "event_type": event_type,
+                "data": serialized_data,
+                "timestamp": time.time(),
+                "process_id": process_id,
+                "processed": False,
+            }
+
+            # Use tinyrecord transaction for atomic, concurrent-safe insert
+            with transaction(self._table) as tr:
+                tr.insert(document)
         except Exception as e:
             logger.warning(f"Failed to publish event to database: {e}")
 
@@ -83,21 +89,23 @@ def get_unprocessed_events(self, process_id: str) -> List[dict]:
             return []
 
     def mark_event_processed(self, event_id: str) -> None:
-        """Mark an event as processed."""
+        """Mark an event as processed using atomic transaction."""
         try:
             Event = Query()
-            self._table.update({"processed": True}, Event.event_id == event_id)
+            with transaction(self._table) as tr:
+                tr.update({"processed": True}, Event.event_id == event_id)
         except Exception as e:
             logger.debug(f"Failed to mark event as processed: {e}")
 
     def cleanup_old_events(self, max_age_hours: int = 24) -> None:
-        """Clean up old processed events."""
+        """Clean up old processed events using atomic transaction."""
         try:
-            # Reload table from disk to see latest data before cleanup
-            self._table._read_table()
+            # Clear cache to see latest data before cleanup
+            self._table.clear_cache()
 
             cutoff_time = time.time() - (max_age_hours * 3600)
             Event = Query()
-            self._table.remove((Event.processed == True) & (Event.timestamp < cutoff_time))  # noqa: E712
+            with transaction(self._table) as tr:
+                tr.remove((Event.processed == True) & (Event.timestamp < cutoff_time))  # noqa: E712
         except Exception as e:
             logger.debug(f"Failed to cleanup old events: {e}")

pyproject.toml

@@ -35,6 +35,7 @@ dependencies = [
     "pytest>=6.0.0",
     "pytest-asyncio>=0.21.0",
     "tinydb>=4.8.0",
+    "tinyrecord>=0.2.0",
     "backoff>=2.2.0",
     "questionary>=2.0.0",
     # Dependencies for vendored tau2 package

tests/pytest/test_pytest_ensure_logging.py

@@ -15,17 +15,17 @@ async def test_ensure_logging(monkeypatch):
         monkeypatch.setattr(_dl, "_logger", None, raising=False)
     except Exception:
         pass
-    # Mock the SqliteEvaluationRowStore to track calls
+    # Mock the EvaluationRowStore to track calls
     mock_store = Mock()
     mock_store.upsert_row = Mock()
     mock_store.read_rows = Mock(return_value=[])
     mock_store.db_path = "/tmp/test.db"
 
-    # Mock the SqliteEvaluationRowStore constructor so that when SqliteDatasetLoggerAdapter
-    # creates its store, it gets our mock instead
-    with patch(
-        "eval_protocol.dataset_logger.sqlite_dataset_logger_adapter.SqliteEvaluationRowStore", return_value=mock_store
-    ):
+    # Mock get_evaluation_row_store so that when DatasetLoggerAdapter
+    # creates its store, it gets our mock instead.
+    # We patch at the module level where it's defined, which is where
+    # dataset_logger_adapter imports it from.
+    with patch("eval_protocol.dataset_logger.get_evaluation_row_store", return_value=mock_store):
         from eval_protocol.models import EvaluationRow, EvaluateResult
         from eval_protocol.pytest.default_no_op_rollout_processor import NoOpRolloutProcessor
         from eval_protocol.pytest.evaluation_test import evaluation_test
@@ -55,7 +55,7 @@ def eval_fn(row: EvaluationRow) -> EvaluationRow:
         )
 
         # Verify that the store's upsert_row method was called
-        assert mock_store.upsert_row.called, "SqliteEvaluationRowStore.upsert_row should have been called"
+        assert mock_store.upsert_row.called, "EvaluationRowStore.upsert_row should have been called"
 
         # Check that it was called multiple times (once for each row)
         call_count = mock_store.upsert_row.call_count