mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-28 14:10:37 +00:00
fix(runtime): preserve per-run files in JSONL batch writes (#4938)
* fix(runtime): preserve per-run JSONL batch files * fix(events): clarify JSONL rollback failures
This commit is contained in:
parent
e5bf3ccf45
commit
c88c2975c9
@ -157,13 +157,17 @@ class JsonlRunEventStore(RunEventStore):
|
||||
return record
|
||||
|
||||
async def put_batch(self, events):
|
||||
"""Persist a batch of events atomically per-thread.
|
||||
"""Persist a batch of events under a per-thread write lock.
|
||||
|
||||
All seq numbers for the batch are reserved under a single per-thread
|
||||
write lock and every record is appended in one file write so a
|
||||
mid-batch failure cannot leave a partial set of records on disk that
|
||||
a retry would then duplicate. Callers (e.g. worker.py's flush-retry
|
||||
path) may safely re-buffer the entire batch on failure.
|
||||
write lock. Records are grouped by run_id and appended to their own
|
||||
run files while that lock is held. If a write fails and rollback
|
||||
succeeds, already-appended groups for the current thread are restored
|
||||
so callers (e.g. worker.py's flush-retry path) may safely re-buffer
|
||||
that thread's batch. When a batch contains multiple thread IDs, thread
|
||||
groups are processed sequentially, so a later failure does not roll
|
||||
back earlier thread groups. This rollback does not make a multi-file
|
||||
batch crash-atomic.
|
||||
"""
|
||||
if not events:
|
||||
return []
|
||||
@ -226,10 +230,11 @@ class JsonlRunEventStore(RunEventStore):
|
||||
"created_at": ev.get("created_at") or datetime.now(UTC).isoformat(),
|
||||
}
|
||||
records.append(record)
|
||||
path = self._run_file(thread_id, batch[0]["run_id"])
|
||||
# Single append/write per thread. If this raises, no records were
|
||||
# persisted; the caller's re-buffer reproduces no duplicates.
|
||||
await asyncio.to_thread(self._append_records, path, records)
|
||||
records_by_run: dict[str, list[dict[str, Any]]] = {}
|
||||
for record in records:
|
||||
records_by_run.setdefault(record["run_id"], []).append(record)
|
||||
run_batches = [(self._run_file(thread_id, run_id), run_records) for run_id, run_records in records_by_run.items()]
|
||||
await asyncio.to_thread(self._append_record_groups, run_batches)
|
||||
return records
|
||||
|
||||
def _append_records(self, path: Path, records: list[dict[str, Any]]) -> None:
|
||||
@ -238,6 +243,30 @@ class JsonlRunEventStore(RunEventStore):
|
||||
with open(path, "a", encoding="utf-8") as f:
|
||||
f.write(lines)
|
||||
|
||||
def _append_record_groups(self, groups: list[tuple[Path, list[dict[str, Any]]]]) -> None:
|
||||
"""Append run groups and restore their original sizes if one fails."""
|
||||
original_sizes: dict[Path, int | None] = {}
|
||||
try:
|
||||
for path, records in groups:
|
||||
original_sizes[path] = path.stat().st_size if path.exists() else None
|
||||
self._append_records(path, records)
|
||||
except Exception:
|
||||
for path, original_size in original_sizes.items():
|
||||
try:
|
||||
if original_size is None:
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
else:
|
||||
with open(path, "r+b") as f:
|
||||
f.truncate(original_size)
|
||||
except OSError:
|
||||
logger.error(
|
||||
"Failed to roll back JSONL batch append for %s; retrying the batch may create duplicate records",
|
||||
path,
|
||||
exc_info=True,
|
||||
)
|
||||
raise
|
||||
|
||||
async def list_messages(self, thread_id, *, limit=50, before_seq=None, after_seq=None, user_id: str | None | _AutoSentinel = AUTO):
|
||||
all_events = await asyncio.to_thread(self._read_thread_events, thread_id)
|
||||
messages = [e for e in all_events if e.get("category") == "message"]
|
||||
|
||||
@ -10,6 +10,7 @@ Verifies:
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
@ -88,7 +89,7 @@ async def test_concurrent_puts_different_threads_independent_seqs():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# put_batch: delegates to put() and preserves order
|
||||
# put_batch: assigns monotonic seqs and preserves per-run files
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@ -103,6 +104,24 @@ async def test_put_batch_seqs_are_monotonic():
|
||||
assert len(set(seqs)) == 5
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_put_batch_writes_mixed_run_ids_to_their_run_files():
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = _make_store(Path(tmp))
|
||||
events = [
|
||||
{"thread_id": "t1", "run_id": "r1", "event_type": "human_message", "category": "message", "content": "r1-first"},
|
||||
{"thread_id": "t1", "run_id": "r2", "event_type": "human_message", "category": "message", "content": "r2-only"},
|
||||
{"thread_id": "t1", "run_id": "r1", "event_type": "ai_message", "category": "message", "content": "r1-last"},
|
||||
]
|
||||
records = await store.put_batch(events)
|
||||
r1_messages = await store.list_messages_by_run("t1", "r1")
|
||||
r2_messages = await store.list_messages_by_run("t1", "r2")
|
||||
|
||||
assert [(record["run_id"], record["seq"]) for record in records] == [("r1", 1), ("r2", 2), ("r1", 3)]
|
||||
assert [message["content"] for message in r1_messages] == ["r1-first", "r1-last"]
|
||||
assert [message["content"] for message in r2_messages] == ["r2-only"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _ensure_seq_loaded: recovers max_seq from disk after fresh store init
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -148,23 +167,15 @@ async def test_put_offloads_write_via_to_thread():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# put_batch atomicity: a failed append must not leave partial records so a
|
||||
# caller re-buffering the batch on retry does not produce duplicates.
|
||||
# put_batch failure rollback: a failed append must not leave partial records
|
||||
# so a caller re-buffering the batch on retry does not produce duplicates.
|
||||
# Regression for deer-flow PR #4082 (review feedback from willem-bd).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_put_batch_failure_rolls_back_no_partial_records(monkeypatch):
|
||||
"""If the disk write inside ``put_batch`` raises after partial output,
|
||||
no records should remain on disk because the seq counter is reserved
|
||||
under the write lock but the seqs were NOT written. A subsequent retry
|
||||
therefore reproduces no duplicates.
|
||||
|
||||
Concretely: the implementation uses a single ``open().write()`` so on
|
||||
failure the file is either empty or has the prior batch's records —
|
||||
never a partial slice of the new batch.
|
||||
"""
|
||||
"""A failed append is rolled back before the re-buffered batch is retried."""
|
||||
import json
|
||||
|
||||
from deerflow.runtime.events.store import jsonl as jsonl_mod
|
||||
@ -194,26 +205,109 @@ async def test_put_batch_failure_rolls_back_no_partial_records(monkeypatch):
|
||||
}
|
||||
for i in range(4)
|
||||
]
|
||||
# First attempt — fails mid-batch; expect raise; the file may have
|
||||
# partial lines but the in-memory seq counter has been advanced
|
||||
# (because seq reservation happened under the lock).
|
||||
# First attempt — fails after partial output; expect raise. The
|
||||
# in-memory seq counter is advanced because reservation happens under
|
||||
# the lock, but the partial file contents must be rolled back.
|
||||
with pytest.raises(OSError):
|
||||
await store.put_batch(events)
|
||||
|
||||
# Now retry with the real append (no failure): only the unreserved
|
||||
# records will be written — but our implementation appends the whole
|
||||
# batch again, so what we really verify here is that after a failure
|
||||
# the seq counter is monotonic and consistent with the recovered
|
||||
# disk state (no half-batch leftover gets accidentally re-numbered).
|
||||
# Retry the full batch with the real append, matching worker.py's
|
||||
# re-buffer path. Verify persisted contents below, not only return
|
||||
# values, so a duplicate or partial disk write cannot pass unnoticed.
|
||||
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", real_append)
|
||||
# Retry the full batch — the re-buffer pattern from worker.py.
|
||||
records = await store.put_batch(events)
|
||||
persisted_events = await store.list_events("t1", "r1")
|
||||
|
||||
# The batch succeeded on retry, every event ended up exactly once in the
|
||||
# file (no duplicates), and seqs are still strictly monotonic.
|
||||
# The batch succeeded on retry, every event ended up exactly once on disk,
|
||||
# and seqs are still strictly monotonic.
|
||||
assert len(records) == 4, f"Expected 4 records, got {len(records)}"
|
||||
seqs = [r["seq"] for r in records]
|
||||
assert seqs == sorted(seqs) and len(set(seqs)) == 4, f"seqs not unique monotonic: {seqs}"
|
||||
assert len(persisted_events) == 4
|
||||
assert [event["content"] for event in persisted_events] == [f"event-{i}" for i in range(4)]
|
||||
assert [event["seq"] for event in persisted_events] == seqs
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_mixed_run_batch_failure_restores_all_run_files(monkeypatch):
|
||||
"""A failed mixed-run append restores prior bytes in every touched file."""
|
||||
from deerflow.runtime.events.store import jsonl as jsonl_mod
|
||||
|
||||
real_append = jsonl_mod.JsonlRunEventStore._append_records
|
||||
append_calls = 0
|
||||
|
||||
def failing_append(self, path, records):
|
||||
nonlocal append_calls
|
||||
append_calls += 1
|
||||
if append_calls == 2:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(path, "a", encoding="utf-8") as f:
|
||||
f.write("partial\n")
|
||||
raise OSError("simulated second-run write failure")
|
||||
real_append(self, path, records)
|
||||
|
||||
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", failing_append)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = _make_store(Path(tmp))
|
||||
await store.put(thread_id="t1", run_id="r1", event_type="trace", category="trace", content="existing")
|
||||
r1_path = store._run_file("t1", "r1")
|
||||
r2_path = store._run_file("t1", "r2")
|
||||
original_r1 = r1_path.read_bytes()
|
||||
|
||||
events = [
|
||||
{"thread_id": "t1", "run_id": "r1", "event_type": "trace", "category": "trace", "content": "new-r1"},
|
||||
{"thread_id": "t1", "run_id": "r2", "event_type": "trace", "category": "trace", "content": "new-r2"},
|
||||
]
|
||||
with pytest.raises(OSError, match="second-run"):
|
||||
await store.put_batch(events)
|
||||
|
||||
assert r1_path.read_bytes() == original_r1
|
||||
assert not r2_path.exists()
|
||||
assert [event["content"] for event in await store.list_events("t1", "r1")] == ["existing"]
|
||||
assert await store.list_events("t1", "r2") == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_mixed_run_batch_logs_error_when_rollback_fails(monkeypatch, caplog):
|
||||
"""A rollback failure must make possible retry duplicates visible to operators."""
|
||||
from deerflow.runtime.events.store import jsonl as jsonl_mod
|
||||
|
||||
real_append = jsonl_mod.JsonlRunEventStore._append_records
|
||||
real_unlink = Path.unlink
|
||||
append_calls = 0
|
||||
|
||||
def failing_append(self, path, records):
|
||||
nonlocal append_calls
|
||||
append_calls += 1
|
||||
if append_calls == 2:
|
||||
real_append(self, path, records)
|
||||
raise OSError("simulated second-run write failure")
|
||||
real_append(self, path, records)
|
||||
|
||||
monkeypatch.setattr(jsonl_mod.JsonlRunEventStore, "_append_records", failing_append)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = _make_store(Path(tmp))
|
||||
r2_path = store._run_file("t1", "r2")
|
||||
|
||||
def failing_unlink(path, missing_ok=False):
|
||||
if path == r2_path:
|
||||
raise OSError("simulated rollback failure")
|
||||
return real_unlink(path, missing_ok=missing_ok)
|
||||
|
||||
monkeypatch.setattr(Path, "unlink", failing_unlink)
|
||||
events = [
|
||||
{"thread_id": "t1", "run_id": "r1", "event_type": "trace", "category": "trace"},
|
||||
{"thread_id": "t1", "run_id": "r2", "event_type": "trace", "category": "trace"},
|
||||
]
|
||||
with caplog.at_level(logging.ERROR, logger=jsonl_mod.__name__), pytest.raises(OSError, match="second-run"):
|
||||
await store.put_batch(events)
|
||||
|
||||
assert r2_path.exists()
|
||||
|
||||
rollback_errors = [record for record in caplog.records if record.levelno >= logging.ERROR]
|
||||
assert any("duplicate records" in record.getMessage() for record in rollback_errors)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user