Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,13 @@ jobs:
# job smoke-tests, now in the blocking suite.
run: cargo test --test integration_cli --release

- name: Run measurement-overhead invariants (no container)
# `--lib --bins` does NOT build `tests/*.rs` integration targets, and no
# job named this one, so INV-2/INV-3 — the source-shape tripwires that
# keep the per-query timed loops lock-free and the percentiles unbiased —
# had never run in CI. INV-4/INV-4b (#214) join them here.
run: cargo test --test overhead_invariants --release

integration-redis:
name: Integration tests (Redis)
runs-on: ubuntu-latest
Expand Down
66 changes: 29 additions & 37 deletions src/bin/vector_db_benchmark/engine/chroma.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ use crate::engine::{Engine, SearchResults, UploadStats};
use vector_db_benchmark::parsers::datetime_to_epoch_secs;
use vector_db_benchmark::query_filter::QueryFilter;
use vector_db_benchmark::readers::metadata::{MetadataItem, MetadataValue};
use vector_db_benchmark::start_gate::WorkerPool;

const DEFAULT_COLLECTION: &str = "benchmark";

Expand Down Expand Up @@ -681,16 +682,13 @@ impl Engine for ChromaEngine {
let query_idx = Arc::new(AtomicUsize::new(0));
let pb = self.progress_bar(num_to_run);

// Barrier-synchronized start so per-worker HTTP client construction AND the
// cold first query fall OUTSIDE the measured window (mirrors redis/vertex).
// Every worker builds its client + primes, then blocks on `ready`; the main
// thread stamps the shared start instant into `start_cell` and releases `go`,
// so the measurement clock starts only once all workers are warm. A worker
// that fails to build its client MUST still pass both barriers before
// returning, or the run would deadlock.
let ready = Arc::new(std::sync::Barrier::new(parallel + 1));
let go = Arc::new(std::sync::Barrier::new(parallel + 1));
let start_cell = Arc::new(std::sync::OnceLock::<Instant>::new());
// Gate-synchronized start so connection setup AND the cold first query
// fall OUTSIDE the measured window. Every worker connects + primes, then
// parks at the gate; `WorkerPool::start` stamps the shared start instant and
// releases everyone, so the measurement clock starts only once all workers
// are warm and poised. The gate is count-agnostic: a worker that fails to
// set up, panics, or is never started by the OS settles its ticket and turns
// the run into a hard error instead of a hang (#214).

let mut times: Vec<f64> = Vec::with_capacity(num_to_run);
let mut precs: Vec<f64> = Vec::with_capacity(num_to_run);
Expand All @@ -700,20 +698,18 @@ impl Engine for ChromaEngine {

let query_url = format!("{}/collections/{}/query", self.api_base, self.collection_id);

std::thread::scope(|s| {
let mut handles = Vec::with_capacity(parallel);
let measured_start = std::thread::scope(|s| -> Result<Instant, String> {
let mut pool = WorkerPool::new(s, "chroma-search", parallel);
for _ in 0..parallel {
let query_url = query_url.clone();
let timeout = self.timeout;
let neighbors = &neighbors;
let tops = &tops;
let bodies = &bodies;
let query_idx = Arc::clone(&query_idx);
let ready = Arc::clone(&ready);
let go = Arc::clone(&go);
let pb = &pb;

handles.push(s.spawn(move || {
pool.spawn(move |ticket| {
let mut t = Vec::new();
let mut p = Vec::new();
let mut r = Vec::new();
Expand All @@ -725,10 +721,11 @@ impl Engine for ChromaEngine {
.build()
{
Ok(c) => c,
Err(_) => {
// Still cross both barriers so peers aren't stranded.
ready.wait();
go.wait();
Err(e) => {
// A worker that cannot set itself up would leave the run at a
// lower real concurrency than the `parallel` it reports. Settle
// the ticket with the reason; the coordinator makes it an error.
ticket.fail(format!("chroma-search worker setup failed: {e}"));
return (t, p, r, mr, nd);
}
};
Expand All @@ -747,10 +744,11 @@ impl Engine for ChromaEngine {
.and_then(|resp| resp.text());
}

// Signal "ready + primed", then block until the main thread stamps
// Signal "ready + primed", then block until the coordinator stamps
// the shared measurement start and releases everyone.
ready.wait();
go.wait();
if ticket.arrive_and_wait().is_none() {
return (t, p, r, mr, nd);
}

loop {
let idx = query_idx.fetch_add(1, Ordering::Relaxed);
Expand Down Expand Up @@ -794,33 +792,27 @@ impl Engine for ChromaEngine {
pb.inc(1);
}
(t, p, r, mr, nd)
}));
})?;
}

// All workers spawned: wait until every one has built its client and
// primed, stamp the shared measurement start, then release them together.
ready.wait();
let st = Instant::now();
start_cell.set(st).ok();
go.wait();
// Every worker is connected + primed and parked at the gate.
// Stamp the shared measurement start and release them together.
let (per_worker, measured_start) = pool.start()?;

for h in handles {
let (t, p, r, mr, nd) = h.join().unwrap();
for (t, p, r, mr, nd) in per_worker {
times.extend(t);
precs.extend(p);
recs.extend(r);
mrr_vals.extend(mr);
ndcg_vals.extend(nd);
}
});
Ok(measured_start)
})?;

pb.finish_and_clear();
// total_time excludes connection setup and the cold first query: it is
// measured from the barrier release (start_cell), not a pre-scope instant.
let total_time = start_cell
.get()
.map(|st| st.elapsed().as_secs_f64())
.unwrap_or(0.0);
// measured from the gate release, not a pre-scope instant.
let total_time = measured_start.elapsed().as_secs_f64();
let top = explicit_top.unwrap_or_else(|| neighbors.first().map(|n| n.len()).unwrap_or(10));
crate::engine::compute_search_stats(
&times, &precs, &recs, &mrr_vals, &ndcg_vals, total_time, top, parallel, num_to_run,
Expand Down
64 changes: 28 additions & 36 deletions src/bin/vector_db_benchmark/engine/dragonfly.rs
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ use crate::metrics::compute_metrics;
use vector_db_benchmark::parsers::{datetime_to_epoch_secs, doc_key_to_id, doc_key_to_id_opt};
use vector_db_benchmark::query_filter::QueryFilter;
use vector_db_benchmark::readers::metadata::MetadataItem;
use vector_db_benchmark::start_gate::WorkerPool;

/// Dragonfly engine configuration.
#[derive(Clone)]
Expand Down Expand Up @@ -1014,25 +1015,22 @@ impl Engine for DragonflyEngine {

let pb = self.create_progress_bar(num_to_run);

// Barrier-synchronized start so connection setup AND the cold first query
// fall OUTSIDE the measured window (mirrors redis.rs/vertex.rs). Every
// worker connects + primes, then blocks on `ready`; the main thread stamps
// the shared start instant into `start_cell` and releases `go`, so the
// measurement clock starts only once all workers are warm and poised. A
// worker that fails to connect MUST still pass both barriers before
// returning, or the run would deadlock.
let ready = Arc::new(std::sync::Barrier::new(parallel + 1));
let go = Arc::new(std::sync::Barrier::new(parallel + 1));
let start_cell = Arc::new(std::sync::OnceLock::<Instant>::new());
// Gate-synchronized start so connection setup AND the cold first query
// fall OUTSIDE the measured window. Every worker connects + primes, then
// parks at the gate; `WorkerPool::start` stamps the shared start instant and
// releases everyone, so the measurement clock starts only once all workers
// are warm and poised. The gate is count-agnostic: a worker that fails to
// set up, panics, or is never started by the OS settles its ticket and turns
// the run into a hard error instead of a hang (#214).

let mut times: Vec<f64> = Vec::with_capacity(num_to_run);
let mut precs: Vec<f64> = Vec::with_capacity(num_to_run);
let mut recs: Vec<f64> = Vec::with_capacity(num_to_run);
let mut mrr_vals: Vec<f64> = Vec::with_capacity(num_to_run);
let mut ndcg_vals: Vec<f64> = Vec::with_capacity(num_to_run);

std::thread::scope(|s| {
let mut handles = Vec::with_capacity(parallel);
let measured_start = std::thread::scope(|s| -> Result<Instant, String> {
let mut pool = WorkerPool::new(s, "dragonfly-search", parallel);
for _ in 0..parallel {
let host = self.host.clone();
let port = self.port;
Expand All @@ -1043,11 +1041,9 @@ impl Engine for DragonflyEngine {
let algorithm = algorithm.as_str();
let index_name = index_name.as_str();
let query_idx = Arc::clone(&query_idx);
let ready = Arc::clone(&ready);
let go = Arc::clone(&go);
let pb = &pb;

handles.push(s.spawn(move || {
pool.spawn(move |ticket| {
let mut t = Vec::new();
let mut p = Vec::new();
let mut r = Vec::new();
Expand All @@ -1057,10 +1053,11 @@ impl Engine for DragonflyEngine {

let mut conn = match DragonflyEngine::connect(&host, port) {
Ok(c) => c,
Err(_) => {
// Still cross both barriers so peers aren't stranded.
ready.wait();
go.wait();
Err(e) => {
// A worker that cannot set itself up would leave the run at a
// lower real concurrency than the `parallel` it reports. Settle
// the ticket with the reason; the coordinator makes it an error.
ticket.fail(format!("dragonfly-search worker setup failed: {e}"));
return (t, p, r, mr, nd);
}
};
Expand All @@ -1083,10 +1080,11 @@ impl Engine for DragonflyEngine {
);
}

// Signal "connected + primed", then block until the main thread
// Signal "connected + primed", then block until the coordinator
// stamps the shared measurement start and releases everyone.
ready.wait();
go.wait();
if ticket.arrive_and_wait().is_none() {
return (t, p, r, mr, nd);
}

loop {
let idx = query_idx.fetch_add(1, Ordering::Relaxed);
Expand Down Expand Up @@ -1144,32 +1142,26 @@ impl Engine for DragonflyEngine {
pb.inc(pb_pending);
}
(t, p, r, mr, nd)
}));
})?;
}

// All workers are connected + primed once they clear `ready`; stamp the
// shared measurement start and release them together via `go`.
ready.wait();
let st = Instant::now();
start_cell.set(st).ok();
go.wait();
// Every worker is connected + primed and parked at the gate.
// Stamp the shared measurement start and release them together.
let (per_worker, measured_start) = pool.start()?;

for h in handles {
let (t, p, r, mr, nd) = h.join().unwrap();
for (t, p, r, mr, nd) in per_worker {
times.extend(t);
precs.extend(p);
recs.extend(r);
mrr_vals.extend(mr);
ndcg_vals.extend(nd);
}
});
Ok(measured_start)
})?;

pb.finish_and_clear();
// total_time excludes connection setup and the cold first query.
let total_time = start_cell
.get()
.map(|st| st.elapsed().as_secs_f64())
.unwrap_or(0.0);
let total_time = measured_start.elapsed().as_secs_f64();

if times.is_empty() {
return Err("No searches completed".to_string());
Expand Down
Loading
Loading