From 37313e7b6487544e899afce9338f3f987361222f Mon Sep 17 00:00:00 2001 From: Ivan Sopin Date: Fri, 31 Jul 2026 20:58:26 -0700 Subject: [PATCH] Fix two memory leaks in shard_ivf_index_centroids (#5495) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5495 `handle_ivf` / `handle_binary_ivf` (behind `shard_ivf_index_centroids` / `shard_binary_ivf_index_centroids`) leaked memory in two places; this fixes both: 1. **Leaked `clone`.** Each function clones the input index into a `clone` whose quantizer is reset, then clones `clone` once per shard. The per-shard clones are deleted, but `clone` itself was never freed — one leaked `IndexIVF` (with its sub-objects) per call. Added the missing `delete clone;` in both helpers. 2. **Leaked quantizer in the `generate_ids=true` path.** That path replaces the per-shard clone's quantizer with `new IndexIDMap2(quantizer)`, but the IDMap did not own the wrapped quantizer, so the cloned quantizer (and its codes) leaked per shard. Set `own_fields = true` on the IDMap so it is freed together with the shard index. Surfaced by a LeakSanitizer failure in a new Galileo unit test that shards a trained IVF model; the fix also benefits the existing LaserKNN sharded-training caller. Reviewed By: mnorris11 Differential Revision: D112336282 fbshipit-source-id: 080657355299a45c2f1dbb36cab9304d825daa5d --- faiss/IVFlib.cpp | 17 ++++- tests/test_ivf_sharding.cpp | 122 ++++++++++++++++++++++++++++++++++++ 2 files changed, 136 insertions(+), 3 deletions(-) create mode 100644 tests/test_ivf_sharding.cpp 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