mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 21:48:04 +00:00
perf(runtime): index MemoryRunEventStore events by run_id to avoid O(n) scans (#3686)
* perf(runtime): index MemoryRunEventStore events by run_id to avoid O(n) scans list_messages was already served from a thread-wide messages projection (#3531), but list_events and list_messages_by_run still scanned the whole thread's event log (every run, every category) to return one run's events -- O(N_thread) on every run-scoped /messages page-load and /events request. Add _events_by_run / _messages_by_run projections (same dict objects, kept in lockstep in _put_one / delete_by_run / delete_by_thread), so both reads are O(M_run), with bisect cursor pagination for messages. Semantics are unchanged, pinned by a brute-force parity test over interleaved traces and both cursors. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * style: ruff format test_run_event_store_by_run_index.py Clears the lint-backend (ruff format --check) failure on the PR; the original commit ran `ruff check` but not `ruff format`. No behavior change (test fixture formatting only). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: ly-wang19 <ly-wang19@users.noreply.github.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
4a8f94eb28
commit
820560e576
@ -19,6 +19,14 @@ class MemoryRunEventStore(RunEventStore):
|
||||
# kept in seq order so message pagination is O(log m + page) via bisect
|
||||
# instead of re-scanning every event on each request.
|
||||
self._messages: dict[str, list[dict]] = {} # thread_id -> seq-sorted message list
|
||||
# Run-keyed projections of the two lists above (same dict objects, no
|
||||
# copies), kept in seq order. Per-run reads then cost O(events-in-run)
|
||||
# instead of O(events-in-thread): without these, ``list_events`` and
|
||||
# ``list_messages_by_run`` re-scan the whole thread's event log on every
|
||||
# request even though one run holds only a handful of events. This is
|
||||
# the per-run analogue of the thread-wide ``_messages`` projection.
|
||||
self._events_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted events
|
||||
self._messages_by_run: dict[str, dict[str, list[dict]]] = {} # thread_id -> run_id -> seq-sorted messages
|
||||
self._seq_counters: dict[str, int] = {} # thread_id -> last assigned seq
|
||||
|
||||
def _next_seq(self, thread_id: str) -> int:
|
||||
@ -50,8 +58,10 @@ class MemoryRunEventStore(RunEventStore):
|
||||
"created_at": created_at or datetime.now(UTC).isoformat(),
|
||||
}
|
||||
self._events.setdefault(thread_id, []).append(record)
|
||||
self._events_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record)
|
||||
if category == "message":
|
||||
self._messages.setdefault(thread_id, []).append(record)
|
||||
self._messages_by_run.setdefault(thread_id, {}).setdefault(run_id, []).append(record)
|
||||
return record
|
||||
|
||||
async def put(
|
||||
@ -100,23 +110,27 @@ class MemoryRunEventStore(RunEventStore):
|
||||
return messages[-limit:]
|
||||
|
||||
async def list_events(self, thread_id, run_id, *, event_types=None, limit=500):
|
||||
all_events = self._events.get(thread_id, [])
|
||||
filtered = [e for e in all_events if e["run_id"] == run_id]
|
||||
# ``_events_by_run`` is already scoped to this run and seq-ordered, so we
|
||||
# touch only this run's events instead of scanning the whole thread.
|
||||
run_events = self._events_by_run.get(thread_id, {}).get(run_id, [])
|
||||
if event_types is not None:
|
||||
filtered = [e for e in filtered if e["event_type"] in event_types]
|
||||
return filtered[:limit]
|
||||
run_events = [e for e in run_events if e["event_type"] in event_types]
|
||||
return run_events[:limit]
|
||||
|
||||
async def list_messages_by_run(self, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None):
|
||||
all_events = self._events.get(thread_id, [])
|
||||
filtered = [e for e in all_events if e["run_id"] == run_id and e["category"] == "message"]
|
||||
if before_seq is not None:
|
||||
filtered = [e for e in filtered if e["seq"] < before_seq]
|
||||
# Per-run, messages-only, seq-sorted: the seq window is a contiguous
|
||||
# slice located with bisect (O(log m_run)) over only this run's
|
||||
# messages, instead of re-scanning the whole thread's event log.
|
||||
messages = self._messages_by_run.get(thread_id, {}).get(run_id, [])
|
||||
lo = 0 if after_seq is None else bisect.bisect_right(messages, after_seq, key=lambda e: e["seq"])
|
||||
hi = len(messages) if before_seq is None else bisect.bisect_left(messages, before_seq, key=lambda e: e["seq"])
|
||||
window = messages[lo:hi]
|
||||
# An ``after_seq`` cursor pages forward (first ``limit``); otherwise
|
||||
# return the last ``limit`` (the latest page, or the page ending just
|
||||
# before ``before_seq``). Matches the prior filter-based semantics.
|
||||
if after_seq is not None:
|
||||
filtered = [e for e in filtered if e["seq"] > after_seq]
|
||||
if after_seq is not None:
|
||||
return filtered[:limit]
|
||||
else:
|
||||
return filtered[-limit:] if len(filtered) > limit else filtered
|
||||
return window[:limit]
|
||||
return window[-limit:]
|
||||
|
||||
async def count_messages(self, thread_id):
|
||||
return len(self._messages.get(thread_id, []))
|
||||
@ -124,6 +138,8 @@ class MemoryRunEventStore(RunEventStore):
|
||||
async def delete_by_thread(self, thread_id):
|
||||
events = self._events.pop(thread_id, [])
|
||||
self._messages.pop(thread_id, None)
|
||||
self._events_by_run.pop(thread_id, None)
|
||||
self._messages_by_run.pop(thread_id, None)
|
||||
self._seq_counters.pop(thread_id, None)
|
||||
return len(events)
|
||||
|
||||
@ -136,4 +152,7 @@ class MemoryRunEventStore(RunEventStore):
|
||||
self._events[thread_id] = remaining
|
||||
# Keep the message projection in lockstep (same surviving dict objects).
|
||||
self._messages[thread_id] = [e for e in remaining if e["category"] == "message"]
|
||||
# Drop the deleted run from the run-keyed projections.
|
||||
self._events_by_run.get(thread_id, {}).pop(run_id, None)
|
||||
self._messages_by_run.get(thread_id, {}).pop(run_id, None)
|
||||
return removed
|
||||
|
||||
130
backend/tests/test_run_event_store_by_run_index.py
Normal file
130
backend/tests/test_run_event_store_by_run_index.py
Normal file
@ -0,0 +1,130 @@
|
||||
"""Regression tests for MemoryRunEventStore's run-keyed event/message index.
|
||||
|
||||
``list_events`` and ``list_messages_by_run`` are served from per-run
|
||||
projections so a single run's reads cost O(events-in-run) instead of
|
||||
re-scanning O(events-in-thread). These tests pin the indexed implementation to
|
||||
the exact semantics of a brute-force full-thread scan -- including interleaved
|
||||
trace events (non-contiguous message seqs), both cursors supplied at once, and
|
||||
index upkeep after ``delete_by_run`` -- so the optimization can never silently
|
||||
drift from the reference behavior.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.runtime.events.store.memory import MemoryRunEventStore
|
||||
|
||||
|
||||
def _ref_messages_by_run(records, thread_id, run_id, *, limit=50, before_seq=None, after_seq=None):
|
||||
"""Brute-force reference: the pre-index full-thread scan it replaced."""
|
||||
filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id and e["category"] == "message"]
|
||||
if before_seq is not None:
|
||||
filtered = [e for e in filtered if e["seq"] < before_seq]
|
||||
if after_seq is not None:
|
||||
filtered = [e for e in filtered if e["seq"] > after_seq]
|
||||
if after_seq is not None:
|
||||
return filtered[:limit]
|
||||
return filtered[-limit:] if len(filtered) > limit else filtered
|
||||
|
||||
|
||||
def _ref_events(records, thread_id, run_id, *, event_types=None, limit=500):
|
||||
filtered = [e for e in records if e["thread_id"] == thread_id and e["run_id"] == run_id]
|
||||
if event_types is not None:
|
||||
filtered = [e for e in filtered if e["event_type"] in event_types]
|
||||
return filtered[:limit]
|
||||
|
||||
|
||||
async def _seed(store):
|
||||
"""Two runs interleaved within one thread; messages and traces mixed so
|
||||
each run's message seqs are non-contiguous (the bisect must handle gaps)."""
|
||||
plan = [
|
||||
("run-a", "message"),
|
||||
("run-a", "trace"),
|
||||
("run-b", "message"),
|
||||
("run-a", "message"),
|
||||
("run-b", "trace"),
|
||||
("run-b", "message"),
|
||||
("run-a", "trace"),
|
||||
("run-a", "message"),
|
||||
("run-b", "message"),
|
||||
("run-a", "message"),
|
||||
("run-b", "message"),
|
||||
("run-a", "message"),
|
||||
]
|
||||
records = []
|
||||
for i, (run_id, category) in enumerate(plan):
|
||||
rec = await store.put(thread_id="t1", run_id=run_id, event_type=f"e{i}", category=category, content=str(i))
|
||||
records.append(rec)
|
||||
return records
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_messages_by_run_matches_reference_across_cursors():
|
||||
store = MemoryRunEventStore()
|
||||
records = await _seed(store)
|
||||
seqs = [r["seq"] for r in records]
|
||||
cursors = [None, 0, *seqs, max(seqs) + 1]
|
||||
for run_id in ("run-a", "run-b", "run-missing"):
|
||||
for limit in (1, 2, 3, 50):
|
||||
for before_seq in cursors:
|
||||
for after_seq in cursors:
|
||||
got = await store.list_messages_by_run("t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq)
|
||||
want = _ref_messages_by_run(records, "t1", run_id, limit=limit, before_seq=before_seq, after_seq=after_seq)
|
||||
assert got == want, (run_id, limit, before_seq, after_seq)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_events_matches_reference_with_filters():
|
||||
store = MemoryRunEventStore()
|
||||
records = await _seed(store)
|
||||
all_types = sorted({r["event_type"] for r in records})
|
||||
for run_id in ("run-a", "run-b", "run-missing"):
|
||||
assert await store.list_events("t1", run_id) == _ref_events(records, "t1", run_id)
|
||||
assert await store.list_events("t1", run_id, limit=2) == _ref_events(records, "t1", run_id, limit=2)
|
||||
for et in all_types:
|
||||
assert await store.list_events("t1", run_id, event_types=[et]) == _ref_events(records, "t1", run_id, event_types=[et])
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_keyed_index_partitions_every_event():
|
||||
"""Every stored event is filed under exactly its (thread, run), each run's
|
||||
list is seq-ordered, and the union reconstructs the flat event log."""
|
||||
store = MemoryRunEventStore()
|
||||
records = await _seed(store)
|
||||
indexed = [e for run_events in store._events_by_run["t1"].values() for e in run_events]
|
||||
assert sorted(e["seq"] for e in indexed) == sorted(r["seq"] for r in records)
|
||||
for run_id, run_events in store._events_by_run["t1"].items():
|
||||
assert all(e["run_id"] == run_id for e in run_events)
|
||||
assert [e["seq"] for e in run_events] == sorted(e["seq"] for e in run_events)
|
||||
for run_id, run_msgs in store._messages_by_run["t1"].items():
|
||||
assert all(e["run_id"] == run_id and e["category"] == "message" for e in run_msgs)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_run_index_stays_in_lockstep_after_delete_by_run():
|
||||
store = MemoryRunEventStore()
|
||||
await _seed(store)
|
||||
removed = await store.delete_by_run("t1", "run-a")
|
||||
assert removed == 7 # run-a: 5 messages + 2 traces
|
||||
|
||||
# The deleted run vanishes from both per-run reads.
|
||||
assert await store.list_events("t1", "run-a") == []
|
||||
assert await store.list_messages_by_run("t1", "run-a") == []
|
||||
assert "run-a" not in store._events_by_run.get("t1", {})
|
||||
assert "run-a" not in store._messages_by_run.get("t1", {})
|
||||
|
||||
# The surviving run is untouched, and the thread-wide projection agrees.
|
||||
msgs_b = await store.list_messages_by_run("t1", "run-b")
|
||||
assert len(msgs_b) == 4
|
||||
assert all(m["run_id"] == "run-b" for m in msgs_b)
|
||||
assert all(m["run_id"] == "run-b" for m in await store.list_messages("t1"))
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delete_by_thread_clears_run_indexes():
|
||||
store = MemoryRunEventStore()
|
||||
await _seed(store)
|
||||
await store.delete_by_thread("t1")
|
||||
assert "t1" not in store._events_by_run
|
||||
assert "t1" not in store._messages_by_run
|
||||
assert await store.list_events("t1", "run-a") == []
|
||||
assert await store.list_messages_by_run("t1", "run-b") == []
|
||||
Loading…
x
Reference in New Issue
Block a user