perf(index): budget in-flight partition preparation process-wide

[BubbleCal]

Problem

Follow-up to #9213. Each search_partitions call keeps up to prepare_parallelism (the CPU count) partitions in flight ahead of scoring. A query fans out over every index segment on the node (a distributed table has tens: the incident cluster searched 13 segments per plan executor) and queries run concurrently, so the node-wide prepare window is segments * queries * CPUs partitions. With ~7 MiB IVF_RQ partitions on a 32-vCPU node that is 13 * 32 * 7 MiB ≈ 2.9 GiB per query just for the window, before the per-segment scoring chunks, and it grows with concurrency.

Change

• Add a process-wide PreparePartitionBudget (a tokio::sync::Semaphore), default 2 * CPUs partitions, overridable with LANCE_IVF_PREPARE_PARTITION_BUDGET.
• Every partition load on the global-top-k path, the streaming path and the multi-query batch path acquires a permit before loading and releases it the moment the prepared partition is pulled out of the buffered window (.buffered(n).map(release_prepare_permit)), before it enters a scoring chunk or channel. Permits are acquired sequentially in probe order in a then stage ahead of buffered (a named async fn, since the compiler cannot type an async block as a then callback): buffered yields in order, so letting the buffered futures acquire on their own allowed later partitions to hold permits behind a waiting head, and eight concurrent probe-everything searches then held the whole budget with none able to advance (reproduced on the benchmark VM and in the new test). With in-order acquisition a window's permit holders are always its oldest entries, scoring never holds permits, and chunk assembly never waits on permits held by its own partitions, so no combination of budget, chunk and channel sizes can deadlock. The chunks stay bounded per segment search as before (GLOBAL_TOPK_CHUNK_BYTES, STREAMING_SEARCH_BATCH_SIZE).
• Segments share the budget dynamically rather than splitting it up front, so a segment stalled on I/O does not idle the others' share, and a single segment search on an idle node is not throttled below its own window.
• The budget lives on IVFIndex as an ArcSwap pointing at the process-wide instance, so tests can run a search under a tiny budget.

Node-wide in-flight prepared partitions become min(segments * queries * CPUs, 2 * CPUs); for the 32-vCPU, 13-segment case above the window drops from ≈2.9 GiB per query to ≈0.45 GiB shared by all queries.

Validation

• New test_prepare_budget_of_one_bounds_window_without_deadlock: an IVF_PQ index with more partitions than any chunk or channel holds, searched under a budget of one partition through the global-top-k path (two searches concurrently), the streaming path (an early-stop control that never stops) and the multi-query batch path, plus eight concurrent global-top-k searches under a budget of four; all complete, return the same rows as under the default budget, and leave the permits released. The eight-search case deadlocked with in-future permit acquisition.
• cargo test -p lance --lib index::vector::ivf::v2 (114 passed), io::exec::knn (152 passed), scanner::test::test_knn / test_ann (65 passed).
• cargo clippy --all --tests --benches -- -D warnings and RUSTDOCFLAGS="-D warnings" cargo doc -p lance --no-deps clean.
• Performance: A/B on an AWS c7i.16xlarge (64 vCPU, so the default budget is 128 partitions), release profile, baseline = main 5d27099, this PR = 9932db3. Ad-hoc harness (not committed): 1M x 768-d random vectors, IVF_RQ8 with 1024 partitions (~780 KiB each), seeded queries after warm-up, 3 repetitions interleaved baseline/PR, means of the per-repetition means. conc is the number of queries in flight in one process; peak RSS is ru_maxrss. Lower latency / higher QPS / lower RSS is better.

Scenario / metric	Baseline	This PR	Benefit
nprobes=20, k=10, warm, conc=1, mean latency	8.63 ms	8.68 ms	1.00x
nprobes=1024, k=10, warm, conc=1, mean latency	26.25 ms	25.94 ms	1.01x
nprobes=1024, k=10, 64 MiB cache, conc=1, mean latency	339.6 ms	339.7 ms	1.00x
nprobes=20, k=10, warm, conc=16, throughput	373 q/s	520 q/s	1.39x speedup
nprobes=20, k=10, warm, conc=16, peak RSS	1,227 MiB	1,059 MiB	1.16x less memory
nprobes=1024, k=10, warm, conc=8, throughput	327 q/s	318 q/s	0.97x
nprobes=1024, k=10, 64 MiB cache, conc=8, throughput	16.3 q/s	15.8 q/s	0.97x
nprobes=1024, k=10, 64 MiB cache, conc=8, peak RSS	1,739 MiB	1,639 MiB	1.06x less memory

The harness has a single index segment, so the budget only bites under concurrency: 8 or 16 concurrent searches want 8 * 64 / 16 * 64 partitions in flight against 128 permits. Cutting the cold-read concurrency 4x costs 3% throughput because the reads saturate the NVMe either way; the 16-query warm case gets faster because far fewer prepared partitions compete for the CPU pool at once. The memory difference is small here because ~780 KiB partitions are 50 MiB per window; on the multi-MiB partitions of a billion-row RQ index the same permit count is GiBs. The repo's vector_throughput bench (IVF_PQ, 1M x 768-d, k=50, nprobes=20, refine 10) run interleaved baseline/PR/baseline/PR: 1 thread 1.867 s / 1.974 s vs 1.883 s / 1.941 s, 16 threads 516 ms / 510 ms vs 512 ms / 516 ms, i.e. no difference beyond its run-to-run spread.

🤖 Generated with Claude Code

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The process-wide memory contract needs to cover partition-parallel execution, not only the sequential search_partitions pipeline. Applying the same shared budget at the direct per-partition load/score boundary is the minimal complete fix; otherwise supported search modes retain the node-wide fan-out this change is meant to remove.

[lance-gatekeeper]

This process-wide cap is bypassed by partition-parallel searches, so the node can still retain segments * queries * parallelism loaded partitions and hit the same OOM mode. ANNIvfSubIndexExec calls search_partitions only when effective parallelism is at most 1; documented query_parallelism >= 2 instead buffers direct search_in_partition calls, and non-global-heap sub-indexes select that branch automatically. search_in_partition loads and then queues/scores the partition without acquiring this budget. Please acquire the same budget around the direct load/score path (no downstream chunk waits there) and add that branch to the regression.

Reproducer

Against 32b512bb891c8a62495936540eef48371e3345e6, I appended this to test_prepare_budget_of_one_bounds_window_without_deadlock after installing its one-permit budget:

let held_permit = budget.acquire().await.unwrap();
let part_id = probes[0].0.value(0) as usize;
let direct_search = index.search_in_partition(
    part_id,
    &queries[0],
    Arc::new(NoFilter),
    &NoOpMetricsCollector,
);
assert!(
    tokio::time::timeout(std::time::Duration::from_secs(2), direct_search)
        .await
        .is_err(),
    "partition-parallel search bypassed the process-wide prepare budget"
);
drop(held_permit);

Then I ran:

cargo test -p lance --lib index::vector::ivf::v2::tests::test_prepare_budget_of_one_bounds_window_without_deadlock -- --exact --nocapture

Expected the direct search to remain blocked while the only permit was held. It completed instead, and the assertion failed with partition-parallel search bypassed the process-wide prepare budget.

[lance-gatekeeper]

❌ Gate recommendation: request changes.

The permit-ordering rewrite fixes the contended-search deadlock, but the process-wide coverage finding remains: documented partition-parallel searches still load and queue partitions through search_in_partition without acquiring the shared budget, retaining segments * queries * parallelism memory scaling. Apply the same budget at that direct load/score boundary to complete the node-wide cap.