scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-05 07:39:47 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 10

Fix initial scan retry scope and stale batches

Browse source
This commit is contained in:
mikemikimike 2026-10-22 02:12:54 +08:00
parent ae01c5eb9c
commit 970b7a3b46
No known key found for this signature in database
4 changed files with 71 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
reme/steps/index/init_changes.py
View file

@ -73,11 +73,17 @@ class InitChangesStep(BaseStep):
self.context["changes"] = changes
if changes:
self.logger.info(f"[{self.name}] scan {self.monitor_type}:{self.monitor_name}: {counts}")
await self.dispatch_steps(
self.dispatch_step_specs,
changes=changes,
retry_source_changes_before_watch=True,
)
retry_key = "retry_source_changes_before_watch"
had_retry_value = retry_key in self.context
previous_retry_value = self.context.get(retry_key)
self.context[retry_key] = True
try:
await self.dispatch_steps(self.dispatch_step_specs, changes=changes)
finally:
if had_retry_value:
self.context[retry_key] = previous_retry_value
else:
del self.context[retry_key]
else:
self.logger.info(f"[{self.name}] {self.monitor_type}:{self.monitor_name} is up to date")
self.context.response.metadata["counts"] = counts

2
reme/steps/index/update_changes.py
View file

@ -189,13 +189,13 @@ class ChangeApplyStep(BaseStep):
memory_budget = self._batch_memory_budget()
items.append(item)
del item
ok_paths.append(path)
estimated_bytes += item_bytes
if self._batch_is_full(len(items), estimated_bytes, 0, memory_budget):
await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results)
estimated_bytes = 0
memory_budget = self._batch_memory_budget()
del item
if items:
await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results)
return results

50
tests/unit/test_background_steps.py
View file

