From bf56bc52d08fab7b020ef26f2b9e1a6add541afa Mon Sep 17 00:00:00 2001 From: generall Date: Thu, 3 Sep 2026 17:04:49 +0200 Subject: [PATCH] Interleave serverless requests across collections MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both `serverless upload` and `serverless query` drained collections sequentially: the whole budget of collection 0, then collection 1, and so on. With bounded parallelism that means only one or two collections ever see concurrent traffic, which is not the multi-collection workload the mode is supposed to produce. Pre-allocate per-collection budgets as before, then drain them in random order (`drain_budgets`) so in-flight requests span many collections while the distribution still comes from the budgets. Sampling while draining is uniform on purpose — the uniform/Zipf skew already lives in the budgets, so re-applying the picker there would skew twice. Query used to pick a collection per request, which made the actual per-collection query counts a random walk rather than the requested distribution; it now walks the same pre-planned batches. Requests are no longer all `batch_size` items, so `Processor` gains `request_count`/`request_size` (defaulting to the old arithmetic) and the progress bar asks the processor instead of assuming a fixed size. Co-Authored-By: Claude Opus 5 (1M context) --- src/processor.rs | 11 ++++++ src/serverless/distribution.rs | 67 ++++++++++++++++++++++++++++++++++ src/serverless/query.rs | 23 +++++++++--- src/serverless/upload.rs | 51 +++++++++++--------------- src/stats.rs | 11 ++---- 5 files changed, 120 insertions(+), 43 deletions(-) diff --git a/src/processor.rs b/src/processor.rs index b9d831d..009d0be 100644 --- a/src/processor.rs +++ b/src/processor.rs @@ -33,4 +33,15 @@ pub trait Processor { } fn get_batch_size(&self) -> usize; + + /// Number of requests needed for `total_items`. Processors with variable + /// request sizes can override this to expose a precomputed schedule. + fn request_count(&self, total_items: usize) -> usize { + total_items.div_ceil(self.get_batch_size()) + } + + /// Number of items represented by one request (for progress accounting). + fn request_size(&self, _req_id: usize) -> usize { + self.get_batch_size() + } } diff --git a/src/serverless/distribution.rs b/src/serverless/distribution.rs index 31a9330..b0143ed 100644 --- a/src/serverless/distribution.rs +++ b/src/serverless/distribution.rs @@ -20,6 +20,56 @@ pub struct CollectionPicker { zipf: Option>, } +/// One request produced while draining pre-allocated collection budgets. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct BudgetBatch { + pub collection: usize, + /// Global item offset. Each collection owns one contiguous range, even + /// though requests from those ranges are emitted in random order. + pub offset: usize, + pub count: usize, +} + +/// Randomly drain collection budgets in batches of at most `batch_size`. +/// +/// Sampling is uniform among collections that still have budget. The desired +/// uniform/Zipf skew is already represented by `budgets`; applying it again +/// while draining would skew the workload twice. +pub fn drain_budgets(budgets: &[usize], batch_size: usize, rng: &mut impl Rng) -> Vec { + assert!(batch_size > 0, "batch size must be positive"); + + let mut remaining = budgets.to_vec(); + let mut next_offset = Vec::with_capacity(budgets.len()); + let mut offset = 0usize; + for &budget in budgets { + next_offset.push(offset); + offset += budget; + } + let mut active: Vec<_> = remaining + .iter() + .enumerate() + .filter_map(|(i, &budget)| (budget > 0).then_some(i)) + .collect(); + let mut batches = Vec::new(); + + while !active.is_empty() { + let active_idx = rng.random_range(0..active.len()); + let collection = active[active_idx]; + let count = remaining[collection].min(batch_size); + batches.push(BudgetBatch { + collection, + offset: next_offset[collection], + count, + }); + next_offset[collection] += count; + remaining[collection] -= count; + if remaining[collection] == 0 { + active.swap_remove(active_idx); + } + } + batches +} + impl CollectionPicker { pub fn new(n: usize, distribution: Distribution) -> anyhow::Result { anyhow::ensure!(n > 0, "collections-count must be > 0"); @@ -92,4 +142,21 @@ mod tests { // Rank 0 should get strictly more than the last rank on average. assert!(counts[0] > counts[9]); } + + #[test] + fn draining_budgets_preserves_limits_offsets_and_totals() { + let mut rng = StdRng::seed_from_u64(7); + let batches = drain_budgets(&[5, 0, 7], 3, &mut rng); + assert!(batches.iter().all(|batch| batch.count <= 3)); + let mut totals = [0usize; 3]; + let mut offsets = [Vec::new(), Vec::new(), Vec::new()]; + for batch in batches { + totals[batch.collection] += batch.count; + offsets[batch.collection].push(batch.offset); + } + assert_eq!(totals, [5, 0, 7]); + offsets.iter_mut().for_each(|values| values.sort_unstable()); + assert_eq!(offsets[0], [0, 3]); + assert_eq!(offsets[2], [5, 8, 11]); + } } diff --git a/src/serverless/query.rs b/src/serverless/query.rs index f90292e..2192c1e 100644 --- a/src/serverless/query.rs +++ b/src/serverless/query.rs @@ -23,7 +23,7 @@ use qdrant_client::serverless::{CollectionConfig, QdrantServerless}; use super::args::ServerlessQueryArgs; use super::client::create_clients; use super::collections::list_matching; -use super::distribution::CollectionPicker; +use super::distribution::{BudgetBatch, CollectionPicker, drain_budgets}; use crate::args::Args; use crate::client::retry_with_clients; use crate::config::examples::{ExampleKind, resolve}; @@ -48,7 +48,7 @@ struct ServerlessQueryProcessor { stopped: Arc, clients: Vec, names: Vec, - picker: CollectionPicker, + batches: Vec, generator: ConfigSearchGenerator, start_timestamp_millis: f64, start_time: Instant, @@ -97,11 +97,12 @@ impl ServerlessQueryProcessor { return Ok(()); } + let batch = self.batches[req_id]; let mut rng = rand::rng(); - let collection = &self.names[self.picker.pick(&mut rng)]; + let collection = &self.names[batch.collection]; let template_idx = self.generator.random_template_idx(&mut rng); - let queries = (0..self.args.search_batch_size) + let queries = (0..batch.count) .map(|_| { let generated = self .generator @@ -183,6 +184,14 @@ impl Processor for ServerlessQueryProcessor { fn get_batch_size(&self) -> usize { self.args.search_batch_size } + + fn request_count(&self, _total_items: usize) -> usize { + self.batches.len() + } + + fn request_size(&self, req_id: usize) -> usize { + self.batches[req_id].count + } } /// Derive a search config from a live collection: one random-query template @@ -292,11 +301,15 @@ pub async fn run(args: &Args, query: ServerlessQueryArgs, stopped: Arc Vec { - let mut batches = Vec::new(); - let mut next_id = offset as u64; - for (coll_idx, &count) in per_collection.iter().enumerate() { - let mut remaining = count; - while remaining > 0 { - let n = min(batch_size, remaining); - batches.push(Batch { - coll_idx, - start_id: next_id, - count: n, - }); - next_id += n as u64; - remaining -= n; - } - } - batches -} - /// Request pacing: `--rps` fires at a fixed interval regardless of how many /// requests are in flight (missed ticks are skipped, not burst); otherwise /// the regular `--throttle` stream. @@ -201,7 +178,17 @@ pub async fn run( println!(" …"); } - let batches = plan_batches(&per_collection, args.batch_size, args.offset); + // Budgets determine the distribution; request order is randomized among + // collections that still have budget so bounded parallelism interleaves + // traffic instead of draining collections sequentially. + let batches: Vec<_> = drain_budgets(&per_collection, args.batch_size, &mut rand::rng()) + .into_iter() + .map(|batch| Batch { + coll_idx: batch.collection, + start_id: (args.offset + batch.offset) as u64, + count: batch.count, + }) + .collect(); let generator: Box = Box::new(ConfigGenerator::new(&upload_config)?); let logger = env_logger::Builder::from_default_env().build(); @@ -302,13 +289,17 @@ pub async fn run( #[cfg(test)] mod tests { use super::*; + use rand::SeedableRng; #[test] fn batches_cover_every_point_with_unique_ids() { - let batches = plan_batches(&[5, 0, 7], 3, 100); - let counts: Vec<_> = batches.iter().map(|b| (b.coll_idx, b.count)).collect(); - assert_eq!(counts, vec![(0, 3), (0, 2), (2, 3), (2, 3), (2, 1)]); - let ids: Vec<_> = batches.iter().map(|b| b.start_id).collect(); - assert_eq!(ids, vec![100, 103, 105, 108, 111]); + let mut rng = rand::rngs::StdRng::seed_from_u64(7); + let batches = drain_budgets(&[5, 0, 7], 3, &mut rng); + let mut ids: Vec<_> = batches + .iter() + .flat_map(|b| b.offset..b.offset + b.count) + .collect(); + ids.sort_unstable(); + assert_eq!(ids, (0..12).collect::>()); } } diff --git a/src/stats.rs b/src/stats.rs index 648f627..0e2657f 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -77,8 +77,7 @@ pub async fn process( processor: P, ) -> Result { let started = Instant::now(); - let batch_size = processor.get_batch_size(); - let batch_count = args.num_vectors_or_default().div_ceil(batch_size); + let batch_count = processor.request_count(args.num_vectors_or_default()); let multiprogress = MultiProgress::new(); let progress_bar = multiprogress.add(ProgressBar::new(args.num_vectors_or_default() as u64)); @@ -98,7 +97,6 @@ pub async fn process( &processor, &progress_bar, batch_count, - batch_size, target_rps, ) .await?; @@ -109,7 +107,6 @@ pub async fn process( &processor, &progress_bar, batch_count, - batch_size, ) .await?; } @@ -183,13 +180,12 @@ async fn process_with_parallel( processor: &P, progress_bar: &ProgressBar, batch_count: usize, - batch_size: usize, ) -> Result<()> { let query_stream = (0..batch_count) .take_while(|_| !stopped.load(Ordering::Relaxed)) .map(|n| { let future = processor.make_request(n, args, progress_bar); - progress_bar.inc(batch_size as u64); + progress_bar.inc(processor.request_size(n) as u64); future }); @@ -219,7 +215,6 @@ async fn process_with_rps( processor: &P, progress_bar: &ProgressBar, batch_count: usize, - batch_size: usize, target_rps: f64, ) -> Result<()> { use futures::stream::FuturesUnordered; @@ -255,7 +250,7 @@ async fn process_with_rps( let req_id = requests_sent; requests_sent += 1; - progress_bar.inc(batch_size as u64); + progress_bar.inc(processor.request_size(req_id) as u64); let future = processor.make_request(req_id, args, progress_bar); in_flight.push(future);