Repository navigation
Enforce the warmup memory limit while warmup data is downloaded - #6990
fulmicoton-dd wants to merge 2 commits into
Conversation
1037b1f to
15a6da3
Compare
Until now, a split search only checked the warmup memory budget when it was granted a permit, using a pessimistic estimate. The memory actually used by the warmup was only known after the fact and could go over the budget. The split's CachingDirectory now reserves memory before each cache miss, through a new MemoryReserver trait in quickwit-directories. SplitWarmupMemory implements it: the permit's initial allocation is used first, then the reservation grows by 1.7x, or by the exact need when 1.7x does not fit. If the reservation cannot grow without going over warmup_memory_budget, the read fails with OutOfMemoryError, and the split search fails. Failed splits are still retried for now. SplitWarmupMemory is shared by the SearchPermit and the CachingDirectory, so its memory is freed when the cache holding the warmup data is dropped. After opening the index, the reservation shrinks to the index size, and after the warmup, to the memory in use. WarmupMemoryGuard::grow_unchecked is removed. List terms uses the same mechanism. Its initial permit size is now clamped to warmup_memory_budget, like leaf search, so that it can always be granted. A new leaf_search_split_warmup_out_of_memory_total counter counts the split searches that ran out of warmup memory.
15a6da3 to
ba5893c
Compare
datadog-official
Bot
added
area:search
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ba5893c641
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| self.memory_reserver | ||
| .try_reserve_memory(byte_range.len() as u64)?; |
There was a problem hiding this comment.
When a requested range partially overlaps an existing cached range—or two concurrent reads miss on the same range—this reserves the full request length even though FileByteRangeCache::put_slice stores only the newly covered union (and may add nothing). SplitWarmupMemory::in_use is never reduced to the cache's actual size, so later reads can incorrectly fail with OutOfMemory while the warmup cache remains below the configured budget. Reserve the incremental cache growth or reconcile the reservation after insertion.
Useful? React with 👍 / 👎.
| &split, | ||
| Some(ctx.doc_mapper.tokenizer_manager()), | ||
| Some(byte_range_cache.clone()), | ||
| Some(search_permit.warmup_memory()), |
There was a problem hiding this comment.
When the newly attached reserver rejects a warmup read, the error is converted through io::Error/anyhow into SearchError::Internal; leaf_search_single_split_wrapper then emits it with retryable_error: true, and LeafSearchRetryPolicy retries every response containing failed splits without inspecting that flag. Under memory pressure this immediately repeats the same memory-intensive warmup, contrary to SCHEDULING_AND_MEMORY_MANAGEMENT.md:63-64, which specifies that these OOM failures are non-retryable; preserve an identifiable OOM classification and make the retry policy exclude it.
AGENTS.md reference: AGENTS.md:L23-L24
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
| For this reason, we also have a warmup_max_split_initial_allocation that will fail the request if the | ||
| single split allocation is greater than a given amount; hence increasing the chance that the faulty query is the one that | ||
| will be OOMed first. |
There was a problem hiding this comment.
This documents warmup_max_split_initial_allocation as an available setting that rejects oversized requests, but a repo-wide rg finds no implementation or configuration field outside this paragraph; SearcherConfig exposes only warmup_memory_budget and warmup_single_split_initial_allocation (quickwit-config/src/node_config/mod.rs:435-436), and scheduling clamps an initial reservation to the total budget rather than applying the described per-split rejection. Operators therefore cannot configure the behavior promised here; either implement the setting or describe the limits that actually exist.
AGENTS.md reference: AGENTS.md:L23-L24
Useful? React with 👍 / 👎.
| self.memory_reserver | ||
| .try_reserve_memory(byte_range.len() as u64)?; |
| fn try_reserve_memory(&self, num_bytes: u64) -> Result<(), OutOfMemoryError> { | ||
| self.reserved | ||
| .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |reserved| { | ||
| let new_reserved = reserved + num_bytes; |
There was a problem hiding this comment.
Even if an overflow with u64 is very unlikely, we can use checked_add and safely return OutOfMemoryError.
reserved
.checked_add(num_bytes)
.filter(|&new_reserved| new_reserved <= self.limit)| &split, | ||
| Some(ctx.doc_mapper.tokenizer_manager()), | ||
| Some(byte_range_cache.clone()), | ||
| Some(search_permit.warmup_memory()), |
There was a problem hiding this comment.
| } | ||
| self.memory_reserver | ||
| .try_reserve_memory(byte_range.len() as u64)?; | ||
| let owned_bytes = self.underlying_filehandle.read_bytes(byte_range.clone())?; |
There was a problem hiding this comment.
If read_bytes fails then memory is not released.
| let read_bytes = self | ||
| .underlying_filehandle | ||
| .read_bytes_async(byte_range.clone()) | ||
| .await?; |
There was a problem hiding this comment.
If read_bytes fails then memory is not released.
| /// | ||
| /// On error, the read fails with an `io::Error` of kind `OutOfMemory`: the underlying directory | ||
| /// is not called and nothing is cached. | ||
| fn try_reserve_memory(&self, num_bytes: u64) -> Result<(), OutOfMemoryError>; |
There was a problem hiding this comment.
Not a blocker but worth discussing: the trait MemoryReserver returns () upon success. However, to release memory the user calls release_unused_memory() on a search_permit.
Would it make sense that try_reserve_memory() returns some sort of MemoryGuard that automatically releases the memory when dropped? This also relates to https://github.com/quickwit-oss/quickwit/pull/6990/changes#r4525592085 and https://github.com/quickwit-oss/quickwit/pull/6990/changes#r4525593712.
Context
Second step of the work to prevent searcher OOMs caused by warmup. It builds on #6988, which moved the warmup memory accounting to a shared atomic (
GlobalWarmupMemory).Until now, a split search only checked
searcher.warmup_memory_budgetwhen it was granted a permit, with a pessimistic estimate. The memory actually used by the warmup was only known after the fact, and could go over the budget. Heavy queries (for example leading wildcards) caused OOMs.What
The split's
CachingDirectorynow reserves memory before each cache miss, so the memory is reserved before the data is downloaded.In order to avoid too many "allocation" we allocate at least 1.7x the current amount of memory held by the split.
If not enough memory is available, the allocation fails with an explicit error.
We will improve the treatment of those error in the next PR.
Allocating is simple and does not need much synchronization.
Deallocating happens when WarmupMemoryGuard is dropped. The SearchPermitProvider then get waken up via a message.
Load test
Run on a build with this PR, #7005 and #6993, compared with an earlier build without any warmup memory limit.
Setup: one searcher (64 vCPU, 512 GiB pod), cold caches,
warmup_memory_budget: 100G,warmup_single_split_initial_allocation: 60MiB. Each query is a count over one day of logs: 338 splits, 3.8B docs. For 15 min, 2 QPS of normal counts, plus 2 leading-wildcard queries (message:*ms* OR message:*elapsed*) every 30 s. Each of those needs about 1.4 GB of warmup memory per split.With the default limit, the leading-wildcard queries are stopped by the node-wide budget, not by the per-split limit. The split searches already running keep downloading until then: 7.7 TiB was read from object storage during the 15 min.
With the per-split limit of #7005 set to
warmup_single_split_max_allocation: 1G, below the 1.4 GB these queries need, a single leading-wildcard query on a cold searcher fails all 338 splits in 8.4 s instead of about 20 s, without retries. This was not run under the full load.