scriptoriumadmin/ReMe
1
0
Fork 0
mirror of https://github.com/agentscope-ai/ReMe.git synced 2026-11-15 09:11:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

refactor(index): carry source versions with batch items

Browse source
This commit is contained in:
jinli.yl 2026-11-03 13:30:57 +08:00
parent 24df62fc90
commit bab44fb832
3 changed files with 26 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
reme/components/file_chunker/base_file_chunker.py
View file

@ -44,7 +44,7 @@ class BaseFileChunker(BaseComponent):
with self._worker_state_lock:
self._worker_closing = True
self._worker_generation += 1
await finish_cleanup(asyncio.create_task(self._close_worker()))
await finish_cleanup(self._close_worker())
async def _close_worker(self) -> None:
await self._acquire_worker_lock()
@ -76,18 +76,15 @@ class BaseFileChunker(BaseComponent):
file_path = Path(path)
rel_path = self.to_workspace_relative(path)
await self._acquire_worker_lock()
temporary_worker = False
try:
with self._worker_state_lock:
if self._worker_closing or generation != self._worker_generation:
raise RuntimeError("File chunker closed while waiting to parse")
worker = self._worker
temporary_worker = worker is None
if temporary_worker:
if worker is None:
# Standalone chunk() remains usable without start(), with no persistent threads.
with ThreadPoolExecutor(max_workers=1, thread_name_prefix="reme-chunker") as worker:
return await self._run_chunk(worker, generation, file_path, rel_path, chunk)
assert worker is not None
return await self._run_chunk(worker, generation, file_path, rel_path, chunk)
finally:
self._worker_lock.release()

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

@ -407,33 +407,22 @@ class UpdateIndexStep(ChangeApplyStep):
reads_file_content = True
source_change_retries: int = 1
def __init__(self, *args: Any, **kwargs: Any) -> None:
super().__init__(*args, **kwargs)
self._source_versions: dict[str, os.stat_result] = {}
def _change_path_key(self, path: str) -> str:
"""Coalesce absolute and workspace-relative aliases of the same file."""
return os.path.normcase(os.fspath(self._to_abs_path(path).absolute()))
async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]:
try:
return await super()._apply_existing(buckets)
finally:
# Steps are invocation-scoped; cancelled batches must release their snapshots too.
self._source_versions.clear()
async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]:
async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk], os.stat_result]:
"""Keep each source version with its parsed item until publication."""
before = path.stat()
item = await self.chunk_file(path)
node, chunks = await self.chunk_file(path)
after = path.stat()
check_file_version(path, before, after)
self._source_versions[item[0].path] = after
return item
return node, chunks, after
async def _try_upsert(
self,
change: Change,
items: list[tuple[FileNode, list[FileChunk]]],
items: list[tuple[FileNode, list[FileChunk], os.stat_result]],
ok_paths: list[str],
*,
retry_source_changes: bool = True,
@ -442,15 +431,15 @@ class UpdateIndexStep(ChangeApplyStep):
retry_paths: list[str] = []
result_by_path: dict[str, dict] = {}
async with self.file_store._maintenance_guard(): # pylint: disable=protected-access
current_items: list[tuple[FileNode, list[FileChunk]]] = []
current_items: list[tuple[FileNode, list[FileChunk], os.stat_result]] = []
current_paths: list[str] = []
item: tuple[FileNode, list[FileChunk]] | None = None
item: tuple[FileNode, list[FileChunk], os.stat_result] | None = None
for item, path in zip(items, ok_paths):
node = item[0]
try:
check_file_version(
self._to_abs_path(node.path),
self._source_versions[node.path],
item[2],
self._to_abs_path(node.path).stat(),
)
except (FileChangedDuringChunkingError, FileNotFoundError) as exc:
@ -474,8 +463,6 @@ class UpdateIndexStep(ChangeApplyStep):
else:
current_items.append(item)
current_paths.append(path)
finally:
self._source_versions.pop(node.path)
if current_items:
result_by_path.update(
{
@ -509,8 +496,8 @@ class UpdateIndexStep(ChangeApplyStep):
result_by_path.update({result["path"]: result for result in retry_results})
return [result_by_path[path] for path in ok_paths]
async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None:
await self.file_store.upsert(items)
async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk], os.stat_result]]) -> None:
await self.file_store.upsert([(node, chunks) for node, chunks, _ in items])
async def delete_paths(self, paths: list[str]) -> None:
await self.file_store.delete(paths)
@ -526,7 +513,7 @@ class UpdateIndexStep(ChangeApplyStep):
def estimate_item_memory(
self,
item: tuple[FileNode, list[FileChunk]],
item: tuple[FileNode, list[FileChunk], os.stat_result],
path: Path,
size_bytes: int,
) -> int:

