Commit a86fe61
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
File tree
- .github/workflows
- diskann-benchmark-runner/src
- utils
- diskann-benchmark
- src/index
- diskann-inmem
- integration
- index
- jsons
- support
- src
- integration
- layers
- test
- diskann-vector/src/distance
- rfcs
Some content is hidden
Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
22 | 22 | | |
23 | 23 | | |
24 | 24 | | |
25 | | - | |
| 25 | + | |
26 | 26 | | |
27 | 27 | | |
28 | 28 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
17 | 17 | | |
18 | 18 | | |
19 | 19 | | |
20 | | - | |
| 20 | + | |
21 | 21 | | |
22 | 22 | | |
23 | 23 | | |
| |||
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
18 | 18 | | |
19 | 19 | | |
20 | 20 | | |
| 21 | + | |
21 | 22 | | |
22 | 23 | | |
23 | 24 | | |
| |||
60 | 61 | | |
61 | 62 | | |
62 | 63 | | |
| 64 | + | |
63 | 65 | | |
64 | 66 | | |
65 | 67 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
62 | 62 | | |
63 | 63 | | |
64 | 64 | | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
65 | 71 | | |
66 | 72 | | |
67 | 73 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
379 | 379 | | |
380 | 380 | | |
381 | 381 | | |
| 382 | + | |
| 383 | + | |
| 384 | + | |
| 385 | + | |
| 386 | + | |
| 387 | + | |
| 388 | + | |
| 389 | + | |
| 390 | + | |
| 391 | + | |
| 392 | + | |
| 393 | + | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
| 397 | + | |
| 398 | + | |
| 399 | + | |
| 400 | + | |
| 401 | + | |
| 402 | + | |
| 403 | + | |
| 404 | + | |
| 405 | + | |
| 406 | + | |
| 407 | + | |
| 408 | + | |
| 409 | + | |
| 410 | + | |
| 411 | + | |
| 412 | + | |
| 413 | + | |
| 414 | + | |
| 415 | + | |
| 416 | + | |
| 417 | + | |
| 418 | + | |
| 419 | + | |
| 420 | + | |
| 421 | + | |
| 422 | + | |
| 423 | + | |
| 424 | + | |
| 425 | + | |
| 426 | + | |
| 427 | + | |
| 428 | + | |
| 429 | + | |
| 430 | + | |
| 431 | + | |
| 432 | + | |
| 433 | + | |
| 434 | + | |
| 435 | + | |
| 436 | + | |
| 437 | + | |
| 438 | + | |
| 439 | + | |
| 440 | + | |
| 441 | + | |
| 442 | + | |
| 443 | + | |
| 444 | + | |
| 445 | + | |
| 446 | + | |
| 447 | + | |
| 448 | + | |
| 449 | + | |
| 450 | + | |
| 451 | + | |
| 452 | + | |
| 453 | + | |
| 454 | + | |
| 455 | + | |
| 456 | + | |
| 457 | + | |
| 458 | + | |
| 459 | + | |
| 460 | + | |
| 461 | + | |
| 462 | + | |
| 463 | + | |
| 464 | + | |
| 465 | + | |
| 466 | + | |
| 467 | + | |
| 468 | + | |
| 469 | + | |
| 470 | + | |
| 471 | + | |
| 472 | + | |
| 473 | + | |
| 474 | + | |
| 475 | + | |
| 476 | + | |
| 477 | + | |
| 478 | + | |
| 479 | + | |
| 480 | + | |
| 481 | + | |
| 482 | + | |
| 483 | + | |
| 484 | + | |
| 485 | + | |
| 486 | + | |
| 487 | + | |
| 488 | + | |
| 489 | + | |
| 490 | + | |
| 491 | + | |
| 492 | + | |
| 493 | + | |
| 494 | + | |
| 495 | + | |
| 496 | + | |
| 497 | + | |
| 498 | + | |
| 499 | + | |
| 500 | + | |
| 501 | + | |
| 502 | + | |
| 503 | + | |
| 504 | + | |
| 505 | + | |
| 506 | + | |
| 507 | + | |
| 508 | + | |
| 509 | + | |
| 510 | + | |
382 | 511 | | |
383 | 512 | | |
384 | 513 | | |
| |||
606 | 735 | | |
607 | 736 | | |
608 | 737 | | |
| 738 | + | |
| 739 | + | |
| 740 | + | |
| 741 | + | |
| 742 | + | |
| 743 | + | |
| 744 | + | |
| 745 | + | |
| 746 | + | |
| 747 | + | |
| 748 | + | |
| 749 | + | |
| 750 | + | |
| 751 | + | |
| 752 | + | |
| 753 | + | |
| 754 | + | |
| 755 | + | |
| 756 | + | |
| 757 | + | |
| 758 | + | |
| 759 | + | |
| 760 | + | |
| 761 | + | |
| 762 | + | |
| 763 | + | |
| 764 | + | |
| 765 | + | |
| 766 | + | |
| 767 | + | |
| 768 | + | |
| 769 | + | |
| 770 | + | |
| 771 | + | |
| 772 | + | |
| 773 | + | |
| 774 | + | |
| 775 | + | |
| 776 | + | |
| 777 | + | |
| 778 | + | |
| 779 | + | |
| 780 | + | |
| 781 | + | |
| 782 | + | |
| 783 | + | |
| 784 | + | |
| 785 | + | |
| 786 | + | |
| 787 | + | |
| 788 | + | |
| 789 | + | |
| 790 | + | |
| 791 | + | |
| 792 | + | |
| 793 | + | |
| 794 | + | |
| 795 | + | |
| 796 | + | |
| 797 | + | |
| 798 | + | |
| 799 | + | |
| 800 | + | |
| 801 | + | |
| 802 | + | |
| 803 | + | |
| 804 | + | |
| 805 | + | |
| 806 | + | |
| 807 | + | |
| 808 | + | |
| 809 | + | |
| 810 | + | |
| 811 | + | |
| 812 | + | |
| 813 | + | |
| 814 | + | |
| 815 | + | |
| 816 | + | |
| 817 | + | |
| 818 | + | |
| 819 | + | |
| 820 | + | |
| 821 | + | |
| 822 | + | |
| 823 | + | |
| 824 | + | |
| 825 | + | |
| 826 | + | |
609 | 827 | | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
39 | 39 | | |
40 | 40 | | |
41 | 41 | | |
| 42 | + | |
42 | 43 | | |
43 | 44 | | |
44 | 45 | | |
| |||
67 | 68 | | |
68 | 69 | | |
69 | 70 | | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
70 | 74 | | |
71 | 75 | | |
72 | 76 | | |
| |||
0 commit comments