Skip to content

Commit a86fe61

Browse files
authored
In-Mem 2.0 (microsoft#1206)
Introduce a second in-memory provider with the intention of replacing the current provider. # Why The [RFC](https://github.com/microsoft/DiskANN/blob/9d5f1435fd3a327a8fc2cdf7f96d0ffeb164fa1b/rfcs/01206-inmem2.md) outlines much of the motivation. In short, the goal here is to: * Make the provider safe under concurrent inserts/search/deletes. * Support proper external/internal ID translation. * Improve test coverage * Do so with minimal performance overhead. The concurrency argument comes from the epoch-based-reclamation (EBR) protection scheme for internal slots. See the RFC for more details. # Known Follow-up Items - Perf Parity: This generally has pretty good performance, but could use a little more tuning to bring it fully on-par with our current inmem index. - Quantization: This initial prototype is lacking quantization support. Adding quantization is relatively straightforward in the primary `Store`, but will need some thought on how to add the reranking layer. This shouldn't be an architectural blocker, though. - Hybrid PQ: A larger open question is how to support the `max_fp_vecs` feature of our current PQ implementation, which reads some full-precision and some quantized vectors during prune. Like with quantization in general, I think this is not a fundamental issue. - Support for non-uniform sized items in slots: For this, I'm mainly thinking of multi-vectors where the number of vectors within each multi-vector can vary. In the context of multi-vectors, fast element access is less important than for traditional vectors as distance computations take considerably longer. # Suggested Reviewing Order The majority of this PR is in a new `diskann-inmem` crate. To facilitate testing, this crate has an "integration-test" feature, which enables the code in `diskann-inmem/src/integration`. This code is an unstable public reexport of internal types meant only for consumption in the `diskann-inmem/integration` integration test binary. The integration test binary is powered by `diskann-benchmark-runner`. ## `diskann-inmem` ### Independent Low-Level Utilities * `num.rs`: Strong type utilities for byte and alignment representations. * `buffer.rs`: A miri-compliant version of [`AlignedMemoryVectorStore`](https://github.com/microsoft/DiskANN/blob/b17240a610ba51aeb2b8bbfebab98d0c8240e12d/diskann-providers/src/model/graph/provider/async_/common.rs#L89-L95). This type allows vectors/neighbors to be stored in a single larger allocation. The use of `RawSlice` allows slots within the `Buffer` to be inspected and manipulated without forcing reference materialization (which is important to prevent aliasing). * `neighbors.rs`: The new version of [`SimpleNeighborVectorProviderAsync`](https://github.com/microsoft/DiskANN/blob/main/diskann-providers/src/model/graph/provider/async_/simple_neighbor_provider.rs#L18-L26). This reuses the sharded-lock idea, but provides additional utility, including the ability to perform read-modify-write operations on adjacency lists. * `counters.rs`: Event counters. When the "integration-test" feature is not enabled, counters become a no-op. These are enabled for testing to monitor changes. * `sharded.rs`: An external-to-internal ID translation utility. The main trick with this struct is to provide utilities like `Sharded::occupied_entry`, which locks and returns an `external/internal` mapping. The proxy `Entry` struct is important as it verifies that such a mapping exists and provides an infallible way of deleting the mapping. This is chained with higher level operations (e.g. `Provider::delete`) to delete both the ID-mapping and the internal data-slot in lock-step. ### Concurrency Protocol The concurrency protocol is built upon three main layers: * `tag.rs`: An atomic slot tag for controlling access to data. * `epoch.rs`: The central registry where readers register and deregister. This is the crux of this PR and probably the most important file. * `store.rs`: A binary blob store built on top of `epoch.rs` to provide the safe concurrent store for data. This provides the following operations: - Storage of binary data in "slots". - Reading of data in slots (provided by `Reader`). - Tracking on the valid/invalid state slots. - Finding available slots into which new data can be inserted. - Safe retirement of slots and eventual reclamation. The `Store` in `store.rs` has some help from `freelist.rs` to accelerate locating available slots internally. **Testing**: `epoch.rs` has unit tests with injectable delays to set up known pathological orderings. The sequencing is helped by `test/sequencer.rs`. In addition, `test/epoch.rs` includes a direct stress test for the `Registry`. This is particularly helpful when run under Miri, which has the ability to detect some race conditions. A larger concurrency stress test lives in the integration-test binary. This directly tests `store.rs` by spinning up readers, writers, and retirers and hammers a single `Store`. Data is read and written into the store in a knowable pattern, allowing readers to detect torn reads, implying a race condition. For this PR, I ran the following stress test file ```json { "search_directories": [ ], "output_directory": null, "jobs": [ { "type": "store-stress", "content": { "capacity": 8192, "duration_secs": 600, "entry_bytes": 256, "epoch_guard_slots": 256, "freelist_recycle_capacity": 1024, "low_watermark": 4096, "max_ops": 50000000000, "readers": 32, "retirers": 16, "seed": 11935966405698895599, "writers": 32 } } ] } ``` with the command ``` cargo run --package diskann-inmem \ --bin integration-test \ --features integration-test \ --release -- \ run --input-file stress.json --output-file temp.json ``` The generated output was ``` readers: 32 writers: 32 retirers: 16 capacity: 8192 entry_bytes: 256 low_watermark: 4096 duration_secs: 600 max_ops: 50000000000 seed: 11935966405698895599 elapsed_secs: 600.001182962 reads: 110206469888 acquires_ok: 248086767 acquires_fail: 619145 retires_ok: 248082573 retires_fail: 243312644 reclaims: 108847566 transitions: 42006520 peak_live: 8131 ``` While not a proof of correctness, this is a pretty decent stress test. ### Providers The implementation of the data provider is split into two logical pieces. The first lives in `layers/` and is focused on computing distances. With this approach, I am trying to avoid the need to replicate `Accessor`s and `Strategy`s for each future quantization type. `layers/full.rs` is the full-precision implementation. One thing to note is the use of the `FullPrecision` marker trait from which the implementations of the whole `layers` API is derived for `layers::Full`. This allows users to include just a `T: FullPrecision` trait bound and **really** simplifies the generics upstream. Within `provider.rs` - my goal here is to minimize the use of generics as much as possible. In particular for search, I use a trait object for `ExpandBeam`. When coupled with the `layers::QueryDistance` API, we can create implementations where the distance function is inlined *and* the number of prefetch instructions can be tailored to the data length. Fully optimizing this is still a work-in-progress. Another thing to call out in `provider.rs` is the care needed for data insertion and deletion. Since external/internal ID translation is supported, we need to ensure that the translation table stays in-sync with the internal store. On insert, if we allocate an internal slot only to find the external ID already exists, we need to abort the operation rather than publish the internal slot. Similarly on delete, we first need to establish if the external/internal ID mapping exists. If so, then we can try to retire the slot. If slot retiring fails (it shouldn't, but bugs can happen) - we need to not commit the ID mapping deletion and instead return an error. #### Testing Like the concurrency stress tests, testing uses the `integration-test` binary. Here, the 10k YFCC dataset is used for non-trivial runs. Using the A/B functionality in `diskann-benchmark-runner`, we can compare against checked-in baselines. This allows us to capture rich metrics for recall, number of operations, etc. and update easily. The main logic for checking and reporting baseline mismatches is in `integration/support/check.rs`. The goal is to summarize all such mismatches for presentation to provide the highest signal possible. ## `diskann-benchmark` Integration into the benchmarks is straightforward. I elected to put everything in a single file to minimize disruption. I'm trying an approach of using `diskann_benchmark_runner::Input::from_raw` to separate out the deserialization types from the actual inputs, allowing richer types (e.g., a full `diskann_benchmark_core::streaming::bigann::RunBook`) to be loaded. Also note how relatively simple the streaming benchmark integration is. Since the new inmem provider supports ID translation and internal slot allocation, it does not need the same level of hand-holding that the current inmem provider needs.
1 parent a10433d commit a86fe61

49 files changed

Lines changed: 13234 additions & 3 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/ci.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ env:
2222
CARGO_TERM_COLOR: always
2323
# The features we want to explicitly test. For example, the `flatbuffers-build` feature
2424
# of `diskann-quantization` requires additional setup and so must not be included by default.
25-
DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree"
25+
DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2"
2626

2727
# Intel SDE version used for baseline and AVX-512 emulation jobs.
2828
SDE_VERSION: "sde-external-10.8.0-2026-03-15-lin"

.github/workflows/nightly.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ env:
1717
RUST_CONFIG: 'build.rustflags=["-Dwarnings"]'
1818
RUST_BACKTRACE: 1
1919
CARGO_TERM_COLOR: always
20-
DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree"
20+
DISKANN_FEATURES: "virtual_storage,spherical-quantization,product-quantization,tracing,experimental_diversity_search,disk-index,flatbuffers,linalg,codegen,multi-vector,bftree,inmem2"
2121

2222
defaults:
2323
run:

Cargo.lock

Lines changed: 25 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ members = [
1818
"diskann-disk",
1919
"diskann-label-filter",
2020
"diskann-garnet",
21+
"diskann-inmem",
2122
# Infrastructure
2223
"diskann-benchmark-runner",
2324
"diskann-benchmark-core",
@@ -60,6 +61,7 @@ diskann-quantization = { path = "diskann-quantization", default-features = false
6061
diskann = { path = "diskann", version = "0.55.0" }
6162
# Providers
6263
diskann-providers = { path = "diskann-providers", default-features = false, version = "0.55.0" }
64+
diskann-inmem = { path = "diskann-inmem", default-features = false, version = "0.55.0" }
6365
diskann-disk = { path = "diskann-disk", version = "0.55.0" }
6466
diskann-label-filter = { path = "diskann-label-filter", version = "0.55.0" }
6567
# Infra

diskann-benchmark-runner/src/files.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,12 @@ impl std::ops::Deref for InputFile {
6262
}
6363
}
6464

65+
impl std::fmt::Display for InputFile {
66+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67+
write!(f, "{}", self.display())
68+
}
69+
}
70+
6571
///////////
6672
// Tests //
6773
///////////

diskann-benchmark-runner/src/utils/fmt.rs

Lines changed: 218 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,6 +379,135 @@ where
379379
}
380380
}
381381

382+
//////////////
383+
// KeyValue //
384+
//////////////
385+
386+
enum MaybeLazy<'a> {
387+
Lazy(&'a dyn std::fmt::Display),
388+
Eager(String),
389+
}
390+
391+
impl std::fmt::Display for MaybeLazy<'_> {
392+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
393+
match self {
394+
Self::Lazy(lazy) => write!(f, "{}", lazy),
395+
Self::Eager(s) => f.write_str(s),
396+
}
397+
}
398+
}
399+
400+
impl std::fmt::Debug for MaybeLazy<'_> {
401+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
402+
struct AsDisplay<'a>(&'a dyn std::fmt::Display);
403+
impl std::fmt::Debug for AsDisplay<'_> {
404+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
405+
self.0.fmt(f)
406+
}
407+
}
408+
409+
match self {
410+
Self::Lazy(o) => {
411+
let as_display = AsDisplay(o);
412+
f.debug_tuple("MaybeLazy::Lazy").field(&as_display).finish()
413+
}
414+
Self::Eager(s) => f.debug_tuple("MaybeLazy::Eager").field(s).finish(),
415+
}
416+
}
417+
}
418+
419+
/// Display a dynamic list of key-value pairs in a YAML-like style.
420+
///
421+
/// Keys are left-aligned and single-line values are aligned into a common column
422+
/// just past the longest key. A value that renders to multiple lines (for example
423+
/// a nested [`KeyValue`] or any other multi-line block) is placed on the lines
424+
/// following its key, indented by two spaces. This keeps nested structures visibly
425+
/// subordinate to their key regardless of whether the value is itself a key-value
426+
/// list or an opaque block.
427+
///
428+
/// # Examples
429+
///
430+
/// ```
431+
/// use diskann_benchmark_runner::utils::fmt::KeyValue;
432+
///
433+
/// let mut kv = KeyValue::new();
434+
/// kv.push("a", &1);
435+
/// kv.push("hello", &"world");
436+
///
437+
/// let expected = "a: 1\nhello: world";
438+
///
439+
/// assert_eq!(kv.to_string(), expected);
440+
/// ```
441+
///
442+
/// Multi-line values are indented beneath their key:
443+
///
444+
/// ```
445+
/// use diskann_benchmark_runner::utils::fmt::KeyValue;
446+
///
447+
/// let mut inner = KeyValue::new();
448+
/// inner.push("x", &1);
449+
/// inner.push("yy", &2);
450+
/// let inner = inner.to_string();
451+
///
452+
/// let mut kv = KeyValue::new();
453+
/// kv.push("name", &"example");
454+
/// kv.push("nested", &inner);
455+
///
456+
/// let expected = "name: example\nnested:\n x: 1\n yy: 2";
457+
///
458+
/// assert_eq!(kv.to_string(), expected);
459+
/// ```
460+
#[derive(Debug, Default)]
461+
pub struct KeyValue<'a> {
462+
kv: Vec<(&'a str, MaybeLazy<'a>)>,
463+
max_key_length: usize,
464+
}
465+
466+
impl<'a> KeyValue<'a> {
467+
/// Create a new empty [`KeyValue`] formatter.
468+
pub fn new() -> Self {
469+
Self {
470+
kv: Vec::new(),
471+
max_key_length: 0,
472+
}
473+
}
474+
475+
/// Push the key-value pair to `self` for formatting.
476+
pub fn push(&mut self, key: &'a str, value: &'a dyn std::fmt::Display) {
477+
self.max_key_length = self.max_key_length.max(key.len());
478+
self.kv.push((key, MaybeLazy::Lazy(value)))
479+
}
480+
481+
/// Push the key-value pair to `self` for formatting - eagerly formatting `value`.
482+
pub fn push_eager<D>(&mut self, key: &'a str, value: D)
483+
where
484+
D: std::fmt::Display,
485+
{
486+
self.max_key_length = self.max_key_length.max(key.len());
487+
self.kv.push((key, MaybeLazy::Eager(value.to_string())))
488+
}
489+
}
490+
491+
impl std::fmt::Display for KeyValue<'_> {
492+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
493+
let width = self.max_key_length;
494+
let mut prefix = "";
495+
for (k, v) in self.kv.iter() {
496+
let rendered = v.to_string();
497+
if rendered.contains('\n') {
498+
write!(f, "{}{}:\n{}", prefix, k, Indent::new(&rendered, 2))?
499+
} else {
500+
// Left-align the key and pad so that all single-line values line up in a
501+
// column one space past the longest key's colon.
502+
let pad = (width + 1).saturating_sub(k.len());
503+
write!(f, "{}{}:{:pad$}{rendered}", prefix, k, "")?;
504+
}
505+
prefix = "\n";
506+
}
507+
Ok(())
508+
}
509+
}
510+
382511
///////////
383512
// Tests //
384513
///////////
@@ -606,4 +735,93 @@ string, , string
606735
.with_pair(" and ");
607736
assert_eq!(d.to_string(), "\"topk\" and \"range\"");
608737
}
738+
739+
//----------//
740+
// KeyValue //
741+
//----------//
742+
743+
// Strip a preceding newline if it exists.
744+
fn process(x: &str) -> &str {
745+
let x = x.strip_prefix('\n').unwrap_or(x);
746+
x.strip_suffix('\n').unwrap_or(x)
747+
}
748+
749+
#[test]
750+
fn test_key_value_empty() {
751+
let kv = KeyValue::new();
752+
assert_eq!(kv.to_string(), "");
753+
}
754+
755+
#[test]
756+
fn test_key_value_single_pair() {
757+
let mut kv = KeyValue::new();
758+
kv.push("a", &1);
759+
assert_eq!(kv.to_string(), "a: 1");
760+
}
761+
762+
#[test]
763+
fn test_key_value_aligns_values() {
764+
let mut kv = KeyValue::new();
765+
kv.push("a", &1);
766+
kv.push("hello", &"world");
767+
let expected = process(
768+
r#"
769+
a: 1
770+
hello: world
771+
"#,
772+
);
773+
assert_eq!(kv.to_string(), expected);
774+
}
775+
776+
#[test]
777+
fn test_key_value_push_eager() {
778+
let mut kv = KeyValue::new();
779+
kv.push_eager("a", 1);
780+
kv.push_eager("hello", "world");
781+
782+
let expected = process(
783+
r#"
784+
a: 1
785+
hello: world
786+
"#,
787+
);
788+
789+
assert_eq!(kv.to_string(), expected);
790+
}
791+
792+
#[test]
793+
fn test_key_value_multiline_value_is_indented() {
794+
let mut inner = KeyValue::new();
795+
inner.push("x", &1);
796+
inner.push("yy", &2);
797+
let inner = inner.to_string();
798+
799+
let mut kv = KeyValue::new();
800+
kv.push("name", &"example");
801+
kv.push("nested", &inner);
802+
kv.push("another line", &1);
803+
804+
let expected = process(
805+
r#"
806+
name: example
807+
nested:
808+
x: 1
809+
yy: 2
810+
another line: 1
811+
"#,
812+
);
813+
814+
assert_eq!(kv.to_string(), expected);
815+
}
816+
817+
#[test]
818+
fn maybe_lazy_debug() {
819+
let x = MaybeLazy::Lazy(&1);
820+
let dbg = format!("{:?}", x);
821+
assert_eq!(dbg, "MaybeLazy::Lazy(1)");
822+
823+
let x = MaybeLazy::Eager("hello".into());
824+
let dbg = format!("{:?}", x);
825+
assert_eq!(dbg, "MaybeLazy::Eager(\"hello\")");
826+
}
609827
}

diskann-benchmark/Cargo.toml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ opentelemetry_sdk = { workspace = true, optional = true }
3939
scopeguard = { version = "1.2", optional = true }
4040
diskann-benchmark-core = { workspace = true, features = ["bigann"] }
4141
itertools.workspace = true
42+
diskann-inmem = { workspace = true, optional = true }
4243

4344
[lints]
4445
clippy.undocumented_unsafe_blocks = "warn"
@@ -67,6 +68,9 @@ minmax-quantization = []
6768
# Enable multi-vector MaxSim distance benchmarks
6869
multi-vector = []
6970

71+
# Enable inmem 2.0
72+
inmem2 = ["dep:diskann-inmem"]
73+
7074
# Enable bftree backend
7175
bftree = ["dep:diskann-bftree"]
7276

0 commit comments

Comments
 (0)