mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-06 08:16:00 +00:00
feat(index): add bounded memory-aware batch processing (#381)
* test(background_steps): add comprehensive tests for batch processing and memory management - Add test for catalog upserts in batches of at most 100 files - Add test for catalog deletes in batches of at most 100 paths - Add test for index memory budget reducing batches to one file - Add test for memory target limiting cumulative batch size - Add test for invalid batch memory settings rejection - Add test for continuing after one batch fails - Add test for yielding to event loop while building batch - Add test for modified file reusing unchanged embedding - Add test for reporting memory estimation failure without aborting feat(update_changes): implement bounded batch processing with memory management - Add configurable batch parameters with default values - Implement memory budget calculation based on available system memory - Add file inspection and memory estimation before processing - Implement batch flushing when limits are reached - Add proper error handling for batch operations - Support async yielding during batch building - Add comprehensive validation for batch configuration parameters - Implement memory estimation for indexing operations - Add batch size limiting for delete operations * test(steps): add tests for memory estimation failure handling - Add test case for isolated file processing when memory estimation fails - Add test case for proper release of flushed items before building next file - Implement weak reference tracking to verify payload lifetime management - Create parametrized tests for both source and item memory estimation methods - Add assertions to verify single-item batch behavior on estimation failures - Include comprehensive error handling verification for memory budget calculations * chore(version): bump version to 0.4.1.3 - Update __version__ from 0.4.1.2 to 0.4.1.3 in __init__.py * feat(index): support batch settings from environment * refactor(index): use direct batch defaults * refactor(index): configure memory estimates through step args * ci: simplify Windows smoke dependencies
This commit is contained in:
jinliyl
GitHub
parent
55ef4bd6ad
commit
b4333fbef8
4 changed files with 644 additions and 37 deletions
3
.github/workflows/windows-smoke.yml
vendored
3
.github/workflows/windows-smoke.yml
vendored
|
|
@ -32,7 +32,8 @@ jobs:
|
|||
- name: Install package
|
||||
run: |
|
||||
python -m pip install --upgrade pip setuptools wheel
|
||||
pip install -e ".[core,benchmark]"
|
||||
pip install agentscope
|
||||
pip install -e .
|
||||
|
||||
- name: Run version job
|
||||
run: reme start service.backend=cli job=version
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
"""ReMe CLI package."""
|
||||
|
||||
__version__ = "0.4.1.2"
|
||||
__version__ = "0.4.1.3"
|
||||
|
||||
from . import config
|
||||
from . import constants
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
"""Apply file change batches to file_catalog or file_store."""
|
||||
"""Apply bounded file change batches to file_catalog or file_store."""
|
||||
|
||||
import asyncio
|
||||
import math
|
||||
from abc import abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import Any, Callable
|
||||
|
||||
import psutil
|
||||
from watchfiles import Change
|
||||
|
||||
from ._change_batch import bucket_changes
|
||||
|
|
@ -20,18 +23,50 @@ class ChangeApplyStep(BaseStep):
|
|||
target_name = "target"
|
||||
reads_file_content = False
|
||||
|
||||
def __init__(self, persist: bool | None = None, **kwargs):
|
||||
def __init__(
|
||||
self,
|
||||
persist: bool | None = None,
|
||||
batch_max_files: int = 100,
|
||||
batch_available_memory_ratio: float = 0.10,
|
||||
batch_memory_target_bytes: int = 512 * 1024 * 1024,
|
||||
batch_memory_expansion_factor: float = 8.0,
|
||||
file_memory_overhead_bytes: int = 32 * 1024,
|
||||
chunk_memory_overhead_bytes: int = 4 * 1024,
|
||||
estimated_chunk_bytes: int = 10_000,
|
||||
float16_bytes: int = 2,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.persist = persist
|
||||
self.batch_max_files = int(batch_max_files)
|
||||
self.batch_available_memory_ratio = float(batch_available_memory_ratio)
|
||||
self.batch_memory_target_bytes = int(batch_memory_target_bytes)
|
||||
self.batch_memory_expansion_factor = float(batch_memory_expansion_factor)
|
||||
self.file_memory_overhead_bytes = int(file_memory_overhead_bytes)
|
||||
self.chunk_memory_overhead_bytes = int(chunk_memory_overhead_bytes)
|
||||
self.estimated_chunk_bytes = int(estimated_chunk_bytes)
|
||||
self.float16_bytes = int(float16_bytes)
|
||||
if self.batch_max_files <= 0:
|
||||
raise ValueError("batch_max_files must be greater than zero")
|
||||
if not math.isfinite(self.batch_available_memory_ratio) or not 0 < self.batch_available_memory_ratio <= 1:
|
||||
raise ValueError("batch_available_memory_ratio must be in the interval (0, 1]")
|
||||
if self.batch_memory_target_bytes <= 0:
|
||||
raise ValueError("batch_memory_target_bytes must be greater than zero")
|
||||
if not math.isfinite(self.batch_memory_expansion_factor) or self.batch_memory_expansion_factor <= 0:
|
||||
raise ValueError("batch_memory_expansion_factor must be finite and greater than zero")
|
||||
if self.file_memory_overhead_bytes < 0:
|
||||
raise ValueError("file_memory_overhead_bytes must be non-negative")
|
||||
if self.chunk_memory_overhead_bytes < 0:
|
||||
raise ValueError("chunk_memory_overhead_bytes must be non-negative")
|
||||
if self.estimated_chunk_bytes <= 0:
|
||||
raise ValueError("estimated_chunk_bytes must be greater than zero")
|
||||
if self.float16_bytes <= 0:
|
||||
raise ValueError("float16_bytes must be greater than zero")
|
||||
|
||||
@abstractmethod
|
||||
async def build_item(self, path: Path) -> Any:
|
||||
"""Parse one existing file into the target item shape."""
|
||||
|
||||
@abstractmethod
|
||||
def item_path(self, item: Any) -> str:
|
||||
"""Return the target-relative path for an upsert item."""
|
||||
|
||||
@abstractmethod
|
||||
async def upsert_items(self, items: list[Any]) -> None:
|
||||
"""Upsert parsed items into the target."""
|
||||
|
|
@ -59,26 +94,79 @@ class ChangeApplyStep(BaseStep):
|
|||
|
||||
async def _apply_existing(self, buckets: dict[Change, list[str]]) -> list[dict]:
|
||||
results: list[dict] = []
|
||||
for change, action in ((Change.added, "Adding"), (Change.modified, "Updating")):
|
||||
for change in (Change.added, Change.modified):
|
||||
paths = buckets[change]
|
||||
if not paths:
|
||||
continue
|
||||
self.logger.info(f"Detected {len(paths)} {change.name} file(s)")
|
||||
items, ok_paths = [], []
|
||||
items: list[Any] = []
|
||||
ok_paths: list[str] = []
|
||||
estimated_bytes = 0
|
||||
memory_budget = self._batch_memory_budget()
|
||||
for path in paths:
|
||||
item = await self._try_build_item(change, action, path, results)
|
||||
if item is not None:
|
||||
items.append(item)
|
||||
ok_paths.append(path)
|
||||
await asyncio.sleep(0)
|
||||
inspected = await self._inspect_path(change, path, results)
|
||||
if inspected is None:
|
||||
continue
|
||||
abs_path, size_bytes = inspected
|
||||
preliminary_bytes = self._try_estimate_memory(
|
||||
path,
|
||||
self.estimate_source_memory,
|
||||
abs_path,
|
||||
size_bytes,
|
||||
)
|
||||
force_single_item_batch = preliminary_bytes is None
|
||||
if preliminary_bytes is None:
|
||||
preliminary_bytes = memory_budget + 1
|
||||
if items and self._batch_is_full(len(items), estimated_bytes, preliminary_bytes, memory_budget):
|
||||
await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results)
|
||||
estimated_bytes = 0
|
||||
memory_budget = self._batch_memory_budget()
|
||||
|
||||
item = await self._try_build_item(change, path, abs_path, results)
|
||||
if item is None:
|
||||
continue
|
||||
item_bytes = self._try_estimate_memory(
|
||||
path,
|
||||
self.estimate_item_memory,
|
||||
item,
|
||||
abs_path,
|
||||
size_bytes,
|
||||
)
|
||||
if item_bytes is None:
|
||||
force_single_item_batch = True
|
||||
item_bytes = memory_budget + 1
|
||||
elif force_single_item_batch:
|
||||
item_bytes = max(item_bytes, memory_budget + 1)
|
||||
if items and self._batch_is_full(len(items), estimated_bytes, item_bytes, memory_budget):
|
||||
await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results)
|
||||
estimated_bytes = 0
|
||||
memory_budget = self._batch_memory_budget()
|
||||
|
||||
items.append(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:
|
||||
results.extend(await self._try_upsert(change, items, ok_paths))
|
||||
await self._flush_upsert_batch(change, items, ok_paths, estimated_bytes, results)
|
||||
return results
|
||||
|
||||
async def _try_build_item(self, change: Change, action: str, path: str, results: list[dict]):
|
||||
async def _inspect_path(self, change: Change, path: str, results: list[dict]) -> tuple[Path, int] | None:
|
||||
abs_path = self._to_abs_path(path)
|
||||
try:
|
||||
if not abs_path.is_file():
|
||||
results.append({"change": change.name, "path": path, "success": False, "error": "not a file"})
|
||||
results.append(
|
||||
{
|
||||
"change": change.name,
|
||||
"path": path,
|
||||
"success": False,
|
||||
"error": "not a file",
|
||||
},
|
||||
)
|
||||
return None
|
||||
size_bytes = abs_path.stat().st_size
|
||||
max_file_bytes = self.max_file_bytes()
|
||||
|
|
@ -99,16 +187,83 @@ class ChangeApplyStep(BaseStep):
|
|||
},
|
||||
)
|
||||
return None
|
||||
self.logger.info(f"{action} file: {path}")
|
||||
return abs_path, size_bytes
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to process {path}")
|
||||
results.append({"change": change.name, "path": path, "success": False, "error": str(e)})
|
||||
return None
|
||||
|
||||
async def _try_build_item(
|
||||
self,
|
||||
change: Change,
|
||||
path: str,
|
||||
abs_path: Path,
|
||||
results: list[dict],
|
||||
):
|
||||
try:
|
||||
self.logger.debug(f"Processing {change.name} file: {path}")
|
||||
return await self.build_item(abs_path)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to process {path}")
|
||||
results.append({"change": change.name, "path": path, "success": False, "error": str(e)})
|
||||
return None
|
||||
|
||||
def _try_estimate_memory(
|
||||
self,
|
||||
path: str,
|
||||
estimate: Callable[..., int],
|
||||
*args: Any,
|
||||
) -> int | None:
|
||||
"""Return an advisory estimate, or ``None`` so the caller can isolate the item."""
|
||||
try:
|
||||
return max(0, int(estimate(*args)))
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Failed to estimate memory for {path}; processing it alone: {e}")
|
||||
return None
|
||||
|
||||
def _batch_memory_budget(self) -> int:
|
||||
"""Return the current estimated-memory budget for a new batch."""
|
||||
try:
|
||||
available = max(0, int(psutil.virtual_memory().available))
|
||||
except (AttributeError, OSError, psutil.Error) as e:
|
||||
self.logger.warning(f"Failed to read available memory, using batch memory target: {e}")
|
||||
return self.batch_memory_target_bytes
|
||||
dynamic_budget = int(available * self.batch_available_memory_ratio)
|
||||
return max(1, min(dynamic_budget, self.batch_memory_target_bytes))
|
||||
|
||||
def estimate_source_memory(self, path: Path, size_bytes: int) -> int:
|
||||
"""Estimate one item before reading it so the current batch can flush first."""
|
||||
del path, size_bytes
|
||||
return self.file_memory_overhead_bytes
|
||||
|
||||
def estimate_item_memory(self, item: Any, path: Path, size_bytes: int) -> int:
|
||||
"""Estimate the retained memory for one built target item."""
|
||||
del item
|
||||
return self.estimate_source_memory(path, size_bytes)
|
||||
|
||||
def _batch_is_full(self, item_count: int, estimated_bytes: int, next_bytes: int, memory_budget: int) -> bool:
|
||||
return item_count >= self.batch_max_files or estimated_bytes + next_bytes > memory_budget
|
||||
|
||||
async def _flush_upsert_batch(
|
||||
self,
|
||||
change: Change,
|
||||
items: list[Any],
|
||||
ok_paths: list[str],
|
||||
estimated_bytes: int,
|
||||
results: list[dict],
|
||||
) -> None:
|
||||
if not items:
|
||||
return
|
||||
self.logger.info(
|
||||
f"Applying {change.name} batch to {self.target_name}: "
|
||||
f"files={len(items)} estimated_bytes={estimated_bytes}",
|
||||
)
|
||||
results.extend(await self._try_upsert(change, items, ok_paths))
|
||||
items.clear()
|
||||
ok_paths.clear()
|
||||
|
||||
async def _try_upsert(self, change: Change, items: list[Any], ok_paths: list[str]) -> list[dict]:
|
||||
try:
|
||||
await self.delete_paths([self.item_path(item) for item in items])
|
||||
await self.upsert_items(items)
|
||||
return [{"change": change.name, "path": p, "success": True} for p in ok_paths]
|
||||
except Exception as e:
|
||||
|
|
@ -119,12 +274,17 @@ class ChangeApplyStep(BaseStep):
|
|||
if not deleted:
|
||||
return []
|
||||
self.logger.info(f"Detected {len(deleted)} deleted file(s)")
|
||||
try:
|
||||
await self.delete_paths([self.to_workspace_relative(p) for p in deleted])
|
||||
return [{"change": "deleted", "path": p, "success": True} for p in deleted]
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to delete {len(deleted)} file(s) from {self.target_name}")
|
||||
return [{"change": "deleted", "path": p, "success": False, "error": str(e)} for p in deleted]
|
||||
results: list[dict] = []
|
||||
for start in range(0, len(deleted), self.batch_max_files):
|
||||
await asyncio.sleep(0)
|
||||
batch = deleted[start : start + self.batch_max_files]
|
||||
try:
|
||||
await self.delete_paths([self.to_workspace_relative(p) for p in batch])
|
||||
results.extend({"change": "deleted", "path": p, "success": True} for p in batch)
|
||||
except Exception as e:
|
||||
self.logger.exception(f"Failed to delete {len(batch)} file(s) from {self.target_name}")
|
||||
results.extend({"change": "deleted", "path": p, "success": False, "error": str(e)} for p in batch)
|
||||
return results
|
||||
|
||||
def _to_abs_path(self, path: str | Path) -> Path:
|
||||
p = Path(path)
|
||||
|
|
@ -141,9 +301,6 @@ class UpdateCatalogStep(ChangeApplyStep):
|
|||
stat = path.stat()
|
||||
return FileNode(path=self.to_workspace_relative(path), st_mtime=stat.st_mtime)
|
||||
|
||||
def item_path(self, item: FileNode) -> str:
|
||||
return item.path
|
||||
|
||||
async def upsert_items(self, items: list[FileNode]) -> None:
|
||||
if self.file_catalog is None:
|
||||
raise RuntimeError("file_catalog is not initialized!")
|
||||
|
|
@ -169,9 +326,6 @@ class UpdateIndexStep(ChangeApplyStep):
|
|||
async def build_item(self, path: Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
return await self.chunk_file(path)
|
||||
|
||||
def item_path(self, item: tuple[FileNode, list[FileChunk]]) -> str:
|
||||
return item[0].path
|
||||
|
||||
async def upsert_items(self, items: list[tuple[FileNode, list[FileChunk]]]) -> None:
|
||||
await self.file_store.upsert(items)
|
||||
|
||||
|
|
@ -181,6 +335,33 @@ class UpdateIndexStep(ChangeApplyStep):
|
|||
async def dump_target(self) -> None:
|
||||
await self.file_store.dump()
|
||||
|
||||
def estimate_source_memory(self, path: Path, size_bytes: int) -> int:
|
||||
chunker = self._resolve_chunker(path)
|
||||
chunk_bytes = max(1, int(getattr(chunker, "chunk_byte_size", self.estimated_chunk_bytes)))
|
||||
estimated_chunks = max(1, (size_bytes + chunk_bytes - 1) // chunk_bytes)
|
||||
return self._estimate_index_memory(size_bytes, estimated_chunks)
|
||||
|
||||
def estimate_item_memory(
|
||||
self,
|
||||
item: tuple[FileNode, list[FileChunk]],
|
||||
path: Path,
|
||||
size_bytes: int,
|
||||
) -> int:
|
||||
del path
|
||||
return self._estimate_index_memory(size_bytes, len(item[1]))
|
||||
|
||||
def _estimate_index_memory(self, size_bytes: int, chunk_count: int) -> int:
|
||||
embedding_bytes = 0
|
||||
embedding_store = getattr(self.file_store, "embedding_store", None)
|
||||
if embedding_store is not None:
|
||||
try:
|
||||
embedding_bytes = max(0, int(embedding_store.dimensions)) * self.float16_bytes
|
||||
except (AttributeError, TypeError, ValueError):
|
||||
embedding_bytes = 0
|
||||
expanded_content = int(size_bytes * self.batch_memory_expansion_factor)
|
||||
per_chunk = self.chunk_memory_overhead_bytes + embedding_bytes
|
||||
return expanded_content + self.file_memory_overhead_bytes + max(0, chunk_count) * per_chunk
|
||||
|
||||
async def chunk_file(self, path: str | Path) -> tuple[FileNode, list[FileChunk]]:
|
||||
"""Chunk a file into (node, chunks)."""
|
||||
if self.app_context is None:
|
||||
|
|
|
|||
|
|
@ -15,9 +15,12 @@ import datetime
|
|||
import os
|
||||
import tempfile
|
||||
import warnings
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from watchfiles import Change
|
||||
|
||||
from reme.components.agent_wrapper import BaseAgentWrapper
|
||||
|
|
@ -462,6 +465,401 @@ def test_update_catalog_relative_path_uses_workspace():
|
|||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_upserts_in_batches_of_at_most_100_files():
|
||||
"""Large change sets are committed incrementally instead of retained as one list."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index:03d}.md", "x") for index in range(205)]
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
batch_sizes = []
|
||||
original_upsert = catalog.upsert
|
||||
|
||||
async def record_upsert(nodes):
|
||||
batch_sizes.append(len(nodes))
|
||||
await original_upsert(nodes)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(file_catalog=catalog, persist=False, app_context=_make_app_context(cwd))
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
available = MagicMock(available=8 * 1024 * 1024 * 1024)
|
||||
with (
|
||||
patch.object(catalog, "upsert", side_effect=record_upsert),
|
||||
patch("reme.steps.index.update_changes.psutil.virtual_memory", return_value=available),
|
||||
):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert batch_sizes == [100, 100, 5]
|
||||
assert len(await catalog.get_nodes()) == 205
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_deletes_in_batches_of_at_most_100_paths():
|
||||
"""Large delete sets use the same bounded failure domain as upserts."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
deleted = [cwd / "daily" / f"{index:03d}.md" for index in range(205)]
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
batch_sizes = []
|
||||
original_delete = catalog.delete
|
||||
|
||||
async def record_delete(paths):
|
||||
batch_sizes.append(len(paths))
|
||||
await original_delete(paths)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(file_catalog=catalog, persist=False, app_context=_make_app_context(cwd))
|
||||
changes = [{"change": "deleted", "path": str(path)} for path in deleted]
|
||||
with patch.object(catalog, "delete", side_effect=record_delete):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert batch_sizes == [100, 100, 5]
|
||||
assert len(response.answer) == 205
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_index_memory_budget_can_reduce_batches_to_one_file():
|
||||
"""A file larger than the target batch budget is still processed alone."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index}.md", f"# Note {index}\nbody\n") for index in range(3)]
|
||||
fs = LocalFileStore(name="default", embedding_store="")
|
||||
chunker = DefaultFileChunker(supported_extensions=["md"])
|
||||
await fs.start()
|
||||
await chunker.start()
|
||||
batch_sizes = []
|
||||
original_upsert = fs.upsert
|
||||
|
||||
async def record_upsert(items):
|
||||
batch_sizes.append(len(items))
|
||||
await original_upsert(items)
|
||||
|
||||
try:
|
||||
app_ctx = _make_app_context(cwd)
|
||||
app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
|
||||
step = UpdateIndexStep(file_store=fs, persist=False, app_context=app_ctx)
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
available = MagicMock(available=1_000)
|
||||
with (
|
||||
patch.object(fs, "upsert", side_effect=record_upsert),
|
||||
patch("reme.steps.index.update_changes.psutil.virtual_memory", return_value=available),
|
||||
):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert batch_sizes == [1, 1, 1]
|
||||
assert len(await fs.get_nodes()) == 3
|
||||
finally:
|
||||
await chunker.close()
|
||||
await fs.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_memory_target_limits_cumulative_batch_size():
|
||||
"""The memory target applies to the accumulated batch, not only individual files."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index}.md", "x") for index in range(3)]
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
batch_sizes = []
|
||||
original_upsert = catalog.upsert
|
||||
|
||||
async def record_upsert(nodes):
|
||||
batch_sizes.append(len(nodes))
|
||||
await original_upsert(nodes)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(
|
||||
file_catalog=catalog,
|
||||
persist=False,
|
||||
batch_memory_target_bytes=64 * 1024,
|
||||
app_context=_make_app_context(cwd),
|
||||
)
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
available = MagicMock(available=8 * 1024 * 1024 * 1024)
|
||||
with (
|
||||
patch.object(catalog, "upsert", side_effect=record_upsert),
|
||||
patch("reme.steps.index.update_changes.psutil.virtual_memory", return_value=available),
|
||||
):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert batch_sizes == [2, 1]
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "message"),
|
||||
[
|
||||
({"batch_available_memory_ratio": float("nan")}, "batch_available_memory_ratio"),
|
||||
({"batch_memory_target_bytes": 0}, "batch_memory_target_bytes"),
|
||||
({"batch_memory_expansion_factor": float("inf")}, "batch_memory_expansion_factor"),
|
||||
({"file_memory_overhead_bytes": -1}, "file_memory_overhead_bytes"),
|
||||
({"chunk_memory_overhead_bytes": -1}, "chunk_memory_overhead_bytes"),
|
||||
({"estimated_chunk_bytes": 0}, "estimated_chunk_bytes"),
|
||||
({"float16_bytes": 0}, "float16_bytes"),
|
||||
],
|
||||
)
|
||||
def test_update_step_rejects_invalid_batch_memory_settings(kwargs, message):
|
||||
"""Invalid batch-memory settings fail during step construction."""
|
||||
with pytest.raises(ValueError, match=message):
|
||||
UpdateCatalogStep(**kwargs)
|
||||
|
||||
|
||||
def test_update_step_accepts_configured_batch_memory_settings():
|
||||
"""Step configuration can override all batching and memory-estimation defaults."""
|
||||
step = UpdateCatalogStep(
|
||||
batch_max_files=25,
|
||||
batch_available_memory_ratio=0.25,
|
||||
batch_memory_target_bytes=1024 * 1024,
|
||||
batch_memory_expansion_factor=4.5,
|
||||
file_memory_overhead_bytes=1024,
|
||||
chunk_memory_overhead_bytes=512,
|
||||
estimated_chunk_bytes=5000,
|
||||
float16_bytes=4,
|
||||
)
|
||||
|
||||
assert step.batch_max_files == 25
|
||||
assert step.batch_available_memory_ratio == 0.25
|
||||
assert step.batch_memory_target_bytes == 1024 * 1024
|
||||
assert step.batch_memory_expansion_factor == 4.5
|
||||
assert step.file_memory_overhead_bytes == 1024
|
||||
assert step.chunk_memory_overhead_bytes == 512
|
||||
assert step.estimated_chunk_bytes == 5000
|
||||
assert step.float16_bytes == 4
|
||||
|
||||
|
||||
@pytest.mark.parametrize("estimate_method", ["estimate_source_memory", "estimate_item_memory"])
|
||||
def test_update_catalog_memory_estimate_failure_processes_each_file_alone(estimate_method):
|
||||
"""Advisory estimate failures isolate files without skipping their updates."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index}.md", "x") for index in range(2)]
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
batch_sizes = []
|
||||
original_upsert = catalog.upsert
|
||||
|
||||
async def record_upsert(nodes):
|
||||
batch_sizes.append(len(nodes))
|
||||
await original_upsert(nodes)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(file_catalog=catalog, persist=False, app_context=_make_app_context(cwd))
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
with (
|
||||
patch.object(step, estimate_method, side_effect=RuntimeError("estimate failed")),
|
||||
patch.object(catalog, "upsert", side_effect=record_upsert),
|
||||
):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert batch_sizes == [1, 1]
|
||||
assert len(await catalog.get_nodes()) == 2
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_releases_flushed_item_before_building_next_file():
|
||||
"""A one-file batch does not retain its payload while the next payload is built."""
|
||||
|
||||
class TrackedItem:
|
||||
"""Weak-referenceable payload used to observe the local item lifetime."""
|
||||
|
||||
class TrackingCatalogStep(UpdateCatalogStep):
|
||||
"""Catalog step that records whether the prior payload was released."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.previous_item = None
|
||||
self.release_observations = []
|
||||
|
||||
async def build_item(self, path):
|
||||
"""Build a payload and inspect the preceding payload reference."""
|
||||
del path
|
||||
if self.previous_item is not None:
|
||||
self.release_observations.append(self.previous_item() is None)
|
||||
item = TrackedItem()
|
||||
self.previous_item = weakref.ref(item)
|
||||
return item
|
||||
|
||||
async def upsert_items(self, items):
|
||||
"""Discard the batch so only the apply loop could retain its payload."""
|
||||
del items
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index}.md", "x") for index in range(2)]
|
||||
step = TrackingCatalogStep(
|
||||
persist=False,
|
||||
batch_max_files=1,
|
||||
app_context=_make_app_context(cwd),
|
||||
)
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is True
|
||||
assert step.release_observations == [True]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_continues_after_one_batch_fails():
|
||||
"""An upsert failure is isolated to its bounded batch and later files continue."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
paths = [write_file(cwd / "daily" / f"{index}.md", "x") for index in range(3)]
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
original_upsert = catalog.upsert
|
||||
call_count = 0
|
||||
|
||||
async def fail_first_batch(nodes):
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
if call_count == 1:
|
||||
raise RuntimeError("batch failed")
|
||||
await original_upsert(nodes)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(
|
||||
file_catalog=catalog,
|
||||
persist=False,
|
||||
batch_max_files=2,
|
||||
app_context=_make_app_context(cwd),
|
||||
)
|
||||
changes = [{"change": "added", "path": str(path)} for path in paths]
|
||||
available = MagicMock(available=8 * 1024 * 1024 * 1024)
|
||||
with (
|
||||
patch.object(catalog, "upsert", side_effect=fail_first_batch),
|
||||
patch("reme.steps.index.update_changes.psutil.virtual_memory", return_value=available),
|
||||
):
|
||||
response = await step(RuntimeContext(changes=changes))
|
||||
|
||||
assert response.success is False
|
||||
assert [result["success"] for result in response.answer] == [False, False, True]
|
||||
assert [node.path for node in await catalog.get_nodes()] == ["daily/2.md"]
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_catalog_yields_to_event_loop_while_building_batch():
|
||||
"""Synchronous file inspection does not monopolize the application event loop."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
source = write_file(cwd / "daily" / "a.md", "alpha")
|
||||
catalog = LocalFileCatalog(name="test_catalog")
|
||||
await catalog.start()
|
||||
yielded = False
|
||||
observed = []
|
||||
original_upsert = catalog.upsert
|
||||
|
||||
def mark_yielded():
|
||||
nonlocal yielded
|
||||
yielded = True
|
||||
|
||||
async def observe_upsert(nodes):
|
||||
observed.append(yielded)
|
||||
await original_upsert(nodes)
|
||||
|
||||
try:
|
||||
step = UpdateCatalogStep(file_catalog=catalog, persist=False, app_context=_make_app_context(cwd))
|
||||
asyncio.get_running_loop().call_soon(mark_yielded)
|
||||
with patch.object(catalog, "upsert", side_effect=observe_upsert):
|
||||
response = await step(RuntimeContext(changes=[{"change": "added", "path": str(source)}]))
|
||||
|
||||
assert response.success is True
|
||||
assert observed == [True]
|
||||
finally:
|
||||
await catalog.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
class _CountingEmbeddingStore:
|
||||
dimensions = 2
|
||||
max_batch_size = 10
|
||||
|
||||
def __init__(self):
|
||||
self.calls = 0
|
||||
|
||||
async def get_node_embeddings(self, nodes, **_kwargs):
|
||||
"""Record one embedding call and populate deterministic vectors."""
|
||||
self.calls += 1
|
||||
for node in nodes:
|
||||
node.embedding = np.array([1.0, 0.0], dtype=np.float16)
|
||||
return nodes
|
||||
|
||||
|
||||
def test_update_index_modified_file_reuses_unchanged_embedding():
|
||||
"""The step relies on replace-aware upsert instead of deleting reusable chunks first."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
source = write_file(cwd / "daily" / "a.md", "alpha")
|
||||
fs = LocalFileStore(name="default", embedding_store="")
|
||||
chunker = DefaultFileChunker(supported_extensions=["md"])
|
||||
await fs.start()
|
||||
await chunker.start()
|
||||
embedding_store = _CountingEmbeddingStore()
|
||||
fs.embedding_store = embedding_store
|
||||
try:
|
||||
app_ctx = _make_app_context(cwd)
|
||||
app_ctx.components = {ComponentEnum.FILE_CHUNKER: {"default": chunker}}
|
||||
await fs.upsert([await chunker.chunk(source)])
|
||||
assert embedding_store.calls == 1
|
||||
|
||||
step = UpdateIndexStep(file_store=fs, persist=False, app_context=app_ctx)
|
||||
with patch.object(fs, "delete", wraps=fs.delete) as delete_mock:
|
||||
response = await step(
|
||||
RuntimeContext(changes=[{"change": "modified", "path": str(source)}]),
|
||||
)
|
||||
|
||||
assert response.success is True
|
||||
assert embedding_store.calls == 1
|
||||
delete_mock.assert_not_awaited()
|
||||
finally:
|
||||
await chunker.close()
|
||||
await fs.close()
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_index_skips_oversized_file_and_clears_stale_index():
|
||||
"""Oversized content is not read and any previous index entry is removed."""
|
||||
|
||||
|
|
@ -513,17 +911,16 @@ def test_update_index_handles_file_removed_before_stat():
|
|||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
step = UpdateIndexStep(app_context=_make_app_context(Path.cwd()))
|
||||
results = []
|
||||
step = UpdateIndexStep(persist=False, app_context=_make_app_context(Path.cwd()))
|
||||
|
||||
with (
|
||||
patch.object(Path, "is_file", return_value=True),
|
||||
patch.object(Path, "stat", side_effect=FileNotFoundError("file disappeared")),
|
||||
):
|
||||
item = await step._try_build_item(Change.modified, "Updating", "daily/a.md", results)
|
||||
response = await step(RuntimeContext(changes=[{"change": "modified", "path": "daily/a.md"}]))
|
||||
|
||||
assert item is None
|
||||
assert results == [
|
||||
assert response.success is False
|
||||
assert response.answer == [
|
||||
{
|
||||
"change": "modified",
|
||||
"path": "daily/a.md",
|
||||
|
|
@ -535,6 +932,34 @@ def test_update_index_handles_file_removed_before_stat():
|
|||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_update_index_reports_memory_estimation_failure_without_aborting():
|
||||
"""A missing chunker is isolated to the affected file and returned as a normal failure."""
|
||||
|
||||
async def run():
|
||||
with tempfile.TemporaryDirectory() as tmpdir, temp_chdir(tmpdir):
|
||||
cwd = Path.cwd()
|
||||
source = write_file(cwd / "daily" / "a.unknown", "alpha")
|
||||
app_ctx = _make_app_context(cwd)
|
||||
app_ctx.components = {ComponentEnum.FILE_CHUNKER: {}}
|
||||
step = UpdateIndexStep(persist=False, app_context=app_ctx)
|
||||
|
||||
response = await step(RuntimeContext(changes=[{"change": "added", "path": str(source)}]))
|
||||
|
||||
assert response.success is False
|
||||
assert response.answer == [
|
||||
{
|
||||
"change": "added",
|
||||
"path": str(source),
|
||||
"success": False,
|
||||
"error": (
|
||||
f"No file chunker supports {source} (suffix='unknown') and no default chunker is configured"
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
|
||||
def test_index_update_loop_init_dispatch_updates_store_across_batches():
|
||||
"""index_update_loop init scan dispatches to update_index_step and preserves final store state."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue