Fix two memory leaks in shard_ivf_index_centroids (#5495)

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

@@ -592,8 +592,12 @@ void handle_ivf(
                 static_cast<faiss::IndexIVF*>(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 <typename IndexType>

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 <cstdint>
+#include <cstdio>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+
+#include <faiss/IVFlib.h>
+#include <faiss/IndexBinaryFlat.h>
+#include <faiss/IndexBinaryIVF.h>
+#include <faiss/IndexFlat.h>
+#include <faiss/IndexIVFFlat.h>
+#include <faiss/index_io.h>
+
+#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 ("<unique-prefix>.%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<float> centroids(nlist * d);
+    for (size_t i = 0; i < centroids.size(); i++) {
+        centroids[i] = static_cast<float>(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<uint8_t> centroids(nlist * (d / 8));
+    for (size_t i = 0; i < centroids.size(); i++) {
+        centroids[i] = static_cast<uint8_t>(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