diff --git a/faiss/IVFlib.cpp b/faiss/IVFlib.cpp index ecc1f23ff..c5a182c73 100644 --- a/faiss/IVFlib.cpp +++ b/faiss/IVFlib.cpp @@ -592,8 +592,12 @@ void handle_ivf( static_cast(faiss::clone_index(clone)); if (generate_ids) { // Assume the quantizer does not natively support add_with_ids. - sharded_index->quantizer = - new IndexIDMap2(sharded_index->quantizer); + // The IDMap wraps the cloned quantizer; own_fields makes it delete + // that quantizer when sharded_index (which owns the IDMap) is + // freed. + auto* id_quantizer = new IndexIDMap2(sharded_index->quantizer); + id_quantizer->own_fields = true; + sharded_index->quantizer = id_quantizer; sharded_index->quantizer->add_with_ids( sharded_centroids[i].size() / index->quantizer->d, sharded_centroids[i].data(), @@ -615,6 +619,7 @@ void handle_ivf( faiss::write_index(sharded_index, fname); delete sharded_index; } + delete clone; } void handle_binary_ivf( @@ -650,8 +655,13 @@ void handle_binary_ivf( faiss::clone_binary_index(clone)); if (generate_ids) { // Assume the quantizer does not natively support add_with_ids. - sharded_index->quantizer = + // The IDMap wraps the cloned quantizer; own_fields makes it delete + // that quantizer when sharded_index (which owns the IDMap) is + // freed. + auto* id_quantizer = new IndexBinaryIDMap2(sharded_index->quantizer); + id_quantizer->own_fields = true; + sharded_index->quantizer = id_quantizer; sharded_index->quantizer->add_with_ids( sharded_centroids[i].size() / reconstruction_size, sharded_centroids[i].data(), @@ -673,6 +683,7 @@ void handle_binary_ivf( faiss::write_index_binary(sharded_index, fname); delete sharded_index; } + delete clone; } template diff --git a/tests/test_ivf_sharding.cpp b/tests/test_ivf_sharding.cpp new file mode 100644 index 000000000..68f6d7057 --- /dev/null +++ b/tests/test_ivf_sharding.cpp @@ -0,0 +1,122 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include "test_util.h" + +// Regression guard for two leaks in handle_ivf / handle_binary_ivf +// (faiss/IVFlib.cpp), both invisible to normal runs but caught under +// AddressSanitizer/LeakSanitizer (fbcode @mode/dev-asan, run in CI): +// 1. the input `clone` created once per call and never freed; +// 2. (generate_ids=true only) the cloned quantizer swapped out for an +// IndexIDMap2 wrapper that previously did not own/free it. +// Each test exercises both generate_ids paths so both leaks are covered. + +namespace { + +pthread_mutex_t temp_file_mutex = PTHREAD_MUTEX_INITIALIZER; + +// Builds a shard filename template (".%d.index") from a fresh +// temp path so parallel test shards don't collide. +std::string shardTemplate(Tempfilename& prefix) { + return std::string(prefix.c_str()) + ".%d.index"; +} + +void removeShardFiles(const std::string& tmpl, int64_t shardCount) { + char fname[256]; + for (int64_t i = 0; i < shardCount; i++) { + snprintf(fname, sizeof(fname), tmpl.c_str(), i); + std::remove(fname); + } +} + +TEST(TestIvfSharding, ShardIvfCentroidsNoLeak) { + constexpr int d = 16; + constexpr int nlist = 32; + constexpr int64_t shardCount = 3; + + std::vector centroids(nlist * d); + for (size_t i = 0; i < centroids.size(); i++) { + centroids[i] = static_cast(i % 100) * 0.01f; + } + + for (bool generate_ids : {false, true}) { + faiss::IndexFlatL2 quantizer(d); + faiss::IndexIVFFlat index(&quantizer, d, nlist); + index.quantizer->add(nlist, centroids.data()); + + Tempfilename prefix(&temp_file_mutex, "/tmp/faiss_ivf_shard_XXXXXX"); + const std::string tmpl = shardTemplate(prefix); + + faiss::ivflib::shard_ivf_index_centroids( + &index, + shardCount, + tmpl, + /*sharding_function=*/nullptr, + generate_ids); + + char fname[256]; + for (int64_t i = 0; i < shardCount; i++) { + snprintf(fname, sizeof(fname), tmpl.c_str(), i); + faiss::Index* shard = faiss::read_index(fname); + EXPECT_NE(shard, nullptr); + delete shard; + } + removeShardFiles(tmpl, shardCount); + } +} + +TEST(TestIvfSharding, ShardBinaryIvfCentroidsNoLeak) { + constexpr int d = 64; // bits; d / 8 bytes per centroid + constexpr int nlist = 32; + constexpr int64_t shardCount = 3; + + std::vector centroids(nlist * (d / 8)); + for (size_t i = 0; i < centroids.size(); i++) { + centroids[i] = static_cast(i); + } + + for (bool generate_ids : {false, true}) { + faiss::IndexBinaryFlat quantizer(d); + faiss::IndexBinaryIVF index(&quantizer, d, nlist); + index.quantizer->add(nlist, centroids.data()); + + Tempfilename prefix(&temp_file_mutex, "/tmp/faiss_binivf_shard_XXXXXX"); + const std::string tmpl = shardTemplate(prefix); + + faiss::ivflib::shard_binary_ivf_index_centroids( + &index, + shardCount, + tmpl, + /*sharding_function=*/nullptr, + generate_ids); + + char fname[256]; + for (int64_t i = 0; i < shardCount; i++) { + snprintf(fname, sizeof(fname), tmpl.c_str(), i); + faiss::IndexBinary* shard = faiss::read_index_binary(fname); + EXPECT_NE(shard, nullptr); + delete shard; + } + removeShardFiles(tmpl, shardCount); + } +} + +} // namespace