Skip to content

Commit 11221fb

Browse files
author
Dylan Huang
committed
Enhance delete_row method in TinyDBEvaluationRowStore to return accurate removal count and clear cache for fresh reads
1 parent 46da8d3 commit 11221fb

File tree

1 file changed

+10
-3
lines changed

1 file changed

+10
-3
lines changed

eval_protocol/dataset_logger/tinydb_evaluation_row_store.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -63,10 +63,17 @@ def read_rows(self, rollout_id: Optional[str] = None) -> List[dict]:
6363

6464
def delete_row(self, rollout_id: str) -> int:
6565
Row = Query()
66+
condition = Row.execution_metadata.rollout_id == rollout_id
67+
6668
with transaction(self._table) as tr:
67-
tr.remove(Row.execution_metadata.rollout_id == rollout_id)
68-
# Return count after removal (we don't have access to removed count in transaction)
69-
return 1
69+
# Clear cache to ensure fresh read in multi-process scenarios
70+
self._table.clear_cache()
71+
# Check if document exists before removal to get accurate count
72+
existing = self._table.search(condition)
73+
if existing:
74+
tr.remove(condition)
75+
return len(existing)
76+
return 0
7077

7178
def delete_all_rows(self) -> int:
7279
count = len(self._table)

0 commit comments

Comments
 (0)