mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-09 16:08:31 +00:00
perf(runtime): index MemoryRunStore by thread_id to avoid O(n) scans (#3562)
* perf(runtime): index MemoryRunStore by thread_id to avoid O(n) scans MemoryRunStore is the default run backend (database.backend=memory) and backs RunManager.list_by_thread, which calls it on every thread-runs query to hydrate persisted runs. list_by_thread scanned every run in the store (O(total runs)) to filter by thread_id, so listing one thread's runs got linearly slower as unrelated runs accumulated across all threads. Add a thread_id -> insertion-ordered run_id set secondary index, maintained in lockstep with _runs in put()/delete(), and use it in list_by_thread for an O(runs-in-thread) lookup. This mirrors the index RunManager already keeps over its own in-memory records (#3499); the store extraction — whose docstring notes it is "Equivalent to the original RunManager._runs dict behavior" — did not carry the index across. Behavior is unchanged: same user_id filtering, newest-first ordering, and limit. The store runs each method without awaits on the event loop, so the index and _runs stay consistent without a lock. Extends tests/test_persistence_scaffold.py::TestMemoryRunStore with coverage for unknown-thread, newest-first ordering, limit, and index cleanup on delete (including empty-bucket removal). * perf(runtime): route aggregate_tokens_by_thread through the thread index too list_by_thread already uses the _runs_by_thread index this PR adds, but aggregate_tokens_by_thread (the /token-usage endpoint) still scanned every run in the process to pick out one thread's runs. Route it through the same index for an O(runs-in-thread) lookup, completing the thread-scoped read coverage. 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:
ly-wang19
GitHub
parent
435edbd8c6
commit
346164d740
2 changed files with 122 additions and 3 deletions
|
|
@ -14,6 +14,23 @@ from deerflow.runtime.runs.store.base import RunStore
|
|||
class MemoryRunStore(RunStore):
|
||||
def __init__(self) -> None:
|
||||
self._runs: dict[str, dict[str, Any]] = {}
|
||||
# Secondary index: thread_id -> insertion-ordered run_id set (a dict is
|
||||
# used as an ordered set), maintained in lockstep with ``_runs`` so
|
||||
# per-thread queries avoid O(total in-memory runs) full scans. Mirrors
|
||||
# the index ``RunManager`` keeps over its own in-memory records.
|
||||
self._runs_by_thread: dict[str, dict[str, None]] = {}
|
||||
|
||||
def _index_run(self, run_id: str, thread_id: str) -> None:
|
||||
"""Register *run_id* under *thread_id* in the secondary index."""
|
||||
self._runs_by_thread.setdefault(thread_id, {})[run_id] = None
|
||||
|
||||
def _unindex_run(self, run_id: str, thread_id: str) -> None:
|
||||
"""Drop *run_id* from the *thread_id* bucket, removing the bucket when empty."""
|
||||
bucket = self._runs_by_thread.get(thread_id)
|
||||
if bucket is not None:
|
||||
bucket.pop(run_id, None)
|
||||
if not bucket:
|
||||
self._runs_by_thread.pop(thread_id, None)
|
||||
|
||||
async def put(
|
||||
self,
|
||||
|
|
@ -45,6 +62,7 @@ class MemoryRunStore(RunStore):
|
|||
"created_at": created_at or now,
|
||||
"updated_at": now,
|
||||
}
|
||||
self._index_run(run_id, thread_id)
|
||||
|
||||
async def get(self, run_id, *, user_id=None):
|
||||
run = self._runs.get(run_id)
|
||||
|
|
@ -55,7 +73,13 @@ class MemoryRunStore(RunStore):
|
|||
return run
|
||||
|
||||
async def list_by_thread(self, thread_id, *, user_id=None, limit=100):
|
||||
results = [r for r in self._runs.values() if r["thread_id"] == thread_id and (user_id is None or r.get("user_id") == user_id)]
|
||||
# Use the thread index for an O(runs-in-thread) lookup instead of
|
||||
# scanning every run. ``self._runs.get`` is defense-in-depth: it drops a
|
||||
# stale id still in the index but already gone from ``_runs``.
|
||||
run_ids = self._runs_by_thread.get(thread_id)
|
||||
if not run_ids:
|
||||
return []
|
||||
results = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and (user_id is None or run.get("user_id") == user_id)]
|
||||
results.sort(key=lambda r: r["created_at"], reverse=True)
|
||||
return results[:limit]
|
||||
|
||||
|
|
@ -74,7 +98,9 @@ class MemoryRunStore(RunStore):
|
|||
self._runs[run_id]["updated_at"] = datetime.now(UTC).isoformat()
|
||||
|
||||
async def delete(self, run_id):
|
||||
self._runs.pop(run_id, None)
|
||||
run = self._runs.pop(run_id, None)
|
||||
if run is not None:
|
||||
self._unindex_run(run_id, run["thread_id"])
|
||||
|
||||
async def update_run_completion(self, run_id, *, status, **kwargs):
|
||||
if run_id in self._runs:
|
||||
|
|
@ -107,7 +133,10 @@ class MemoryRunStore(RunStore):
|
|||
|
||||
async def aggregate_tokens_by_thread(self, thread_id: str, *, include_active: bool = False) -> dict[str, Any]:
|
||||
statuses = ("success", "error", "running") if include_active else ("success", "error")
|
||||
completed = [r for r in self._runs.values() if r["thread_id"] == thread_id and r.get("status") in statuses]
|
||||
# Use the thread index for an O(runs-in-thread) lookup instead of
|
||||
# scanning every run in the process (mirrors ``list_by_thread``).
|
||||
run_ids = self._runs_by_thread.get(thread_id) or ()
|
||||
completed = [run for run_id in run_ids if (run := self._runs.get(run_id)) is not None and run.get("status") in statuses]
|
||||
by_model: dict[str, dict] = {}
|
||||
for r in completed:
|
||||
usage_by_model = r.get("token_usage_by_model") or {}
|
||||
|
|
|
|||
|
|
@ -131,6 +131,96 @@ class TestMemoryRunStore:
|
|||
async def test_delete_nonexistent_is_noop(self, store):
|
||||
await store.delete("nope") # should not raise
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_by_thread_unknown_thread_is_empty(self, store):
|
||||
await store.put("r1", thread_id="t1")
|
||||
assert await store.list_by_thread("missing") == []
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_by_thread_newest_first(self, store):
|
||||
await store.put("r1", thread_id="t1", created_at="2024-01-01T00:00:00+00:00")
|
||||
await store.put("r2", thread_id="t1", created_at="2024-01-03T00:00:00+00:00")
|
||||
await store.put("r3", thread_id="t1", created_at="2024-01-02T00:00:00+00:00")
|
||||
rows = await store.list_by_thread("t1")
|
||||
assert [r["run_id"] for r in rows] == ["r2", "r3", "r1"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_by_thread_respects_limit(self, store):
|
||||
for i in range(5):
|
||||
await store.put(f"r{i}", thread_id="t1", created_at=f"2024-01-0{i + 1}T00:00:00+00:00")
|
||||
rows = await store.list_by_thread("t1", limit=2)
|
||||
assert [r["run_id"] for r in rows] == ["r4", "r3"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delete_keeps_thread_index_consistent(self, store):
|
||||
await store.put("r1", thread_id="t1")
|
||||
await store.put("r2", thread_id="t1")
|
||||
await store.delete("r1")
|
||||
rows = await store.list_by_thread("t1")
|
||||
assert [r["run_id"] for r in rows] == ["r2"]
|
||||
# deleting the last run in a thread drops the now-empty index bucket
|
||||
await store.delete("r2")
|
||||
assert await store.list_by_thread("t1") == []
|
||||
assert "t1" not in store._runs_by_thread
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aggregate_tokens_by_thread_scopes_to_thread(self, store):
|
||||
await store.put("r1", thread_id="t1")
|
||||
await store.update_run_completion("r1", status="success", model_name="m-a", total_tokens=100)
|
||||
await store.put("r2", thread_id="t1")
|
||||
await store.update_run_completion("r2", status="error", model_name="m-a", total_tokens=20)
|
||||
await store.put("r3", thread_id="t2")
|
||||
await store.update_run_completion("r3", status="success", model_name="m-b", total_tokens=999)
|
||||
|
||||
agg = await store.aggregate_tokens_by_thread("t1")
|
||||
assert agg["total_tokens"] == 120 # the other thread's run is excluded
|
||||
assert agg["total_runs"] == 2
|
||||
assert agg["by_model"]["m-a"] == {"tokens": 120, "runs": 2}
|
||||
assert "m-b" not in agg["by_model"]
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aggregate_tokens_by_thread_excludes_active_unless_requested(self, store):
|
||||
await store.put("r1", thread_id="t1")
|
||||
await store.update_run_completion("r1", status="success", total_tokens=10)
|
||||
await store.put("r2", thread_id="t1")
|
||||
await store.update_run_completion("r2", status="running", total_tokens=5)
|
||||
|
||||
assert (await store.aggregate_tokens_by_thread("t1"))["total_tokens"] == 10
|
||||
assert (await store.aggregate_tokens_by_thread("t1", include_active=True))["total_tokens"] == 15
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aggregate_tokens_by_thread_unknown_thread_is_zero(self, store):
|
||||
await store.put("r1", thread_id="t1")
|
||||
await store.update_run_completion("r1", status="success", total_tokens=10)
|
||||
agg = await store.aggregate_tokens_by_thread("missing")
|
||||
assert agg["total_tokens"] == 0
|
||||
assert agg["total_runs"] == 0
|
||||
assert agg["by_model"] == {}
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_aggregate_tokens_by_thread_matches_full_scan_reference(self, store):
|
||||
plan = [
|
||||
("r0", "t1", "success", "m-a", 10),
|
||||
("r1", "t1", "error", "m-b", 20),
|
||||
("r2", "t1", "running", "m-a", 7),
|
||||
("r3", "t2", "success", "m-a", 999),
|
||||
("r4", "t1", "pending", "m-a", 3),
|
||||
]
|
||||
for run_id, thread_id, status, model, tokens in plan:
|
||||
await store.put(run_id, thread_id=thread_id)
|
||||
await store.update_run_completion(run_id, status=status, model_name=model, total_tokens=tokens)
|
||||
|
||||
def _reference(thread_id, include_active):
|
||||
statuses = ("success", "error", "running") if include_active else ("success", "error")
|
||||
completed = [r for r in store._runs.values() if r["thread_id"] == thread_id and r.get("status") in statuses]
|
||||
return len(completed), sum(r.get("total_tokens", 0) for r in completed)
|
||||
|
||||
for thread_id in ("t1", "t2", "missing"):
|
||||
for include_active in (False, True):
|
||||
agg = await store.aggregate_tokens_by_thread(thread_id, include_active=include_active)
|
||||
ref_runs, ref_tokens = _reference(thread_id, include_active)
|
||||
assert (agg["total_runs"], agg["total_tokens"]) == (ref_runs, ref_tokens), (thread_id, include_active)
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_list_pending(self, store):
|
||||
await store.put("r1", thread_id="t1", status="pending")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue