From f238e5a2b3b2b438fa703148db35440ae6b7aaa4 Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Thu, 22 Jan 2026 22:59:13 -0800 Subject: [PATCH 1/2] Fixing flaky test_compute_relative_contrast_with_random_data --- diskann-tools/src/bin/subsample_bin.rs | 4 ++-- diskann-tools/src/utils/relative_contrast.rs | 7 ++++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/diskann-tools/src/bin/subsample_bin.rs b/diskann-tools/src/bin/subsample_bin.rs index b43536d66..ad6f5ddc7 100644 --- a/diskann-tools/src/bin/subsample_bin.rs +++ b/diskann-tools/src/bin/subsample_bin.rs @@ -15,7 +15,7 @@ use rand_distr::{Distribution, StandardUniform}; use diskann::utils::VectorRepr; use diskann_providers::storage::FileStorageProvider; use diskann_providers::storage::StorageWriteProvider; -use diskann_providers::utils::{write_metadata, SampleVectorReader, SamplingDensity}; +use diskann_providers::utils::{random, write_metadata, SampleVectorReader, SamplingDensity}; use diskann_tools::utils::DataType; /// Subsamples vectors from a DiskANN style binary file. @@ -43,7 +43,7 @@ fn create_rng(seed: Option) -> StdRng { match seed { Some(seed) => StdRng::seed_from_u64(seed), None => { - let mut system_rng = rand::rng(); + let mut system_rng = random::create_rnd(); StdRng::from_rng(&mut system_rng) } } diff --git a/diskann-tools/src/utils/relative_contrast.rs b/diskann-tools/src/utils/relative_contrast.rs index 65a30fac7..c67e0cf48 100644 --- a/diskann-tools/src/utils/relative_contrast.rs +++ b/diskann-tools/src/utils/relative_contrast.rs @@ -5,6 +5,7 @@ use diskann::{utils::VectorRepr, ANNError}; use diskann_providers::storage::StorageReadProvider; +use diskann_providers::utils::random; use diskann_providers::{model::graph::traits::GraphDataType, utils::file_util::load_bin}; use rand::Rng; @@ -33,7 +34,7 @@ fn average_squared_distance( base: &[Vec], num_random_samples: usize, ) -> CMDResult { - let mut rng = rand::rng(); + let mut rng = random::create_rnd(); let n = base.len(); let mut sum_dist = 0.0; for _ in 0..num_random_samples { @@ -125,7 +126,7 @@ mod relative_contrast_tests { // Generate 1000 random vectors of fp16 data type with 384 dimensions let num_vectors = 1000; let dim = 384; - let mut rng = rand::rng(); + let mut rng = random::create_rnd_in_tests(); let base: Vec = (0..num_vectors * dim) .map(|_| f16::from_f32(rng.random_range(0.0..1.0))) .collect(); @@ -213,7 +214,7 @@ mod relative_contrast_tests { let num_queries = 10; let dim = 128; - let mut rng = rand::rng(); + let mut rng = random::create_rnd_in_tests(); let query: Vec = (0..num_queries * dim) .map(|_| f16::from_f32(rng.random_range(0.0..1.0))) .collect(); From 94048377bf7db4555c0abf5484bb3f80981108df Mon Sep 17 00:00:00 2001 From: "Alex Razumov (from Dev Box)" Date: Fri, 23 Jan 2026 12:28:49 -0800 Subject: [PATCH 2/2] Addressing comments --- diskann-tools/src/bin/subsample_bin.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/diskann-tools/src/bin/subsample_bin.rs b/diskann-tools/src/bin/subsample_bin.rs index ad6f5ddc7..891d71ec6 100644 --- a/diskann-tools/src/bin/subsample_bin.rs +++ b/diskann-tools/src/bin/subsample_bin.rs @@ -9,7 +9,7 @@ use std::path::PathBuf; use anyhow::{ensure, Result}; use clap::Parser; use half::f16; -use rand::{rngs::StdRng, SeedableRng}; +use rand::rngs::StdRng; use rand_distr::{Distribution, StandardUniform}; use diskann::utils::VectorRepr; @@ -41,11 +41,8 @@ struct Args { fn create_rng(seed: Option) -> StdRng { match seed { - Some(seed) => StdRng::seed_from_u64(seed), - None => { - let mut system_rng = random::create_rnd(); - StdRng::from_rng(&mut system_rng) - } + Some(seed) => random::create_rnd_from_seed(seed), + None => random::create_rnd(), } }