@ -1150,6 +1150,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches():
first = await init_step(ctx)
assert first.metadata["counts"] == {"added": 2, "modified": 0, "deleted": 0}
assert "retry_source_changes_before_watch" not in ctx
nodes = {n.path: n for n in await fs.get_nodes()}
assert set(nodes) == {"daily/a.md", "digest/report.md"}
assert all(nodes[p].chunk_ids for p in nodes)
@ -1161,6 +1162,7 @@ def test_index_update_loop_init_dispatch_updates_store_across_batches():
second = await init_step(ctx)
assert second.metadata["counts"] == {"added": 1, "modified": 1, "deleted": 1}
assert "retry_source_changes_before_watch" not in ctx
assert {(c["change"], Path(c["path"]).name) for c in ctx["changes"]} == {
("modified", "a.md"),
("deleted", "report.md"),
@ -1186,9 +1188,11 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
cwd = Path.cwd()
path = write_file(cwd / "daily" / "source.md", "initial contents")
parse_started = [threading.Event(), threading.Event()]
resume_parse = [threading.Event(), threading.Event()]
parse_started = [threading.Event() for _ in range(5)]
resume_parse = [threading.Event() for _ in range(5)]
watcher_started = asyncio.Event()
watcher_update_requested = asyncio.Event()
watcher_update_completed = asyncio.Event()
chunk_calls: list[Path] = []
watcher_contents: list[str] = []
@ -1227,16 +1231,12 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
result = original_chunk(file_path, rel_path)
chunk_calls.append(file_path)
attempt = len(chunk_calls) - 1
if attempt < len(parse_started):
if attempt in {0, 1, 3, 4}:
parse_started[attempt].set()
if not resume_parse[attempt].wait(timeout=5):
raise TimeoutError(f"initial parse {attempt + 1} was not resumed")
raise TimeoutError(f"parse attempt {attempt + 1} was not resumed")
return result
from reme.steps.index import watch_changes as watch_changes_module
original_awatch = watch_changes_module.awatch
async def track_watcher_start(
*paths: str | Path,
watch_filter: Callable[[Change, str], bool],
@ -1247,24 +1247,20 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
poll_delay_ms: int,
stop_event: asyncio.Event,
) -> AsyncGenerator[set[tuple[Change, str]], None]:
del paths, watch_filter, recursive, force_polling, debounce, step, poll_delay_ms
nodes = await store.get_nodes()
watcher_contents.append(
"".join(store.file_chunks[chunk_id].text for node in nodes for chunk_id in node.chunk_ids),
)
watcher_started.set()
async for changes in original_awatch(
*paths,
watch_filter=watch_filter,
recursive=recursive,
force_polling=force_polling,
debounce=debounce,
step=step,
poll_delay_ms=poll_delay_ms,
stop_event=stop_event,
):
yield changes
await watcher_update_requested.wait()
yield {(Change.modified, str(path))}
watcher_update_completed.set()
await stop_event.wait()
monkeypatch.setattr(chunker, "_chunk_sync", count_chunk)
from reme.steps.index import watch_changes as watch_changes_module
monkeypatch.setattr(watch_changes_module, "awatch", track_watcher_start)
try:
@ -1286,6 +1282,22 @@ def test_index_update_loop_retries_source_changed_during_initial_scan(monkeypatc
assert "".join(chunk.text for chunk in chunks) == latest
assert watcher_contents == [latest]
assert len(chunk_calls) == 3
path.write_text("watcher update before first parse", encoding="utf-8")
watcher_update_requested.set()
assert await asyncio.to_thread(parse_started[3].wait, 5)
path.write_text("watcher update before retry parse", encoding="utf-8")
resume_parse[3].set()
assert await asyncio.to_thread(parse_started[4].wait, 5)
path.write_text("latest watcher update source", encoding="utf-8")
resume_parse[4].set()
await asyncio.wait_for(watcher_update_completed.wait(), timeout=5)
assert len(chunk_calls) == 5
assert job._task is not None and not job._task.done()
nodes = await store.get_nodes()
chunks = [store.file_chunks[chunk_id] for chunk_id in nodes[0].chunk_ids]
assert "".join(chunk.text for chunk in chunks) == latest
finally:
for event in resume_parse:
event.set()

54
tests/unit/test_file_chunker_workers.py
View file

@ -7,6 +7,7 @@ import json
import os
import threading
import weakref
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Callable
from unittest.mock import patch
@ -582,27 +583,25 @@ async def test_stale_batch_chunks_are_released_before_source_retry(
) -> None:
"""The retry parse does not overlap the retained stale chunk batch in memory."""
monkeypatch.chdir(tmp_path)
first_path, second_path = tmp_path / "first.txt", tmp_path / "second.txt"
first_path.write_text("original first")
second_path.write_text("original second")
source_path = tmp_path / "source.txt"
source_path.write_text("original source")
context = ApplicationContext(workspace_dir=str(tmp_path))
chunker = DefaultFileChunker(app_context=context)
context.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
async with LocalFileStore(embedding_store="") as store, chunker:
await store.upsert([await chunker.chunk(first_path), await chunker.chunk(second_path)])
first_path.write_text("parsed first")
second_path.write_text("parsed second")
entered, release = threading.Event(), threading.Event()
original_chunk = chunker.chunk
await store.upsert([await chunker.chunk(source_path)])
source_path.write_text("parsed source")
flush_waiting = asyncio.Event()
original_guard = store._maintenance_guard
async def pause_second(file_path: Path) -> tuple[FileNode, list[FileChunk]]:
if file_path == second_path:
entered.set()
while not release.is_set():
await asyncio.sleep(0.01)
return await original_chunk(file_path)
@asynccontextmanager
async def track_store_guard():
flush_waiting.set()
async with original_guard():
yield
step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=2, app_context=context)
monkeypatch.setattr(store, "_maintenance_guard", track_store_guard)
step = UpdateIndexStep(file_store=store, persist=False, batch_max_files=1, app_context=context)
original_build = step._try_build_item
stale_chunk_released: threading.Event | None = None
@ -614,35 +613,38 @@ async def test_stale_batch_chunks_are_released_before_source_retry(
results: list[dict],
) -> tuple[tuple[FileNode, list[FileChunk]], int] | None:
nonlocal stale_chunk_released
if Path(path) == first_path and stale_chunk_released is not None:
if stale_chunk_released is not None:
gc.collect()
assert stale_chunk_released.is_set()
built = await original_build(change, path, abs_path, size_bytes, results)
if Path(path) == first_path and stale_chunk_released is None and built is not None:
if stale_chunk_released is None and built is not None:
stale_chunk_released = threading.Event()
weakref.finalize(built[0][1][0], stale_chunk_released.set)
return built
monkeypatch.setattr(step, "_try_build_item", capture_and_check_stale_chunk)
with patch.object(chunker, "chunk", side_effect=pause_second):
await store._maintenance_lock.acquire()
try:
task = asyncio.create_task(
step(
RuntimeContext(
changes=[
{"change": "modified", "path": str(first_path)},
{"change": "modified", "path": str(second_path)},
],
changes=[{"change": "modified", "path": str(source_path)}],
retry_source_changes_before_watch=True,
),
),
)
await _wait_started(entered)
first_path.write_text("latest first")
release.set()
response = await task
await asyncio.wait_for(flush_waiting.wait(), timeout=5)
source_path.write_text("latest source with a different size")
finally:
store._maintenance_lock.release()
response = await task
assert response.success is True
assert stale_chunk_released is not None and stale_chunk_released.is_set()
assert not step._source_versions
nodes = await store.get_nodes()
assert len(nodes) == 1
assert store.file_chunks[nodes[0].chunk_ids[0]].text == "latest source with a different size"
@pytest.mark.asyncio