mirror of
https://github.com/facebookresearch/faiss.git
synced 2026-10-12 02:13:36 +00:00
Ivan Sopin
37313e7b64
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
123 lines
3.9 KiB
C++
123 lines
3.9 KiB
C++
/*
|
|
* 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
|