Repository navigation
fix: embed documents in batches with a per-batch budget and fair scheduling - #140
Merged
Merged
Conversation
…duling A chunked document went to the embedding provider in one request, so one timeout bounded the whole document and a large text could not be stored at any allowed EMBEDDING_TIMEOUT_S. Each batch of chunks is now one provider request with its own retries and total timeout, so a retry re-sends only the batch that failed. EMBEDDING_BATCH_SIZE sizes the batches, by default from the measured seconds per chunk toward EMBEDDING_BATCH_TARGET_S, or as a fixed count. Each batch holds its own embedding semaphore slot, so a caller that arrives during a long document is served after the current batch instead of after the whole document. EMBEDDING_TIMEOUT_S bounds one attempt of one request and accepts up to 3600 seconds, and the error messages name the failing batch and the chunks embedded so far.
A 280-chunk document at batch 3, chunks 9-15, projects 3 + ceil(265 / 7) = 41 batches, so the examples now read batch 3 of 41.
The progress in a failed batch's error message counted the failing batch's own attempts, backoff and timeout, so a batch that exhausted its retries reported the chunks already embedded as taking at least the whole per-batch timeout. The time now runs from the document's start to the arrival of the last successful batch's vectors.
The auto sizer learned seconds per chunk, so a one-chunk short text such as a one-line message made the full-size chunks of the next large document look cheap and its first batch ran far past EMBEDDING_BATCH_TARGET_S. The sizer now keeps decayed sums of seconds and characters, and fills each batch with the document's next chunks while their predicted time stays within the target. The settings, server.json and documentation describe the per-character sizing.
No test or fixture called _reset_batch_sizer, and tests that need a particular sizer patch _batch_sizer directly. The test of the lazy build now patches _batch_sizer to None and checks that the fixed and the auto sizer come from the settings.
… three batches With the shared auto sizer the twelve-chunk document could go out in one request once earlier short stores had measured a fast provider, so the check did not exercise batching. Every harness server now runs with EMBEDDING_BATCH_SIZE=4, and the check's docstring and recorded detail state what it verifies: the store completes and semantic search finds the entry.
Each embedding batch holds one EMBEDDING_MAX_CONCURRENT slot, and a slot kept after a batch fails would leave every later store waiting without a timeout. The new tests run a provider that raises, one that outlasts the batch timeout, and a batch whose task is cancelled, and assert the semaphore is free afterwards.
The comment above the offload of the chunk split states what the code does: the splitter's fast path at or below chunk_size decides whether the split runs on a worker thread.
The batching loop sends one embed_documents request per batch, so an INFO line per request repeated once per batch at the default log level without the batch position. Each provider now logs its request at DEBUG beside the batching loop's per-batch line, and the document's start and completion stay at INFO as the semantic search guide states.
Coverage report
Click to see where and how coverage changed
This report was generated by python-coverage-comment-action |
||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
The embedding leg sent every chunk of a document to the provider in one request under one
EMBEDDING_TIMEOUT_Swindow, and each retry re-sent the whole document. On a CPU Ollama deployment a text of a few hundred kilobytes therefore could not be stored or updated at all once its chunks needed more than the 300 s ceiling, and a small store issued while another client's document was embedding waited for that whole run, because the embedding slot was held for the entire document.What changes
EMBEDDING_TIMEOUT_Snow bound one request (one batch on the write path, one query on search), so a retry never re-sends chunks that are already embedded and the document's size no longer decides whether it can be stored.EMBEDDING_TIMEOUT_Saccepts up to 3600 s.EMBEDDING_BATCH_SIZE(defaultauto) sizes batches from the measured seconds per character towardEMBEDDING_BATCH_TARGET_S(default 5 s), bounded between 1 and 64 chunks after a first batch of 4; an integer from 1 to 512 fixes the size for operators who know their provider.server.jsonand the project guidance file describe the batches, the per-request timeout and the two settings.Verification
server.jsonparity.qwen3-embedding:0.6b, default embedding settings) a 400 KB Markdown document stores and then updates on SQLite and on PostgreSQL: its 341 chunks embed in about 30 batches in about 135 s when the embedding model has the CPU to itself, and in about 70 smaller batches in about 300 s while a summary is generated beside it, the sizer following the measured throughput either way. A 3-chunk store issued during that run gets its embeddings in under 5 s. The per-batch slot bounds the embedding wait only: when summaries run on the same single-request Ollama model, a concurrent store's summary still waits for the running summary to finish.