Repository navigation
fix(retain): write a held batch when the producer is blocked on the memory budget - #5620
Merged
Merged
Conversation
…emory budget (#5296) A consumer holding a batch under half the budget waited for the next queue item, while the producer waited for room that only that batch could free. The budget now signals when a reservation is blocked; a consumer with an open batch and an empty queue writes it on that signal. With items queued, batching is unchanged.
Review fixes: - next_item_or_producer_waiting is now a RetainMemoryBudget method instead of a module function reaching into the budget's private event - drop the duplicated sentinel check in the consumer loop - explain why the sentinel is an enum
This was referenced Oct 16, 2026
Merged
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.
Fixes #5296.
Problem
The streaming retain pipeline could hang for good. The writer held a small batch (under half the memory budget) and waited for the next chunk. The extractor's next chunk did not fit next to that batch, so it waited for room. Room only frees up when the writer saves. Each side waited on the other until the retain's time limit.
Example: 1 MB budget, a 0.4 MB chunk, then a 0.7 MB chunk, batch size > 1.
Fix
RetainMemoryBudgetnow tracks when a reservation is blocked.next_item_or_producer_waiting(queue, budget): if the queue is empty and the extractor is blocked, the writer gets a "producer waiting" signal instead of waiting for an item.The signal is only checked when the writer holds a batch and the queue is empty. If items are queued, batching works exactly as before, so batches don't get smaller in normal runs. Abandoning the pending
Queue.getis safe because an item only leaves the queue when the getter resumes.Tests (
tests/test_retain_memory_budget.py)[40, 70]with budget 100 → batches[[40], [70]], no hang._streaming_retain_batch(0.4 MB + 0.7 MB on a 1 MB budget). It times out without the orchestrator change and passes with it.test_retain.py,test_retain_pipeline_cancellation.pyandtest_retain_chunks_embedding_batching.pyalso pass (88 passed).