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);