Docker下SPTAG的安装与测试

2019-11-15 本文已影响0人叁公子

关于SPTAG

SPTAG是微软开发的一款近似最近邻搜索（ approximate nearest neighbor search）的库，可以用它来做dense vector的索引。

最常见的就是图像搜索这样的应用，当然文本检索做语义匹配也可以用到。

Docker 环境安装

因为SPTAG目前不支持mac版本，所以安装在docker里面就好了。我试了一下官方的dockerfile写的有点儿问题，我没运行起来。于是打算直接建一个docker镜像安装。

docker run -it ubuntu:18.04 "/bin/bash"

进入docker后运行：

apt update
apt -y install wget build-essential openjdk-8-jdk python3-pip swig 
apt -y install software-properties-common git

安装boost1.67

add-apt-repository ppa:mhier/libboost-latest
apt update
apt install libboost1.67 libboost1.67-dev

安装cmake3.15.5

wget https://cmake.org/files/v3.15/cmake-3.15.5.tar.gz
tar zvxf cmake-3.15.5.tar.gz
cd cmake-3.15.5
./bootstrap
make -j2 && make install

编译SPTAG

git clone https://github.com/microsoft/SPTAG.git
cd SPTAG && mkdir build && cd build
cmake ..
make -j2

到这里环境就算搞好了。把编译好的结果放到/app下

mkdir /app
mv Release /app

装一下测试用到的python库

pip3 install numpy rpyc

Docker镜像的保存

这里算是装好了基本的环境，这里再把环境保存一下。

先查看一下自己的container id。

docker ps

docker ps
我这里是：46b0c72411dc

docker commit -m "create SPTAG env" 46b0c72411dc nladuo/sptag-rpc-server:latest

再看下当前的镜像。

docker images

编写rpc服务

因为SPTAG不支持Mac，所以为了能在Mac上访问，这里可以编写一个简单的Rpc Demo服务，将接口稍微封装一下即可。

这里代码放到了：https://github.com/nladuo/SPTAG-rpc-service

SPTAG_rpc_demo_server.py需要放到docker中，SPTAG_rpc_demo_client.py则直接import到自己的包里即可。

这里我们先把之前的镜像停掉，重新开一个带端口映射的容器（我这里用的8888端口）。

docker run -p 8888:8888 -t -i nladuo/sptag-rpc-server:latest "/bin/bash"

这里先把SPTAG_rpc_demo_server.py拷贝到新的docker容器中（注意容器的id的变化）。

docker cp SPTAG_rpc_demo_server.py 25042d741f07:/app/Release/

然后通过python运行起来：

到这里SPTAG的rpc服务算是搞好了，我们可以摁下Ctrl+p 然后再摁下Ctrl+q把服务放到后台运行。

测试Demo API

添加索引测试

import numpy as np
from SPTAG_rpc_demo_client import SPTAG_RpcDemoClient, DataBean

client = SPTAG_RpcDemoClient("127.0.0.1", "8888")
beans = []
for i in range(5):
    vec = i * np.ones((10,), dtype=np.float32)
    beans.append(DataBean(_id=f"s{i}", vec=vec))

index_name = "test"
print("Adding Data:", client.add_data(index_name, beans))

这里添加了5个向量，分别是10个0，10个1，..., 10个4。

搜索测试

print("Test Search")
q = DataBean(_id=f"s{0}", vec=0 * np.ones((10,), dtype=np.float32))
print(client.search(index_name, [q], 3))

然后测试下搜索，我们搜索10个0的向量，可以看到返回的10个0（本身）的距离是0，10个1的距离为 $10*(1-0)^2$ =10，10个2的距离为 $10*(2-0)^2$ =40。没有问题

删除数据测试

print("*"*100)
print("Test Delete:", client.delete_data(index_name, [q]))

print("*"*100)
print("Test Search After Deletion")
print(client.search(index_name, [q], 3))

删除之后，本身不在了。第三近的变成了10个3， $10*(3-0)^2$ =90

删除索引测试

print("*"*100)
print("Test Delete Index:", client.delete_index(index_name))

最后是删除索引，可以看到返回结果为True，删除成功。

真实场景下的接口

在真实场景中，

1 . 索引文件可能非常大，我们不会一批一批的添加数据，也不会使用网络开销很大的数据传输。
2 . 每次搜索时候不会每次都重新导入index，然后再调用搜索。

基于以上两个缺陷，这里我想到了以下解决方法：

1 . 针对索引的很大的问题，直接使用SPTAG提供的indexbuilder工具建立索引。
2 . 对于搜索问题，直接编写一个专门的搜索服务接口，在刚启动的时候就导入index。（如有需要可以对索引做定期更新，而非每次都重新加载）

通过indexbuilder建立索引

这里先导出一批测试数据到test_index_input.txt中

import numpy as np

with open("test_index_input.txt", "w") as f:
    for i in range(5):
        _id = f"s{i}"
        vec = i * np.ones((10,), dtype=np.float32)
        vec_str = "|".join([str(i) for i in vec.tolist()])
        f.write(f"{_id}\t{vec_str}\n")

然后把test_index_input.txt放到docker里面

然后进入到/app/Release目录建立索引

cd /app/Release
./indexbuilder -d 10 -v Float  -i ./test_index_input.txt -o data/test_index -a BKT -t 2 Index.DistCalcMethod=L2

测试搜索服务

然后和上面的demo API类似，这里我编写了个search API。还是在这个项目里：https://github.com/nladuo/SPTAG-rpc-service

SPTAG_rpc_search_server.py需要放到docker中，SPTAG_rpc_search_client.py则直接import到自己的包里即可。

这里把SPTAG_rpc_search_server.py放到/app/Release目录后，启动起来。

python3 SPTAG_rpc_search_server.py -i test_index

然后再测一下搜索客户端。

和之前的一样，没问题。

到这里SPTAG的安装和测试就结束了。