From a11c9f87c1e08ecb4f61dc9c5297d5ea2fa89026 Mon Sep 17 00:00:00 2001 From: Jack Moffitt Date: Mon, 20 Jul 2026 11:42:51 -0500 Subject: [PATCH] [vectorset] Runbook support --- vectorset/.gitignore | 1 + vectorset/Cargo.lock | 486 ++++++++++++++++- vectorset/Cargo.toml | 10 + vectorset/README.md | 22 +- vectorset/datasets/wikipedia-100K.yaml | 9 + vectorset/datasets/wikipedia-1M.yaml | 9 + vectorset/src/catalog.rs | 66 +++ vectorset/src/dataset.rs | 579 ++++++++++++++++++++ vectorset/src/driver.rs | 68 +++ vectorset/src/garnet.rs | 380 +++++++++++++ vectorset/src/main.rs | 249 ++++++++- vectorset/src/report.rs | 63 +++ vectorset/src/runbook.rs | 272 ++++++++++ vectorset/src/runner.rs | 704 +++++++++++++++++++++++++ vectorset/src/test_utils.rs | 18 + vectorset/test_data/f32_good.bin | Bin 0 -> 408 bytes vectorset/test_data/f32_large_dim.bin | Bin 0 -> 8 bytes vectorset/test_data/f32_too_long.bin | Bin 0 -> 608 bytes vectorset/test_data/f32_too_short.bin | Bin 0 -> 208 bytes vectorset/test_data/f32_zero_count.bin | Bin 0 -> 8 bytes vectorset/test_data/f32_zero_dim.bin | Bin 0 -> 8 bytes vectorset/test_data/u32_good.bin | Bin 0 -> 408 bytes 22 files changed, 2902 insertions(+), 34 deletions(-) create mode 100644 vectorset/datasets/wikipedia-100K.yaml create mode 100644 vectorset/datasets/wikipedia-1M.yaml create mode 100644 vectorset/src/catalog.rs create mode 100644 vectorset/src/dataset.rs create mode 100644 vectorset/src/driver.rs create mode 100644 vectorset/src/garnet.rs create mode 100644 vectorset/src/report.rs create mode 100644 vectorset/src/runbook.rs create mode 100644 vectorset/src/runner.rs create mode 100644 vectorset/src/test_utils.rs create mode 100644 vectorset/test_data/f32_good.bin create mode 100644 vectorset/test_data/f32_large_dim.bin create mode 100644 vectorset/test_data/f32_too_long.bin create mode 100644 vectorset/test_data/f32_too_short.bin create mode 100644 vectorset/test_data/f32_zero_count.bin create mode 100644 vectorset/test_data/f32_zero_dim.bin create mode 100644 vectorset/test_data/u32_good.bin diff --git a/vectorset/.gitignore b/vectorset/.gitignore index 5b6c0960c6..2ab4e3d926 100644 --- a/vectorset/.gitignore +++ b/vectorset/.gitignore @@ -1 +1,2 @@ config.toml +reports \ No newline at end of file diff --git a/vectorset/Cargo.lock b/vectorset/Cargo.lock index 974b5808b5..40bd7c0e8f 100644 --- a/vectorset/Cargo.lock +++ b/vectorset/Cargo.lock @@ -8,6 +8,26 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "annotate-snippets" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f211a51805bc641f3ad5b7664c77d2547af685cc33b4cd8d31964027a46f13f1" +dependencies = [ + "anstyle", + "memchr", + "unicode-width", +] + [[package]] name = "anstream" version = "1.0.0" @@ -70,6 +90,12 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" +[[package]] +name = "arraydeque" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d902e3d592a523def97af8f317b08ce16b7ab854c1985a0c671e6f15cebc236" + [[package]] name = "async-compression" version = "0.4.42" @@ -213,6 +239,20 @@ name = "bytemuck" version = "1.25.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95832e849adfb21180ccb6826a99da14e5d266ae5c2e668e1602cf234f153797" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f65693059b6b9c588b9f62fed1cedbf0a8b805631457ea162d68f0de186f3de5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "bytes" @@ -252,7 +292,21 @@ checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" dependencies = [ "cfg-if", "cpufeatures", - "rand_core", + "rand_core 0.10.1", +] + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", ] [[package]] @@ -396,12 +450,37 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "deranged" version = "0.5.8" @@ -411,6 +490,95 @@ dependencies = [ "serde_core", ] +[[package]] +name = "diskann" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31a7680a192e1abc16cb70415cb6c1aa4c4cf3269294ce2583c5d81763ee9762" +dependencies = [ + "anyhow", + "bytemuck", + "diskann-utils", + "diskann-vector", + "diskann-wide", + "futures-util", + "half", + "hashbrown 0.16.1", + "num-traits", + "rand 0.9.5", + "thiserror", + "tokio", + "tracing", +] + +[[package]] +name = "diskann-benchmark-core" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e9ecfa86aa7fcab045157644b0d06bfefb76b44e8d3decded140c9baed282a" +dependencies = [ + "anyhow", + "diskann", + "diskann-benchmark-runner", + "diskann-utils", + "futures-util", + "thiserror", + "tokio", +] + +[[package]] +name = "diskann-benchmark-runner" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d52165452b429d84f8f61682fc16c2190edc1d16cddb514c8cd5f49b85971ce" +dependencies = [ + "anyhow", + "clap", + "half", + "indicatif", + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "diskann-utils" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e2d755436dbe7141db9d2311ba99e885126003251d9392d5e41768548df777a" +dependencies = [ + "bytemuck", + "cfg-if", + "diskann-vector", + "diskann-wide", + "half", + "rand 0.9.5", + "rand_distr", + "rayon", + "thiserror", +] + +[[package]] +name = "diskann-vector" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf19284b35e697b4c49e71306425490b526ebe70b5602eb06d19c3013fb5138c" +dependencies = [ + "cfg-if", + "diskann-wide", + "half", +] + +[[package]] +name = "diskann-wide" +version = "0.55.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d3f814c2b66a98968630b7ad31afbb3c5e62b94b26f80542878eb7440aa8e5d" +dependencies = [ + "cfg-if", + "half", +] + [[package]] name = "displaydoc" version = "0.2.6" @@ -434,12 +602,36 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + [[package]] name = "encode_unicode" version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "encoding_rs_io" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cc3c5651fb62ab8aa3103998dade57efdd028544bd300516baa31840c252a83" +dependencies = [ + "encoding_rs", +] + [[package]] name = "equivalent" version = "1.0.2" @@ -499,6 +691,12 @@ dependencies = [ "miniz_oxide", ] +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + [[package]] name = "foreign-types" version = "0.3.2" @@ -630,6 +828,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.3" @@ -639,11 +849,45 @@ dependencies = [ "cfg-if", "js-sys", "libc", - "r-efi", - "rand_core", + "r-efi 6.0.0", + "rand_core 0.10.1", "wasm-bindgen", ] +[[package]] +name = "granit-parser" +version = "1.0.0-rc.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a8a3d7c6a0b0e077d7f8d21df7ec74ce97870ef25f40af4336df33b6cfdf1d4" +dependencies = [ + "arraydeque", + "smallvec", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "bytemuck", + "cfg-if", + "crunchy", + "num-traits", + "rand 0.9.5", + "rand_distr", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "foldhash", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -753,6 +997,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -863,7 +1131,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown", + "hashbrown 0.17.1", ] [[package]] @@ -973,6 +1241,12 @@ version = "0.2.188" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22053b6a34f84abc97f9129e61334f40174659a1b9bd18c970b83db6a9a6348b" +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1041,6 +1315,12 @@ dependencies = [ "tempfile", ] +[[package]] +name = "nohash-hasher" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2bf50223579dc7cdcfb3bfcacf7069ff68243f8c363f62ffa99cf000a6b9c451" + [[package]] name = "num-bigint" version = "0.4.8" @@ -1073,6 +1353,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -1195,6 +1476,15 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "proc-macro2" version = "1.0.107" @@ -1234,7 +1524,7 @@ dependencies = [ "bytes", "getrandom 0.4.3", "lru-slab", - "rand", + "rand 0.10.2", "rand_pcg", "ring", "rustc-hash", @@ -1270,12 +1560,28 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.2" @@ -1284,7 +1590,26 @@ checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ "chacha20", "getrandom 0.4.3", - "rand_core", + "rand_core 0.10.1", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -1293,13 +1618,43 @@ version = "0.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.5", +] + [[package]] name = "rand_pcg" version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" dependencies = [ - "rand_core", + "rand_core 0.10.1", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", ] [[package]] @@ -1555,6 +1910,23 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "serde-saphyr" +version = "1.0.0-rc.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f585dc5cd9a358fd20811c238205b1bf009f061de1d0477263159631dd4cc92" +dependencies = [ + "annotate-snippets", + "base64", + "encoding_rs_io", + "granit-parser", + "nohash-hasher", + "num-traits", + "serde_core", + "smallvec", + "zmij", +] + [[package]] name = "serde_core" version = "1.0.229" @@ -2013,7 +2385,7 @@ dependencies = [ "dyn-clone", "futures", "pin-project", - "rand", + "rand 0.10.2", "reqwest", "serde", "serde_json", @@ -2111,11 +2483,19 @@ dependencies = [ "azure_core", "azure_identity", "bytemuck", + "chrono", "clap", + "diskann-benchmark-core", + "diskann-benchmark-runner", + "diskann-utils", "indicatif", - "rand", + "rand 0.10.2", "redis", "serde", + "serde-saphyr", + "serde_json", + "tempfile", + "thiserror", "tokio", "toml", ] @@ -2145,6 +2525,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.126" @@ -2251,12 +2640,65 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" @@ -2345,6 +2787,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "writeable" version = "0.6.3" @@ -2380,6 +2828,26 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zerofrom" version = "0.1.8" diff --git a/vectorset/Cargo.toml b/vectorset/Cargo.toml index 2c5630a5c1..cb54f3ce93 100644 --- a/vectorset/Cargo.toml +++ b/vectorset/Cargo.toml @@ -13,6 +13,10 @@ azure_core = "1.1" azure_identity = "1.0" bytemuck = "1.23.0" clap = { version = "4.5.45", features = ["derive"] } +chrono = { version = "0.4.45", features = ["serde"] } +diskann-benchmark-core = "0.55.0" +diskann-benchmark-runner = "0.55.0" +diskann-utils = "0.55.0" indicatif = "0.18.3" rand = "0.10.2" redis = { version = "1.4.1", features = [ @@ -20,6 +24,9 @@ redis = { version = "1.4.1", features = [ "vector-sets", ] } serde = { version = "1.0.228", features = ["derive"] } +serde-saphyr = "1.0.0-rc.1" +serde_json = "1.0.151" +thiserror = "2.0.19" tokio = { version = "1.48.0", features = [ "rt-multi-thread", "net", @@ -28,3 +35,6 @@ tokio = { version = "1.48.0", features = [ "sync", ] } toml = "1.1.3" + +[dev-dependencies] +tempfile = "3.27.0" diff --git a/vectorset/README.md b/vectorset/README.md index c292c1d88e..2492f2320d 100644 --- a/vectorset/README.md +++ b/vectorset/README.md @@ -39,8 +39,10 @@ Benchmarks](https://github.com/harsha-simhadri/big-ann-benchmarks)'s `create_dataset.py` to download datasets. Once you have a dataset and Garnet is running, copy config.toml.example to config.toml and modify as necessary. -Since vectorset has its own workspace in the DiskANN repo, the following examples should be run from the `vectorset` directory in the repo. Running `cargo run --release -- --help` will enumerate -all the various subcommands and arguments. +Since vectorset has its own workspace in the DiskANN repo, the following +examples should be run from the `vectorset` directory in the repo. Running +`cargo run --release -- --help` and `cargo run --release -- --help` +will enumerate all the various subcommands and arguments. ### Ingest @@ -53,3 +55,19 @@ The `ingest` subcommand will ingest vectors and build the index. The following e The `query` subcommand queries the database while calculating recall against precomputed ground truth. By default, it will measure 10-recall@10. The following example queries the index created above and reports queries per second (QPS) and recall results. `cargo run --release -- --config path/to/config.toml --quantizer no-quant query --tasks 32 -k 100 -n 100 --l-search 192 path/to/data/wikipedia_cohere/wikipedia_query.bin path/to/data/wikipedia_cohere/wikipedia-10M` + +### Run + +The `run` subcommand executes [Big ANN +Benchmark](https://github.com/harsha-simhadri/big-ann-benchmarks) style +runbooks. It requires a dataset catalog, describing metadata about available +datasets. You can find an example in the `datasets` directory. Aside from +console output, it will write reports to disk in the `reports` directory by +default. + +Unlike the `query` command, searches executed from runbooks are not pipelined in +order to get more accurate latency metrics. + +The following example command executes a runbook: + +`cargo run --release -- --config path/to/config.toml --quantizer no-quant --threads 8 run --tasks 8 --degree 48 --l-build 300 --l-search 100 path/to/runbook.yaml path/to/datasets` diff --git a/vectorset/datasets/wikipedia-100K.yaml b/vectorset/datasets/wikipedia-100K.yaml new file mode 100644 index 0000000000..8d0c3b4f75 --- /dev/null +++ b/vectorset/datasets/wikipedia-100K.yaml @@ -0,0 +1,9 @@ +vectors: 100000 +queries: 5000 +dim: 768 +data-type: "float32" +metric: "innerproduct" +base-path: "wikipedia_cohere/wikipedia_base.bin.crop_nb_100000" +query-path: "wikipedia_cohere/wikipedia_query.bin" +gt-path: "wikipedia_cohere/wikipedia-100K" +step-gt-dir: "wikipedia_cohere/100000" diff --git a/vectorset/datasets/wikipedia-1M.yaml b/vectorset/datasets/wikipedia-1M.yaml new file mode 100644 index 0000000000..2b4c78e6ef --- /dev/null +++ b/vectorset/datasets/wikipedia-1M.yaml @@ -0,0 +1,9 @@ +vectors: 1000000 +queries: 5000 +dim: 768 +data-type: "float32" +metric: "innerproduct" +base-path: "wikipedia_cohere/wikipedia_base.bin.crop_nb_1000000" +query-path: "wikipedia_cohere/wikipedia_query.bin" +gt-path: "wikipedia_cohere/wikipedia-1M" +step-gt-dir: "wikipedia_cohere/1000000" diff --git a/vectorset/src/catalog.rs b/vectorset/src/catalog.rs new file mode 100644 index 0000000000..faf5d01b6e --- /dev/null +++ b/vectorset/src/catalog.rs @@ -0,0 +1,66 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + collections::HashMap, + ffi::OsStr, + path::{Path, PathBuf}, +}; + +use thiserror::Error; + +use crate::dataset::{Dataset, DatasetError}; + +#[derive(Debug, Error)] +pub enum CatalogError { + #[error("dataset key (stem of {0}) is not valid")] + BadName(PathBuf), + #[error("dataset error: {0}")] + Dataset(#[from] DatasetError), + #[error("io error: {0}")] + Io(#[from] std::io::Error), +} + +pub struct Catalog { + datasets: HashMap, +} + +impl Catalog { + pub fn load_directory, S: AsRef>( + path: P, + search_paths: Option<&[S]>, + ) -> Result { + let mut datasets = HashMap::new(); + + for spec_path in std::fs::read_dir(path)? + .filter_map(|de| de.ok()) + .filter_map(|de| { + if let Ok(ft) = de.file_type() + && ft.is_file() + && de.path().extension().unwrap_or(OsStr::new("")) == "yaml" + { + Some(de.path()) + } else { + None + } + }) + { + let name = spec_path.file_stem().and_then(|s| s.to_str()).unwrap_or(""); + if name.is_empty() { + return Err(CatalogError::BadName(spec_path)); + } + + let ds = Dataset::from_path(spec_path.clone(), search_paths)?; + + datasets.insert(name.to_string(), ds); + } + + Ok(Self { datasets }) + } + + pub fn dataset(&self, name: &str) -> Option<&Dataset> { + self.datasets.get(name) + } +} diff --git a/vectorset/src/dataset.rs b/vectorset/src/dataset.rs new file mode 100644 index 0000000000..dcba0ab5bb --- /dev/null +++ b/vectorset/src/dataset.rs @@ -0,0 +1,579 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use diskann_benchmark_core::recall::Rows; +use diskann_utils::views::Matrix; +use serde::Deserialize; +use std::{ + fs::File, + io::{Read, Seek}, + ops::Deref, + path::{Path, PathBuf}, +}; +use thiserror::Error; + +use crate::{DistanceMetric, Element, ElementType}; + +const MAX_DIM: usize = 16_384; + +#[derive(Debug, Clone, PartialEq, Deserialize)] +pub enum DataType { + #[serde( + rename = "float32", + alias = "f32", + alias = "Float32", + alias = "FLOAT32" + )] + F32, + #[serde( + rename = "int8", + alias = "i8", + alias = "Int8", + alias = "INT8", + alias = "I8" + )] + I8, + #[serde( + rename = "uint8", + alias = "u8", + alias = "Uint8", + alias = "UINT8", + alias = "U8" + )] + U8, +} + +impl DataType { + pub fn size(&self) -> usize { + match self { + DataType::F32 => std::mem::size_of::(), + DataType::I8 => std::mem::size_of::(), + DataType::U8 => std::mem::size_of::(), + } + } +} + +#[derive(Debug, Error)] +pub enum DatasetSpecError { + #[error("spec has no vectors")] + NoVectors, + #[error("spec has no queries")] + NoQueries, + #[error("bad dimension of {0}: must be positive and less than {MAX_DIM}")] + BadDim(usize), + #[error("io error: {0}")] + Io(#[from] std::io::Error), + #[error("bad yaml: {0}")] + BadYaml(#[from] serde_saphyr::Error), +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "kebab-case")] +struct DatasetSpec { + vectors: usize, + queries: usize, + dim: usize, + data_type: DataType, + metric: DistanceMetric, + base_path: PathBuf, + query_path: PathBuf, + gt_path: PathBuf, + step_gt_dir: PathBuf, +} + +impl DatasetSpec { + pub fn from_path>(path: P) -> Result { + let f = File::open(path)?; + let spec = serde_saphyr::from_reader(f)?; + Self::validate(&spec)?; + Ok(spec) + } + + fn validate(&self) -> Result<(), DatasetSpecError> { + if self.vectors == 0 { + return Err(DatasetSpecError::NoVectors); + } + if self.queries == 0 { + return Err(DatasetSpecError::NoQueries); + } + if self.dim == 0 || self.dim > MAX_DIM { + return Err(DatasetSpecError::BadDim(self.dim)); + } + + Ok(()) + } +} + +impl From for ElementType { + fn from(value: DataType) -> Self { + match value { + DataType::F32 => ElementType::F32, + DataType::I8 => ElementType::I8, + DataType::U8 => ElementType::U8, + } + } +} + +#[derive(Debug, Error)] +pub enum RowBufError { + #[error("error creating rowbuf")] + Create, +} + +#[derive(Debug)] +pub struct RowBuf { + start: usize, + buf: Matrix, +} + +impl RowBuf { + fn from_buffer( + start: usize, + dim: usize, + count: usize, + buffer: Vec, + ) -> Result { + let buf = Matrix::try_from(buffer.into_boxed_slice(), count, dim) + .map_err(|_| RowBufError::Create)?; + Ok(Self { start, buf }) + } + + pub fn row(&self, index: usize) -> &[T] { + assert!(index >= self.start && index - self.start < self.buf.nrows()); + + self.buf.row(index - self.start) + } + + pub fn start(&self) -> usize { + self.start + } +} + +impl Deref for RowBuf { + type Target = Matrix; + + fn deref(&self) -> &Self::Target { + &self.buf + } +} + +impl Rows for RowBuf { + fn nrows(&self) -> usize { + self.buf.nrows() + } + + fn ncols(&self) -> Option { + Some(self.buf.ncols()) + } + + fn row(&self, i: usize) -> &[T] { + RowBuf::row(self, i) + } +} + +#[derive(Debug, Error)] +pub enum DatasetError { + #[error("io error: {0}")] + Io(#[from] std::io::Error), + #[error("invalid dataset spec: {0}")] + BadSpec(#[from] DatasetSpecError), + #[error("missing path: {0}")] + MissingPath(PathBuf), + #[error("dimension specified ({1}) doesn't match dimension in {0} data file ({2})")] + DimMismatch(String, usize, usize), + #[error("{0} data file had wrong size")] + WrongSize(String), + #[error("{0} data file had different count ({1}) than expected ({2})")] + WrongCount(String, usize, usize), + #[error("found {0} ground truth entries, but expected {1}")] + GtMismatch(usize, usize), + #[error("gt dimension ({0}) must be positive and less than {MAX_DIM}")] + BadGtDim(usize), + #[error("index range ({0}..{1}) greater than max bound ({2})")] + IndexRangeOutOfBounds(usize, usize, usize), + #[error("rowbuf: {0}")] + RowBuf(#[from] RowBufError), + #[error("element type mismatch (expected: {0:?}; got {1:?})")] + ElementType(ElementType, ElementType), +} + +#[derive(Debug)] +pub struct Dataset { + data_type: DataType, + + base_path: PathBuf, + query_path: PathBuf, + step_gt_dir: PathBuf, + + dim: usize, + metric: DistanceMetric, + vector_count: usize, + query_count: usize, + + search_paths: Option>, +} + +impl Dataset { + pub fn from_path, S: AsRef>( + path: P, + search_paths: Option<&[S]>, + ) -> Result { + let spec = DatasetSpec::from_path(&path)?; + + let base_path = Self::resolve_path(&spec.base_path, search_paths)?; + let query_path = Self::resolve_path(&spec.query_path, search_paths)?; + let gt_path = Self::resolve_path(&spec.gt_path, search_paths)?; + let step_gt_dir = Self::resolve_path(&spec.step_gt_dir, search_paths)?; + + let dim = spec.dim; + let metric = spec.metric; + + let (vector_count, base_dim, base_len) = Self::metadata(&base_path)?; + if vector_count != spec.vectors { + return Err(DatasetError::WrongCount( + "base".to_string(), + vector_count, + spec.vectors, + )); + } + if dim != base_dim { + return Err(DatasetError::DimMismatch("base".to_string(), dim, base_dim)); + } + let base_len_check = vector_count * base_dim * spec.data_type.size() + 8; // 8 bytes extra for the u32 count and dimension + if base_len != base_len_check { + return Err(DatasetError::WrongSize("base".to_string())); + } + + let (query_count, query_dim, query_len) = Self::metadata(&query_path)?; + if query_count != spec.queries { + return Err(DatasetError::WrongCount( + "query".to_string(), + query_count, + spec.queries, + )); + } + if dim != query_dim { + return Err(DatasetError::DimMismatch( + "query".to_string(), + dim, + query_dim, + )); + } + let query_len_check = query_count * query_dim * spec.data_type.size() + 8; // 8 bytes extra for the u32 count and dimension + if query_len != query_len_check { + return Err(DatasetError::WrongSize("query".to_string())); + } + + let (gt_count, gt_dim, gt_len) = Self::metadata(>_path)?; + if gt_dim == 0 || gt_dim > MAX_DIM { + return Err(DatasetError::BadGtDim(gt_dim)); + } + if gt_count != query_count { + return Err(DatasetError::GtMismatch(gt_count, query_count)); + } + let gt_len_check = + gt_count * gt_dim * (std::mem::size_of::() + std::mem::size_of::()) + 8; // 8 bytes extra for the u32 count and dimension + if gt_len != gt_len_check { + return Err(DatasetError::WrongSize("gt".to_string())); + } + + let search_paths = search_paths.map(|sp| { + sp.iter() + .map(|p| p.as_ref().to_owned()) + .collect::>() + }); + + Ok(Dataset { + data_type: spec.data_type, + base_path, + query_path, + step_gt_dir, + dim, + metric, + vector_count, + query_count, + search_paths, + }) + } + + pub fn vector_count(&self) -> usize { + self.vector_count + } + + pub fn query_count(&self) -> usize { + self.query_count + } + + pub fn metric(&self) -> DistanceMetric { + self.metric + } + + fn read_rows( + &self, + path: &Path, + index: usize, + count: usize, + dim: usize, + ) -> Result, DatasetError> { + let mut f = File::open(path)?; + + let offset = index * dim * std::mem::size_of::() + 8; + + let mut row = vec![T::default(); dim * count]; + f.seek(std::io::SeekFrom::Start(offset as u64))?; + f.read_exact(bytemuck::cast_slice_mut(&mut row))?; + + Ok(RowBuf::from_buffer(index, dim, count, row)?) + } + + pub fn vectors( + &self, + index: usize, + count: usize, + ) -> Result, DatasetError> { + if T::ELEMENT_TYPE != self.data_type.clone().into() { + return Err(DatasetError::ElementType( + self.data_type.clone().into(), + T::ELEMENT_TYPE, + )); + } + + if index >= self.vector_count || index + count > self.vector_count { + return Err(DatasetError::IndexRangeOutOfBounds( + index, + count, + self.vector_count, + )); + } + + self.read_rows::(&self.base_path, index, count, self.dim) + } + + pub fn queries( + &self, + index: usize, + count: usize, + ) -> Result, DatasetError> { + if T::ELEMENT_TYPE != self.data_type.clone().into() { + return Err(DatasetError::ElementType( + self.data_type.clone().into(), + T::ELEMENT_TYPE, + )); + } + + if index >= self.query_count || index + count > self.query_count { + return Err(DatasetError::IndexRangeOutOfBounds( + index, + count, + self.query_count, + )); + } + + self.read_rows::(&self.query_path, index, count, self.dim) + } + + fn resolve_path, S: AsRef>( + path: P, + search_paths: Option<&[S]>, + ) -> Result { + if path.as_ref().exists() { + return Ok(path.as_ref().to_path_buf()); + } + + if !path.as_ref().is_relative() { + return Err(DatasetError::MissingPath(path.as_ref().to_path_buf())); + } + + if let Some(search_paths) = search_paths { + for search_path in search_paths { + let p = search_path.as_ref().join(&path); + if p.exists() { + return Ok(p); + } + } + } + + Err(DatasetError::MissingPath(path.as_ref().to_path_buf())) + } + + fn metadata>(path: P) -> Result<(usize, usize, usize), DatasetError> { + let mut f = File::open(&path)?; + + let len = f.metadata()?.len() as usize; + + let mut count = 0u32; + let mut dim = 0u32; + + f.read_exact(bytemuck::bytes_of_mut(&mut count))?; + f.read_exact(bytemuck::bytes_of_mut(&mut dim))?; + + Ok((count as usize, dim as usize, len)) + } + + pub fn step_gt( + &self, + runbook_name: &str, + step: usize, + ) -> Result<(RowBuf, RowBuf), DatasetError> { + let step_gt_path = PathBuf::new() + .join(&self.step_gt_dir) + .join(runbook_name) + .join(format!("step{step}.gt100")); + + let resolved = Self::resolve_path(step_gt_path, self.search_paths.as_deref())?; + + // Step ground truth files carry their own dimension, which need not match `gt_path`'s. + let (gt_count, gt_dim, _) = Self::metadata(&resolved)?; + if gt_dim == 0 || gt_dim > MAX_DIM { + return Err(DatasetError::BadGtDim(gt_dim)); + } + + let nrows = self.read_rows::(&resolved, 0, self.query_count, gt_dim)?; + let drows = self.read_rows::(&resolved, gt_count, self.query_count, gt_dim)?; + + Ok((nrows, drows)) + } +} + +#[cfg(test)] +mod test { + use std::assert_matches; + + use crate::{ + dataset::{Dataset, DatasetError, DatasetSpec, DatasetSpecError}, + test_utils::create_test_yaml, + }; + + #[test] + fn really_invalid_yaml() { + let (_tmpdir, path) = create_test_yaml( + r#" + this is not valid + "#, + ); + + let res = DatasetSpec::from_path(&path); + assert_matches!(res, Err(DatasetSpecError::BadYaml(_))); + } + + #[test] + fn zero_dim() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 10 + queries: 10 + dim: 0 + data-type: float32 + metric: l2 + base-path: base + query-path: query + gt-path: gt + step-gt-dir: . + "#, + ); + + let res = DatasetSpec::from_path(&path); + assert_matches!(res, Err(DatasetSpecError::BadDim(0))); + } + #[test] + fn dim_too_big() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 10 + queries: 10 + dim: 120039102938123 + data-type: float32 + metric: l2 + base-path: base + query-path: query + gt-path: gt + step-gt-dir: . + "#, + ); + + let res = DatasetSpec::from_path(&path); + assert_matches!(res, Err(DatasetSpecError::BadDim(120039102938123))); + } + + #[test] + fn spec_vector_count_zero() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 0 + queries: 10 + dim: 10 + data-type: "float32" + metric: l2 + base-path: f32_zero_count.bin + query-path: f32_good.bin + gt-path: u32_good.bin + step-gt-dir: . + "#, + ); + + let res = Dataset::from_path(&path, Some(&["test_data"])); + assert_matches!(res, Err(DatasetError::BadSpec(DatasetSpecError::NoVectors))); + } + + #[test] + fn vector_count_zero() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 10 + queries: 10 + dim: 10 + data-type: "float32" + metric: l2 + base-path: f32_zero_count.bin + query-path: f32_good.bin + gt-path: u32_good.bin + step-gt-dir: . + "#, + ); + + let res = Dataset::from_path(&path, Some(&["test_data"])); + assert_matches!(res, Err(DatasetError::WrongCount(_, 0, 10))); + } + + #[test] + fn spec_query_count_zero() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 10 + queries: 0 + dim: 10 + data-type: "float32" + metric: l2 + base-path: f32_good.bin + query-path: f32_zero_count.bin + gt-path: u32_good.bin + step-gt-dir: . + "#, + ); + + let res = Dataset::from_path(&path, Some(&["test_data"])); + assert_matches!(res, Err(DatasetError::BadSpec(DatasetSpecError::NoQueries))); + } + + #[test] + fn query_count_zero() { + let (_tmpdir, path) = create_test_yaml( + r#" + vectors: 10 + queries: 10 + dim: 10 + data-type: "float32" + metric: l2 + base-path: f32_good.bin + query-path: f32_zero_count.bin + gt-path: u32_good.bin + step-gt-dir: . + "#, + ); + + let res = Dataset::from_path(&path, Some(&["test_data"])); + assert_matches!(res, Err(DatasetError::WrongCount(_, 0, 10))); + } +} diff --git a/vectorset/src/driver.rs b/vectorset/src/driver.rs new file mode 100644 index 0000000000..dc091a7bba --- /dev/null +++ b/vectorset/src/driver.rs @@ -0,0 +1,68 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::time::Duration; + +use thiserror::Error; +use tokio::sync::mpsc; + +use crate::{DistanceMetric, Element, dataset::RowBuf}; + +#[derive(Debug, Error)] +#[error("controller: {0}")] +pub struct ControllerError(pub Box); + +pub type Timings = Vec<(usize, Duration)>; +pub type SearchResults = (Timings, usize, Vec>); + +pub trait Driver { + type Connection: Send; + type Error: std::error::Error + Send + Sync + 'static + From; + type Data: Element + Send; + + fn name(&self) -> String; + fn parallelism(&self) -> usize; + fn get_connection(&self) -> impl Future> + Send; + fn prepare( + &self, + con: Self::Connection, + ) -> impl Future> + Send; + fn finish(&self, con: Self::Connection) + -> impl Future> + Send; + + fn insert( + &self, + con: Self::Connection, + metric: DistanceMetric, + vectors: RowBuf, + reporter: mpsc::UnboundedSender, + ) -> impl Future> + Send; + + fn delete( + &self, + con: Self::Connection, + start: usize, + end: usize, + reporter: mpsc::UnboundedSender, + ) -> impl Future> + Send; + + fn replace( + &self, + con: Self::Connection, + metric: DistanceMetric, + tags_start: usize, + tags_end: usize, + vectors: RowBuf, + reporter: mpsc::UnboundedSender, + ) -> impl Future> + Send; + + fn search( + &self, + con: Self::Connection, + queries: RowBuf, + recall_n: usize, + reporter: mpsc::UnboundedSender, + ) -> impl Future> + Send; +} diff --git a/vectorset/src/garnet.rs b/vectorset/src/garnet.rs new file mode 100644 index 0000000000..d83d69c25c --- /dev/null +++ b/vectorset/src/garnet.rs @@ -0,0 +1,380 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{marker::PhantomData, time::Instant}; + +use redis::{AsyncCommands, Pipeline}; +use thiserror::Error; +use tokio::sync::mpsc; + +use crate::{ + DataType, DistanceMetric, Element, ExpiringCredential, ExpiringCredentialError, Quantizer, + VectorId, + dataset::RowBuf, + driver::{ControllerError, Driver, SearchResults, Timings}, +}; + +#[derive(Debug, Error)] +pub enum GarnetError { + #[error("redis error: {0}")] + Redis(#[from] redis::RedisError), + #[error("credential error: {0}")] + Credential(#[from] ExpiringCredentialError), + #[error("controller error: {0}")] + Controller(#[from] ControllerError), + #[error("replace tags/ids mismatch (ids: {0}, tags: {1})")] + ReplaceMismatch(usize, usize), +} + +pub struct Garnet { + client: redis::Client, + cred: Option, + + vset: String, + pipeline_size: usize, + parallelism: usize, + data_type: DataType, + + degree: usize, + l_build: usize, + l_search: usize, + quantizer: Quantizer, + + _phantom: PhantomData, +} + +impl Garnet { + #[allow(clippy::too_many_arguments)] + pub fn new( + client: redis::Client, + cred: Option, + vset: String, + pipeline_size: usize, + parallelism: usize, + data_type: DataType, + degree: usize, + l_build: usize, + l_search: usize, + quantizer: Quantizer, + ) -> Self { + Self { + client, + cred, + vset, + pipeline_size, + parallelism, + data_type, + degree, + l_build, + l_search, + quantizer, + _phantom: PhantomData, + } + } +} + +impl Driver for Garnet { + type Connection = redis::aio::MultiplexedConnection; + type Error = GarnetError; + type Data = T; + + fn name(&self) -> String { + "Garnet".to_string() + } + + fn parallelism(&self) -> usize { + self.parallelism + } + + async fn get_connection(&self) -> Result { + Ok(self + .client + .get_multiplexed_async_connection_with_config(&crate::connection_config()) + .await?) + } + + async fn prepare(&self, mut con: Self::Connection) -> Result<(), Self::Error> { + let _: usize = con.del(self.vset.as_bytes()).await?; + Ok(()) + } + + async fn finish(&self, mut con: Self::Connection) -> Result<(), Self::Error> { + let _: usize = con.del(self.vset.as_bytes()).await?; + Ok(()) + } + + async fn insert( + &self, + mut con: Self::Connection, + metric: DistanceMetric, + vectors: RowBuf, + reporter: mpsc::UnboundedSender, + ) -> Result { + let mut pipeline = Pipeline::with_capacity(self.pipeline_size); + let mut cred = self.cred.clone(); + + let start = vectors.start(); + let end = start + vectors.nrows(); + + let mut timings = Vec::new(); + let mut id = start; + while id < end { + if let Some(c) = cred { + cred = Some(c.refresh_if_needed(&mut con).await?); + } + + while id < end { + pipeline.clear(); + + let count = (end - id).min(self.pipeline_size); + for i in 0..count { + let element = VectorId((id + i) as u32); + + pipeline.cmd("VADD").arg(&self.vset); + + match self.data_type { + DataType::Float32 => { + pipeline.arg(b"FP32"); + } + DataType::Int8 => { + pipeline.arg(b"XI8"); + } + DataType::Uint8 => { + pipeline.arg(b"XU8"); + } + } + + pipeline + .arg(bytemuck::cast_slice::(vectors.row(id + i))) + .arg(element); + + pipeline.arg(self.quantizer); + + pipeline.arg(b"XDISTANCE_METRIC").arg(metric); + + pipeline + .arg(b"EF") + .arg(self.l_build.to_string().as_bytes()) + .arg(b"M") + .arg(self.degree.to_string().as_bytes()); + } + + id += count; + + let started = Instant::now(); + + pipeline.exec_async(&mut con).await?; + + let duration = Instant::now().duration_since(started); + + timings.push((count, duration)); + + // A closed channel just means nobody is collecting stats. + let _ = reporter.send(count); + } + } + + Ok(timings) + } + + async fn delete( + &self, + mut con: Self::Connection, + start: usize, + end: usize, + reporter: mpsc::UnboundedSender, + ) -> Result { + let mut pipeline = Pipeline::with_capacity(self.pipeline_size); + let mut cred = self.cred.clone(); + + let mut timings = Vec::new(); + let mut id = start; + while id < end { + if let Some(c) = cred { + cred = Some(c.refresh_if_needed(&mut con).await?); + } + + while id < end { + pipeline.clear(); + + let count = (end - id).min(self.pipeline_size); + for i in 0..count { + let element = VectorId((id + i) as u32); + + pipeline.cmd("VREM").arg(&self.vset).arg(element); + } + + id += count; + + let started = Instant::now(); + + pipeline.exec_async(&mut con).await?; + + let duration = Instant::now().duration_since(started); + + timings.push((count, duration)); + + // A closed channel just means nobody is collecting stats. + let _ = reporter.send(count); + } + } + + Ok(timings) + } + + async fn replace( + &self, + mut con: Self::Connection, + metric: DistanceMetric, + tags_start: usize, + tags_end: usize, + vectors: RowBuf, + reporter: mpsc::UnboundedSender, + ) -> Result { + let mut pipeline = Pipeline::with_capacity(self.pipeline_size); + let mut cred = self.cred.clone(); + + let start = vectors.start(); + let end = start + vectors.nrows(); + + if vectors.nrows() != tags_end - tags_start { + return Err(GarnetError::ReplaceMismatch( + vectors.nrows(), + tags_end - tags_start, + )); + } + + let mut timings = Vec::new(); + let mut id = start; + let mut tag = tags_start; + while id < end { + if let Some(c) = cred { + cred = Some(c.refresh_if_needed(&mut con).await?); + } + + while id < end { + pipeline.clear(); + + let count = (end - id).min(self.pipeline_size); + for i in 0..count { + let element = VectorId((tag + i) as u32); + + pipeline.cmd("VADD").arg(&self.vset); + + match self.data_type { + DataType::Float32 => { + pipeline.arg(b"FP32"); + } + DataType::Int8 => { + pipeline.arg(b"XI8"); + } + DataType::Uint8 => { + pipeline.arg(b"XU8"); + } + } + + pipeline + .arg(bytemuck::cast_slice::(vectors.row(id + i))) + .arg(element); + + pipeline.arg(self.quantizer); + + pipeline.arg(b"XDISTANCE_METRIC").arg(metric); + + pipeline + .arg(b"EF") + .arg(self.l_build.to_string().as_bytes()) + .arg(b"M") + .arg(self.degree.to_string().as_bytes()); + } + + id += count; + tag += count; + + let started = Instant::now(); + + pipeline.exec_async(&mut con).await?; + + let duration = Instant::now().duration_since(started); + + timings.push((count, duration)); + + // A closed channel just means nobody is collecting stats. + let _ = reporter.send(count); + } + } + + Ok(timings) + } + + async fn search( + &self, + mut con: Self::Connection, + queries: RowBuf, + recall_n: usize, + reporter: mpsc::UnboundedSender, + ) -> Result { + let mut cred = self.cred.clone(); + + let start = queries.start(); + let end = start + queries.nrows(); + + let mut timings = Vec::new(); + let mut results = Vec::new(); + + let mut id = start; + while id < end { + if let Some(c) = cred { + cred = Some(c.refresh_if_needed(&mut con).await?); + } + + let mut cmd = redis::cmd("VSIM"); + cmd.arg(&self.vset); + + match self.data_type { + DataType::Float32 => { + cmd.arg(b"FP32"); + } + DataType::Int8 => { + cmd.arg(b"XI8"); + } + DataType::Uint8 => { + cmd.arg(b"XU8"); + } + } + + cmd.arg(bytemuck::cast_slice::(queries.row(id))) + .arg(b"COUNT") + .arg(recall_n.to_string().as_bytes()) + .arg(b"EF") + .arg(self.l_search.to_string().as_bytes()); + + id += 1; + + let started = Instant::now(); + + let all_results: Vec<[u8; 4]> = cmd.query_async(&mut con).await?; + + let duration = Instant::now().duration_since(started); + + timings.push((1, duration)); + results.push( + all_results + .into_iter() + .map(|b| { + let mut id = 0u32; + bytemuck::bytes_of_mut(&mut id).copy_from_slice(&b); + id + }) + .collect::>(), + ); + + // A closed channel just means nobody is collecting stats. + let _ = reporter.send(1); + } + + Ok((timings, start, results)) + } +} diff --git a/vectorset/src/main.rs b/vectorset/src/main.rs index 5d6a9fd9bb..5eb8ec3c00 100644 --- a/vectorset/src/main.rs +++ b/vectorset/src/main.rs @@ -10,9 +10,10 @@ use clap::{Args, Parser, Subcommand, ValueEnum}; use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; use loader::DatasetLoader; use redis::{ - AsyncTypedCommands, IntoConnectionInfo, Pipeline, ToRedisArgs, aio::MultiplexedConnection, + AsyncConnectionConfig, AsyncTypedCommands, IntoConnectionInfo, Pipeline, ToRedisArgs, + aio::MultiplexedConnection, }; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use std::{ collections::HashSet, path::PathBuf, @@ -20,6 +21,7 @@ use std::{ thread, time::{Duration, Instant}, }; +use thiserror::Error; use tokio::{ fs::File, io::AsyncReadExt, @@ -27,23 +29,70 @@ use tokio::{ task::{JoinHandle, JoinSet}, }; +use crate::{ + catalog::Catalog, + garnet::Garnet, + runbook::Runbook, + runner::{Filter, Runner}, +}; + +mod catalog; +mod dataset; +mod driver; +mod garnet; mod loader; +mod report; +mod runbook; +mod runner; +#[cfg(test)] +mod test_utils; const DEFAULT_PORT: u16 = 6379; -trait Element: bytemuck::Pod + std::fmt::Debug + Send + Sync + 'static {} +/// redis-rs defaults to a 500ms response timeout, which a pipelined batch of searches or +/// inserts will always exceed. +fn connection_config() -> AsyncConnectionConfig { + AsyncConnectionConfig::new() + .set_response_timeout(None) + .set_connection_timeout(Some(Duration::from_secs(30))) +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ElementType { + F32, + U8, + I8, + U32, +} + +pub trait Element: bytemuck::Pod + Default + std::fmt::Debug + Send + Sync + 'static { + const ELEMENT_TYPE: ElementType; +} + +impl Element for f32 { + const ELEMENT_TYPE: ElementType = ElementType::F32; +} + +impl Element for u8 { + const ELEMENT_TYPE: ElementType = ElementType::U8; +} -impl Element for u8 {} -impl Element for i8 {} -impl Element for f32 {} +impl Element for i8 { + const ELEMENT_TYPE: ElementType = ElementType::I8; +} +impl Element for u32 { + const ELEMENT_TYPE: ElementType = ElementType::U32; +} #[derive(Deserialize)] +#[serde(rename_all = "kebab-case")] struct Config { ips: Vec, port: Option, secure: bool, scope: Option, username: Option, + dataset_search_paths: Option>, } #[derive(Parser)] @@ -74,6 +123,8 @@ enum Commands { Delete(DeleteArgs), /// Run queries and calculate recall Query(QueryArgs), + /// Run runbook + Run(RunArgs), } #[derive(Args)] @@ -143,11 +194,11 @@ struct QueryArgs { #[arg(long, default_value = "15")] l_search: usize, - /// Number of search results to return + /// Number of ground truth neighbors to score against (the k in k-recall@n) #[arg(short, long, default_value = "10")] k: usize, - /// Number of ground vectors to consider + /// Number of search results to return (the n in k-recall@n) #[arg(short, long, default_value = "10")] n: usize, @@ -162,6 +213,71 @@ struct QueryArgs { gt_path: PathBuf, } +#[derive(Args)] +struct RunArgs { + /// Vector set key prefix + #[arg(short, long, value_name = "VECTOR_SET", default_value = "vs0")] + set: String, + + /// Number of parallel search tasks + #[arg( + short, + long, + value_parser = clap::builder::RangedU64ValueParser::::new().range(1..) + )] + tasks: Option, + + /// Number of pipelined commands to the server (searches are never pipelined) + #[arg( + long, + default_value_t = 64, + value_parser = clap::builder::RangedU64ValueParser::::new().range(1..) + )] + pipeline_size: usize, + + /// Graph degree + #[arg(long, default_value = "16")] + degree: usize, + + /// Candidate list size during build + #[arg(long, default_value = "15")] + l_build: usize, + + /// Candidate list size during search + #[arg(long, default_value = "15")] + l_search: usize, + + /// Number of ground truth neighbors to score against (the k in k-recall@n) + #[arg(short, long, default_value = "10")] + k: usize, + + /// Number of search results to return (the n in k-recall@n) + #[arg(short, long, default_value = "10")] + n: usize, + + /// Include dataset filter + #[arg(long)] + include: Vec, + + /// Exclude dataset filter + #[arg(long)] + exclude: Vec, + + /// Repeat search steps + #[arg(long, default_value_t = 5, value_parser = clap::builder::RangedU64ValueParser::::new().range(1..))] + search_repetitions: usize, + + /// Output directory for reports + #[arg(long, default_value = "reports")] + report_path: PathBuf, + + /// Runbook to execute + runbook: PathBuf, + + /// Dataset catalog directory + catalog: PathBuf, +} + #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] enum DataType { Uint8, @@ -170,7 +286,7 @@ enum DataType { } #[allow(non_camel_case_types)] -#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Default)] +#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Default, Serialize)] enum Quantizer { /// f32 vectors; no quantization #[default] @@ -232,11 +348,22 @@ impl std::fmt::Display for Quantizer { } } -#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)] +#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum, Serialize, Deserialize)] enum DistanceMetric { + #[serde(rename = "l2", alias = "L2")] L2, + #[serde(rename = "cosine", alias = "COSINE")] Cosine, + #[serde(rename = "cosine_normalized", alias = "COSINE_NORMALIZED")] CosineNormalized, + #[serde( + rename = "innerproduct", + alias = "InnerProduct", + alias = "ip", + alias = "INNERPRODUCT", + alias = "inner_product", + alias = "INNER_PRODUCT" + )] InnerProduct, } @@ -266,6 +393,14 @@ impl ToRedisArgs for VectorId { } } +#[derive(Debug, Error)] +pub enum ExpiringCredentialError { + #[error("redis error: {0}")] + Redis(#[from] redis::RedisError), + #[error("azure error: {0}")] + Azure(#[from] azure_core::Error), +} + #[derive(Clone)] struct ExpiringCredential { scope: String, @@ -289,7 +424,10 @@ impl ExpiringCredential { } } - async fn refresh_if_needed(mut self, con: &mut MultiplexedConnection) -> Result { + async fn refresh_if_needed( + mut self, + con: &mut MultiplexedConnection, + ) -> std::result::Result { if self.expires - OffsetDateTime::now_utc() < Duration::from_secs(300) { let res = self.cred.get_token(&[&self.scope], None).await?; @@ -387,31 +525,35 @@ async fn async_main(opts: Options) -> Result<()> { }; match opts.quantizer.data_type() { - DataType::Uint8 => dispatch::(&opts.command, &opts, infos, cred).await, - DataType::Int8 => dispatch::(&opts.command, &opts, infos, cred).await, - DataType::Float32 => dispatch::(&opts.command, &opts, infos, cred).await, + DataType::Uint8 => dispatch::(&config, &opts.command, &opts, infos, cred).await, + DataType::Int8 => dispatch::(&config, &opts.command, &opts, infos, cred).await, + DataType::Float32 => dispatch::(&config, &opts.command, &opts, infos, cred).await, } } async fn dispatch( + config: &Config, command: &Commands, opts: &Options, infos: Vec, cred: Option, ) -> Result<()> { match command { - Commands::Ping => ping::(infos[0].clone()).await?, + Commands::Ping => ping(infos[0].clone()).await?, Commands::Ingest(args) => ingest::(opts, args, infos[0].clone(), cred).await?, - Commands::Delete(args) => delete::(args, infos[0].clone()).await?, + Commands::Delete(args) => delete(args, infos[0].clone()).await?, Commands::Query(args) => query::(opts, args, infos, cred).await?, + Commands::Run(args) => run::(config, opts, args, infos, cred).await?, } Ok(()) } -async fn ping(info: redis::ConnectionInfo) -> Result<()> { +async fn ping(info: redis::ConnectionInfo) -> Result<()> { let client = redis::Client::open(info).unwrap(); - let mut con = client.get_multiplexed_async_connection().await?; + let mut con = client + .get_multiplexed_async_connection_with_config(&connection_config()) + .await?; println!("PING..."); let result = con.ping().await?; @@ -454,7 +596,9 @@ async fn ingest( // Insert base vectors for _ in 0..parallelism { let client = redis::Client::open(info.clone())?; - let mut con = client.get_multiplexed_async_connection().await?; + let mut con = client + .get_multiplexed_async_connection_with_config(&connection_config()) + .await?; let ds = ds.clone(); let pipeline_size = args.pipeline_size; let vset = vset.clone(); @@ -556,12 +700,13 @@ async fn ingest( Ok(()) } -async fn delete(args: &DeleteArgs, info: redis::ConnectionInfo) -> Result<()> { +async fn delete(args: &DeleteArgs, info: redis::ConnectionInfo) -> Result<()> { let client = redis::Client::open(info).unwrap(); - let mut con = client.get_multiplexed_async_connection().await?; + let mut con = client + .get_multiplexed_async_connection_with_config(&connection_config()) + .await?; con.del(&args.set).await?; - con.flushdb().await?; Ok(()) } @@ -600,7 +745,9 @@ async fn query( let time_start = Instant::now(); for task_idx in 0..parallelism { let client = redis::Client::open(infos[task_idx % infos.len()].clone())?; - let mut con = client.get_multiplexed_async_connection().await?; + let mut con = client + .get_multiplexed_async_connection_with_config(&connection_config()) + .await?; let pipeline_size = args.pipeline_size; let tx = tx.clone(); let queries = queries.clone(); @@ -722,3 +869,59 @@ async fn query( Ok(()) } + +async fn run( + config: &Config, + opts: &Options, + args: &RunArgs, + infos: Vec, + cred: Option, +) -> Result<()> { + let book = Runbook::from_path(&args.runbook)?; + let cat = Catalog::load_directory(&args.catalog, config.dataset_search_paths.as_deref())?; + + let parallelism = args.tasks.unwrap_or(thread::available_parallelism()?.get()); + let client = redis::Client::open(infos[0].clone())?; + + let filter = if !args.include.is_empty() || !args.exclude.is_empty() { + let mut filter = Filter::default(); + for included in &args.include { + filter.include(included); + } + for excluded in &args.exclude { + filter.exclude(excluded); + } + Some(filter) + } else { + None + }; + + // Execute the runbook + let driver = Garnet::::new( + client, + cred, + args.set.clone(), + args.pipeline_size, + parallelism, + opts.quantizer.data_type(), + args.degree, + args.l_build, + args.l_search, + opts.quantizer, + ); + let runner = Runner::new(driver); + runner + .run( + &book, + &cat, + &args.report_path, + args.k, + args.n, + filter, + args, + opts, + ) + .await?; + + Ok(()) +} diff --git a/vectorset/src/report.rs b/vectorset/src/report.rs new file mode 100644 index 0000000000..c1b4f9ce1e --- /dev/null +++ b/vectorset/src/report.rs @@ -0,0 +1,63 @@ +use std::collections::HashMap; + +use chrono::{DateTime, Utc}; +use serde::Serialize; + +use crate::Quantizer; + +#[derive(Serialize)] +pub struct Report { + pub date: DateTime, + + pub num_threads: usize, + pub quantizer: Quantizer, + + pub num_tasks: usize, + pub pipeline_size: usize, + pub search_repetitions: usize, + pub degree: usize, + pub l_build: usize, + pub l_search: usize, + pub k: usize, + pub n: usize, + + pub runbook: String, + pub dataset: HashMap>, +} + +#[derive(Clone, Serialize)] +pub enum StepReport { + Insert(OpReport), + Delete(OpReport), + Replace(OpReport), + Search(SearchReport), +} + +#[derive(Clone, Serialize)] +pub struct OpReport { + pub parallelism: usize, + pub count: usize, + pub wall_time_s: f64, + pub busy_time_s: f64, + pub latency_us_mean: f64, + pub latency_us_p90: f64, + pub latency_us_p99: f64, +} + +impl OpReport { + pub fn throughput(&self) -> f64 { + self.count as f64 / self.wall_time_s + } + + pub fn utilization(&self) -> f64 { + self.busy_time_s / (self.wall_time_s * self.parallelism as f64) + } +} + +#[derive(Clone, Serialize)] +pub struct SearchReport { + pub op_reports: Vec, + pub k: usize, + pub n: usize, + pub recall: f64, +} diff --git a/vectorset/src/runbook.rs b/vectorset/src/runbook.rs new file mode 100644 index 0000000000..9644f2694c --- /dev/null +++ b/vectorset/src/runbook.rs @@ -0,0 +1,272 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + collections::HashMap, + fmt, + fs::File, + path::{Path, PathBuf}, +}; + +use serde::{ + Deserialize, Deserializer, + de::{MapAccess, Visitor}, +}; +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum RunbookError { + #[error("io error: {0}")] + Io(#[from] std::io::Error), + #[error("parsing error: {0}")] + Parsing(#[from] serde_saphyr::Error), + #[error("runbook missing steps for {0}")] + NoSteps(String), +} + +#[derive(Debug, Deserialize)] +#[serde(transparent)] +pub struct Runbook { + #[serde(skip)] + path: PathBuf, + recipes: HashMap, +} + +#[derive(Debug, Deserialize)] +pub struct Recipe { + #[serde(rename = "max_pts")] + max_points: usize, + #[serde(flatten)] + steps: Steps, +} + +#[derive(Debug)] +struct Steps(Vec); + +impl<'de> Deserialize<'de> for Steps { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct StepsVisitor; + + impl<'de> Visitor<'de> for StepsVisitor { + type Value = Steps; + + fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.write_str("a mapping of step numbers to operations") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut steps = Vec::new(); + while let Some((idx, operation)) = map.next_entry::()? { + let idx = idx + .parse::() + .map_err(|_| serde::de::Error::custom("operation index not an integer"))?; + if idx != steps.len() + 1 { + return Err(serde::de::Error::custom(format!( + "operation index {idx} out of sequence" + ))); + } + match operation { + Operation::Insert { start, end } if start >= end => { + return Err(serde::de::Error::custom( + "start/end ranges invalid".to_string(), + )); + } + Operation::Delete { start, end } if start >= end => { + return Err(serde::de::Error::custom( + "start/end ranges invalid".to_string(), + )); + } + Operation::Replace { + tags_start, + tags_end, + ids_start, + ids_end, + } => { + if tags_start >= tags_end || ids_start >= ids_end { + return Err(serde::de::Error::custom( + "tag or id ranges are invalid".to_string(), + )); + } + if tags_end.wrapping_sub(tags_start) != ids_end.wrapping_sub(ids_start) + { + return Err(serde::de::Error::custom(format!( + "replace operation {idx} has mismatched tag and id range sizes" + ))); + } + } + _ => {} + } + steps.push(operation); + } + Ok(Steps(steps)) + } + } + + deserializer.deserialize_map(StepsVisitor) + } +} + +impl Recipe { + pub fn steps(&self) -> impl Iterator { + self.steps.0.iter() + } + + pub fn max_points(&self) -> usize { + self.max_points + } +} + +#[derive(Debug, Deserialize)] +#[serde(tag = "operation", rename_all = "lowercase")] +pub enum Operation { + Insert { + start: usize, + end: usize, + }, + Delete { + start: usize, + end: usize, + }, + Search, + Replace { + tags_start: usize, + tags_end: usize, + ids_start: usize, + ids_end: usize, + }, +} + +impl Operation { + pub fn name(&self) -> &str { + match self { + Operation::Insert { .. } => "insert", + Operation::Delete { .. } => "delete", + Operation::Search => "search", + Operation::Replace { .. } => "replace", + } + } +} + +impl Runbook { + pub fn from_path>(path: P) -> Result { + let f = File::open(&path)?; + + let mut rb: Runbook = serde_saphyr::from_reader(f)?; + if let Some((name, _)) = rb.recipes.iter().find(|(_, r)| r.steps.0.is_empty()) { + return Err(RunbookError::NoSteps(name.clone())); + } + rb.path = path.as_ref().into(); + Ok(rb) + } + + pub fn recipes(&self) -> impl Iterator { + self.recipes.keys() + } + + pub fn recipe(&self, key: &str) -> Option<&Recipe> { + self.recipes.get(key) + } + + pub fn name(&self) -> Option<&str> { + self.path.file_name().and_then(|s| s.to_str()) + } +} + +#[cfg(test)] +mod test { + use std::assert_matches; + + use crate::{ + runbook::{Operation, Runbook, RunbookError}, + test_utils::create_test_yaml, + }; + + #[test] + fn basic() { + let (_tmpdir, path) = create_test_yaml( + r#" + dataset: + max_pts: 1000 + 1: + operation: "insert" + start: 0 + end: 1000 + 2: + operation: "search" + "#, + ); + + let rb = Runbook::from_path(&path).expect("runbook failed parsing"); + assert_eq!(rb.recipes().count(), 1); + let r = rb.recipe("dataset").expect("missing recipe"); + assert_eq!(r.max_points, 1000); + assert_eq!(r.steps.0.len(), 2); + assert_matches!( + r.steps.0[0], + Operation::Insert { + start: 0, + end: 1000 + }, + ); + assert_matches!(r.steps.0[1], Operation::Search); + } + + #[test] + fn bad_op_index() { + let (_tmpdir, path) = create_test_yaml( + r#" + dataset: + max_pts: 1000 + a: + operation: "insert" + start: 0 + end: 1000 + 2: + operation: "search" + "#, + ); + + let rb = Runbook::from_path(&path); + assert_matches!(rb, Err(RunbookError::Parsing(_))); + } + + #[test] + fn bad_op_index_seq() { + let (_tmpdir, path) = create_test_yaml( + r#" + dataset: + max_pts: 1000 + 1: + operation: "insert" + start: 0 + end: 1000 + 5: + operation: "search" + "#, + ); + + let rb = Runbook::from_path(&path); + assert_matches!(rb, Err(RunbookError::Parsing(_))); + } + + #[test] + fn missing_steps() { + let (_tmpdir, path) = create_test_yaml( + r#" + dataset: + max_pts: 1000 + "#, + ); + + let rb = Runbook::from_path(&path); + assert_matches!(rb, Err(RunbookError::NoSteps(name)) if name == "dataset"); + } +} diff --git a/vectorset/src/runner.rs b/vectorset/src/runner.rs new file mode 100644 index 0000000000..9cb4f491ae --- /dev/null +++ b/vectorset/src/runner.rs @@ -0,0 +1,704 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{ + collections::HashMap, + fs::{self, File}, + path::Path, + sync::Arc, + thread, + time::{Duration, Instant}, +}; + +use chrono::Utc; +use diskann_benchmark_core::recall::{self, ComputeRecallError}; +use diskann_benchmark_runner::utils::percentiles::{self, CannotBeEmpty, Percentiles}; +use diskann_utils::views::Matrix; +use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle}; +use thiserror::Error; +use tokio::{ + sync::mpsc, + task::{JoinError, JoinSet}, +}; + +use crate::{ + Options, RunArgs, + catalog::Catalog, + dataset::{Dataset, DatasetError, RowBuf}, + driver::{ControllerError, Driver}, + report::{OpReport, Report, SearchReport, StepReport}, + runbook::{Operation, Runbook}, +}; + +#[derive(Debug)] +pub struct Filter { + includes: Vec, + excludes: Vec, +} + +impl Filter { + pub fn include(&mut self, key: &str) { + self.includes.push(key.to_string()); + } + + pub fn exclude(&mut self, key: &str) { + self.excludes.push(key.to_string()); + } + + pub fn matches(&self, key: &str) -> bool { + let included = if !self.includes.is_empty() { + let mut included = false; + for in_rule in &self.includes { + if key.starts_with(in_rule) { + included = true; + break; + } + } + included + } else { + true + }; + + for ex_rule in &self.excludes { + if key.starts_with(ex_rule) { + return false; + } + } + + included + } +} + +impl Default for Filter { + fn default() -> Self { + let includes = vec![]; + let excludes = vec![]; + Self { includes, excludes } + } +} + +pub struct Runner { + driver: Arc, +} + +#[derive(Debug, Error)] +pub enum RunnerError { + #[error("Nothing to do; empty runbook or nothing matched filter")] + NothingToDo, + #[error("data set missing: {0}")] + DatasetMissing(String), + #[error("recipe missing: {0}")] + RecipeMissing(String), + #[error("recipe max points ({0}) bigger than dataset size ({1})")] + RecipeMaxPoints(usize, usize), + #[error("bad or missing runbook name")] + BadName, + #[error("redis error: {0}")] + Redis(#[from] redis::RedisError), + #[error("dataset error: {0}")] + Dataset(#[from] DatasetError), + #[error("driver error: {0}")] + Driver(Box), + #[error("compute recall error: {0}")] + ComputeRecall(#[from] ComputeRecallError), + #[error("progress task failed: {0}")] + ProgressTask(#[from] JoinError), + #[error("progress was empty")] + EmptyProgress(#[from] CannotBeEmpty), + #[error("file i/o error: {0}")] + FileIo(#[from] std::io::Error), + #[error("json error: {0}")] + Json(#[from] serde_json::Error), +} + +impl Runner { + pub fn new(driver: D) -> Self { + Self { + driver: Arc::new(driver), + } + } + + #[allow(clippy::too_many_arguments)] + pub async fn run( + &self, + runbook: &Runbook, + data_manager: &Catalog, + report_path: &Path, + recall_k: usize, + recall_n: usize, + filter: Option, + args: &RunArgs, + opts: &Options, + ) -> Result<(), RunnerError> { + println!( + "Executing runbook: {}...", + runbook.name().unwrap_or("unknown") + ); + + // Filter recipes + let recipes: Vec<&String> = if let Some(filter) = filter { + runbook.recipes().filter(|ds| filter.matches(ds)).collect() + } else { + runbook.recipes().collect() + }; + + if recipes.is_empty() { + println!("error: Runbook is empty."); + return Err(RunnerError::NothingToDo); + } + + let mut step_reports = HashMap::new(); + + // Run each recipe + for recipe_name in recipes { + println!("Running dataset {recipe_name}..."); + + // Make sure dataset <-> recipe makes sense + let ds = data_manager + .dataset(recipe_name) + .ok_or_else(|| RunnerError::DatasetMissing(recipe_name.clone()))?; + + let recipe = runbook + .recipe(recipe_name) + .ok_or_else(|| RunnerError::RecipeMissing(recipe_name.clone()))?; + + if recipe.max_points() > ds.vector_count() { + return Err(RunnerError::RecipeMaxPoints( + recipe.max_points(), + ds.vector_count(), + )); + } + + let con = self + .driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + self.driver + .prepare(con) + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + + let mut error = None; + for (step, op) in recipe.steps().enumerate() { + println!("Step {step}: {}", op.name()); + + let step_report = match op { + Operation::Insert { start, end } => self.op_insert(ds, *start, *end).await, + + Operation::Delete { start, end } => self.op_delete(*start, *end).await, + Operation::Replace { + tags_start, + tags_end, + ids_start, + .. + } => { + let count = tags_end - tags_start; + self.op_replace(ds, *tags_start, *ids_start, count).await + } + Operation::Search => { + match runbook + .name() + .ok_or(RunnerError::BadName) + .and_then(|name| ds.step_gt(name, step + 1).map_err(|e| e.into())) + .map(Arc::new) + { + Ok(step_gt) => { + let mut step_reports = Vec::new(); + for _ in 0..args.search_repetitions { + let step_report = self + .op_search(ds, step_gt.clone(), recall_k, recall_n) + .await; + + match step_report { + Ok(sr) => { + step_reports.push(sr); + } + Err(e) => { + error = Some(e); + break; + } + } + } + + if let Some(e) = error.take() { + Err(e) + } else { + let step_report = consolidate_search_reports(&step_reports); + Ok(step_report) + } + } + Err(e) => Err(e), + } + } + }; + + let step_report = match step_report { + Ok(sr) => sr, + Err(e) => { + error = Some(e); + break; + } + }; + + step_reports + .entry(recipe_name.clone()) + .or_insert_with(Vec::new) + .push(step_report.clone()); + + match step_report { + StepReport::Insert(op_report) => { + println!( + " Inserted {} vectors in {:0.3}s ({:0.3} inserts/sec; utilization was {:0.3}).", + op_report.count, + op_report.wall_time_s, + op_report.throughput(), + op_report.utilization(), + ); + } + StepReport::Search(search_report) => { + println!( + " Searches completed with {}-recall@{} of {:0.3}.", + search_report.k, search_report.n, search_report.recall + ); + + for (i, or) in search_report.op_reports.iter().enumerate() { + println!( + " Repetition {}: Queried {} vectors in {:0.3}s ({:0.3} qps; utilization was {:0.3}).", + i + 1, + or.count, + or.wall_time_s, + or.throughput(), + or.utilization(), + ); + } + } + StepReport::Delete(op_report) => { + println!( + " Deleted {} vectors in {:0.3}s ({:0.3} deletes/sec; utilization was {:0.3}).", + op_report.count, + op_report.wall_time_s, + op_report.throughput(), + op_report.utilization(), + ); + } + StepReport::Replace(op_report) => { + println!( + " Replaced {} vectors in {:0.3}s ({:0.3} replaces/sec; utilization was {:0.3}).", + op_report.count, + op_report.wall_time_s, + op_report.throughput(), + op_report.utilization(), + ); + } + } + } + + let result: Result<(), RunnerError> = async { + let con = self + .driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + self.driver + .finish(con) + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + Ok(()) + } + .await; + + if let Some(e) = error { + if let Err(e) = result { + eprintln!("cleanup failed: {e}"); + } + return Err(e); + } + + result?; + } + + let available_parallelism = thread::available_parallelism()?.get(); + let runbook_name = runbook.name().unwrap_or("unknown").to_string(); + let date = Utc::now(); + let report = Report { + date, + num_threads: opts.threads.unwrap_or(available_parallelism), + quantizer: opts.quantizer, + num_tasks: args.tasks.unwrap_or(available_parallelism), + pipeline_size: args.pipeline_size, + search_repetitions: args.search_repetitions, + degree: args.degree, + l_build: args.l_build, + l_search: args.l_search, + k: args.k, + n: args.n, + runbook: runbook_name.clone(), + dataset: step_reports, + }; + + let report_path = report_path.join(self.driver.name()).join(format!( + "{runbook_name}-{}.json", + date.format("%Y%m%dT%H%M%S%.3fZ") + )); + fs::create_dir_all(report_path.parent().unwrap_or(&report_path))?; + + let mut f = File::create(report_path)?; + serde_json::to_writer(&mut f, &report)?; + + Ok(()) + } + + async fn op_insert( + &self, + dataset: &Dataset, + start: usize, + end: usize, + ) -> Result { + let (tx, rx) = mpsc::unbounded_channel::(); + let mut tasks = JoinSet::, D::Error>>::new(); + + let progress_handle = tokio::spawn(progress_task::(rx, end - start)); + + let chunks = chunk_range(start, end, self.driver.parallelism()); + let num_tasks = chunks.len(); + let mut inputs = Vec::new(); + + for (chunk_start, chunk_end) in chunks { + inputs.push(( + Arc::clone(&self.driver), + self.driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?, + dataset.metric(), + dataset.vectors::(chunk_start, chunk_end - chunk_start)?, + tx.clone(), + )); + } + + let wall_time_start = Instant::now(); + + for (driver, con, metric, data, tx) in inputs { + tasks.spawn(async move { driver.insert(con, metric, data, tx).await }); + } + + let mut timings = Vec::new(); + for result in tasks.join_all().await { + timings.extend(result.map_err(|e| RunnerError::Driver(Box::new(e)))?); + } + + let wall_time_s = Instant::now().duration_since(wall_time_start).as_secs_f64(); + + drop(tx); + progress_handle + .await? + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + + let report = StepReport::Insert(make_op_report(timings, wall_time_s, num_tasks)?); + + Ok(report) + } + + async fn op_delete(&self, start: usize, end: usize) -> Result { + let (tx, rx) = mpsc::unbounded_channel::(); + let mut tasks = JoinSet::, D::Error>>::new(); + + let progress_handle = tokio::spawn(progress_task::(rx, end - start)); + + let chunks = chunk_range(start, end, self.driver.parallelism()); + let num_tasks = chunks.len(); + let mut inputs = Vec::new(); + + for _ in 0..num_tasks { + inputs.push(( + Arc::clone(&self.driver), + self.driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?, + tx.clone(), + )); + } + + let wall_time_start = Instant::now(); + + for ((chunk_start, chunk_end), (driver, con, tx)) in chunks.zip(inputs) { + tasks.spawn(async move { driver.delete(con, chunk_start, chunk_end, tx).await }); + } + + let mut timings = Vec::new(); + for result in tasks.join_all().await { + timings.extend(result.map_err(|e| RunnerError::Driver(Box::new(e)))?); + } + + let wall_time_s = Instant::now().duration_since(wall_time_start).as_secs_f64(); + + drop(tx); + progress_handle + .await? + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + + let report = StepReport::Delete(make_op_report(timings, wall_time_s, num_tasks)?); + + Ok(report) + } + + async fn op_search( + &self, + dataset: &Dataset, + step_gt: Arc<(RowBuf, RowBuf)>, + recall_k: usize, + recall_n: usize, + ) -> Result { + let (tx, rx) = mpsc::unbounded_channel::(); + let mut tasks = + JoinSet::, usize, Vec>), D::Error>>::new(); + + let progress_handle = tokio::spawn(progress_task::(rx, dataset.query_count())); + + let chunks = chunk_range(0, dataset.query_count(), self.driver.parallelism()); + let num_tasks = chunks.len(); + let mut inputs = Vec::new(); + + for (chunk_start, chunk_end) in chunks { + inputs.push(( + Arc::clone(&self.driver), + self.driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?, + dataset.queries::(chunk_start, chunk_end - chunk_start)?, + tx.clone(), + )); + } + + let wall_time_start = Instant::now(); + + for (driver, con, data, tx) in inputs { + tasks.spawn(async move { driver.search(con, data, recall_n, tx).await }); + } + + let results = tasks.join_all().await; + let wall_time_s = Instant::now().duration_since(wall_time_start).as_secs_f64(); + + let mut timings = Vec::new(); + let mut query_results = Matrix::::new(u32::MAX, dataset.query_count(), recall_n); + for result in results { + let (batch_timings, start_idx, results) = + result.map_err(|e| RunnerError::Driver(Box::new(e)))?; + timings.extend(batch_timings); + + for (i, row) in results.into_iter().enumerate() { + let n = row.len().min(recall_n); + query_results.row_mut(start_idx + i)[..n].copy_from_slice(&row[..n]); + } + } + + drop(tx); + progress_handle + .await? + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + + let op_report = make_op_report(timings, wall_time_s, num_tasks)?; + + let recall_metrics = recall::knn( + &step_gt.0, + Some(step_gt.1.as_view().into()), + &query_results, + recall_k, + recall_n, + recall::GroundTruthMode::Fixed, + )?; + + let op_reports = vec![op_report]; + let report = StepReport::Search(SearchReport { + op_reports, + k: recall_metrics.recall_k, + n: recall_metrics.recall_n, + recall: recall_metrics.average, + }); + + Ok(report) + } + + async fn op_replace( + &self, + dataset: &Dataset, + tags_start: usize, + ids_start: usize, + count: usize, + ) -> Result { + let (tx, rx) = mpsc::unbounded_channel::(); + let mut tasks = JoinSet::, D::Error>>::new(); + + let progress_handle = tokio::spawn(progress_task::(rx, count)); + + let chunks = chunk_range(0, count, self.driver.parallelism()); + let num_tasks = chunks.len(); + let mut inputs = Vec::new(); + + for (chunk_start, chunk_end) in chunks.clone() { + inputs.push(( + Arc::clone(&self.driver), + self.driver + .get_connection() + .await + .map_err(|e| RunnerError::Driver(Box::new(e)))?, + dataset.metric(), + dataset.vectors::(chunk_start + ids_start, chunk_end - chunk_start)?, + tx.clone(), + )); + } + + let wall_time_start = Instant::now(); + + for ((chunk_start, chunk_end), (driver, con, metric, data, tx)) in chunks.zip(inputs) { + tasks.spawn(async move { + driver + .replace( + con, + metric, + chunk_start + tags_start, + chunk_end + tags_start, + data, + tx, + ) + .await + }); + } + + let mut timings = Vec::new(); + for result in tasks.join_all().await { + timings.extend(result.map_err(|e| RunnerError::Driver(Box::new(e)))?); + } + + let wall_time_s = Instant::now().duration_since(wall_time_start).as_secs_f64(); + + drop(tx); + progress_handle + .await? + .map_err(|e| RunnerError::Driver(Box::new(e)))?; + + let report = StepReport::Replace(make_op_report(timings, wall_time_s, num_tasks)?); + + Ok(report) + } +} + +fn make_op_report( + results: Vec<(usize, Duration)>, + wall_time_s: f64, + parallelism: usize, +) -> Result { + let count = results.iter().fold(0, |acc, (cnt, _dur)| acc + *cnt); + let busy_time_s = results + .iter() + .fold(Duration::default(), |acc, (_cnt, dur)| acc + *dur) + .as_secs_f64(); + + let mut times: Vec = results + .iter() + .map(|(cnt, dur)| (dur.as_nanos() / *cnt as u128) as u64) + .collect(); + let Percentiles { mean, p90, p99, .. } = percentiles::compute_percentiles(&mut times)?; + + Ok(OpReport { + parallelism, + count, + wall_time_s, + busy_time_s, + latency_us_mean: mean / 1000.0, + latency_us_p90: p90 as f64 / 1000.0, + latency_us_p99: p99 as f64 / 1000.0, + }) +} + +/// Consolidates individual search reports into a single one folded under op_reports. +/// Only the first recall metrics are preserved. Other StepReport variants are passed +/// through, discarding all but the first. +fn consolidate_search_reports(reports: &[StepReport]) -> StepReport { + if reports.len() == 1 { + return reports[0].clone(); + } + + let mut report = reports[0].clone(); + let op_reports = if let StepReport::Search(sr) = &mut report { + &mut sr.op_reports + } else { + return report; + }; + + for r in &reports[1..] { + if let StepReport::Search(sr) = r { + op_reports.push(sr.op_reports[0].clone()); + } else { + return report; + } + } + + report +} + +async fn progress_task( + mut rx: mpsc::UnboundedReceiver, + count: usize, +) -> Result<(), D::Error> { + let progress = + ProgressBar::with_draw_target(Some(count as u64), ProgressDrawTarget::stderr_with_hz(1)); + progress.set_style( + ProgressStyle::with_template("{wide_bar} {pos}/{len} {elapsed}/{eta} {per_sec}") + .map_err(|e| ControllerError(Box::new(e)))?, + ); + + while let Some(count) = rx.recv().await { + progress.inc(count as u64); + } + + progress.finish_and_clear(); + + Ok(()) +} + +#[derive(Debug, Clone)] +pub struct ChunkIterator { + size: usize, + start: usize, + end: usize, +} + +impl ChunkIterator { + fn new(start: usize, end: usize, count: usize) -> Self { + let size = (end - start).div_ceil(count); + Self { size, start, end } + } + + fn len(&self) -> usize { + if self.size == 0 { + return 0; + } + + (self.end - self.start).div_ceil(self.size) + } +} + +impl Iterator for ChunkIterator { + type Item = (usize, usize); + + fn next(&mut self) -> Option { + if self.start >= self.end { + return None; + } + + let end = (self.start + self.size).min(self.end); + let item = (self.start, end); + self.start = end; + Some(item) + } +} + +fn chunk_range(start: usize, end: usize, count: usize) -> ChunkIterator { + ChunkIterator::new(start, end, count) +} diff --git a/vectorset/src/test_utils.rs b/vectorset/src/test_utils.rs new file mode 100644 index 0000000000..8503aa4c17 --- /dev/null +++ b/vectorset/src/test_utils.rs @@ -0,0 +1,18 @@ +/* + * Copyright (c) Microsoft Corporation. + * Licensed under the MIT license. + */ + +use std::{fs::File, io::Write, path::PathBuf}; + +pub fn create_test_yaml(s: &str) -> (tempfile::TempDir, PathBuf) { + let tmpdir = tempfile::tempdir().unwrap(); + let path = tmpdir.path().join("test.yaml"); + + { + let mut f = File::create(&path).unwrap(); + f.write_all(s.as_bytes()).unwrap(); + } + + (tmpdir, path) +} diff --git a/vectorset/test_data/f32_good.bin b/vectorset/test_data/f32_good.bin new file mode 100644 index 0000000000000000000000000000000000000000..2c29683a689d178e3fed46eee5a6e665ad9f68cf GIT binary patch literal 408 zcma*fK@9*P6vNP-{|QHcGo&B|DM(=s=#qD9RzQW{>R>SPDU3A6N@whJ#z|+~bOsZQ HmhSHXMTI&) literal 0 HcmV?d00001 diff --git a/vectorset/test_data/f32_large_dim.bin b/vectorset/test_data/f32_large_dim.bin new file mode 100644 index 0000000000000000000000000000000000000000..5bb8ae72c220f88aa29d671194f038a94b84e569 GIT binary patch literal 8 Pcmd;LU|?`?;%5K=0z3em literal 0 HcmV?d00001 diff --git a/vectorset/test_data/f32_too_long.bin b/vectorset/test_data/f32_too_long.bin new file mode 100644 index 0000000000000000000000000000000000000000..5a0f50e947a56bacaa59561d9df878cb4a36c41f GIT binary patch literal 608 zcma*fK@os36a&$mdqM&daKNvFN~}&>1J4anTt?XAC;yq1*cbn!QtT literal 0 HcmV?d00001 diff --git a/vectorset/test_data/f32_too_short.bin b/vectorset/test_data/f32_too_short.bin new file mode 100644 index 0000000000000000000000000000000000000000..359968f48ba0622949cb13b3ba9ccd256a296cce GIT binary patch literal 208 lcmd;LU|`?^V&XxAJp%(Vn85*PCK1>HsD>Ea08~d90|3|57K;D? literal 0 HcmV?d00001 diff --git a/vectorset/test_data/f32_zero_count.bin b/vectorset/test_data/f32_zero_count.bin new file mode 100644 index 0000000000000000000000000000000000000000..b41157351ae065612ceab973e57922d0cba8f00c GIT binary patch literal 8 NcmZQzU|`?^VgLXz01E&B literal 0 HcmV?d00001 diff --git a/vectorset/test_data/f32_zero_dim.bin b/vectorset/test_data/f32_zero_dim.bin new file mode 100644 index 0000000000000000000000000000000000000000..5de48411147e09c99d309bca0cc1a7f433fc6720 GIT binary patch literal 8 Kcmd;LfB*mhSO5zE literal 0 HcmV?d00001 diff --git a/vectorset/test_data/u32_good.bin b/vectorset/test_data/u32_good.bin new file mode 100644 index 0000000000000000000000000000000000000000..a20d19e1a704e4c5718b96e640f67e0cdd0f3dae GIT binary patch literal 408 ycma*d$q@h`5QM<-*3Ca)2?475M