mirror of
https://github.com/facebookresearch/faiss.git
synced 2025-06-03 21:54:02 +08:00
a996a4a052
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2582 A few more or less cosmetic improvements * Index::idx_t was in the Index object, which does not make much sense, this diff moves it to faiss::idx_t * replace multiprocessing.dummy with multiprocessing.pool * add Alexandr as a core contributor of Faiss in the README ;-) ``` for i in $( find . -name \*.cu -o -name \*.cuh -o -name \*.h -o -name \*.cpp ) ; do sed -i s/Index::idx_t/idx_t/ $i done ``` For the fbcode deps: ``` for i in $( fbgs Index::idx_t --exclude fbcode/faiss -l ) ; do sed -i s/Index::idx_t/idx_t/ $i done ``` Reviewed By: algoriddle Differential Revision: D41437507 fbshipit-source-id: 8300f2a3ae97cace6172f3f14a9be3a83999fb89
92 lines
2.7 KiB
Python
Executable File
92 lines
2.7 KiB
Python
Executable File
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
from multiprocessing.pool import ThreadPool
|
|
import faiss
|
|
from typing import List, Tuple
|
|
|
|
from . import rpc
|
|
|
|
############################################################
|
|
# Server implementation
|
|
############################################################
|
|
|
|
|
|
class SearchServer(rpc.Server):
|
|
""" Assign version that can be exposed via RPC """
|
|
|
|
def __init__(self, s: int, index: faiss.Index):
|
|
rpc.Server.__init__(self, s)
|
|
self.index = index
|
|
self.index_ivf = faiss.extract_index_ivf(index)
|
|
|
|
def set_nprobe(self, nprobe: int) -> int:
|
|
""" set nprobe field """
|
|
self.index_ivf.nprobe = nprobe
|
|
|
|
def get_ntotal(self) -> int:
|
|
return self.index.ntotal
|
|
|
|
def __getattr__(self, f):
|
|
# all other functions get forwarded to the index
|
|
return getattr(self.index, f)
|
|
|
|
|
|
def run_index_server(index: faiss.Index, port: int, v6: bool = False):
|
|
""" serve requests for that index forerver """
|
|
rpc.run_server(
|
|
lambda s: SearchServer(s, index),
|
|
port, v6=v6)
|
|
|
|
|
|
############################################################
|
|
# Client implementation
|
|
############################################################
|
|
|
|
class ClientIndex:
|
|
"""manages a set of distance sub-indexes. The sub_indexes search a
|
|
subset of the inverted lists. Searches are merged afterwards
|
|
"""
|
|
|
|
def __init__(self, machine_ports: List[Tuple[str, int]], v6: bool = False):
|
|
""" connect to a series of (host, port) pairs """
|
|
self.sub_indexes = []
|
|
for machine, port in machine_ports:
|
|
self.sub_indexes.append(rpc.Client(machine, port, v6))
|
|
|
|
self.ni = len(self.sub_indexes)
|
|
# pool of threads. Each thread manages one sub-index.
|
|
self.pool = ThreadPool(self.ni)
|
|
# test connection...
|
|
self.ntotal = self.get_ntotal()
|
|
self.verbose = False
|
|
|
|
def set_nprobe(self, nprobe: int) -> None:
|
|
self.pool.map(
|
|
lambda idx: idx.set_nprobe(nprobe),
|
|
self.sub_indexes
|
|
)
|
|
|
|
def set_omp_num_threads(self, nt: int) -> None:
|
|
self.pool.map(
|
|
lambda idx: idx.set_omp_num_threads(nt),
|
|
self.sub_indexes
|
|
)
|
|
|
|
def get_ntotal(self) -> None:
|
|
return sum(self.pool.map(
|
|
lambda idx: idx.get_ntotal(),
|
|
self.sub_indexes
|
|
))
|
|
|
|
def search(self, x, k: int):
|
|
|
|
rh = faiss.ResultHeap(x.shape[0], k)
|
|
|
|
for Di, Ii in self.pool.imap(lambda idx: idx.search(x, k), self.sub_indexes):
|
|
rh.add_result(Di, Ii)
|
|
rh.finalize()
|
|
return rh.D, rh.I
|