Skip to content

fix(_memory): make interrupted reindex resumable, verified by model identity - #1792

Open
omar-nahhas wants to merge 2 commits into
agent0ai:mainfrom
agent-zero-operator:upstream-fix/memory-rebuild-resume-model-check
Open

fix(_memory): make interrupted reindex resumable, verified by model identity#1792
omar-nahhas wants to merge 2 commits into
agent0ai:mainfrom
agent-zero-operator:upstream-fix/memory-rebuild-resume-model-check

Conversation

@omar-nahhas

Copy link
Copy Markdown

Summary

When the configured embedding model no longer matches the one a memory store was built with, Memory.initialize() triggers a full reindex — but it embedded every document in a single pass and only persisted the result at the very end. Interrupting that pass (a crash, a restart, a slow/shared embeddings backend timing out mid-batch) threw away every already-embedded document, so the next attempt started over from zero. Against a slow or shared embeddings backend this can cost hours of redone work on every interruption — hit this in production on a small local GPU box.

  • Reindexing now embeds in batches (Memory._REBUILD_BATCH_SIZE, 20) and checkpoints to disk after each batch (index.rebuilding.faiss/.pkl), so an interrupted rebuild resumes where it left off.
  • The checkpoint is tagged with the embedding model it was built under (index.rebuilding.json) and is only trusted if that matches the model we're about to (re)index with. If the target model changed again while a rebuild was interrupted, the stale checkpoint is discarded and a fresh rebuild starts — otherwise resuming would add new-model vectors into an old-model FAISS index, failing a dimension assertion on the next batch.
  • On successful completion the checkpoint files are cleaned up.

No change to the existing reindex-trigger logic (comparing embedding.json's recorded provider/name against the current config) — this only changes how the reindex, once triggered, is executed.

Test plan

  • New tests/test_memory_rebuild_resume.py — 9 tests covering the checkpoint save/load/clear helpers directly, plus Memory.initialize() end-to-end for both the resume-under-same-model and discard-under-changed-model cases. All pass.
  • Ran the broader suite — no regressions (two pre-existing, unrelated failures already present on main untouched: a stale import in email_parser_test.py, and rate_limiter_test.py requiring live OpenRouter credentials).

🤖 Generated with Claude Code

…dentity

Memory.initialize() used to embed every document in one pass when a
reindex was needed (embedding model changed) and only persist the result
at the very end. Interrupting that pass -- a crash, a restart, a slow or
shared embeddings backend timing out mid-batch -- threw away every
already-embedded document; the next attempt started over from zero. On a
large memory store against a slow backend this can cost hours of redone
work every time the process is interrupted.

Reindexing now embeds in batches (Memory._REBUILD_BATCH_SIZE) and
checkpoints to disk after each one (index.rebuilding.faiss/.pkl), so an
interrupted rebuild resumes from where it left off instead of restarting.

The checkpoint is not trusted blindly: it's tagged with the embedding
model it was built under (index.rebuilding.json), and a resume is only
used if that matches the model we're about to (re)index with. If the
target model changed again while a rebuild was interrupted, the stale
checkpoint is discarded and a fresh rebuild starts -- otherwise resuming
would add new-model vectors into an old-model FAISS index, which fails a
dimension assertion on the very next batch.

Adds tests/test_memory_rebuild_resume.py covering the checkpoint
save/load/clear helpers directly and Memory.initialize() end-to-end for
both the resume-under-same-model and discard-under-changed-model cases.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: c4dca4d4ed

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread plugins/_memory/helpers/memory.py Outdated
Comment on lines +115 to +118
if (
meta.get("model_provider") != model_config.provider
or meta.get("model_name") != model_config.name
):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Include embedding kwargs in checkpoint identity

When a rebuild is resumed after the user changes an embedding option while keeping the same provider/name (for example OpenAI dimensions, or api_base pointing to a different compatible backend), this check still accepts the old checkpoint even though initialize() constructs the embedder from model_config.build_kwargs(). The next batch can then append vectors from a different dimension/space to the checkpoint, causing a FAISS dimension assertion or a mixed index instead of the fresh rebuild intended for the new configuration; include the effective non-secret embedding signature (or at least the vector dimension) in the checkpoint validation.

AGENTS.md reference: plugins/_memory/AGENTS.md:L17-L19

Useful? React with 👍 / 👎.

Comment thread plugins/_memory/helpers/memory.py Outdated
PrintStyle.standard(f"Indexed {done}/{total} memories...")
Memory._write_rebuild_checkpoint(db, db_dir, model_config)

Memory._clear_rebuild_checkpoint(db_dir)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear the rebuild checkpoint only after the final save

Because the checkpoint is removed before the rebuilt index.* files and the new embedding.json are durably written, a kill/restart during _save_db_file() or in the gap just before it leaves no index.rebuilding.* files to resume from while the persisted index is still stale or only partially written. For large stores this reintroduces the full re-embed-on-retry failure mode that the checkpoint is meant to avoid; move the cleanup after the final index and metadata writes succeed.

Useful? React with 👍 / 👎.

…e ordering

Two gaps flagged by review on the resumable-reindex fix:

1. The rebuild checkpoint's model-identity check only compared
   provider/name. A kwarg that changes the embedder's actual output (e.g.
   an OpenAI-style `dimensions` override, or an api_base pointed at a
   different backend) could change while provider/name stayed the same,
   letting a checkpoint built under the old kwargs get resumed into under
   the new ones. The signature now includes build_kwargs() (minus api_key,
   which is deliberately excluded since it's persisted to disk in the
   memory dir and a credential rotation alone isn't "a different model").

2. The rebuild checkpoint was cleared right after the batch loop, before
   the final index/embedding.json were durably written. A crash in that
   gap left neither a valid final index nor a resumable checkpoint,
   reintroducing the full re-embed-from-scratch failure mode the checkpoint
   exists to prevent. Clearing now happens strictly after the final save
   succeeds.

Adds 3 tests: api_key exclusion, kwarg-change rejection, and
checkpoint-survives-a-crash-during-final-save.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants