diff --git a/Cargo.lock b/Cargo.lock index ff39960..fd6dc49 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1472,11 +1472,10 @@ dependencies = [ [[package]] name = "qdrant-client" version = "1.16.1-dev" -source = "git+https://github.com/qdrant/rust-client?branch=dev#738227709dab2a6473f3d60cc845bf2313f9dfea" +source = "git+https://github.com/qdrant/rust-client?branch=dev#b6bacf7cdc50dc0a64ed6ff7e6d159bc186e903a" dependencies = [ "anyhow", "derive_builder", - "futures", "parking_lot", "prost", "prost-types", diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 99e3f57..3edfe2e 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -9,6 +9,7 @@ src/ ├── collection/ # collection (re)creation: from CLI flags / from YAML config ├── generators/ # data generation: points (legacy & config), queries, random primitives ├── search/ # search benchmarking processors (flag-driven & config-driven) +├── serverless/ # multi-collection mode for Qdrant Serverless (upload/clear/query) ├── upload.rs # upload pipeline (parallelism, batching, progress) ├── upsert.rs # upsert request execution + timing ├── scroll.rs # scroll benchmarking processor diff --git a/README.md b/README.md index d8ff8bb..8f3ab0a 100644 --- a/README.md +++ b/README.md @@ -366,6 +366,77 @@ and search configs. The CLI still controls *how* the benchmark runs (`-n`, `-p`, The flag-driven path (`bfb --scroll --keywords 100 …`) is still available when you prefer flat CLI flags over a YAML file. +### `serverless` — multi-collection benchmarks against Qdrant Serverless + +Serverless uses a **collection per tenant**, so the interesting workload is +traffic spread across many collections rather than one shared collection with +tenant payload. `bfb serverless` talks to the space through the +[`QdrantServerless`](https://github.com/qdrant/rust-client) client. It takes +the same connection flags as the rest of bfb (`--uri`, `--connections`, +`--timeout`, API key via `QDRANT_API_KEY`); a `--uri` without an explicit port +defaults to 443. + +```bash +# Upload 10M points across 100 collections (created lazily on first upsert) +bfb serverless upload \ + --uri https://serverless.example.cloud.qdrant.io \ + --collection-prefix benchmark- \ + --collections-count 100 \ + --distribution uniform \ + --total-points 10M \ + --example serverless-upload \ + -b 64 -p 8 + +# Query existing collections; query shape is derived from a collection's config +bfb serverless query \ + --uri https://serverless.example.cloud.qdrant.io \ + --collection-prefix benchmark- \ + --distribution zipf \ + -n 10k --rps 200 --json results.json + +# Same, but with a search YAML (filters, sparse queries, dataset query sets …) +bfb serverless query \ + --uri https://serverless.example.cloud.qdrant.io \ + --collection-prefix benchmark- \ + --file search.yaml -n 10k -p 8 + +# Show what is there: one line per collection with its point count +bfb serverless list \ + --uri https://serverless.example.cloud.qdrant.io \ + --collection-prefix benchmark- + +# Tear down everything with that prefix +bfb serverless clear \ + --uri https://serverless.example.cloud.qdrant.io \ + --collection-prefix benchmark- +``` + +| Flag | Meaning | +|------|---------| +| `--collection-prefix` | Shared name prefix (`benchmark-` → `benchmark-0` … `benchmark-99`). Optional for `list`, which then shows the whole space | +| `--collections-count` | Upload only: how many collection slots to spread points across | +| `--distribution` | `uniform` or `zipf` — how points/queries are routed across slots | +| `--total-points` | Upload only: total points across all collections (falls back to `-n`) | +| `--file` / `--example` | Upload: collection shape (`bfb upload` schema, required). Query: request shape (`bfb search` schema, optional) | + +The run itself is controlled by the usual global flags: `-p` / `--rps` / +`--throttle`, `-b`, `--search-batch-size`, `--search-limit`, `--retries`, +`--ignore-errors`, `--jsonl-updates` / `--jsonl-searches` / `--jsonl-rps`, and +`--json` for the unified results document. + +Collections are **not** created up front. On first upsert to a slot, bfb creates +it from the upload YAML (only the tenant-facing subset: dense/sparse vectors and +payload indexes — HNSW/quantization/on-disk knobs are ignored and the precision +tier is left at the service default). A slot that already exists is reused +after checking its vectors match the YAML. Point ids are contiguous across +collections (starting at `--offset`), so with a dataset source every +collection receives a different slice of the data. A registry prints how many +collections were preexisting / created / used after the run. + +`query` only targets collections that hold points. Without a search YAML it +issues random dense and sparse queries matching the vectors it finds in the +first collection's config. + ### `schema` — print the upload-config file schema Print an annotated YAML reference enumerating every option accepted by an diff --git a/examples/serverless-upload.yaml b/examples/serverless-upload.yaml new file mode 100644 index 0000000..9c7a588 --- /dev/null +++ b/examples/serverless-upload.yaml @@ -0,0 +1,26 @@ +# Minimal upload shape for `bfb serverless upload`. +# +# bfb serverless upload \ +# --uri https://serverless.example.cloud.qdrant.io \ +# --collection-prefix benchmark- \ +# --collections-count 100 \ +# --distribution uniform \ +# --total-points 1M \ +# --example serverless-upload +# +# Only the tenant-facing subset is sent to serverless (dense/sparse vectors + +# payload indexes). HNSW / quantization / on-disk knobs are ignored by the +# converter — keep them if you also reuse a file with `bfb upload`. + +collection: + vectors: + - size: 128 + distance: cosine + source: random + + fields: + - name: color + type: keyword + source: + type: random + cardinality: 100 diff --git a/src/args/mod.rs b/src/args/mod.rs index 9e5a19f..88d221d 100644 --- a/src/args/mod.rs +++ b/src/args/mod.rs @@ -550,6 +550,12 @@ pub enum Command { #[clap(value_enum)] shell: clap_complete::Shell, }, + + /// Benchmark Qdrant Serverless (multi-collection / per-tenant mode). + /// + /// Spreads upload and query traffic across a range of collections instead + /// of a single shared collection. See `bfb serverless --help`. + Serverless(crate::serverless::ServerlessArgs), } /// `--file` or `--example` (exactly one required). Flattened into upload / @@ -603,7 +609,7 @@ impl Args { } } -fn parse_number(n: &str) -> Result { +pub(crate) fn parse_number(n: &str) -> Result { parse_number_impl(n) .and_then(|v| v.try_into().ok()) .ok_or_else(|| format!("Invalid number: {n}")) diff --git a/src/client.rs b/src/client.rs index 495a532..b88028c 100644 --- a/src/client.rs +++ b/src/client.rs @@ -56,10 +56,18 @@ pub fn create_clients(args: &Args) -> Result> { /// Try the request on every client in random order, retrying `args.retries` /// times with `args.retry_interval` between rounds. -pub async fn retry_with_clients<'a, R, T: std::future::Future>>( - clients: &'a [Qdrant], +/// +/// Generic over the client type so the regular [`Qdrant`] client and the +/// serverless client share one retry policy. +pub async fn retry_with_clients< + 'a, + C, + R, + T: std::future::Future>, +>( + clients: &'a [C], args: &Args, - mut call: impl FnMut(&'a Qdrant) -> T, + mut call: impl FnMut(&'a C) -> T, ) -> anyhow::Result { let mut rng = rand::rng(); let mut permutation = (0..clients.len()).collect::>(); diff --git a/src/config/examples.rs b/src/config/examples.rs index 2585045..458952b 100644 --- a/src/config/examples.rs +++ b/src/config/examples.rs @@ -90,6 +90,11 @@ pub static EXAMPLES: &[Example] = &[ Upload, "Full LAION-400M corpus (~410 parts, streamed with cache: evict)" ), + example!( + "serverless-upload", + Upload, + "Minimal upload shape for `bfb serverless upload` (dense + keyword)" + ), example!( "search-config", Search, diff --git a/src/main.rs b/src/main.rs index 1da7930..de2f03f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -23,6 +23,7 @@ mod save_jsonl; mod scroll; mod search; mod self_update; +mod serverless; mod stats; mod upload; mod upsert; @@ -129,6 +130,9 @@ async fn run_benchmark(args: Args, stopped: Arc) -> Result<()> { Some(Command::Search(search_args)) => return run_search(args, search_args, stopped).await, Some(Command::Upload(upload_args)) => return run_upload(args, upload_args, stopped).await, Some(Command::Scroll(scroll_args)) => return run_scroll(args, scroll_args, stopped).await, + Some(Command::Serverless(serverless_args)) => { + return serverless::run(args, serverless_args, stopped).await; + } // `Schema` / `SelfUpdate` / `Completions` are handled before the runtime // starts; `None` falls through. Some( diff --git a/src/serverless/args.rs b/src/serverless/args.rs new file mode 100644 index 0000000..6bd4b66 --- /dev/null +++ b/src/serverless/args.rs @@ -0,0 +1,119 @@ +//! CLI for `bfb serverless {upload,clear,query}`. + +use clap::{Args as ClapArgs, Subcommand, ValueEnum}; + +use super::distribution::Distribution; +use crate::args::ConfigArgs; + +/// `bfb serverless` — multi-collection benchmarks against Qdrant Serverless. +#[derive(ClapArgs, Debug, Clone)] +pub struct ServerlessArgs { + #[command(subcommand)] + pub command: ServerlessCommand, +} + +#[derive(Subcommand, Debug, Clone)] +pub enum ServerlessCommand { + /// Upload points across a range of collections (created lazily on first use). + Upload(ServerlessUploadArgs), + + /// Delete every collection whose name starts with `--collection-prefix`. + Clear(ServerlessClearArgs), + + /// List collections and their point counts. + List(ServerlessListArgs), + + /// Run queries routed across existing collections matching the prefix. + Query(ServerlessQueryArgs), +} + +#[derive(ClapArgs, Debug, Clone)] +pub struct ServerlessUploadArgs { + /// Prefix shared by every collection name (`benchmark-` → `benchmark-0` …). + #[clap(long)] + pub collection_prefix: String, + + /// How many collections to spread points across. + #[clap(long, value_parser = crate::args::parse_number)] + pub collections_count: usize, + + /// How points are allocated across collections. + #[clap(long, value_enum, default_value_t = DistributionArg::Uniform)] + pub distribution: DistributionArg, + + /// Total number of points to upload across all collections. + /// Falls back to the global `-n` / `--num-vectors` when omitted. + #[clap(long, value_parser = crate::args::parse_number)] + pub total_points: Option, + + /// Upload-shape YAML: `--file ` or `--example ` (same schema + /// as `bfb upload`). + #[clap(flatten)] + pub config: ConfigArgs, +} + +#[derive(ClapArgs, Debug, Clone)] +pub struct ServerlessClearArgs { + /// Delete every collection whose name starts with this prefix. + #[clap(long)] + pub collection_prefix: String, +} + +#[derive(ClapArgs, Debug, Clone)] +pub struct ServerlessListArgs { + /// Only list collections whose name starts with this prefix. + /// Omit to list every collection in the space. + #[clap(long, default_value = "")] + pub collection_prefix: String, +} + +#[derive(ClapArgs, Debug, Clone)] +pub struct ServerlessQueryArgs { + /// Query every existing collection whose name starts with this prefix. + #[clap(long)] + pub collection_prefix: String, + + /// How queries are routed across matching collections. + #[clap(long, value_enum, default_value_t = DistributionArg::Uniform)] + pub distribution: DistributionArg, + + /// Optional search-shape YAML (same schema as `bfb search`). When omitted, + /// one random dense or sparse query template is derived from a matching + /// collection's config. + #[clap(flatten)] + pub config: OptionalConfigArgs, +} + +/// `--file` or `--example`, both optional (at most one). +#[derive(ClapArgs, Debug, Clone)] +#[group(multiple = false)] +pub struct OptionalConfigArgs { + /// Path to a YAML config file + #[clap(long)] + pub file: Option, + + /// Built-in example name (`bfb examples` lists them) + #[clap(long, value_name = "NAME")] + pub example: Option, +} + +impl OptionalConfigArgs { + pub fn is_some(&self) -> bool { + self.file.is_some() || self.example.is_some() + } +} + +#[derive(Debug, Clone, Copy, ValueEnum, PartialEq, Eq)] +pub enum DistributionArg { + Uniform, + Zipf, +} + +impl From for Distribution { + fn from(value: DistributionArg) -> Self { + match value { + DistributionArg::Uniform => Distribution::Uniform, + DistributionArg::Zipf => Distribution::Zipf, + } + } +} diff --git a/src/serverless/clear.rs b/src/serverless/clear.rs new file mode 100644 index 0000000..232127f --- /dev/null +++ b/src/serverless/clear.rs @@ -0,0 +1,44 @@ +//! `bfb serverless clear` — delete collections matching a prefix. + +use anyhow::Result; + +use super::args::ServerlessClearArgs; +use super::client::single_client; +use super::collections::list_matching; +use crate::args::Args; + +pub async fn run(args: &Args, clear: ServerlessClearArgs) -> Result<()> { + let client = single_client(args)?; + let names: Vec = list_matching(&client, &clear.collection_prefix) + .await? + .into_iter() + .map(|c| c.collection_name) + .collect(); + + if names.is_empty() { + println!( + "No collections matched prefix {:?}", + clear.collection_prefix + ); + return Ok(()); + } + + println!( + "Deleting {} collection(s) with prefix {:?}", + names.len(), + clear.collection_prefix + ); + + let mut deleted = 0usize; + for name in &names { + let ok = client.delete_collection(name).await?; + if ok { + deleted += 1; + println!(" deleted {name}"); + } else { + println!(" skipped {name} (already gone)"); + } + } + println!("Deleted {deleted}/{} collections", names.len()); + Ok(()) +} diff --git a/src/serverless/client.rs b/src/serverless/client.rs new file mode 100644 index 0000000..07cad72 --- /dev/null +++ b/src/serverless/client.rs @@ -0,0 +1,27 @@ +//! Build [`QdrantServerless`] clients from shared BFB [`Args`]. +//! +//! Reuses the regular client configuration (`--uri` × `--connections`, +//! `--timeout`, `QDRANT_API_KEY`) so both modes are configured the same way. + +use anyhow::Result; +use qdrant_client::serverless::QdrantServerless; + +use crate::args::Args; +use crate::client::get_config; + +/// One serverless client per (`uri` × `connections`) pair, matching regular BFB. +pub fn create_clients(args: &Args) -> Result> { + get_config(args) + .into_iter() + .map(|config| Ok(QdrantServerless::new(config)?)) + .collect() +} + +/// A single client for one-off administrative calls (`clear`, listing). +pub fn single_client(args: &Args) -> Result { + let config = get_config(args) + .into_iter() + .next() + .ok_or_else(|| anyhow::anyhow!("no --uri given"))?; + Ok(QdrantServerless::new(config)?) +} diff --git a/src/serverless/collections.rs b/src/serverless/collections.rs new file mode 100644 index 0000000..a48a669 --- /dev/null +++ b/src/serverless/collections.rs @@ -0,0 +1,220 @@ +//! Track which serverless collections existed before a run and which were +//! created during upload, creating missing ones lazily on first upsert. + +use std::collections::HashSet; +use std::sync::Mutex; + +use anyhow::{Context, Result, bail}; +use qdrant_client::serverless::{CollectionConfig, CollectionSummary, QdrantServerless}; +use tokio::sync::Mutex as AsyncMutex; + +/// Name of the `i`-th collection under `prefix` (`benchmark-` + `0` → `benchmark-0`). +pub fn collection_name(prefix: &str, index: usize) -> String { + format!("{prefix}{index}") +} + +/// All collections currently in the space whose name starts with `prefix`, +/// sorted by name. +pub async fn list_matching( + client: &QdrantServerless, + prefix: &str, +) -> Result> { + let mut summaries: Vec = client + .list_collections() + .await + .context("list_collections")? + .into_iter() + .filter(|c| c.collection_name.starts_with(prefix)) + .collect(); + summaries.sort_by(|a, b| a.collection_name.cmp(&b.collection_name)); + Ok(summaries) +} + +/// Check that every vector the upload config will send exists in `existing` +/// with the same shape. Extra vectors or indexes on the existing collection +/// are fine; a missing or mismatched one would make every upsert fail, so +/// report it up front instead. +pub fn check_vector_shape( + name: &str, + existing: &CollectionConfig, + wanted: &CollectionConfig, +) -> Result<()> { + for (vector, want) in &wanted.dense_vectors { + let label = if vector.is_empty() { + "default dense vector".to_string() + } else { + format!("dense vector {vector:?}") + }; + let Some(have) = existing.dense_vectors.get(vector) else { + bail!( + "collection {name} has no {label}; delete it (`bfb serverless clear`) or use a matching config" + ); + }; + if have.size != want.size || have.distance != want.distance { + bail!( + "collection {name}: {label} is {}d/{:?} but the config wants {}d/{:?}", + have.size, + have.distance, + want.size, + want.distance + ); + } + if have.multivector != want.multivector { + bail!( + "collection {name}: {label} multivector={} but the config wants {}", + have.multivector, + want.multivector + ); + } + } + for vector in wanted.sparse_vectors.keys() { + if !existing.sparse_vectors.contains_key(vector) { + bail!("collection {name} has no sparse vector {vector:?}"); + } + } + Ok(()) +} + +/// Registry for one serverless upload experiment. +pub struct CollectionRegistry { + prefix: String, + /// Present in the space before the experiment started. + preexisting: HashSet, + /// Known to exist and match the config: created by this process, or + /// preexisting and verified on first use. + ready: Mutex>, + /// Created by this process during upload. + created: Mutex>, + /// Config used when lazily creating a missing collection. + create_config: CollectionConfig, + /// One lock per slot, so parallel workers on the same collection do not + /// race `create_collection` while other slots proceed unblocked. + slot_locks: Vec>, +} + +impl CollectionRegistry { + pub async fn bootstrap( + client: &QdrantServerless, + prefix: &str, + collections_count: usize, + create_config: CollectionConfig, + ) -> Result { + let preexisting: HashSet = list_matching(client, prefix) + .await? + .into_iter() + .map(|c| c.collection_name) + .collect(); + + println!( + "Serverless registry: prefix={prefix:?} slots={collections_count} preexisting={}", + preexisting.len() + ); + let mut shown: Vec<&String> = preexisting.iter().collect(); + shown.sort(); + for name in shown.iter().take(5) { + println!(" preexisting: {name}"); + } + if preexisting.len() > 5 { + println!(" … {} more", preexisting.len() - 5); + } + + Ok(Self { + prefix: prefix.to_string(), + preexisting, + ready: Mutex::new(HashSet::new()), + created: Mutex::new(HashSet::new()), + create_config, + slot_locks: (0..collections_count) + .map(|_| AsyncMutex::new(())) + .collect(), + }) + } + + pub fn name(&self, index: usize) -> String { + collection_name(&self.prefix, index) + } + + /// Ensure collection `index` exists with a compatible vector shape, + /// creating it lazily on first use. Returns the collection name. + pub async fn ensure(&self, client: &QdrantServerless, index: usize) -> Result { + let name = self.name(index); + + if self.ready.lock().unwrap().contains(&name) { + return Ok(name); + } + + let _guard = self.slot_locks[index].lock().await; + + // Re-check under the slot lock. + if self.ready.lock().unwrap().contains(&name) { + return Ok(name); + } + + let info = client + .get_collection(&name) + .await + .with_context(|| format!("get_collection {name}"))?; + + if info.exists { + // Preexisting, or appeared between bootstrap and now (another + // process / prior run). Do not claim we created it, but make sure + // the upload config fits it before sending points. + if let Some(existing) = &info.config { + check_vector_shape(&name, existing, &self.create_config)?; + } + if !self.preexisting.contains(&name) { + println!("Collection {name} appeared during the run; reusing it"); + } + } else { + client + .create_collection(&name, self.create_config.clone()) + .await + .with_context(|| format!("create_collection {name}"))?; + self.created.lock().unwrap().insert(name.clone()); + println!("Created collection {name}"); + } + + self.ready.lock().unwrap().insert(name.clone()); + Ok(name) + } + + pub fn summary(&self) { + let created = self.created.lock().unwrap(); + let ready = self.ready.lock().unwrap(); + println!( + "Serverless registry summary: preexisting={} created={} used={}", + self.preexisting.len(), + created.len(), + ready.len() + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use qdrant_client::serverless::{DenseVectorConfig, Distance, SparseVectorConfig}; + + fn dense(size: u64) -> CollectionConfig { + CollectionConfig::new().dense_vector(DenseVectorConfig::new(size, Distance::Cosine)) + } + + #[test] + fn matching_shape_passes_and_extra_vectors_are_fine() { + let existing = dense(128).named_sparse_vector("text", SparseVectorConfig::new()); + check_vector_shape("c", &existing, &dense(128)).unwrap(); + } + + #[test] + fn size_mismatch_is_reported() { + let err = check_vector_shape("c", &dense(128), &dense(256)).unwrap_err(); + assert!(err.to_string().contains("128d"), "{err}"); + } + + #[test] + fn missing_sparse_vector_is_reported() { + let wanted = dense(128).named_sparse_vector("text", SparseVectorConfig::new()); + let err = check_vector_shape("c", &dense(128), &wanted).unwrap_err(); + assert!(err.to_string().contains("sparse vector \"text\""), "{err}"); + } +} diff --git a/src/serverless/convert.rs b/src/serverless/convert.rs new file mode 100644 index 0000000..5dc10eb --- /dev/null +++ b/src/serverless/convert.rs @@ -0,0 +1,112 @@ +//! Convert a BFB [`UploadConfig`] into a serverless [`CollectionConfig`]. +//! +//! Serverless only accepts the tenant-facing shape (dense/sparse vectors + +//! payload indexes). Storage knobs from the upload YAML (HNSW, quantization, +//! on-disk placement, datatype, …) are ignored — the serverless manager +//! decides those. The serverless `precision_tier` is deliberately left unset +//! (service default): there is no faithful mapping from BFB's quantization / +//! datatype settings onto the tiers. + +use anyhow::{Result, bail}; +use qdrant_client::serverless::{ + BoolIndex, CollectionConfig, DenseVectorConfig, Distance, FloatIndex, GeoIndex, IntegerIndex, + KeywordIndex, PayloadIndex, SparseVectorConfig, TextIndex, Tokenizer, UuidIndex, +}; + +use crate::config::payload::PayloadType; +use crate::config::vector::{DistanceKind, ModifierKind}; +use crate::config::{TokenizerKind, UploadConfig}; + +/// Map an upload-shape YAML into the serverless create-collection config. +pub fn to_serverless_config(upload: &UploadConfig) -> Result { + let c = &upload.collection; + if c.vectors.is_empty() && c.sparse_vectors.is_empty() { + bail!("config must define at least one dense or sparse vector"); + } + + let mut config = CollectionConfig::new(); + + for v in &c.vectors { + let dense = DenseVectorConfig::new(v.size, map_distance(v.distance)) + .multivector(v.multivector.is_some()); + match &v.name { + Some(name) => config = config.named_dense_vector(name.clone(), dense), + None => config = config.dense_vector(dense), + } + } + + for s in &c.sparse_vectors { + let sparse = SparseVectorConfig::new().use_idf(s.modifier == ModifierKind::Idf); + config = config.named_sparse_vector(s.name.clone(), sparse); + } + + for field in &c.fields { + if !field.index { + continue; + } + let index: PayloadIndex = match field.kind { + PayloadType::Keyword => KeywordIndex.into(), + PayloadType::Integer => IntegerIndex::new() + .lookup(true) + .range(field.range_index) + .into(), + PayloadType::Float => FloatIndex.into(), + PayloadType::Bool => BoolIndex.into(), + PayloadType::Uuid => UuidIndex.into(), + PayloadType::Geo => GeoIndex.into(), + PayloadType::Text => { + let mut text = TextIndex::new(); + if let Some(tok) = field.tokenizer { + text = text.tokenizer(map_tokenizer(tok)); + } + text.into() + } + PayloadType::Datetime => qdrant_client::serverless::DatetimeIndex.into(), + }; + config = config.payload_index(field.name.clone(), index); + } + + Ok(config) +} + +fn map_distance(d: DistanceKind) -> Distance { + match d { + DistanceKind::Cosine => Distance::Cosine, + DistanceKind::Dot => Distance::Dot, + DistanceKind::Euclid => Distance::Euclid, + DistanceKind::Manhattan => Distance::Manhattan, + } +} + +fn map_tokenizer(t: TokenizerKind) -> Tokenizer { + match t { + TokenizerKind::Word => Tokenizer::Word, + TokenizerKind::Whitespace => Tokenizer::Whitespace, + TokenizerKind::Prefix => Tokenizer::Prefix, + TokenizerKind::Multilingual => Tokenizer::Multilingual, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn converts_basic_upload_config() { + let yaml = r#" +collection: + vectors: + - size: 128 + distance: cosine + fields: + - name: color + type: keyword +"#; + let upload = crate::config::parse(yaml, "test").unwrap(); + let cfg = to_serverless_config(&upload).unwrap(); + assert_eq!(cfg.dense_vectors.len(), 1); + assert!(cfg.dense_vectors.contains_key("")); + assert_eq!(cfg.dense_vectors[""].size, 128); + assert!(cfg.payload_indexes.contains_key("color")); + } +} diff --git a/src/serverless/distribution.rs b/src/serverless/distribution.rs new file mode 100644 index 0000000..31a9330 --- /dev/null +++ b/src/serverless/distribution.rs @@ -0,0 +1,95 @@ +//! Uniform / Zipf sampling over a fixed set of collection slots. + +use rand::Rng; +use rand::RngExt; +use rand::prelude::Distribution as RandDistribution; +use rand_distr::Zipf; + +/// How work (points or queries) is spread across collections. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Distribution { + Uniform, + /// Zipf with exponent 1.03 over ranks `1..=n` (same default as payload text). + Zipf, +} + +/// Pick a collection index in `0..n` according to [`Distribution`]. +#[derive(Debug, Clone)] +pub struct CollectionPicker { + n: usize, + zipf: Option>, +} + +impl CollectionPicker { + pub fn new(n: usize, distribution: Distribution) -> anyhow::Result { + anyhow::ensure!(n > 0, "collections-count must be > 0"); + let zipf = match distribution { + Distribution::Uniform => None, + Distribution::Zipf => Some( + Zipf::new(n as f64, 1.03) + .map_err(|e| anyhow::anyhow!("failed to build Zipf(n={n}): {e}"))?, + ), + }; + Ok(Self { n, zipf }) + } + + /// Sample a collection index in `0..n`. + pub fn pick(&self, rng: &mut impl Rng) -> usize { + match &self.zipf { + None => rng.random_range(0..self.n), + // Zipf samples in `1..=n`; convert to 0-based. + Some(zipf) => ((zipf.sample(rng) as usize).saturating_sub(1)).min(self.n - 1), + } + } + + /// Allocate `total` points across `n` collections according to the + /// distribution. Returns a vec of length `n` summing to `total`. + pub fn allocate(&self, total: usize, rng: &mut impl Rng) -> Vec { + if self.n == 1 { + return vec![total]; + } + match self.zipf { + None => { + // Even split; remainder goes to the first slots. + let base = total / self.n; + let rem = total % self.n; + (0..self.n).map(|i| base + usize::from(i < rem)).collect() + } + Some(_) => { + // Monte-Carlo allocation: sample `total` ranks. + let mut counts = vec![0usize; self.n]; + for _ in 0..total { + counts[self.pick(rng)] += 1; + } + counts + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::SeedableRng; + use rand::rngs::StdRng; + + #[test] + fn uniform_allocate_sums() { + let picker = CollectionPicker::new(10, Distribution::Uniform).unwrap(); + let mut rng = StdRng::seed_from_u64(1); + let counts = picker.allocate(1000, &mut rng); + assert_eq!(counts.len(), 10); + assert_eq!(counts.iter().sum::(), 1000); + assert!(counts.iter().all(|&c| c == 100)); + } + + #[test] + fn zipf_allocate_sums_and_skews() { + let picker = CollectionPicker::new(10, Distribution::Zipf).unwrap(); + let mut rng = StdRng::seed_from_u64(42); + let counts = picker.allocate(10_000, &mut rng); + assert_eq!(counts.iter().sum::(), 10_000); + // Rank 0 should get strictly more than the last rank on average. + assert!(counts[0] > counts[9]); + } +} diff --git a/src/serverless/list.rs b/src/serverless/list.rs new file mode 100644 index 0000000..41c600e --- /dev/null +++ b/src/serverless/list.rs @@ -0,0 +1,49 @@ +//! `bfb serverless list` — collections matching a prefix and their point counts. + +use anyhow::Result; + +use super::args::ServerlessListArgs; +use super::client::single_client; +use super::collections::list_matching; +use crate::args::Args; + +pub async fn run(args: &Args, list: ServerlessListArgs) -> Result<()> { + let client = single_client(args)?; + let summaries = list_matching(&client, &list.collection_prefix).await?; + + if summaries.is_empty() { + if list.collection_prefix.is_empty() { + println!("No collections in the space"); + } else { + println!("No collections matched prefix {:?}", list.collection_prefix); + } + return Ok(()); + } + + let width = summaries + .iter() + .map(|c| c.collection_name.len()) + .max() + .unwrap_or(0); + let mut total: u64 = 0; + let mut unknown = 0usize; + for c in &summaries { + match c.point_count { + Some(n) => { + total += n; + println!("{: { + unknown += 1; + println!("{: 0 { + print!(" (+{unknown} with unknown count)"); + } + println!(); + Ok(()) +} diff --git a/src/serverless/mod.rs b/src/serverless/mod.rs new file mode 100644 index 0000000..2258f72 --- /dev/null +++ b/src/serverless/mod.rs @@ -0,0 +1,43 @@ +//! Serverless benchmarking mode. +//! +//! Unlike the regular single-collection workflow, serverless mode spreads work +//! across a *range* of collections (one per tenant). Collections are created +//! lazily on first upsert; the registry tracks which existed before the run +//! and which were created during upload. +//! +//! ```text +//! bfb serverless upload --collection-prefix benchmark- --collections-count 100 \ +//! --distribution uniform --total-points 10M --example serverless-upload +//! bfb serverless list --collection-prefix benchmark- +//! bfb serverless clear --collection-prefix benchmark- +//! bfb serverless query --collection-prefix benchmark- --distribution zipf -n 10k +//! ``` + +mod args; +mod clear; +mod client; +mod collections; +mod convert; +mod distribution; +mod list; +mod query; +mod upload; + +pub use args::{ServerlessArgs, ServerlessCommand}; + +use std::sync::Arc; +use std::sync::atomic::AtomicBool; + +use anyhow::Result; + +use crate::args::Args; + +/// Dispatch a `bfb serverless …` subcommand. +pub async fn run(args: Args, serverless: ServerlessArgs, stopped: Arc) -> Result<()> { + match serverless.command { + ServerlessCommand::Upload(upload_args) => upload::run(&args, upload_args, stopped).await, + ServerlessCommand::Clear(clear_args) => clear::run(&args, clear_args).await, + ServerlessCommand::List(list_args) => list::run(&args, list_args).await, + ServerlessCommand::Query(query_args) => query::run(&args, query_args, stopped).await, + } +} diff --git a/src/serverless/query.rs b/src/serverless/query.rs new file mode 100644 index 0000000..f90292e --- /dev/null +++ b/src/serverless/query.rs @@ -0,0 +1,351 @@ +//! `bfb serverless query` — route searches across existing collections. +//! +//! Requests are generated by the same [`ConfigSearchGenerator`] as +//! `bfb search`, either from a search YAML (`--file` / `--example`) or from a +//! template derived from a matching collection's vector config. Each request +//! picks a collection according to `--distribution` and sends +//! `--search-batch-size` queries to it; the run itself (parallelism, `--rps`, +//! stats, `--jsonl-*`, `--json`) is driven by the shared [`process`] loop. + +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Instant; + +use anyhow::{Context, Result, bail}; +use indicatif::ProgressBar; +use qdrant_client::qdrant::{ + IdfParamsBuilder, Query, QueryBatchPointsBuilder, QueryPointsBuilder, SearchParamsBuilder, + VectorInput, +}; +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 crate::args::Args; +use crate::client::retry_with_clients; +use crate::config::examples::{ExampleKind, resolve}; +use crate::config::search::{SearchCollectionConfig, SearchConfig, SearchRequestConfig}; +use crate::config::{DatatypeKind, SparseSource, VectorSource}; +use crate::generators::ConfigSearchGenerator; +use crate::generators::queries::GeneratedQuery; +use crate::processor::{Processor, Timing}; +use crate::results::BenchmarkResults; +use crate::stats::process; + +#[derive(Default)] +struct QueryStats { + server_timings: Vec, + full_timings: Vec, + qps: Vec, + rps: Vec, +} + +struct ServerlessQueryProcessor { + args: Args, + stopped: Arc, + clients: Vec, + names: Vec, + picker: CollectionPicker, + generator: ConfigSearchGenerator, + start_timestamp_millis: f64, + start_time: Instant, + stats: Mutex, +} + +impl ServerlessQueryProcessor { + fn build_query( + &self, + generated: GeneratedQuery, + collection: &str, + ) -> Result { + let (vector, using) = if let Some((values, indices, name)) = generated.sparse { + (VectorInput::new_sparse(indices.data, values), Some(name)) + } else if let Some((values, name)) = generated.dense { + (VectorInput::new_dense(values), name) + } else { + bail!("search config request must produce a dense or sparse vector"); + }; + + let mut params = SearchParamsBuilder::default().exact(self.args.search_exact); + if let Some(hnsw_ef) = self.args.search_hnsw_ef { + params = params.hnsw_ef(hnsw_ef as u64); + } + if let Some(corpus) = generated.idf_corpus { + params = params.idf(IdfParamsBuilder::default().corpus(corpus)); + } + + let mut builder = QueryPointsBuilder::new(collection) + .query(Query::new_nearest(vector)) + .params(params) + .limit(self.args.search_limit as u64) + .with_payload(self.args.search_with_payload) + .with_vectors(self.args.search_with_vectors); + if let Some(name) = using.filter(|n| !n.is_empty()) { + builder = builder.using(name); + } + if let Some(filter) = generated.filter { + builder = builder.filter(filter); + } + Ok(builder) + } + + async fn search(&self, req_id: usize, progress_bar: &ProgressBar) -> Result<()> { + if self.stopped.load(Ordering::Relaxed) { + return Ok(()); + } + + let mut rng = rand::rng(); + let collection = &self.names[self.picker.pick(&mut rng)]; + let template_idx = self.generator.random_template_idx(&mut rng); + + let queries = (0..self.args.search_batch_size) + .map(|_| { + let generated = self + .generator + .make_query_for(template_idx, req_id, &mut rng); + self.build_query(generated, collection).map(|b| b.build()) + }) + .collect::>>()?; + + let mut request = QueryBatchPointsBuilder::new(collection, queries); + if let Some(timeout) = self.args.timeout { + request = request.timeout(timeout as u64); + } + let request = request.build(); + + let start = Instant::now(); + let res = retry_with_clients(&self.clients, &self.args, |c| { + c.query_batch(request.clone()) + }) + .await + .with_context(|| format!("query on {collection}"))?; + let elapsed = start.elapsed().as_secs_f32(); + let delay_millis = self.start_time.elapsed().as_millis() as u32; + + if res.time > self.args.timing_threshold { + progress_bar.println(format!("Slow query on {collection}: {:?}", res.time)); + } + + let mut stats = self.stats.lock().unwrap(); + stats.server_timings.push(Timing { + delay_millis, + value: res.time as f32, + }); + stats.full_timings.push(Timing { + delay_millis, + value: elapsed, + }); + stats.qps.push(Timing { + delay_millis, + value: progress_bar.per_sec() as f32, + }); + stats.rps.push(Timing { + delay_millis, + value: (progress_bar.per_sec() / self.args.search_batch_size as f64) as f32, + }); + Ok(()) + } +} + +impl Processor for ServerlessQueryProcessor { + async fn make_request( + &self, + req_id: usize, + _args: &Args, + progress_bar: &ProgressBar, + ) -> Result<()> { + self.search(req_id, progress_bar).await + } + + fn start_timestamp_millis(&self) -> f64 { + self.start_timestamp_millis + } + + fn server_timings(&self) -> Vec { + self.stats.lock().unwrap().server_timings.clone() + } + + fn qps(&self) -> Vec { + self.stats.lock().unwrap().qps.clone() + } + + fn rps(&self) -> Vec { + self.stats.lock().unwrap().rps.clone() + } + + fn full_timings(&self) -> Vec { + self.stats.lock().unwrap().full_timings.clone() + } + + fn get_batch_size(&self) -> usize { + self.args.search_batch_size + } +} + +/// Derive a search config from a live collection: one random-query template +/// per dense and sparse vector, no filters. +fn infer_search_config(name: &str, config: &CollectionConfig) -> Result { + let mut requests = Vec::new(); + + let mut dense: Vec<_> = config.dense_vectors.iter().collect(); + dense.sort_by(|a, b| a.0.cmp(b.0)); + for (vector, cfg) in dense { + requests.push(SearchRequestConfig::Dense { + using: Some(vector.clone()).filter(|n| !n.is_empty()), + size: cfg.size, + datatype: DatatypeKind::default(), + source: VectorSource::Random, + filters: Vec::new(), + }); + } + + let mut sparse: Vec<_> = config.sparse_vectors.keys().collect(); + sparse.sort(); + for vector in sparse { + requests.push(SearchRequestConfig::Sparse { + using: vector.clone(), + source: SparseSource::default(), + filters: Vec::new(), + idf_corpus: Vec::new(), + }); + } + + if requests.is_empty() { + bail!("collection {name} has no vectors to query"); + } + Ok(SearchConfig { + collection: SearchCollectionConfig { + name: name.to_string(), + }, + requests, + }) +} + +async fn fetch_config(client: &QdrantServerless, name: &str) -> Result { + let info = client + .get_collection(name) + .await + .with_context(|| format!("get_collection {name}"))?; + if !info.exists { + bail!("collection {name} no longer exists"); + } + info.config.ok_or_else(|| { + anyhow::anyhow!("collection {name} has no config; re-upload or pass a search config") + }) +} + +pub async fn run(args: &Args, query: ServerlessQueryArgs, stopped: Arc) -> Result<()> { + let clients = create_clients(args)?; + let summaries = list_matching(&clients[0], &query.collection_prefix).await?; + if summaries.is_empty() { + bail!( + "no collections matched prefix {:?} — run `bfb serverless upload` first", + query.collection_prefix + ); + } + + let (names, empty): (Vec<_>, Vec<_>) = summaries + .into_iter() + .partition(|c| c.point_count != Some(0)); + let names: Vec = names.into_iter().map(|c| c.collection_name).collect(); + if !empty.is_empty() { + println!( + "Skipping {} empty collection(s) matching the prefix", + empty.len() + ); + } + if names.is_empty() { + bail!( + "every collection matching prefix {:?} is empty — run `bfb serverless upload` first", + query.collection_prefix + ); + } + + println!( + "Querying {} collection(s) with prefix {:?} ({:?})", + names.len(), + query.collection_prefix, + query.distribution + ); + + let (search_config, origin) = if query.config.is_some() { + let resolved = resolve( + query.config.file.as_deref(), + query.config.example.as_deref(), + ExampleKind::Search, + )?; + let config = crate::config::search::parse(&resolved.yaml, &resolved.origin)?; + (config, Some(resolved.origin)) + } else { + let config = fetch_config(&clients[0], &names[0]).await?; + println!( + "No search config given; deriving query shape from {}", + names[0] + ); + (infer_search_config(&names[0], &config)?, None) + }; + + let mut args = args.clone(); + args.collection_name = format!("{}*", query.collection_prefix); + let mut results = BenchmarkResults::new(&args, origin); + + let processor = ServerlessQueryProcessor { + args: args.clone(), + stopped: stopped.clone(), + clients, + picker: CollectionPicker::new(names.len(), query.distribution.into())?, + names, + generator: ConfigSearchGenerator::new(&search_config)?, + start_timestamp_millis: std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as f64, + start_time: Instant::now(), + stats: Mutex::new(QueryStats::default()), + }; + + results.results.search = Some(process(&args, stopped, processor).await?); + results.write_if_requested(&args) +} + +#[cfg(test)] +mod tests { + use super::*; + use qdrant_client::serverless::{DenseVectorConfig, Distance, SparseVectorConfig}; + + #[test] + fn infers_one_template_per_vector() { + let config = CollectionConfig::new() + .dense_vector(DenseVectorConfig::new(64, Distance::Dot)) + .named_dense_vector("image", DenseVectorConfig::new(512, Distance::Cosine)) + .named_sparse_vector("text", SparseVectorConfig::new()); + let search = infer_search_config("c", &config).unwrap(); + assert_eq!(search.collection.name, "c"); + assert_eq!(search.requests.len(), 3); + match &search.requests[0] { + SearchRequestConfig::Dense { using, size, .. } => { + assert_eq!(using.as_deref(), None); + assert_eq!(*size, 64); + } + other => panic!("unexpected {other:?}"), + } + match &search.requests[1] { + SearchRequestConfig::Dense { using, size, .. } => { + assert_eq!(using.as_deref(), Some("image")); + assert_eq!(*size, 512); + } + other => panic!("unexpected {other:?}"), + } + assert!( + matches!(&search.requests[2], SearchRequestConfig::Sparse { using, .. } if using == "text") + ); + } + + #[test] + fn vectorless_collection_is_rejected() { + assert!(infer_search_config("c", &CollectionConfig::new()).is_err()); + } +} diff --git a/src/serverless/upload.rs b/src/serverless/upload.rs new file mode 100644 index 0000000..909377a --- /dev/null +++ b/src/serverless/upload.rs @@ -0,0 +1,314 @@ +//! `bfb serverless upload` — spread points across lazily-created collections. + +use std::cmp::min; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use anyhow::{Context, Result}; +use futures::stream::{Stream, StreamExt}; +use indicatif::{MultiProgress, ProgressBar, ProgressDrawTarget, ProgressStyle}; +use qdrant_client::qdrant::UpsertPointsBuilder; +use qdrant_client::serverless::QdrantServerless; +use tokio::time::sleep; +use tokio_stream::wrappers::IntervalStream; + +use super::args::ServerlessUploadArgs; +use super::client::create_clients; +use super::collections::CollectionRegistry; +use super::convert::to_serverless_config; +use super::distribution::CollectionPicker; +use crate::args::Args; +use crate::client::retry_with_clients; +use crate::config; +use crate::config::examples::{ExampleKind, resolve}; +use crate::dataset; +use crate::generators::{ConfigGenerator, PointGenerator}; +use crate::processor::Timing; +use crate::results::{BenchmarkResults, UploadPhase}; +use crate::save_jsonl::save_timings_as_jsonl; +use crate::stats::{print_stats, throttler}; + +/// One upsert request: `count` points of collection slot `coll_idx`, with +/// point ids starting at `start_id` (unique across the whole upload). +#[derive(Debug, Clone, Copy)] +struct Batch { + coll_idx: usize, + start_id: u64, + count: usize, +} + +/// Split the per-collection allocation into upsert batches. Point ids are +/// laid out contiguously across collections starting at `offset`, so with a +/// dataset source every collection gets a different slice of the data. +fn plan_batches(per_collection: &[usize], batch_size: usize, offset: usize) -> 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. +fn pacer(args: &Args) -> Box + Unpin> { + match args.rps.filter(|rps| *rps > 0.0 && rps.is_finite()) { + Some(rps) => { + let mut interval = tokio::time::interval(Duration::from_secs_f64(1.0 / rps)); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + Box::new(IntervalStream::new(interval).map(|_| ())) + } + None => throttler(args.throttle), + } +} + +/// Records the first failure of a run; later ones are only logged. +#[derive(Default)] +struct FirstError(Mutex>); + +impl FirstError { + fn record(&self, err: anyhow::Error) { + let mut slot = self.0.lock().unwrap(); + if slot.is_none() { + *slot = Some(err); + } + } + + fn into_result(self) -> Result<()> { + match self.0.into_inner().unwrap() { + Some(err) => Err(err), + None => Ok(()), + } + } +} + +struct Uploader { + args: Args, + clients: Vec, + registry: CollectionRegistry, + generator: Box, + bar: ProgressBar, + started: Instant, + /// Server-reported upsert latency per batch. + timings: Mutex>, + /// Points acknowledged by the server. + uploaded: Mutex, +} + +impl Uploader { + async fn upsert(&self, batch_no: usize, batch: Batch) -> Result<()> { + let client = &self.clients[batch_no % self.clients.len()]; + let name = self.registry.ensure(client, batch.coll_idx).await?; + + let points: Vec<_> = (0..batch.count as u64) + .map(|i| self.generator.make_point(batch.start_id + i)) + .collect(); + + let mut request = + UpsertPointsBuilder::new(name.clone(), points).wait(self.args.wait_on_upsert); + if let Some(timeout) = self.args.timeout { + request = request.timeout(timeout as u64); + } + let request = request.build(); + + let resp = retry_with_clients(&self.clients, &self.args, |c| { + c.upsert_points(request.clone()) + }) + .await + .with_context(|| format!("upsert into {name}"))?; + + self.timings.lock().unwrap().push(Timing { + delay_millis: self.started.elapsed().as_millis() as u32, + value: resp.time as f32, + }); + *self.uploaded.lock().unwrap() += batch.count; + if resp.time > self.args.timing_threshold { + self.bar + .println(format!("Slow upsert on {name}: {:?}", resp.time)); + } + self.bar.inc(batch.count as u64); + + if let Some(delay_millis) = self.args.delay { + sleep(Duration::from_millis(delay_millis as u64)).await; + } + Ok(()) + } +} + +pub async fn run( + args: &Args, + upload: ServerlessUploadArgs, + stopped: Arc, +) -> Result<()> { + anyhow::ensure!( + upload.collections_count > 0, + "--collections-count must be > 0" + ); + + let resolved = resolve( + upload.config.file.as_deref(), + upload.config.example.as_deref(), + ExampleKind::Upload, + )?; + let upload_config = config::parse(&resolved.yaml, &resolved.origin)?; + let serverless_config = to_serverless_config(&upload_config)?; + + // Dataset-backed configs cap the total at what the dataset holds. + let total_points = dataset::resolve_num_vectors( + upload.total_points.or(args.num_vectors), + args.offset, + &upload_config, + &dataset::default_datasets_dir(), + )?; + + let mut args = args.clone(); + args.num_vectors = Some(total_points); + args.collection_name = format!("{}*", upload.collection_prefix); + let mut results = BenchmarkResults::new(&args, Some(resolved.origin)); + + let clients = create_clients(&args)?; + let registry = CollectionRegistry::bootstrap( + &clients[0], + &upload.collection_prefix, + upload.collections_count, + serverless_config, + ) + .await?; + + let picker = CollectionPicker::new(upload.collections_count, upload.distribution.into())?; + let per_collection = picker.allocate(total_points, &mut rand::rng()); + + println!( + "Uploading {total_points} points across {} collections ({:?})", + upload.collections_count, upload.distribution + ); + for (i, &n) in per_collection.iter().enumerate().take(5) { + println!(" {} → {n} points", registry.name(i)); + } + if per_collection.len() > 5 { + println!(" …"); + } + + let batches = plan_batches(&per_collection, args.batch_size, args.offset); + let generator: Box = Box::new(ConfigGenerator::new(&upload_config)?); + + let logger = env_logger::Builder::from_default_env().build(); + let multiprogress = MultiProgress::new(); + indicatif_log_bridge::LogWrapper::new(multiprogress.clone(), logger) + .try_init() + .ok(); + + let bar = multiprogress.add(ProgressBar::new(total_points as u64)); + bar.set_style( + ProgressStyle::default_bar() + .template("{msg} [{elapsed_precise}] {wide_bar} [{per_sec:>3}] {pos}/{len} (eta:{eta})") + .expect("progress style"), + ); + bar.set_draw_target(ProgressDrawTarget::stdout_with_hz(2)); + + // `--rps` decides the send rate on its own, so concurrency is unbounded; + // otherwise `-p` bounds the number of in-flight batches. + let concurrency = if args.rps.is_some() { + None + } else { + Some(args.parallel.max(1)) + }; + let pacer = pacer(&args); + let start_timestamp_millis = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() as f64; + + let uploader = Uploader { + args: args.clone(), + clients, + registry, + generator, + bar: bar.clone(), + started: Instant::now(), + timings: Mutex::new(Vec::new()), + uploaded: Mutex::new(0), + }; + let first_error = FirstError::default(); + + futures::stream::iter(batches.into_iter().enumerate()) + .take_while(|_| futures::future::ready(!stopped.load(Ordering::Relaxed))) + .zip(pacer) + .for_each_concurrent(concurrency, |((batch_no, batch), _)| { + let uploader = &uploader; + let first_error = &first_error; + let stopped = &stopped; + let args = &args; + async move { + if stopped.load(Ordering::Relaxed) { + return; + } + if let Err(err) = uploader.upsert(batch_no, batch).await { + uploader.bar.println(format!("Error: {err:?}")); + if !args.ignore_errors { + first_error.record(err); + stopped.store(true, Ordering::Relaxed); + } + } + } + }) + .await; + + let duration_secs = uploader.started.elapsed().as_secs_f64(); + if stopped.load(Ordering::Relaxed) { + bar.abandon(); + } else { + bar.finish(); + } + + let uploaded = *uploader.uploaded.lock().unwrap(); + let phase = UploadPhase::new(duration_secs, uploaded); + println!( + "Uploaded {} points in {:.3} s ({:.0} points/s)", + phase.num_points, phase.duration_secs, phase.points_per_sec + ); + uploader.registry.summary(); + + let mut timings = uploader.timings.into_inner().unwrap(); + println!("--- Upsert timings ---"); + print_stats(&args, &mut timings, "upsert time", true); + if let Some(jsonl_path) = &args.jsonl_updates { + save_timings_as_jsonl( + jsonl_path, + args.absolute_time.unwrap_or(false), + &timings, + start_timestamp_millis, + "upsert_latency", + )?; + } + + first_error.into_result()?; + results.results.upload = Some(phase); + results.write_if_requested(&args) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[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]); + } +}