fix(index): batch IVF streaming partition search off the CPU pool#7680
fix(index): batch IVF streaming partition search off the CPU pool#7680a-agmon wants to merge 4 commits into
Conversation
The streaming/sequential partition search ran its whole recv/search/send loop inside a single `spawn_cpu` closure, using `blocking_recv` and `blocking_send` on capacity-1 channels. Both park the CPU-pool thread, and the producer feeding the input channel does object-store I/O whose partition decode can itself need the CPU pool. On hosts whose CPU pool collapses to a single thread (`<= 3` CPUs), or under enough concurrent queries, the parked search thread and the work that would unblock it can starve each other and deadlock the pool with a silent 0% hang — the same class as lance-format#7423. Move the channel `recv`/`send` back into the async task and hand only the pure-CPU search to `spawn_cpu`, over a batch of `STREAMING_SEARCH_BATCH_SIZE` partitions at a time. No CPU-pool thread ever parks on a channel, so the deadlock is structurally gone, while batching amortizes the per-dispatch overhead the single-worker design in lance-format#6475 was avoiding (fan-out paid N/K times instead of N). `should_stop` is still checked per partition, so early-stop granularity is unchanged. The `PreparedThreadCapturingIndex` test mock is updated to mirror the new non-blocking structure. Adds a regression test that a search spanning multiple batches still prepares and searches every partition in order. Closes lance-format#7642 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Builds on the batched streaming search (lance-format#7673) with review follow-ups: - Form batches greedily: wait for one prepared partition, then drain whatever else is already prepared, up to the batch size. Waiting for a full batch delayed the first search (and the early-stop feedback it produces) behind up to a whole batch of prepare I/O, which is significant when prepare parallelism is low. The prepared channel now holds a full batch so partitions prepared while the previous batch is searching are ready for the next drain. - Stop pulling and searching as soon as the output receiver is dropped (cancelled queries no longer search out their whole batch first). - Make the batch size tunable via LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE (default 16, must be > 0) so the perf sweep lance-format#7642 asks for can run without recompiles. - Correct the comment on the defensive spawn_cpu error arm. - Add an end-to-end smoke test that runs the streaming search path (HNSW sub-index, prefiltered stable-row-id query) in a child process with LANCE_CPU_THREADS=1. The old deadlock itself is a race against prefilter-mask readiness that local filesystems always win (mask I/O is sub-millisecond and local reads bypass the ObjectStore trait, so latency cannot be injected); the test documents that honestly and guarantees the path completes end to end on a single-thread pool.
…ellation - Adapt test_sequential_initial_search_prepares_all_then_searches_on_one_cpu_thread to the configured batch size (min(3, batch)) instead of asserting 3 <= batch, so every valid LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE, including 1, passes. - Probe Sender::is_closed() per partition inside the spawn_cpu search closure (via a sender clone; it is synchronously callable), so a cancelled query stops after at most one more partition search instead of searching out the whole in-flight batch. A select! on closed() would not achieve this: spawn_cpu closures are not cancellable, so abandoning the await leaves the CPU work running.
|
ACTION NEEDED The PR title and description are used as the merge commit message. Please update your PR title and description to match the specification. For details on the error please inspect the "PR Title Check" action. |
a-agmon
changed the title
|
Hi @wjones127 , @BubbleCal - would be happy for your review |
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Overall good to me, we need to verify the performance deeper, could you try this:
- 1M x 1024d f32 vectors
- creating IVF_RQ index with the default index params(just not specify any index params)
- use
prewarmto warm the index before benchmarking - run with concurrency=1 and concurrency=(num_cores-2) to get the numbers
|
Sounds good @BubbleCal num_bits=1, fast rotation), L2prewarm_index before measuringResults (ranges across the 3 rounds)
No measurable difference — per-round deltas are under ~2% and flip sign between rounds. p90/p99 behave the same way. One note for context: unfiltered IVF_RQ queries take the collect-all/global-top-k path, which this PR doesn't modify — so this run is a no-regression check on the mainline query path. The code this PR actually changes (the batched streaming loop) runs for prefiltered late-search and for sub-indices without global-top-k support; the batch-size sweep in the PR description benchmarks that path directly against main and also shows parity. benchmark harness available here : https://gist.github.com/a-agmon/d22e1cf7045112b45bf6f619636af1ed (it's a small standalone example that builds the dataset on first run). |
wjones127
left a comment
wjones127
left a comment
There was a problem hiding this comment.
I think I'd like that test removed, but otherwise this fix looks good.
📝 WalkthroughWalkthroughThis PR reworks the streaming IVF partition search path to batch prepared partitions per spawn_cpu dispatch instead of processing them one at a time via blocking channel operations. A configurable batch size constant/static is added, the prepared-partition channel is resized, and the consumer loop is rewritten to greedily accumulate batches, checking early-stop and channel-closed conditions. Corresponding test mocks and cases are updated to mirror the new batching behavior. ChangesBatched streaming partition search
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Producer
participant PreparedChannel
participant Consumer
participant SpawnCpu
participant BatchTx
Producer->>PreparedChannel: send prepared partition
Consumer->>PreparedChannel: recv first prepared partition
Consumer->>PreparedChannel: try_recv additional partitions up to batch size
Consumer->>SpawnCpu: dispatch single batch search task
SpawnCpu->>SpawnCpu: check search_control.should_stop()
SpawnCpu->>BatchTx: check is_closed() for cancellation
SpawnCpu->>BatchTx: forward per-partition results/errors
Consumer->>Consumer: repeat until producer done and batches handled
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Good point. Thanks @wjones127. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
rust/lance/src/index/vector/ivf/v2.rs (1)
145-158: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAvoid
.expect()/assert!()panics in library init; log and fall back instead.A malformed or non-positive
LANCE_IVF_STREAMING_SEARCH_BATCH_SIZEwill panic on first search rather than degrade gracefully. Note theassert!(batch_size > 0)is load-bearing (a0would reachmpsc::channel(0)at Line 1664, which panics in tokio), so the invariant must stay — but both can be enforced without a panic.As per coding guidelines: "Never use
.unwrap(),.expect(),panic!(), orassert!()in library code for fallible operations."♻️ Log a warning and fall back to the default
pub(crate) static STREAMING_SEARCH_BATCH_SIZE: LazyLock<usize> = LazyLock::new(|| { - let batch_size = std::env::var("LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE") - .map(|value| { - value - .parse() - .expect("failed to parse LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE") - }) - .unwrap_or(DEFAULT_STREAMING_SEARCH_BATCH_SIZE); - assert!( - batch_size > 0, - "LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE must be greater than 0, got {batch_size}" - ); - batch_size + match std::env::var("LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE") { + Err(_) => DEFAULT_STREAMING_SEARCH_BATCH_SIZE, + Ok(value) => match value.parse::<usize>() { + Ok(n) if n > 0 => n, + _ => { + warn!( + "invalid LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE={value:?}; \ + falling back to {DEFAULT_STREAMING_SEARCH_BATCH_SIZE}" + ); + DEFAULT_STREAMING_SEARCH_BATCH_SIZE + } + }, + } });Requires a
warn!import if not already in scope.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rust/lance/src/index/vector/ivf/v2.rs` around lines 145 - 158, Replace the panic-based initialization in STREAMING_SEARCH_BATCH_SIZE with fallible parsing and validation that logs a warning and falls back to DEFAULT_STREAMING_SEARCH_BATCH_SIZE instead of using .expect() or assert!(). Update the LazyLock closure in v2.rs so malformed, missing, zero, or negative LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE values are handled gracefully, while still enforcing the > 0 invariant before the value is used by the search path that leads to mpsc::channel in the IVF v2 flow. Use the existing STREAMING_SEARCH_BATCH_SIZE symbol and add a warn! message that explains the invalid env var and the fallback.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@rust/lance/src/index/vector/ivf/v2.rs`:
- Around line 145-158: Replace the panic-based initialization in
STREAMING_SEARCH_BATCH_SIZE with fallible parsing and validation that logs a
warning and falls back to DEFAULT_STREAMING_SEARCH_BATCH_SIZE instead of using
.expect() or assert!(). Update the LazyLock closure in v2.rs so malformed,
missing, zero, or negative LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE values are
handled gracefully, while still enforcing the > 0 invariant before the value is
used by the search path that leads to mpsc::channel in the IVF v2 flow. Use the
existing STREAMING_SEARCH_BATCH_SIZE symbol and add a warn! message that
explains the invalid env var and the fallback.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 4f91e922-380b-4a0c-88aa-2b598844b53f
📒 Files selected for processing (2)
rust/lance/src/index/vector/ivf/v2.rsrust/lance/src/io/exec/knn.rs
Deliberate, following the predominant pattern for lance's compute knobs — |
3 participants
Closes #7642. Supersedes #7673 — the first commit is @wjones127's original
implementation from that PR, unchanged; the second applies the review
follow-ups and makes the batch size measurable.
Problem
The streaming partition search ran its whole recv/search/send loop inside one
spawn_cpuclosure with blocking channel ops. If the prefilter's row mask wasn'tready when that closure parked the pool thread, the mask's own
spawn_cpuworkqueued behind it and the query hung at 0% CPU — same class as #7423.
Fix
Channel
recv/sendmove to the async task; only the pure-CPU search runs inspawn_cpu, over a batch of partitions per dispatch (from #7673). On top of that:already ready, up to the batch size — instead of waiting for a full batch,
which delayed the first search behind up to a batch of prepare I/O. Batch size
adapts to producer speed: slow producer → batches of ~1 (old latency behavior),
fast producer → full batches.
partition (checked inside the CPU closure via
Sender::is_closed(), sincespawn_cpuwork isn't cancellable from the async side).LANCE_IVF_STREAMING_SEARCH_BATCH_SIZE(default 16).Performance
Swept the batch size with prefiltered IVF_PQ late-search queries (100K×64d,
64 partitions, 0.2% filter selectivity, k=100 — each query streams ~30+
partitions), 512 queries at concurrency 16, on a 4-CPU box (2-thread pool):
No measurable regression vs the blocking baseline, and throughput is insensitive
to the batch size — with greedy draining, actual batches track producer speed
rather than the cap, so the env var is an escape hatch, not a required tuning.
Testing
fix(index): batch IVF streaming partition search off the CPU pool #7673, under any valid batch size (including 1).
process with
LANCE_CPU_THREADS=1, asserting completion under a deadline.is a race against prefilter-mask readiness that local filesystems always win
(mask I/O is sub-millisecond, and local reads bypass the
ObjectStoretraitso latency can't be injected). Real object-storage latency is what loses the
race. The smoke test documents this and guards the path end to end.
Summary by CodeRabbit
New Features
Bug Fixes