Skip to content

fix(retain): write a held batch when the producer is blocked on the memory budget - #5620

Merged
nicoloboschi merged 2 commits into
mainfrom
fix/retain-budget-deadlock-5296
Oct 15, 2026
Merged

nicoloboschi merged 2 commits into
mainfrom
fix/retain-budget-deadlock-5296

Conversation

@nicoloboschi

Copy link
Copy Markdown
Collaborator

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

  • RetainMemoryBudget now 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.
  • On that signal the writer saves its open batch, which frees the room.

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.get is safe because an item only leaves the queue when the getter resumes.

Tests (tests/test_retain_memory_budget.py)

  • The issue's case over the real budget: [40, 70] with budget 100 → batches [[40], [70]], no hang.
  • Normal batching is unchanged while the extractor fits.
  • Queued items still win over the signal; the signal doesn't swallow a queue item; the signal clears once no reservation waits (admitted or cancelled).
  • End to end through the real _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.py and test_retain_chunks_embedding_batching.py also pass (88 passed).

…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
@nicoloboschi
nicoloboschi merged commit a8866f4 into main Oct 15, 2026
117 checks passed
@nicoloboschi
nicoloboschi deleted the fix/retain-budget-deadlock-5296 branch October 15, 2026 13:28
@nicoloboschi nicoloboschi mentioned this pull request Nov 3, 2026
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.

retain: unequal extracted chunk sizes can deadlock the memory-budget producer and DB consumer

1 participant