24
tests/unit/test_file_chunker_workers.py
View file

@ -506,7 +506,6 @@ async def test_existing_content_searchable_during_indexing(
assert recovered.success is True
nodes = await store.get_nodes()
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
assert not step._source_versions
@pytest.mark.asyncio
@ -576,7 +575,6 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file
if change == "oversize":
assert response.answer[0]["skipped"] is True
assert response.answer[0]["reason"] == "file_too_large"
assert not step._source_versions
nodes = {node.path: node for node in await store.get_nodes()}
if change == "oversize":
assert "first.txt" not in nodes
@ -602,7 +600,6 @@ async def test_changed_batched_file_is_rebuilt_before_publication_and_valid_file
assert "first.txt" not in nodes
else:
assert store.file_chunks[nodes["first.txt"].chunk_ids[0]].text == first_path.read_text()
assert not step._source_versions
@pytest.mark.asyncio
@ -640,7 +637,7 @@ async def test_stale_batch_chunks_are_released_before_source_retry(
abs_path: Path,
size_bytes: int,
results: list[dict],
) -> tuple[tuple[FileNode, list[FileChunk]], int] | None:
) -> tuple[tuple[FileNode, list[FileChunk], os.stat_result], int] | None:
nonlocal stale_chunk_released
if stale_chunk_released is not None:
gc.collect()
@ -670,7 +667,6 @@ async def test_stale_batch_chunks_are_released_before_source_retry(
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"
@ -708,7 +704,11 @@ async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate(
estimated_sizes: list[int] = []
original_estimate = step.estimate_item_memory
def record_estimate(item: tuple[FileNode, list[FileChunk]], file_path: Path, size_bytes: int) -> int:
def record_estimate(
item: tuple[FileNode, list[FileChunk], os.stat_result],
file_path: Path,
size_bytes: int,
) -> int:
estimated_sizes.append(size_bytes)
return original_estimate(item, file_path, size_bytes)
@ -736,7 +736,6 @@ async def test_parse_retry_rechecks_size_limit_and_refreshes_memory_estimate(
assert estimated_sizes == [new_size]
nodes = await store.get_nodes()
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
assert not step._source_versions
@pytest.mark.asyncio
@ -767,7 +766,6 @@ async def test_index_updates_deduplicate_absolute_and_relative_paths(
nodes = await store.get_nodes()
assert len(nodes) == 1
assert nodes[0].path == "source.txt"
assert not step._source_versions
@pytest.mark.asyncio
@ -788,11 +786,13 @@ async def test_version_is_checked_after_waiting_for_store_lock(
await store.upsert([await chunker.chunk(path)])
path.write_text("updated source")
built = asyncio.Event()
pending_chunks: list[weakref.ReferenceType[FileChunk]] = []
step = UpdateIndexStep(file_store=store, persist=False, app_context=context)
original = step.build_item
async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk]]:
async def record_build(file_path: Path) -> tuple[FileNode, list[FileChunk], os.stat_result]:
item = await original(file_path)
pending_chunks.append(weakref.ref(item[1][0]))
built.set()
return item
@ -810,7 +810,10 @@ async def test_version_is_checked_after_waiting_for_store_lock(
if outcome == "change":
response = await task
assert response.success is True
assert not step._source_versions
if outcome == "cancel":
await asyncio.sleep(0)
gc.collect()
assert pending_chunks and all(ref() is None for ref in pending_chunks)
nodes = await store.get_nodes()
expected_text = "latest source" if outcome == "change" else "original source"
assert store.file_chunks[nodes[0].chunk_ids[0]].text == expected_text
@ -819,4 +822,3 @@ async def test_version_is_checked_after_waiting_for_store_lock(
assert recovered.success is True
nodes = await store.get_nodes()
assert store.file_chunks[nodes[0].chunk_ids[0]].text == path.read_text()
assert not step._source_versions