milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #192 from youny626/cpu_version 

Pure CPU version for Milvus
pull/209/head
Jin Hai 2019-11-05 22:36:06 +08:00 committed by GitHub
parent af319887d4 ece0b61ac8
commit c806b5d1d7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
111 changed files with 14004 additions and 14540 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
CHANGELOG.md
View File

@ -7,6 +7,7 @@ Please mark all change in change log and use the ticket from JIRA.
## Bug
## Feature
- \#12 - Pure CPU version for Milvus
## Improvement

242
core/CMakeLists.txt
View File

@ -23,9 +23,9 @@ message(STATUS "Building using CMake version: ${CMAKE_VERSION}")
set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
MACRO (GET_CURRENT_TIME CURRENT_TIME)
MACRO(GET_CURRENT_TIME CURRENT_TIME)
execute_process(COMMAND "date" +"%Y-%m-%d %H:%M.%S" OUTPUT_VARIABLE ${CURRENT_TIME})
ENDMACRO (GET_CURRENT_TIME)
ENDMACRO(GET_CURRENT_TIME)
GET_CURRENT_TIME(BUILD_TIME)
string(REGEX REPLACE "\n" "" BUILD_TIME ${BUILD_TIME})
@ -42,23 +42,20 @@ GET_GIT_BRANCH_NAME(GIT_BRANCH_NAME)
message(STATUS "GIT_BRANCH_NAME = ${GIT_BRANCH_NAME}")
if(NOT GIT_BRANCH_NAME STREQUAL "")
string(REGEX REPLACE "\n" "" GIT_BRANCH_NAME ${GIT_BRANCH_NAME})
endif()
endif ()
set(MILVUS_VERSION "0.5.1")
set(MILVUS_VERSION "${GIT_BRANCH_NAME}")
string(REGEX MATCH "[0-9]+\\.[0-9]+\\.[0-9]" MILVUS_VERSION "${MILVUS_VERSION}")
find_package(ClangTools)
set(BUILD_SUPPORT_DIR "${CMAKE_SOURCE_DIR}/build-support")
if(CMAKE_BUILD_TYPE STREQUAL "Release")
if (CMAKE_BUILD_TYPE STREQUAL "Release")
set(BUILD_TYPE "Release")
else()
else ()
set(BUILD_TYPE "Debug")
endif()
endif ()
message(STATUS "Build type = ${BUILD_TYPE}")
project(milvus VERSION "${MILVUS_VERSION}")
project(milvus_engine LANGUAGES CUDA CXX)
project(milvus_engine LANGUAGES CXX)
unset(CMAKE_EXPORT_COMPILE_COMMANDS CACHE)
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
@ -67,15 +64,15 @@ set(MILVUS_VERSION_MAJOR "${milvus_VERSION_MAJOR}")
set(MILVUS_VERSION_MINOR "${milvus_VERSION_MINOR}")
set(MILVUS_VERSION_PATCH "${milvus_VERSION_PATCH}")
if(MILVUS_VERSION_MAJOR STREQUAL ""
if (MILVUS_VERSION_MAJOR STREQUAL ""
OR MILVUS_VERSION_MINOR STREQUAL ""
OR MILVUS_VERSION_PATCH STREQUAL "")
message(WARNING "Failed to determine Milvus version from git branch name")
set(MILVUS_VERSION "0.5.0")
endif()
endif ()
message(STATUS "Build version = ${MILVUS_VERSION}")
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/src/version.h.macro ${CMAKE_CURRENT_SOURCE_DIR}/src/version.h)
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/src/config.h.in ${CMAKE_CURRENT_SOURCE_DIR}/src/config.h @ONLY)
message(STATUS "Milvus version: "
"${MILVUS_VERSION_MAJOR}.${MILVUS_VERSION_MINOR}.${MILVUS_VERSION_PATCH} "
@ -84,46 +81,33 @@ message(STATUS "Milvus version: "
set(CMAKE_CXX_STANDARD 14)
set(CMAKE_CXX_STANDARD_REQUIRED on)
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
message(STATUS "Building milvus_engine on x86 architecture")
set(MILVUS_BUILD_ARCH x86_64)
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)")
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)")
message(STATUS "Building milvus_engine on ppc architecture")
set(MILVUS_BUILD_ARCH ppc64le)
else()
else ()
message(WARNING "Unknown processor type")
message(WARNING "CMAKE_SYSTEM_PROCESSOR=${CMAKE_SYSTEM_PROCESSOR}")
set(MILVUS_BUILD_ARCH unknown)
endif()
find_package (Python COMPONENTS Interpreter Development)
find_package(CUDA)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda")
if(CMAKE_BUILD_TYPE STREQUAL "Release")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3")
else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
endif()
endif ()
# Ensure that a default make is set
if("${MAKE}" STREQUAL "")
if(NOT MSVC)
if ("${MAKE}" STREQUAL "")
if (NOT MSVC)
find_program(MAKE make)
endif()
endif()
endif ()
endif ()
find_path(MYSQL_INCLUDE_DIR
NAMES "mysql.h"
PATH_SUFFIXES "mysql")
NAMES "mysql.h"
PATH_SUFFIXES "mysql")
if (${MYSQL_INCLUDE_DIR} STREQUAL "MYSQL_INCLUDE_DIR-NOTFOUND")
message(FATAL_ERROR "Could not found MySQL include directory")
else()
else ()
include_directories(${MYSQL_INCLUDE_DIR})
endif()
endif ()
set(MILVUS_SOURCE_DIR ${PROJECT_SOURCE_DIR})
set(MILVUS_BINARY_DIR ${PROJECT_BINARY_DIR})
@ -134,26 +118,50 @@ include(DefineOptions)
include(BuildUtils)
include(ThirdPartyPackages)
config_summary()
set(MILVUS_GPU_VERSION false)
if (MILVUS_CPU_VERSION)
message(STATUS "Building Milvus CPU version")
add_compile_definitions("MILVUS_CPU_VERSION")
else ()
message(STATUS "Building Milvus GPU version")
set(MILVUS_GPU_VERSION true)
add_compile_definitions("MILVUS_GPU_VERSION")
enable_language(CUDA)
find_package(CUDA 10 REQUIRED)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda")
endif ()
if (CMAKE_BUILD_TYPE STREQUAL "Release")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
if (MILVUS_GPU_VERSION)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3")
endif ()
else ()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
if (MILVUS_GPU_VERSION)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
endif ()
endif ()
if (CUSTOMIZATION)
add_definitions(-DCUSTOMIZATION)
endif (CUSTOMIZATION)
config_summary()
add_subdirectory(src)
if (BUILD_UNIT_TEST STREQUAL "ON")
if (BUILD_COVERAGE STREQUAL "ON")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage")
endif()
endif ()
add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unittest)
endif()
endif ()
add_custom_target(Clean-All COMMAND ${CMAKE_BUILD_TOOL} clean)
if("${MILVUS_DB_PATH}" STREQUAL "")
if ("${MILVUS_DB_PATH}" STREQUAL "")
set(MILVUS_DB_PATH "/tmp/milvus")
endif()
endif ()
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/conf/server_config.template ${CMAKE_CURRENT_SOURCE_DIR}/conf/server_config.yaml)
configure_file(${CMAKE_CURRENT_SOURCE_DIR}/conf/log_config.template ${CMAKE_CURRENT_SOURCE_DIR}/conf/log_config.conf)
@ -169,19 +177,22 @@ install(FILES
DESTINATION
conf)
find_package(Python COMPONENTS Interpreter Development)
find_package(ClangTools)
set(BUILD_SUPPORT_DIR "${CMAKE_SOURCE_DIR}/build-support")
#
# "make lint" target
#
if(NOT MILVUS_VERBOSE_LINT)
set(MILVUS_LINT_QUIET "--quiet")
endif()
if (NOT MILVUS_VERBOSE_LINT)
set(MILVUS_LINT_QUIET "--quiet")
endif ()
if(NOT LINT_EXCLUSIONS_FILE)
# source files matching a glob from a line in this file
# will be excluded from linting (cpplint, clang-tidy, clang-format)
set(LINT_EXCLUSIONS_FILE ${BUILD_SUPPORT_DIR}/lint_exclusions.txt)
endif()
if (NOT LINT_EXCLUSIONS_FILE)
# source files matching a glob from a line in this file
# will be excluded from linting (cpplint, clang-tidy, clang-format)
set(LINT_EXCLUSIONS_FILE ${BUILD_SUPPORT_DIR}/lint_exclusions.txt)
endif ()
find_program(CPPLINT_BIN NAMES cpplint cpplint.py HINTS ${BUILD_SUPPORT_DIR})
message(STATUS "Found cpplint executable at ${CPPLINT_BIN}")
@ -190,77 +201,76 @@ message(STATUS "Found cpplint executable at ${CPPLINT_BIN}")
# "make lint" targets
#
add_custom_target(lint
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_cpplint.py
--cpplint_binary
${CPPLINT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}
${MILVUS_LINT_QUIET})
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_cpplint.py
--cpplint_binary
${CPPLINT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}
${MILVUS_LINT_QUIET})
#
# "make clang-format" and "make check-clang-format" targets
#
if(${CLANG_FORMAT_FOUND})
# runs clang format and updates files in place.
add_custom_target(clang-format
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_format.py
--clang_format_binary
${CLANG_FORMAT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
--fix
${MILVUS_LINT_QUIET})
if (${CLANG_FORMAT_FOUND})
# runs clang format and updates files in place.
add_custom_target(clang-format
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_format.py
--clang_format_binary
${CLANG_FORMAT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
--fix
${MILVUS_LINT_QUIET})
# runs clang format and exits with a non-zero exit code if any files need to be reformatted
add_custom_target(check-clang-format
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_format.py
--clang_format_binary
${CLANG_FORMAT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
${MILVUS_LINT_QUIET})
endif()
# runs clang format and exits with a non-zero exit code if any files need to be reformatted
add_custom_target(check-clang-format
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_format.py
--clang_format_binary
${CLANG_FORMAT_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
${MILVUS_LINT_QUIET})
endif ()
#
# "make clang-tidy" and "make check-clang-tidy" targets
#
if(${CLANG_TIDY_FOUND})
# runs clang-tidy and attempts to fix any warning automatically
add_custom_target(clang-tidy
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_tidy.py
--clang_tidy_binary
${CLANG_TIDY_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--compile_commands
${CMAKE_BINARY_DIR}/compile_commands.json
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
--fix
${MILVUS_LINT_QUIET})
# runs clang-tidy and exits with a non-zero exit code if any errors are found.
add_custom_target(check-clang-tidy
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_tidy.py
--clang_tidy_binary
${CLANG_TIDY_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--compile_commands
${CMAKE_BINARY_DIR}/compile_commands.json
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
${MILVUS_LINT_QUIET})
endif()
if (${CLANG_TIDY_FOUND})
# runs clang-tidy and attempts to fix any warning automatically
add_custom_target(clang-tidy
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_tidy.py
--clang_tidy_binary
${CLANG_TIDY_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--compile_commands
${CMAKE_BINARY_DIR}/compile_commands.json
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
--fix
${MILVUS_LINT_QUIET})
# runs clang-tidy and exits with a non-zero exit code if any errors are found.
add_custom_target(check-clang-tidy
${PYTHON_EXECUTABLE}
${BUILD_SUPPORT_DIR}/run_clang_tidy.py
--clang_tidy_binary
${CLANG_TIDY_BIN}
--exclude_globs
${LINT_EXCLUSIONS_FILE}
--compile_commands
${CMAKE_BINARY_DIR}/compile_commands.json
--source_dir
${CMAKE_CURRENT_SOURCE_DIR}/src
${MILVUS_LINT_QUIET})
endif ()

18
core/build.sh
View File

@ -12,6 +12,8 @@ USE_JFROG_CACHE="OFF"
RUN_CPPLINT="OFF"
CUSTOMIZATION="OFF" # default use ori faiss
CUDA_COMPILER=/usr/local/cuda/bin/nvcc
CPU_VERSION="OFF"
WITH_MKL="OFF"
CUSTOMIZED_FAISS_URL="${FAISS_URL:-NONE}"
wget -q --method HEAD ${CUSTOMIZED_FAISS_URL}
@ -21,7 +23,7 @@ else
CUSTOMIZATION="OFF"
fi
while getopts "p:d:t:ulrcgjhx" arg
while getopts "p:d:t:ulrcgjhxzm" arg
do
case $arg in
p)
@ -58,6 +60,12 @@ do
x)
CUSTOMIZATION="OFF" # force use ori faiss
;;
z)
CPU_VERSION="ON"
;;
m)
WITH_MKL="ON"
;;
h) # help
echo "
@ -71,10 +79,12 @@ parameter:
-c: code coverage(default: OFF)
-g: profiling(default: OFF)
-j: use jfrog cache build directory(default: OFF)
-z: build pure CPU version(default: OFF)
-m: build with MKL(default: OFF)
-h: help
usage:
./build.sh -p \${INSTALL_PREFIX} -t \${BUILD_TYPE} [-u] [-l] [-r] [-c] [-g] [-j] [-h]
./build.sh -p \${INSTALL_PREFIX} -t \${BUILD_TYPE} [-u] [-l] [-r] [-c] [-g] [-j] [-z] [-m] [-h]
"
exit 0
;;
@ -106,6 +116,8 @@ CMAKE_CMD="cmake \
-DUSE_JFROG_CACHE=${USE_JFROG_CACHE} \
-DCUSTOMIZATION=${CUSTOMIZATION} \
-DFAISS_URL=${CUSTOMIZED_FAISS_URL} \
-DMILVUS_CPU_VERSION=${CPU_VERSION} \
-DBUILD_FAISS_WITH_MKL=${WITH_MKL} \
../"
echo ${CMAKE_CMD}
${CMAKE_CMD}
@ -147,4 +159,4 @@ else
# compile and build
make -j 8 install || exit 1
fi
fi

5
core/cmake/DefineOptions.cmake
View File

@ -40,6 +40,11 @@ macro(define_option_string name description default)
endif()
endmacro()
#----------------------------------------------------------------------
set_option_category("CPU version")
define_option(MILVUS_CPU_VERSION "Build CPU version only" OFF)
#----------------------------------------------------------------------
set_option_category("Thirdparty")

10
core/conf/server_config.template
View File

@ -27,11 +27,14 @@ metric_config:
port: 8080 # port prometheus uses to fetch metrics, must in range [1025, 65534]
cache_config:
cpu_cache_capacity: 16 # GB, CPU memory used for cache, must be a positive integer
cpu_cache_threshold: 0.85 # percentage of data that will be kept when cache cleanup is triggered, must be in range (0.0, 1.0]
cache_insert_data: false # whether to load inserted data into cache, must be a boolean
# Skip the following config if you are using GPU version
gpu_cache_capacity: 4 # GB, GPU memory used for cache, must be a positive integer
gpu_cache_threshold: 0.85 # percentage of data that will be kept when cache cleanup is triggered, must be in range (0.0, 1.0]
cache_insert_data: false # whether to load inserted data into cache, must be a boolean
engine_config:
use_blas_threshold: 20 # if nq < use_blas_threshold, use SSE, faster with fluctuated response times
@ -39,6 +42,7 @@ engine_config:
gpu_search_threshold: 1000 # threshold beyond which the search computation is executed on GPUs only
resource_config:
search_resources: # define the GPUs used for search computation, must be in format: gpux
search_resources: # define the CPU / GPUs used for search computation, must be in format: cpu / gpux
- cpu
- gpu0
index_build_device: gpu0 # GPU used for building index, must be in format: gpux
index_build_device: gpu0 # CPU / GPU used for building index, must be in format: cpu / gpux

42
core/src/CMakeLists.txt
View File

@ -20,11 +20,9 @@
include_directories(${MILVUS_SOURCE_DIR})
include_directories(${MILVUS_ENGINE_SRC})
include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include)
include_directories(${MILVUS_ENGINE_SRC}/grpc/gen-status)
include_directories(${MILVUS_ENGINE_SRC}/grpc/gen-milvus)
#this statement must put here, since the INDEX_INCLUDE_DIRS is defined in code/CMakeList.txt
add_subdirectory(index)
set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE)
@ -109,35 +107,45 @@ set(boost_lib
libboost_serialization.a
)
set(cuda_lib
${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so
cudart
cublas
)
set(third_party_libs
sqlite
${client_grpc_lib}
yaml-cpp
${prometheus_lib}
${cuda_lib}
mysqlpp
zlib
${boost_lib}
)
if (MILVUS_ENABLE_PROFILING STREQUAL "ON")
if (MILVUS_GPU_VERSION)
include_directories(${CUDA_INCLUDE_DIRS})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(cuda_lib
${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so
cudart
cublas
)
set(third_party_libs ${third_party_libs}
gperftools
libunwind
)
${cuda_lib}
)
aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/gpu wrapper_gpu_files)
set(engine_files ${engine_files}
${wrapper_gpu_files}
)
endif ()
if (MILVUS_ENABLE_PROFILING STREQUAL "ON")
set(third_party_libs ${third_party_libs}
gperftools
libunwind
)
endif ()
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(engine_libs
pthread
libgomp.a
libgfortran.a
dl
)
if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64")
@ -147,11 +155,11 @@ if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64")
)
endif ()
cuda_add_library(milvus_engine STATIC ${engine_files})
add_library(milvus_engine STATIC ${engine_files})
target_link_libraries(milvus_engine
knowhere
${engine_libs}
${third_party_libs}
${engine_libs}
)
add_library(metrics STATIC ${metrics_files})
@ -165,8 +173,6 @@ target_link_libraries(metrics ${metrics_lib})
set(server_libs
milvus_engine
pthread
dl
metrics
)

86
core/src/cache/Cache.inl vendored
View File

@ -15,24 +15,18 @@
// specific language governing permissions and limitations
// under the License.
namespace milvus {
namespace cache {
constexpr double DEFAULT_THRESHHOLD_PERCENT = 0.85;
template<typename ItemObj>
template <typename ItemObj>
Cache<ItemObj>::Cache(int64_t capacity, uint64_t cache_max_count)
: usage_(0),
capacity_(capacity),
freemem_percent_(DEFAULT_THRESHHOLD_PERCENT),
lru_(cache_max_count) {
// AGENT_LOG_DEBUG << "Construct Cache with capacity " << std::to_string(mem_capacity)
: usage_(0), capacity_(capacity), freemem_percent_(DEFAULT_THRESHHOLD_PERCENT), lru_(cache_max_count) {
// AGENT_LOG_DEBUG << "Construct Cache with capacity " << std::to_string(mem_capacity)
}
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::set_capacity(int64_t capacity) {
if (capacity > 0) {
@ -41,23 +35,23 @@ Cache<ItemObj>::set_capacity(int64_t capacity) {
}
}
template<typename ItemObj>
template <typename ItemObj>
size_t
Cache<ItemObj>::size() const {
std::lock_guard<std::mutex> lock(mutex_);
return lru_.size();
}
template<typename ItemObj>
template <typename ItemObj>
bool
Cache<ItemObj>::exists(const std::string &key) {
Cache<ItemObj>::exists(const std::string& key) {
std::lock_guard<std::mutex> lock(mutex_);
return lru_.exists(key);
}
template<typename ItemObj>
template <typename ItemObj>
ItemObj
Cache<ItemObj>::get(const std::string &key) {
Cache<ItemObj>::get(const std::string& key) {
std::lock_guard<std::mutex> lock(mutex_);
if (!lru_.exists(key)) {
return nullptr;
@ -66,60 +60,59 @@ Cache<ItemObj>::get(const std::string &key) {
return lru_.get(key);
}
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::insert(const std::string &key, const ItemObj &item) {
Cache<ItemObj>::insert(const std::string& key, const ItemObj& item) {
if (item == nullptr) {
return;
}
// if(item->size() > capacity_) {
// SERVER_LOG_ERROR << "Item size " << item->size()
// << " is too large to insert into cache, capacity " << capacity_;
// return;
// }
// if(item->size() > capacity_) {
// SERVER_LOG_ERROR << "Item size " << item->size()
// << " is too large to insert into cache, capacity " << capacity_;
// return;
// }
//calculate usage
// calculate usage
{
std::lock_guard<std::mutex> lock(mutex_);
//if key already exist, subtract old item size
// if key already exist, subtract old item size
if (lru_.exists(key)) {
const ItemObj &old_item = lru_.get(key);
const ItemObj& old_item = lru_.get(key);
usage_ -= old_item->Size();
}
//plus new item size
// plus new item size
usage_ += item->Size();
}
//if usage exceed capacity, free some items
// if usage exceed capacity, free some items
if (usage_ > capacity_) {
SERVER_LOG_DEBUG << "Current usage " << usage_
<< " exceeds cache capacity " << capacity_
SERVER_LOG_DEBUG << "Current usage " << usage_ << " exceeds cache capacity " << capacity_
<< ", start free memory";
free_memory();
}
//insert new item
// insert new item
{
std::lock_guard<std::mutex> lock(mutex_);
lru_.put(key, item);
SERVER_LOG_DEBUG << "Insert " << key << " size:" << item->Size()
<< " bytes into cache, usage: " << usage_ << " bytes";
SERVER_LOG_DEBUG << "Insert " << key << " size:" << item->Size() << " bytes into cache, usage: " << usage_
<< " bytes";
}
}
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::erase(const std::string &key) {
Cache<ItemObj>::erase(const std::string& key) {
std::lock_guard<std::mutex> lock(mutex_);
if (!lru_.exists(key)) {
return;
}
const ItemObj &old_item = lru_.get(key);
const ItemObj& old_item = lru_.get(key);
usage_ -= old_item->Size();
SERVER_LOG_DEBUG << "Erase " << key << " size: " << old_item->Size();
@ -127,7 +120,7 @@ Cache<ItemObj>::erase(const std::string &key) {
lru_.erase(key);
}
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::clear() {
std::lock_guard<std::mutex> lock(mutex_);
@ -137,15 +130,16 @@ Cache<ItemObj>::clear() {
}
/* free memory space when CACHE occupation exceed its capacity */
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::free_memory() {
if (usage_ <= capacity_) return;
if (usage_ <= capacity_)
return;
int64_t threshhold = capacity_ * freemem_percent_;
int64_t delta_size = usage_ - threshhold;
if (delta_size <= 0) {
delta_size = 1;//ensure at least one item erased
delta_size = 1; // ensure at least one item erased
}
std::set<std::string> key_array;
@ -156,8 +150,8 @@ Cache<ItemObj>::free_memory() {
auto it = lru_.rbegin();
while (it != lru_.rend() && released_size < delta_size) {
auto &key = it->first;
auto &obj_ptr = it->second;
auto& key = it->first;
auto& obj_ptr = it->second;
key_array.emplace(key);
released_size += obj_ptr->Size();
@ -167,14 +161,14 @@ Cache<ItemObj>::free_memory() {
SERVER_LOG_DEBUG << "to be released memory size: " << released_size;
for (auto &key : key_array) {
for (auto& key : key_array) {
erase(key);
}
print();
}
template<typename ItemObj>
template <typename ItemObj>
void
Cache<ItemObj>::print() {
size_t cache_count = 0;
@ -188,7 +182,5 @@ Cache<ItemObj>::print() {
SERVER_LOG_DEBUG << "[Cache capacity]: " << capacity_ << " bytes";
}
} // namespace cache
} // namespace milvus
} // namespace cache
} // namespace milvus

42
core/src/cache/CacheMgr.inl vendored
View File

@ -15,21 +15,18 @@
// specific language governing permissions and limitations
// under the License.
namespace milvus {
namespace cache {
template<typename ItemObj>
template <typename ItemObj>
CacheMgr<ItemObj>::CacheMgr() {
}
template<typename ItemObj>
template <typename ItemObj>
CacheMgr<ItemObj>::~CacheMgr() {
}
template<typename ItemObj>
template <typename ItemObj>
uint64_t
CacheMgr<ItemObj>::ItemCount() const {
if (cache_ == nullptr) {
@ -37,12 +34,12 @@ CacheMgr<ItemObj>::ItemCount() const {
return 0;
}
return (uint64_t) (cache_->size());
return (uint64_t)(cache_->size());
}
template<typename ItemObj>
template <typename ItemObj>
bool
CacheMgr<ItemObj>::ItemExists(const std::string &key) {
CacheMgr<ItemObj>::ItemExists(const std::string& key) {
if (cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
return false;
@ -51,9 +48,9 @@ CacheMgr<ItemObj>::ItemExists(const std::string &key) {
return cache_->exists(key);
}
template<typename ItemObj>
template <typename ItemObj>
ItemObj
CacheMgr<ItemObj>::GetItem(const std::string &key) {
CacheMgr<ItemObj>::GetItem(const std::string& key) {
if (cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
return nullptr;
@ -62,9 +59,9 @@ CacheMgr<ItemObj>::GetItem(const std::string &key) {
return cache_->get(key);
}
template<typename ItemObj>
template <typename ItemObj>
void
CacheMgr<ItemObj>::InsertItem(const std::string &key, const ItemObj &data) {
CacheMgr<ItemObj>::InsertItem(const std::string& key, const ItemObj& data) {
if (cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
return;
@ -74,9 +71,9 @@ CacheMgr<ItemObj>::InsertItem(const std::string &key, const ItemObj &data) {
server::Metrics::GetInstance().CacheAccessTotalIncrement();
}
template<typename ItemObj>
template <typename ItemObj>
void
CacheMgr<ItemObj>::EraseItem(const std::string &key) {
CacheMgr<ItemObj>::EraseItem(const std::string& key) {
if (cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
return;
@ -86,7 +83,7 @@ CacheMgr<ItemObj>::EraseItem(const std::string &key) {
server::Metrics::GetInstance().CacheAccessTotalIncrement();
}
template<typename ItemObj>
template <typename ItemObj>
void
CacheMgr<ItemObj>::PrintInfo() {
if (cache_ == nullptr) {
@ -97,7 +94,7 @@ CacheMgr<ItemObj>::PrintInfo() {
cache_->print();
}
template<typename ItemObj>
template <typename ItemObj>
void
CacheMgr<ItemObj>::ClearCache() {
if (cache_ == nullptr) {
@ -108,7 +105,7 @@ CacheMgr<ItemObj>::ClearCache() {
cache_->clear();
}
template<typename ItemObj>
template <typename ItemObj>
int64_t
CacheMgr<ItemObj>::CacheUsage() const {
if (cache_ == nullptr) {
@ -119,7 +116,7 @@ CacheMgr<ItemObj>::CacheUsage() const {
return cache_->usage();
}
template<typename ItemObj>
template <typename ItemObj>
int64_t
CacheMgr<ItemObj>::CacheCapacity() const {
if (cache_ == nullptr) {
@ -130,7 +127,7 @@ CacheMgr<ItemObj>::CacheCapacity() const {
return cache_->capacity();
}
template<typename ItemObj>
template <typename ItemObj>
void
CacheMgr<ItemObj>::SetCapacity(int64_t capacity) {
if (cache_ == nullptr) {
@ -140,6 +137,5 @@ CacheMgr<ItemObj>::SetCapacity(int64_t capacity) {
cache_->set_capacity(capacity);
}
} // namespace cache
} // namespace milvus
} // namespace cache
} // namespace milvus

8
core/src/version.h.macro → core/src/config.h
View File

@ -15,8 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#define MILVUS_VERSION "@MILVUS_VERSION@"
#define BUILD_TYPE "@BUILD_TYPE@"
#define BUILD_TIME @BUILD_TIME@
#define MILVUS_VERSION "0.5.0"
#define BUILD_TYPE "Debug"
#define BUILD_TIME "2019-11-05 18:49.05"

20
core/src/config.h.in Normal file
View File

@ -0,0 +1,20 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#cmakedefine MILVUS_VERSION "@MILVUS_VERSION@"
#cmakedefine BUILD_TYPE "@BUILD_TYPE@"
#cmakedefine BUILD_TIME @BUILD_TIME@

7
core/src/db/DB.h
View File

@ -67,16 +67,15 @@ class DB {
virtual Status
Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
ResultIds& result_ids, ResultDistances& result_distances) = 0;
QueryResults& results) = 0;
virtual Status
Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) = 0;
const meta::DatesT& dates, QueryResults& results) = 0;
virtual Status
Query(const std::string& table_id, const std::vector<std::string>& file_ids, uint64_t k, uint64_t nq,
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids,
ResultDistances& result_distances) = 0;
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) = 0;
virtual Status
Size(uint64_t& result) = 0;

18
core/src/db/DBImpl.cpp
View File

@ -336,20 +336,20 @@ DBImpl::DropIndex(const std::string& table_id) {
Status
DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
ResultIds& result_ids, ResultDistances& result_distances) {
QueryResults& results) {
if (shutting_down_.load(std::memory_order_acquire)) {
return Status(DB_ERROR, "Milsvus server is shutdown!");
}
meta::DatesT dates = {utils::GetDate()};
Status result = Query(table_id, k, nq, nprobe, vectors, dates, result_ids, result_distances);
Status result = Query(table_id, k, nq, nprobe, vectors, dates, results);
return result;
}
Status
DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) {
const meta::DatesT& dates, QueryResults& results) {
if (shutting_down_.load(std::memory_order_acquire)) {
return Status(DB_ERROR, "Milsvus server is shutdown!");
}
@ -372,15 +372,14 @@ DBImpl::Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t npr
}
cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info before query
status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, result_ids, result_distances);
status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, results);
cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info after query
return status;
}
Status
DBImpl::Query(const std::string& table_id, const std::vector<std::string>& file_ids, uint64_t k, uint64_t nq,
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids,
ResultDistances& result_distances) {
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) {
if (shutting_down_.load(std::memory_order_acquire)) {
return Status(DB_ERROR, "Milsvus server is shutdown!");
}
@ -414,7 +413,7 @@ DBImpl::Query(const std::string& table_id, const std::vector<std::string>& file_
}
cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info before query
status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, result_ids, result_distances);
status = QueryAsync(table_id, file_id_array, k, nq, nprobe, vectors, results);
cache::CpuCacheMgr::GetInstance()->PrintInfo(); // print cache info after query
return status;
}
@ -433,7 +432,7 @@ DBImpl::Size(uint64_t& result) {
///////////////////////////////////////////////////////////////////////////////////////////////////////////////////
Status
DBImpl::QueryAsync(const std::string& table_id, const meta::TableFilesSchema& files, uint64_t k, uint64_t nq,
uint64_t nprobe, const float* vectors, ResultIds& result_ids, ResultDistances& result_distances) {
uint64_t nprobe, const float* vectors, QueryResults& results) {
server::CollectQueryMetrics metrics(nq);
TimeRecorder rc("");
@ -454,8 +453,7 @@ DBImpl::QueryAsync(const std::string& table_id, const meta::TableFilesSchema& fi
}
// step 3: construct results
result_ids = job->GetResultIds();
result_distances = job->GetResultDistances();
results = job->GetResult();
rc.ElapseFromBegin("Engine query totally cost");
return Status::OK();

9
core/src/db/DBImpl.h
View File

@ -91,16 +91,15 @@ class DBImpl : public DB {
Status
Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
ResultIds& result_ids, ResultDistances& result_distances) override;
QueryResults& results) override;
Status
Query(const std::string& table_id, uint64_t k, uint64_t nq, uint64_t nprobe, const float* vectors,
const meta::DatesT& dates, ResultIds& result_ids, ResultDistances& result_distances) override;
const meta::DatesT& dates, QueryResults& results) override;
Status
Query(const std::string& table_id, const std::vector<std::string>& file_ids, uint64_t k, uint64_t nq,
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, ResultIds& result_ids,
ResultDistances& result_distances) override;
uint64_t nprobe, const float* vectors, const meta::DatesT& dates, QueryResults& results) override;
Status
Size(uint64_t& result) override;
@ -108,7 +107,7 @@ class DBImpl : public DB {
private:
Status
QueryAsync(const std::string& table_id, const meta::TableFilesSchema& files, uint64_t k, uint64_t nq,
uint64_t nprobe, const float* vectors, ResultIds& result_ids, ResultDistances& result_distances);
uint64_t nprobe, const float* vectors, QueryResults& results);
void
BackgroundTimerTask();

8
core/src/db/Types.h
View File

@ -19,7 +19,6 @@
#include "db/engine/ExecutionEngine.h"
#include <faiss/Index.h>
#include <stdint.h>
#include <utility>
#include <vector>
@ -27,13 +26,12 @@
namespace milvus {
namespace engine {
using IDNumber = faiss::Index::idx_t;
typedef int64_t IDNumber;
typedef IDNumber* IDNumberPtr;
typedef std::vector<IDNumber> IDNumbers;
typedef std::vector<faiss::Index::idx_t> ResultIds;
typedef std::vector<faiss::Index::distance_t> ResultDistances;
typedef std::vector<std::pair<IDNumber, double>> QueryResult;
typedef std::vector<QueryResult> QueryResults;
struct TableIndex {
int32_t engine_type_ = (int)EngineType::FAISS_IDMAP;

38
core/src/db/engine/ExecutionEngineImpl.cpp
View File

@ -25,6 +25,7 @@
#include "utils/CommonUtil.h"
#include "utils/Exception.h"
#include "utils/Log.h"
#include "wrapper/ConfAdapter.h"
#include "wrapper/ConfAdapterMgr.h"
#include "wrapper/VecImpl.h"
@ -92,11 +93,19 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
break;
}
case EngineType::FAISS_IVFFLAT: {
#ifdef MILVUS_CPU_VERSION
index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_CPU);
#else
index = GetVecIndexFactory(IndexType::FAISS_IVFFLAT_MIX);
#endif
break;
}
case EngineType::FAISS_IVFSQ8: {
#ifdef MILVUS_CPU_VERSION
index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_CPU);
#else
index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_MIX);
#endif
break;
}
case EngineType::NSG_MIX: {
@ -309,13 +318,30 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id, bool hybrid) {
return Status::OK();
}
#endif
try {
index_ = index_->CopyToGpu(device_id);
ENGINE_LOG_DEBUG << "CPU to GPU" << device_id;
} catch (std::exception& e) {
ENGINE_LOG_ERROR << e.what();
return Status(DB_ERROR, e.what());
auto index = std::static_pointer_cast<VecIndex>(cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_));
bool already_in_cache = (index != nullptr);
if (already_in_cache) {
index_ = index;
} else {
if (index_ == nullptr) {
ENGINE_LOG_ERROR << "ExecutionEngineImpl: index is null, failed to copy to gpu";
return Status(DB_ERROR, "index is null");
}
try {
index_ = index_->CopyToGpu(device_id);
ENGINE_LOG_DEBUG << "CPU to GPU" << device_id;
} catch (std::exception& e) {
ENGINE_LOG_ERROR << e.what();
return Status(DB_ERROR, e.what());
}
}
if (!already_in_cache) {
GpuCache(device_id);
}
return Status::OK();
}

671
core/src/db/meta/SqliteMetaImpl.cpp
View File

File diff suppressed because it is too large Load Diff

5107
core/src/external/easyloggingpp/easylogging++.cc vendored
View File

File diff suppressed because it is too large Load Diff

6756
core/src/external/easyloggingpp/easylogging++.h vendored
View File

File diff suppressed because it is too large Load Diff

10958
core/src/external/nlohmann/json.hpp vendored
View File

File diff suppressed because it is too large Load Diff

73
core/src/index/CMakeLists.txt
View File

@ -19,83 +19,96 @@
cmake_minimum_required(VERSION 3.14)
message(STATUS "---------------core--------------")
message(STATUS "------------------------------KNOWHERE-----------------------------------")
message(STATUS "Building using CMake version: ${CMAKE_VERSION}")
set(KNOWHERE_VERSION "0.1.0")
set(KNOWHERE_VERSION "0.5.0")
string(REGEX MATCH "^[0-9]+\\.[0-9]+\\.[0-9]+" KNOWHERE_BASE_VERSION "${KNOWHERE_VERSION}")
project(knowhere VERSION "${KNOWHERE_BASE_VERSION}" LANGUAGES CUDA C CXX)
project(knowhere VERSION "${KNOWHERE_BASE_VERSION}" LANGUAGES C CXX)
set(CMAKE_CXX_STANDARD 14)
set(KNOWHERE_VERSION_MAJOR "${knowhere_VERSION_MAJOR}")
set(KNOWHERE_VERSION_MINOR "${knowhere_VERSION_MINOR}")
set(KNOWHERE_VERSION_PATCH "${knowhere_VERSION_PATCH}")
if(KNOWHERE_VERSION_MAJOR STREQUAL ""
if (KNOWHERE_VERSION_MAJOR STREQUAL ""
OR KNOWHERE_VERSION_MINOR STREQUAL ""
OR KNOWHERE_VERSION_PATCH STREQUAL "")
message(FATAL_ERROR "Failed to determine Knowhere version from '${KNOWHERE_VERSION}'")
endif()
endif ()
message(STATUS "Knowhere version: "
"${KNOWHERE_VERSION_MAJOR}.${KNOWHERE_VERSION_MINOR}.${KNOWHERE_VERSION_PATCH} "
"(full: '${KNOWHERE_VERSION}')")
# if no build build type is specified, default to release builds
if(NOT CMAKE_BUILD_TYPE)
if (NOT CMAKE_BUILD_TYPE)
set(CMAKE_BUILD_TYPE Release)
endif(NOT CMAKE_BUILD_TYPE)
endif (NOT CMAKE_BUILD_TYPE)
if(CMAKE_BUILD_TYPE STREQUAL "Release")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3")
else()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
endif()
MESSAGE(STATUS "CMAKE_CXX_FLAGS" ${CMAKE_CXX_FLAGS})
find_package(CUDA)
if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
if (CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)")
message(STATUS "building milvus_engine on x86 architecture")
set(KNOWHERE_BUILD_ARCH x86_64)
elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)")
elseif (CMAKE_SYSTEM_PROCESSOR MATCHES "(ppc)")
message(STATUS "building milvus_engine on ppc architecture")
set(KNOWHERE_BUILD_ARCH ppc64le)
else()
else ()
message(WARNING "unknown processor type")
message(WARNING "CMAKE_SYSTEM_PROCESSOR=${CMAKE_SYSTEM_PROCESSOR}")
set(KNOWHERE_BUILD_ARCH unknown)
endif()
endif ()
if(CMAKE_BUILD_TYPE STREQUAL "Release")
if (CMAKE_BUILD_TYPE STREQUAL "Release")
set(BUILD_TYPE "release")
else()
else ()
set(BUILD_TYPE "debug")
endif()
endif ()
message(STATUS "Build type = ${BUILD_TYPE}")
set(INDEX_SOURCE_DIR ${PROJECT_SOURCE_DIR})
set(INDEX_BINARY_DIR ${PROJECT_BINARY_DIR})
message(STATUS "Core source dir: ${PROJECT_SOURCE_DIR}")
message(STATUS "Core binary dir: ${PROJECT_BINARY_DIR}")
set(CMAKE_MODULE_PATH ${CMAKE_MODULE_PATH} "${INDEX_SOURCE_DIR}/cmake")
include(ExternalProject)
include(DefineOptionsCore)
include(BuildUtilsCore)
set(KNOWHERE_GPU_VERSION false)
if (MILVUS_CPU_VERSION OR KNOWHERE_CPU_VERSION)
message(STATUS "Building Knowhere CPU version")
add_compile_definitions("MILVUS_CPU_VERSION")
else ()
message(STATUS "Building Knowhere GPU version")
add_compile_definitions("MILVUS_GPU_VERSION")
set(KNOWHERE_GPU_VERSION true)
enable_language(CUDA)
find_package(CUDA 10 REQUIRED)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -Xcompiler -fPIC -std=c++11 -D_FORCE_INLINES --expt-extended-lambda")
endif ()
include(ThirdPartyPackagesCore)
if (CMAKE_BUILD_TYPE STREQUAL "Release")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -DELPP_THREAD_SAFE -fopenmp")
if (KNOWHERE_GPU_VERSION)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O3")
endif ()
else ()
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O0 -g -fPIC -DELPP_THREAD_SAFE -fopenmp")
if (KNOWHERE_GPU_VERSION)
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -O0 -g")
endif ()
endif ()
add_subdirectory(knowhere)
if (BUILD_COVERAGE STREQUAL "ON")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fprofile-arcs -ftest-coverage")
endif()
endif ()
set(INDEX_INCLUDE_DIRS ${INDEX_INCLUDE_DIRS} PARENT_SCOPE)
if(BUILD_UNIT_TEST STREQUAL "ON")
if (KNOWHERE_BUILD_TESTS)
add_subdirectory(unittest)
endif()
endif ()
config_summary()

116
core/src/index/cmake/BuildUtilsCore.cmake
View File

@ -1,50 +1,50 @@
# Define a function that check last file modification
function(Check_Last_Modify cache_check_lists_file_path working_dir last_modified_commit_id)
if(EXISTS "${working_dir}")
if(EXISTS "${cache_check_lists_file_path}")
if (EXISTS "${working_dir}")
if (EXISTS "${cache_check_lists_file_path}")
set(GIT_LOG_SKIP_NUM 0)
set(_MATCH_ALL ON CACHE BOOL "Match all")
set(_LOOP_STATUS ON CACHE BOOL "Whether out of loop")
file(STRINGS ${cache_check_lists_file_path} CACHE_IGNORE_TXT)
while(_LOOP_STATUS)
foreach(_IGNORE_ENTRY ${CACHE_IGNORE_TXT})
if(NOT _IGNORE_ENTRY MATCHES "^[^#]+")
while (_LOOP_STATUS)
foreach (_IGNORE_ENTRY ${CACHE_IGNORE_TXT})
if (NOT _IGNORE_ENTRY MATCHES "^[^#]+")
continue()
endif()
endif ()
set(_MATCH_ALL OFF)
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --name-status --pretty= WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE CHANGE_FILES)
if(NOT CHANGE_FILES STREQUAL "")
if (NOT CHANGE_FILES STREQUAL "")
string(REPLACE "\n" ";" _CHANGE_FILES ${CHANGE_FILES})
foreach(_FILE_ENTRY ${_CHANGE_FILES})
foreach (_FILE_ENTRY ${_CHANGE_FILES})
string(REGEX MATCH "[^ \t]+$" _FILE_NAME ${_FILE_ENTRY})
execute_process(COMMAND sh -c "echo ${_FILE_NAME} | grep ${_IGNORE_ENTRY}" RESULT_VARIABLE return_code)
if (return_code EQUAL 0)
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
set (${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
set(_LOOP_STATUS OFF)
endif()
endforeach()
else()
endif ()
endforeach ()
else ()
set(_LOOP_STATUS OFF)
endif()
endforeach()
endif ()
endforeach ()
if(_MATCH_ALL)
if (_MATCH_ALL)
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
set (${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
set(_LOOP_STATUS OFF)
endif()
endif ()
math(EXPR GIT_LOG_SKIP_NUM "${GIT_LOG_SKIP_NUM} + 1")
endwhile(_LOOP_STATUS)
else()
endwhile (_LOOP_STATUS)
else ()
execute_process(COMMAND git log --no-merges -1 --skip=${GIT_LOG_SKIP_NUM} --pretty=%H WORKING_DIRECTORY ${working_dir} OUTPUT_VARIABLE LAST_MODIFIED_COMMIT_ID)
set (${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
endif()
else()
set(${last_modified_commit_id} ${LAST_MODIFIED_COMMIT_ID} PARENT_SCOPE)
endif ()
else ()
message(FATAL_ERROR "The directory ${working_dir} does not exist")
endif()
endif ()
endfunction()
# Define a function that extracts a cached package
@ -83,15 +83,15 @@ endfunction()
# Define a function that to create a new cached package
function(ExternalProject_Create_Cache project_name package_file install_path cache_username cache_password cache_path)
if(EXISTS ${package_file})
if (EXISTS ${package_file})
message(STATUS "Removing existing package file: ${package_file}")
file(REMOVE ${package_file})
endif()
endif ()
string(REGEX REPLACE "(.+)/.+$" "\\1" package_dir ${package_file})
if(NOT EXISTS ${package_dir})
if (NOT EXISTS ${package_dir})
file(MAKE_DIRECTORY ${package_dir})
endif()
endif ()
message(STATUS "Will create cached package file: ${package_file}")
@ -116,89 +116,89 @@ function(ADD_THIRDPARTY_LIB LIB_NAME)
"${one_value_args}"
"${multi_value_args}"
${ARGN})
if(ARG_UNPARSED_ARGUMENTS)
if (ARG_UNPARSED_ARGUMENTS)
message(SEND_ERROR "Error: unrecognized arguments: ${ARG_UNPARSED_ARGUMENTS}")
endif()
endif ()
if(ARG_STATIC_LIB AND ARG_SHARED_LIB)
if(NOT ARG_STATIC_LIB)
if (ARG_STATIC_LIB AND ARG_SHARED_LIB)
if (NOT ARG_STATIC_LIB)
message(FATAL_ERROR "No static or shared library provided for ${LIB_NAME}")
endif()
endif ()
set(AUG_LIB_NAME "${LIB_NAME}_static")
add_library(${AUG_LIB_NAME} STATIC IMPORTED)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_LOCATION "${ARG_STATIC_LIB}")
if(ARG_DEPS)
if (ARG_DEPS)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
endif()
endif ()
message(STATUS "Added static library dependency ${AUG_LIB_NAME}: ${ARG_STATIC_LIB}")
if(ARG_INCLUDE_DIRECTORIES)
if (ARG_INCLUDE_DIRECTORIES)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
"${ARG_INCLUDE_DIRECTORIES}")
endif()
endif ()
set(AUG_LIB_NAME "${LIB_NAME}_shared")
add_library(${AUG_LIB_NAME} SHARED IMPORTED)
if(WIN32)
if (WIN32)
# Mark the ".lib" location as part of a Windows DLL
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_IMPLIB "${ARG_SHARED_LIB}")
else()
else ()
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_LOCATION "${ARG_SHARED_LIB}")
endif()
if(ARG_DEPS)
endif ()
if (ARG_DEPS)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
endif()
endif ()
message(STATUS "Added shared library dependency ${AUG_LIB_NAME}: ${ARG_SHARED_LIB}")
if(ARG_INCLUDE_DIRECTORIES)
if (ARG_INCLUDE_DIRECTORIES)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
"${ARG_INCLUDE_DIRECTORIES}")
endif()
elseif(ARG_STATIC_LIB)
endif ()
elseif (ARG_STATIC_LIB)
set(AUG_LIB_NAME "${LIB_NAME}_static")
add_library(${AUG_LIB_NAME} STATIC IMPORTED)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_LOCATION "${ARG_STATIC_LIB}")
if(ARG_DEPS)
if (ARG_DEPS)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
endif()
endif ()
message(STATUS "Added static library dependency ${AUG_LIB_NAME}: ${ARG_STATIC_LIB}")
if(ARG_INCLUDE_DIRECTORIES)
if (ARG_INCLUDE_DIRECTORIES)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
"${ARG_INCLUDE_DIRECTORIES}")
endif()
elseif(ARG_SHARED_LIB)
endif ()
elseif (ARG_SHARED_LIB)
set(AUG_LIB_NAME "${LIB_NAME}_shared")
add_library(${AUG_LIB_NAME} SHARED IMPORTED)
if(WIN32)
if (WIN32)
# Mark the ".lib" location as part of a Windows DLL
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_IMPLIB "${ARG_SHARED_LIB}")
else()
else ()
set_target_properties(${AUG_LIB_NAME}
PROPERTIES IMPORTED_LOCATION "${ARG_SHARED_LIB}")
endif()
endif ()
message(STATUS "Added shared library dependency ${AUG_LIB_NAME}: ${ARG_SHARED_LIB}")
if(ARG_DEPS)
if (ARG_DEPS)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_LINK_LIBRARIES "${ARG_DEPS}")
endif()
if(ARG_INCLUDE_DIRECTORIES)
endif ()
if (ARG_INCLUDE_DIRECTORIES)
set_target_properties(${AUG_LIB_NAME}
PROPERTIES INTERFACE_INCLUDE_DIRECTORIES
"${ARG_INCLUDE_DIRECTORIES}")
endif()
else()
endif ()
else ()
message(FATAL_ERROR "No static or shared library provided for ${LIB_NAME}")
endif()
endif ()
endfunction()

81
core/src/index/cmake/DefineOptionsCore.cmake
View File

@ -13,16 +13,16 @@ macro(define_option name description default)
endmacro()
function(list_join lst glue out)
if("${${lst}}" STREQUAL "")
if ("${${lst}}" STREQUAL "")
set(${out} "" PARENT_SCOPE)
return()
endif()
endif ()
list(GET ${lst} 0 joined)
list(REMOVE_AT ${lst} 0)
foreach(item ${${lst}})
foreach (item ${${lst}})
set(joined "${joined}${glue}${item}")
endforeach()
endforeach ()
set(${out} ${joined} PARENT_SCOPE)
endfunction()
@ -35,22 +35,31 @@ macro(define_option_string name description default)
set("${name}_OPTION_ENUM" ${ARGN})
list_join("${name}_OPTION_ENUM" "|" "${name}_OPTION_ENUM")
if(NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
if (NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
set_property(CACHE ${name} PROPERTY STRINGS ${ARGN})
endif()
endif ()
endmacro()
#----------------------------------------------------------------------
set_option_category("CPU version")
if (MILVUS_CPU_VERSION)
define_option(KNOWHERE_CPU_VERSION "Build CPU version only" ON)
else ()
define_option(KNOWHERE_CPU_VERSION "Build CPU version only" OFF)
endif ()
#----------------------------------------------------------------------
set_option_category("Thirdparty")
set(KNOWHERE_DEPENDENCY_SOURCE_DEFAULT "AUTO")
define_option_string(KNOWHERE_DEPENDENCY_SOURCE
"Method to use for acquiring KNOWHERE's build dependencies"
"${KNOWHERE_DEPENDENCY_SOURCE_DEFAULT}"
"AUTO"
"BUNDLED"
"SYSTEM")
"Method to use for acquiring KNOWHERE's build dependencies"
"${KNOWHERE_DEPENDENCY_SOURCE_DEFAULT}"
"AUTO"
"BUNDLED"
"SYSTEM")
define_option(KNOWHERE_VERBOSE_THIRDPARTY_BUILD
"Show output from ExternalProjects rather than just logging to files" ON)
@ -70,10 +79,10 @@ define_option(KNOWHERE_WITH_FAISS "Build with FAISS library" ON)
define_option(KNOWHERE_WITH_FAISS_GPU_VERSION "Build with FAISS GPU version" ON)
define_option(KNOWHERE_WITH_OPENBLAS "Build with OpenBLAS library" ON)
define_option(BUILD_FAISS_WITH_MKL "Build FAISS with MKL" OFF)
#----------------------------------------------------------------------
if(MSVC)
if (MSVC)
set_option_category("MSVC")
define_option(MSVC_LINK_VERBOSE
@ -81,16 +90,16 @@ if(MSVC)
OFF)
define_option(KNOWHERE_USE_STATIC_CRT "Build KNOWHERE with statically linked CRT" OFF)
endif()
endif ()
#----------------------------------------------------------------------
set_option_category("Test and benchmark")
if (BUILD_UNIT_TEST)
define_option(KNOWHERE_BUILD_TESTS "Build the KNOWHERE googletest unit tests" ON)
else()
else ()
define_option(KNOWHERE_BUILD_TESTS "Build the KNOWHERE googletest unit tests" OFF)
endif(BUILD_UNIT_TEST)
endif (BUILD_UNIT_TEST)
#----------------------------------------------------------------------
macro(config_summary)
@ -102,12 +111,12 @@ macro(config_summary)
message(STATUS " Generator: ${CMAKE_GENERATOR}")
message(STATUS " Build type: ${CMAKE_BUILD_TYPE}")
message(STATUS " Source directory: ${CMAKE_CURRENT_SOURCE_DIR}")
if(${CMAKE_EXPORT_COMPILE_COMMANDS})
if (${CMAKE_EXPORT_COMPILE_COMMANDS})
message(
STATUS " Compile commands: ${INDEX_BINARY_DIR}/compile_commands.json")
endif()
endif ()
foreach(category ${KNOWHERE_OPTION_CATEGORIES})
foreach (category ${KNOWHERE_OPTION_CATEGORIES})
message(STATUS)
message(STATUS "${category} options:")
@ -115,50 +124,50 @@ macro(config_summary)
set(option_names ${KNOWHERE_${category}_OPTION_NAMES})
set(max_value_length 0)
foreach(name ${option_names})
foreach (name ${option_names})
string(LENGTH "\"${${name}}\"" value_length)
if(${max_value_length} LESS ${value_length})
if (${max_value_length} LESS ${value_length})
set(max_value_length ${value_length})
endif()
endforeach()
endif ()
endforeach ()
foreach(name ${option_names})
if("${${name}_OPTION_TYPE}" STREQUAL "string")
foreach (name ${option_names})
if ("${${name}_OPTION_TYPE}" STREQUAL "string")
set(value "\"${${name}}\"")
else()
else ()
set(value "${${name}}")
endif()
endif ()
set(default ${${name}_OPTION_DEFAULT})
set(description ${${name}_OPTION_DESCRIPTION})
string(LENGTH ${description} description_length)
if(${description_length} LESS 70)
if (${description_length} LESS 70)
string(
SUBSTRING
" "
${description_length} -1 description_padding)
else()
else ()
set(description_padding "
")
endif()
endif ()
set(comment "[${name}]")
if("${value}" STREQUAL "${default}")
if ("${value}" STREQUAL "${default}")
set(comment "[default] ${comment}")
endif()
endif ()
if(NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
if (NOT ("${${name}_OPTION_ENUM}" STREQUAL ""))
set(comment "${comment} [${${name}_OPTION_ENUM}]")
endif()
endif ()
string(
SUBSTRING "${value} "
0 ${max_value_length} value)
message(STATUS " ${description} ${description_padding} ${value} ${comment}")
endforeach()
endforeach ()
endforeach()
endforeach ()
endmacro()

322
core/src/index/cmake/ThirdPartyPackagesCore.cmake
View File

@ -26,24 +26,24 @@ set(KNOWHERE_THIRDPARTY_DEPENDENCIES
message(STATUS "Using ${KNOWHERE_DEPENDENCY_SOURCE} approach to find dependencies")
# For each dependency, set dependency source to global default, if unset
foreach(DEPENDENCY ${KNOWHERE_THIRDPARTY_DEPENDENCIES})
if("${${DEPENDENCY}_SOURCE}" STREQUAL "")
foreach (DEPENDENCY ${KNOWHERE_THIRDPARTY_DEPENDENCIES})
if ("${${DEPENDENCY}_SOURCE}" STREQUAL "")
set(${DEPENDENCY}_SOURCE ${KNOWHERE_DEPENDENCY_SOURCE})
endif()
endforeach()
endif ()
endforeach ()
macro(build_dependency DEPENDENCY_NAME)
if("${DEPENDENCY_NAME}" STREQUAL "ARROW")
if ("${DEPENDENCY_NAME}" STREQUAL "ARROW")
build_arrow()
elseif("${DEPENDENCY_NAME}" STREQUAL "LAPACK")
elseif ("${DEPENDENCY_NAME}" STREQUAL "LAPACK")
build_lapack()
elseif ("${DEPENDENCY_NAME}" STREQUAL "GTest")
build_gtest()
elseif ("${DEPENDENCY_NAME}" STREQUAL "OpenBLAS")
build_openblas()
elseif("${DEPENDENCY_NAME}" STREQUAL "FAISS")
elseif ("${DEPENDENCY_NAME}" STREQUAL "FAISS")
build_faiss()
else()
else ()
message(FATAL_ERROR "Unknown thirdparty dependency to build: ${DEPENDENCY_NAME}")
endif ()
endmacro()
@ -51,7 +51,7 @@ endmacro()
macro(resolve_dependency DEPENDENCY_NAME)
if (${DEPENDENCY_NAME}_SOURCE STREQUAL "AUTO")
#message(STATUS "Finding ${DEPENDENCY_NAME} package")
#message(STATUS "${DEPENDENCY_NAME} package not found")
#message(STATUS "${DEPENDENCY_NAME} package not found")
build_dependency(${DEPENDENCY_NAME})
elseif (${DEPENDENCY_NAME}_SOURCE STREQUAL "BUNDLED")
build_dependency(${DEPENDENCY_NAME})
@ -64,28 +64,28 @@ endmacro()
# Identify OS
if (UNIX)
if (APPLE)
set (CMAKE_OS_NAME "osx" CACHE STRING "Operating system name" FORCE)
set(CMAKE_OS_NAME "osx" CACHE STRING "Operating system name" FORCE)
else (APPLE)
## Check for Debian GNU/Linux ________________
find_file (DEBIAN_FOUND debian_version debconf.conf
find_file(DEBIAN_FOUND debian_version debconf.conf
PATHS /etc
)
if (DEBIAN_FOUND)
set (CMAKE_OS_NAME "debian" CACHE STRING "Operating system name" FORCE)
set(CMAKE_OS_NAME "debian" CACHE STRING "Operating system name" FORCE)
endif (DEBIAN_FOUND)
## Check for Fedora _________________________
find_file (FEDORA_FOUND fedora-release
find_file(FEDORA_FOUND fedora-release
PATHS /etc
)
if (FEDORA_FOUND)
set (CMAKE_OS_NAME "fedora" CACHE STRING "Operating system name" FORCE)
set(CMAKE_OS_NAME "fedora" CACHE STRING "Operating system name" FORCE)
endif (FEDORA_FOUND)
## Check for RedHat _________________________
find_file (REDHAT_FOUND redhat-release inittab.RH
find_file(REDHAT_FOUND redhat-release inittab.RH
PATHS /etc
)
if (REDHAT_FOUND)
set (CMAKE_OS_NAME "redhat" CACHE STRING "Operating system name" FORCE)
set(CMAKE_OS_NAME "redhat" CACHE STRING "Operating system name" FORCE)
endif (REDHAT_FOUND)
## Extra check for Ubuntu ____________________
if (DEBIAN_FOUND)
@ -94,18 +94,18 @@ if (UNIX)
## a first superficial inspection a system will
## be considered as Debian, which signifies an
## extra check is required.
find_file (UBUNTU_EXTRA legal issue
find_file(UBUNTU_EXTRA legal issue
PATHS /etc
)
if (UBUNTU_EXTRA)
## Scan contents of file
file (STRINGS ${UBUNTU_EXTRA} UBUNTU_FOUND
file(STRINGS ${UBUNTU_EXTRA} UBUNTU_FOUND
REGEX Ubuntu
)
## Check result of string search
if (UBUNTU_FOUND)
set (CMAKE_OS_NAME "ubuntu" CACHE STRING "Operating system name" FORCE)
set (DEBIAN_FOUND FALSE)
set(CMAKE_OS_NAME "ubuntu" CACHE STRING "Operating system name" FORCE)
set(DEBIAN_FOUND FALSE)
endif (UBUNTU_FOUND)
endif (UBUNTU_EXTRA)
endif (DEBIAN_FOUND)
@ -119,17 +119,17 @@ set(THIRDPARTY_DIR "${INDEX_SOURCE_DIR}/thirdparty")
# ----------------------------------------------------------------------
# JFrog
if(NOT DEFINED USE_JFROG_CACHE)
if (NOT DEFINED USE_JFROG_CACHE)
set(USE_JFROG_CACHE "OFF")
endif()
if(USE_JFROG_CACHE STREQUAL "ON")
endif ()
if (USE_JFROG_CACHE STREQUAL "ON")
set(JFROG_ARTFACTORY_CACHE_URL "${JFROG_ARTFACTORY_URL}/milvus/thirdparty/cache/${CMAKE_OS_NAME}/${KNOWHERE_BUILD_ARCH}/${BUILD_TYPE}")
set(THIRDPARTY_PACKAGE_CACHE "${THIRDPARTY_DIR}/cache")
if(NOT EXISTS ${THIRDPARTY_PACKAGE_CACHE})
if (NOT EXISTS ${THIRDPARTY_PACKAGE_CACHE})
message(STATUS "Will create cached directory: ${THIRDPARTY_PACKAGE_CACHE}")
file(MAKE_DIRECTORY ${THIRDPARTY_PACKAGE_CACHE})
endif()
endif()
endif ()
endif ()
macro(resolve_dependency DEPENDENCY_NAME)
if (${DEPENDENCY_NAME}_SOURCE STREQUAL "AUTO")
@ -150,11 +150,11 @@ string(TOUPPER ${CMAKE_BUILD_TYPE} UPPERCASE_BUILD_TYPE)
set(EP_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${CMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}}")
set(EP_C_FLAGS "${CMAKE_C_FLAGS} ${CMAKE_C_FLAGS_${UPPERCASE_BUILD_TYPE}}")
if(NOT MSVC)
if (NOT MSVC)
# Set -fPIC on all external projects
set(EP_CXX_FLAGS "${EP_CXX_FLAGS} -fPIC")
set(EP_C_FLAGS "${EP_C_FLAGS} -fPIC")
endif()
endif ()
# CC/CXX environment variables are captured on the first invocation of the
# builder (e.g make or ninja) instead of when CMake is invoked into to build
@ -164,13 +164,13 @@ endif()
set(EP_COMMON_TOOLCHAIN -DCMAKE_C_COMPILER=${CMAKE_C_COMPILER}
-DCMAKE_CXX_COMPILER=${CMAKE_CXX_COMPILER})
if(CMAKE_AR)
if (CMAKE_AR)
set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_AR=${CMAKE_AR})
endif()
endif ()
if(CMAKE_RANLIB)
if (CMAKE_RANLIB)
set(EP_COMMON_TOOLCHAIN ${EP_COMMON_TOOLCHAIN} -DCMAKE_RANLIB=${CMAKE_RANLIB})
endif()
endif ()
# External projects are still able to override the following declarations.
# cmake command line will favor the last defined variable when a duplicate is
@ -184,18 +184,18 @@ set(EP_COMMON_CMAKE_ARGS
-DCMAKE_CXX_FLAGS=${EP_CXX_FLAGS}
-DCMAKE_CXX_FLAGS_${UPPERCASE_BUILD_TYPE}=${EP_CXX_FLAGS})
if(NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD)
if (NOT KNOWHERE_VERBOSE_THIRDPARTY_BUILD)
set(EP_LOG_OPTIONS LOG_CONFIGURE 1 LOG_BUILD 1 LOG_INSTALL 1 LOG_DOWNLOAD 1)
else()
else ()
set(EP_LOG_OPTIONS)
endif()
endif ()
# Ensure that a default make is set
if("${MAKE}" STREQUAL "")
if(NOT MSVC)
if ("${MAKE}" STREQUAL "")
if (NOT MSVC)
find_program(MAKE make)
endif()
endif()
endif ()
endif ()
set(MAKE_BUILD_ARGS "-j8")
@ -212,32 +212,32 @@ find_package(Threads REQUIRED)
# Read toolchain versions from cpp/thirdparty/versions.txt
file(STRINGS "${THIRDPARTY_DIR}/versions.txt" TOOLCHAIN_VERSIONS_TXT)
foreach(_VERSION_ENTRY ${TOOLCHAIN_VERSIONS_TXT})
foreach (_VERSION_ENTRY ${TOOLCHAIN_VERSIONS_TXT})
# Exclude comments
if(NOT _VERSION_ENTRY MATCHES "^[^#][A-Za-z0-9-_]+_VERSION=")
if (NOT _VERSION_ENTRY MATCHES "^[^#][A-Za-z0-9-_]+_VERSION=")
continue()
endif()
endif ()
string(REGEX MATCH "^[^=]*" _LIB_NAME ${_VERSION_ENTRY})
string(REPLACE "${_LIB_NAME}=" "" _LIB_VERSION ${_VERSION_ENTRY})
# Skip blank or malformed lines
if(${_LIB_VERSION} STREQUAL "")
if (${_LIB_VERSION} STREQUAL "")
continue()
endif()
endif ()
# For debugging
#message(STATUS "${_LIB_NAME}: ${_LIB_VERSION}")
set(${_LIB_NAME} "${_LIB_VERSION}")
endforeach()
endforeach ()
if(CUSTOMIZATION)
if (CUSTOMIZATION)
execute_process(COMMAND wget -q --method HEAD ${FAISS_URL} RESULT_VARIABLE return_code)
message(STATUS "Check the remote cache file ${FAISS_URL}. return code = ${return_code}")
if (NOT return_code EQUAL 0)
MESSAGE(FATAL_ERROR "Can't access to ${FAISS_URL}")
else()
else ()
set(FAISS_SOURCE_URL ${FAISS_URL})
# set(FAISS_MD5 "a589663865a8558205533c8ac414278c")
# set(FAISS_MD5 "57da9c4f599cc8fa4260488b1c96e1cc") # commit-id 6dbdf75987c34a2c853bd172ea0d384feea8358c branch-0.2.0
@ -246,36 +246,35 @@ if(CUSTOMIZATION)
# set(FAISS_MD5 "c89ea8e655f5cdf58f42486f13614714") # commit-id 9c28a1cbb88f41fa03b03d7204106201ad33276b branch-0.2.1
# set(FAISS_MD5 "87fdd86351ffcaf3f80dc26ade63c44b") # commit-id 841a156e67e8e22cd8088e1b58c00afbf2efc30b branch-0.2.1
# set(FAISS_MD5 "f3b2ce3364c3fa7febd3aa7fdd0fe380") # commit-id 694e03458e6b69ce8a62502f71f69a614af5af8f branch-0.3.0
# set(FAISS_MD5 "bb30722c22390ce5f6759ccb216c1b2a") # commit-id d324db297475286afe107847c7fb7a0f9dc7e90e branch-0.3.0
set(FAISS_MD5 "2293cdb209c3718e3b19f3edae8b32b3") # commit-id a13c1205dc52977a9ad3b33a14efa958604a8bff branch-0.3.0
endif()
else()
set(FAISS_MD5 "bb30722c22390ce5f6759ccb216c1b2a") # commit-id d324db297475286afe107847c7fb7a0f9dc7e90e branch-0.3.0
endif ()
else ()
set(FAISS_SOURCE_URL "https://github.com/milvus-io/faiss/archive/1.6.0.tar.gz")
set(FAISS_MD5 "eb96d84f98b078a9eec04a796f5c792e")
endif()
endif ()
message(STATUS "FAISS URL = ${FAISS_SOURCE_URL}")
if(DEFINED ENV{KNOWHERE_ARROW_URL})
if (DEFINED ENV{KNOWHERE_ARROW_URL})
set(ARROW_SOURCE_URL "$ENV{KNOWHERE_ARROW_URL}")
else()
else ()
set(ARROW_SOURCE_URL
"https://github.com/apache/arrow.git"
)
endif()
endif ()
if (DEFINED ENV{KNOWHERE_GTEST_URL})
set(GTEST_SOURCE_URL "$ENV{KNOWHERE_GTEST_URL}")
else ()
set(GTEST_SOURCE_URL
"https://github.com/google/googletest/archive/release-${GTEST_VERSION}.tar.gz")
endif()
endif ()
set(GTEST_MD5 "2e6fbeb6a91310a16efe181886c59596")
if(DEFINED ENV{KNOWHERE_LAPACK_URL})
if (DEFINED ENV{KNOWHERE_LAPACK_URL})
set(LAPACK_SOURCE_URL "$ENV{KNOWHERE_LAPACK_URL}")
else()
else ()
set(LAPACK_SOURCE_URL "https://github.com/Reference-LAPACK/lapack/archive/${LAPACK_VERSION}.tar.gz")
endif()
endif ()
set(LAPACK_MD5 "96591affdbf58c450d45c1daa540dbd2")
if (DEFINED ENV{KNOWHERE_OPENBLAS_URL})
@ -283,7 +282,7 @@ if (DEFINED ENV{KNOWHERE_OPENBLAS_URL})
else ()
set(OPENBLAS_SOURCE_URL
"https://github.com/xianyi/OpenBLAS/archive/${OPENBLAS_VERSION}.tar.gz")
endif()
endif ()
set(OPENBLAS_MD5 "8a110a25b819a4b94e8a9580702b6495")
# ----------------------------------------------------------------------
@ -293,10 +292,10 @@ set(ARROW_PREFIX "${INDEX_BINARY_DIR}/arrow_ep-prefix/src/arrow_ep/cpp")
macro(build_arrow)
message(STATUS "Building Apache ARROW-${ARROW_VERSION} from source")
set(ARROW_STATIC_LIB_NAME arrow)
set(ARROW_STATIC_LIB
set(ARROW_STATIC_LIB
"${ARROW_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}${ARROW_STATIC_LIB_NAME}${CMAKE_STATIC_LIBRARY_SUFFIX}"
)
set(ARROW_INCLUDE_DIR "${ARROW_PREFIX}/include")
set(ARROW_INCLUDE_DIR "${ARROW_PREFIX}/include")
set(ARROW_CMAKE_ARGS
${EP_COMMON_CMAKE_ARGS}
@ -326,10 +325,10 @@ macro(build_arrow)
-DBOOST_SOURCE=AUTO #try to find BOOST in the system default locations and build from source if not found
)
if(USE_JFROG_CACHE STREQUAL "ON")
if (USE_JFROG_CACHE STREQUAL "ON")
execute_process(COMMAND sh -c "git ls-remote --heads --tags ${ARROW_SOURCE_URL} ${ARROW_VERSION} | cut -f 1" OUTPUT_VARIABLE ARROW_LAST_COMMIT_ID)
if(${ARROW_LAST_COMMIT_ID} MATCHES "^[^#][a-z0-9]+")
if (${ARROW_LAST_COMMIT_ID} MATCHES "^[^#][a-z0-9]+")
string(MD5 ARROW_COMBINE_MD5 "${ARROW_LAST_COMMIT_ID}")
set(ARROW_CACHE_PACKAGE_NAME "arrow_${ARROW_COMBINE_MD5}.tar.gz")
set(ARROW_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${ARROW_CACHE_PACKAGE_NAME}")
@ -359,18 +358,18 @@ macro(build_arrow)
)
ExternalProject_Create_Cache(arrow_ep ${ARROW_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/arrow_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${ARROW_CACHE_URL})
else()
else ()
file(DOWNLOAD ${ARROW_CACHE_URL} ${ARROW_CACHE_PACKAGE_PATH} STATUS status)
list(GET status 0 status_code)
message(STATUS "DOWNLOADING FROM ${ARROW_CACHE_URL} TO ${ARROW_CACHE_PACKAGE_PATH}. STATUS = ${status_code}")
if (status_code EQUAL 0)
ExternalProject_Use_Cache(arrow_ep ${ARROW_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR})
endif()
endif()
else()
endif ()
endif ()
else ()
message(FATAL_ERROR "The last commit ID of \"${ARROW_SOURCE_URL}\" repository don't match!")
endif()
else()
endif ()
else ()
externalproject_add(arrow_ep
GIT_REPOSITORY
${ARROW_SOURCE_URL}
@ -390,14 +389,14 @@ macro(build_arrow)
BUILD_BYPRODUCTS
"${ARROW_STATIC_LIB}"
)
endif()
endif ()
file(MAKE_DIRECTORY "${ARROW_PREFIX}/include")
add_library(arrow STATIC IMPORTED)
set_target_properties(arrow
PROPERTIES IMPORTED_LOCATION "${ARROW_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${ARROW_INCLUDE_DIR}")
add_dependencies(arrow arrow_ep)
add_dependencies(arrow arrow_ep)
set(JEMALLOC_PREFIX "${INDEX_BINARY_DIR}/arrow_ep-prefix/src/arrow_ep-build/jemalloc_ep-prefix/src/jemalloc_ep")
@ -408,13 +407,13 @@ macro(build_arrow)
endmacro()
if(KNOWHERE_WITH_ARROW AND NOT TARGET arrow_ep)
if (KNOWHERE_WITH_ARROW AND NOT TARGET arrow_ep)
resolve_dependency(ARROW)
link_directories(SYSTEM ${ARROW_PREFIX}/lib/)
include_directories(SYSTEM ${ARROW_INCLUDE_DIR})
endif()
endif ()
# ----------------------------------------------------------------------
# OpenBLAS
@ -428,7 +427,7 @@ macro(build_openblas)
set(OPENBLAS_REAL_STATIC_LIB
"${OPENBLAS_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}openblas_haswellp-r0.3.6${CMAKE_STATIC_LIBRARY_SUFFIX}")
if(USE_JFROG_CACHE STREQUAL "ON")
if (USE_JFROG_CACHE STREQUAL "ON")
set(OPENBLAS_CACHE_PACKAGE_NAME "openblas_${OPENBLAS_MD5}.tar.gz")
set(OPENBLAS_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${OPENBLAS_CACHE_PACKAGE_NAME}")
set(OPENBLAS_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${OPENBLAS_CACHE_PACKAGE_NAME}")
@ -455,15 +454,15 @@ macro(build_openblas)
${OPENBLAS_STATIC_LIB})
ExternalProject_Create_Cache(openblas_ep ${OPENBLAS_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/openblas_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${OPENBLAS_CACHE_URL})
else()
else ()
file(DOWNLOAD ${OPENBLAS_CACHE_URL} ${OPENBLAS_CACHE_PACKAGE_PATH} STATUS status)
list(GET status 0 status_code)
message(STATUS "DOWNLOADING FROM ${OPENBLAS_CACHE_URL} TO ${OPENBLAS_CACHE_PACKAGE_PATH}. STATUS = ${status_code}")
if (status_code EQUAL 0)
ExternalProject_Use_Cache(openblas_ep ${OPENBLAS_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR})
endif()
endif()
else()
endif ()
endif ()
else ()
externalproject_add(openblas_ep
URL
${OPENBLAS_SOURCE_URL}
@ -481,7 +480,7 @@ macro(build_openblas)
install
BUILD_BYPRODUCTS
${OPENBLAS_STATIC_LIB})
endif()
endif ()
file(MAKE_DIRECTORY "${OPENBLAS_INCLUDE_DIR}")
add_library(openblas STATIC IMPORTED)
@ -510,7 +509,7 @@ macro(build_lapack)
"-DCMAKE_INSTALL_PREFIX=${LAPACK_PREFIX}"
-DCMAKE_INSTALL_LIBDIR=lib)
if(USE_JFROG_CACHE STREQUAL "ON")
if (USE_JFROG_CACHE STREQUAL "ON")
set(LAPACK_CACHE_PACKAGE_NAME "lapack_${LAPACK_MD5}.tar.gz")
set(LAPACK_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${LAPACK_CACHE_PACKAGE_NAME}")
set(LAPACK_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${LAPACK_CACHE_PACKAGE_NAME}")
@ -531,15 +530,15 @@ macro(build_lapack)
${LAPACK_STATIC_LIB})
ExternalProject_Create_Cache(lapack_ep ${LAPACK_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/lapack_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${LAPACK_CACHE_URL})
else()
else ()
file(DOWNLOAD ${LAPACK_CACHE_URL} ${LAPACK_CACHE_PACKAGE_PATH} STATUS status)
list(GET status 0 status_code)
message(STATUS "DOWNLOADING FROM ${LAPACK_CACHE_URL} TO ${LAPACK_CACHE_PACKAGE_PATH}. STATUS = ${status_code}")
if (status_code EQUAL 0)
ExternalProject_Use_Cache(lapack_ep ${LAPACK_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR})
endif()
endif()
else()
endif ()
endif ()
else ()
externalproject_add(lapack_ep
URL
${LAPACK_SOURCE_URL}
@ -551,7 +550,7 @@ macro(build_lapack)
${MAKE_BUILD_ARGS}
BUILD_BYPRODUCTS
${LAPACK_STATIC_LIB})
endif()
endif ()
file(MAKE_DIRECTORY "${LAPACK_INCLUDE_DIR}")
add_library(lapack STATIC IMPORTED)
@ -571,13 +570,13 @@ macro(build_gtest)
set(GTEST_VENDORED TRUE)
set(GTEST_CMAKE_CXX_FLAGS "${EP_CXX_FLAGS}")
if(APPLE)
if (APPLE)
set(GTEST_CMAKE_CXX_FLAGS
${GTEST_CMAKE_CXX_FLAGS}
-DGTEST_USE_OWN_TR1_TUPLE=1
-Wno-unused-value
-Wno-ignored-attributes)
endif()
endif ()
set(GTEST_PREFIX "${INDEX_BINARY_DIR}/googletest_ep-prefix/src/googletest_ep")
set(GTEST_INCLUDE_DIR "${GTEST_PREFIX}/include")
@ -596,10 +595,10 @@ macro(build_gtest)
set(GMOCK_INCLUDE_DIR "${GTEST_PREFIX}/include")
set(GMOCK_STATIC_LIB
"${GTEST_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}gmock${CMAKE_STATIC_LIBRARY_SUFFIX}"
)
)
if(USE_JFROG_CACHE STREQUAL "ON")
if (USE_JFROG_CACHE STREQUAL "ON")
set(GTEST_CACHE_PACKAGE_NAME "googletest_${GTEST_MD5}.tar.gz")
set(GTEST_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${GTEST_CACHE_PACKAGE_NAME}")
set(GTEST_CACHE_PACKAGE_PATH "${THIRDPARTY_PACKAGE_CACHE}/${GTEST_CACHE_PACKAGE_NAME}")
@ -622,15 +621,15 @@ macro(build_gtest)
${EP_LOG_OPTIONS})
ExternalProject_Create_Cache(googletest_ep ${GTEST_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/googletest_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${GTEST_CACHE_URL})
else()
else ()
file(DOWNLOAD ${GTEST_CACHE_URL} ${GTEST_CACHE_PACKAGE_PATH} STATUS status)
list(GET status 0 status_code)
message(STATUS "DOWNLOADING FROM ${GTEST_CACHE_URL} TO ${GTEST_CACHE_PACKAGE_PATH}. STATUS = ${status_code}")
if (status_code EQUAL 0)
ExternalProject_Use_Cache(googletest_ep ${GTEST_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR})
endif()
endif()
else()
endif ()
endif ()
else ()
ExternalProject_Add(googletest_ep
URL
${GTEST_SOURCE_URL}
@ -644,20 +643,20 @@ macro(build_gtest)
CMAKE_ARGS
${GTEST_CMAKE_ARGS}
${EP_LOG_OPTIONS})
endif()
endif ()
# The include directory must exist before it is referenced by a target.
file(MAKE_DIRECTORY "${GTEST_INCLUDE_DIR}")
add_library(gtest STATIC IMPORTED)
set_target_properties(gtest
PROPERTIES IMPORTED_LOCATION "${GTEST_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}")
PROPERTIES IMPORTED_LOCATION "${GTEST_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}")
add_library(gtest_main STATIC IMPORTED)
set_target_properties(gtest_main
PROPERTIES IMPORTED_LOCATION "${GTEST_MAIN_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}")
PROPERTIES IMPORTED_LOCATION "${GTEST_MAIN_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${GTEST_INCLUDE_DIR}")
add_library(gmock STATIC IMPORTED)
set_target_properties(gmock
@ -673,44 +672,88 @@ endmacro()
if (KNOWHERE_BUILD_TESTS AND NOT TARGET googletest_ep)
resolve_dependency(GTest)
if(NOT GTEST_VENDORED)
endif()
if (NOT GTEST_VENDORED)
endif ()
# TODO: Don't use global includes but rather target_include_directories
get_target_property(GTEST_INCLUDE_DIR gtest INTERFACE_INCLUDE_DIRECTORIES)
link_directories(SYSTEM "${GTEST_PREFIX}/lib")
include_directories(SYSTEM ${GTEST_INCLUDE_DIR})
endif()
endif ()
# ----------------------------------------------------------------------
# FAISS
macro(build_faiss)
message(STATUS "Building FAISS-${FAISS_VERSION} from source")
if (NOT DEFINED BUILD_FAISS_WITH_MKL)
set(BUILD_FAISS_WITH_MKL OFF)
endif ()
if (EXISTS "/proc/cpuinfo")
FILE(READ /proc/cpuinfo PROC_CPUINFO)
SET(VENDOR_ID_RX "vendor_id[ \t]*:[ \t]*([a-zA-Z]+)\n")
STRING(REGEX MATCH "${VENDOR_ID_RX}" VENDOR_ID "${PROC_CPUINFO}")
STRING(REGEX REPLACE "${VENDOR_ID_RX}" "\\1" VENDOR_ID "${VENDOR_ID}")
if (NOT ${VENDOR_ID} STREQUAL "GenuineIntel")
set(BUILD_FAISS_WITH_MKL OFF)
endif ()
endif ()
set(FAISS_PREFIX "${INDEX_BINARY_DIR}/faiss_ep-prefix/src/faiss_ep")
set(FAISS_INCLUDE_DIR "${FAISS_PREFIX}/include")
set(FAISS_STATIC_LIB
"${FAISS_PREFIX}/lib/${CMAKE_STATIC_LIBRARY_PREFIX}faiss${CMAKE_STATIC_LIBRARY_SUFFIX}")
set(FAISS_CONFIGURE_ARGS
"--prefix=${FAISS_PREFIX}"
"CFLAGS=${EP_C_FLAGS}"
"CXXFLAGS=${EP_CXX_FLAGS}"
"LDFLAGS=-L${OPENBLAS_PREFIX}/lib -L${LAPACK_PREFIX}/lib -lopenblas -llapack"
--without-python)
if(${KNOWHERE_WITH_FAISS_GPU_VERSION} STREQUAL "ON")
set(FAISS_CFLAGS ${EP_C_FLAGS})
set(FAISS_CXXFLAGS ${EP_CXX_FLAGS})
if (BUILD_FAISS_WITH_MKL)
find_path(MKL_LIB_PATH
NAMES "libmkl_intel_ilp64.a" "libmkl_gnu_thread.a" "libmkl_core.a"
PATH_SUFFIXES "intel/compilers_and_libraries_${MKL_VERSION}/linux/mkl/lib/intel64/")
if (${MKL_LIB_PATH} STREQUAL "MKL_LIB_PATH-NOTFOUND")
message(FATAL_ERROR "Could not find MKL libraries")
endif ()
message(STATUS "Build Faiss with MKL. MKL lib path = ${MKL_LIB_PATH}")
set(MKL_LIBS
${MKL_LIB_PATH}/libmkl_intel_ilp64.a
${MKL_LIB_PATH}/libmkl_gnu_thread.a
${MKL_LIB_PATH}/libmkl_core.a
)
set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS}
"CPPFLAGS=-DFINTEGER=long -DMKL_ILP64 -m64 -I${MKL_LIB_PATH}/../../include"
"LDFLAGS=-L${MKL_LIB_PATH}"
)
else ()
message(STATUS "Build Faiss with OpenBlas/LAPACK")
set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS}
"LDFLAGS=-L${OPENBLAS_PREFIX}/lib -L${LAPACK_PREFIX}/lib")
endif ()
if (KNOWHERE_GPU_VERSION)
set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS}
"--with-cuda=${CUDA_TOOLKIT_ROOT_DIR}"
"--with-cuda-arch=-gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_61,code=sm_61 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75"
)
else()
else ()
set(FAISS_CONFIGURE_ARGS ${FAISS_CONFIGURE_ARGS} --without-cuda)
endif()
endif ()
if(USE_JFROG_CACHE STREQUAL "ON")
if (USE_JFROG_CACHE STREQUAL "ON")
string(MD5 FAISS_COMBINE_MD5 "${FAISS_MD5}${LAPACK_MD5}${OPENBLAS_MD5}")
set(FAISS_CACHE_PACKAGE_NAME "faiss_${FAISS_COMBINE_MD5}.tar.gz")
set(FAISS_CACHE_URL "${JFROG_ARTFACTORY_CACHE_URL}/${FAISS_CACHE_PACKAGE_NAME}")
@ -735,18 +778,20 @@ macro(build_faiss)
BUILD_BYPRODUCTS
${FAISS_STATIC_LIB})
ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep)
if (NOT BUILD_FAISS_WITH_MKL)
ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep)
endif ()
ExternalProject_Create_Cache(faiss_ep ${FAISS_CACHE_PACKAGE_PATH} "${INDEX_BINARY_DIR}/faiss_ep-prefix" ${JFROG_USER_NAME} ${JFROG_PASSWORD} ${FAISS_CACHE_URL})
else()
else ()
file(DOWNLOAD ${FAISS_CACHE_URL} ${FAISS_CACHE_PACKAGE_PATH} STATUS status)
list(GET status 0 status_code)
message(STATUS "DOWNLOADING FROM ${FAISS_CACHE_URL} TO ${FAISS_CACHE_PACKAGE_PATH}. STATUS = ${status_code}")
if (status_code EQUAL 0)
ExternalProject_Use_Cache(faiss_ep ${FAISS_CACHE_PACKAGE_PATH} ${INDEX_BINARY_DIR})
endif()
endif()
else()
endif ()
endif ()
else ()
externalproject_add(faiss_ep
URL
${FAISS_SOURCE_URL}
@ -763,35 +808,54 @@ macro(build_faiss)
BUILD_BYPRODUCTS
${FAISS_STATIC_LIB})
ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep)
endif()
if (NOT BUILD_FAISS_WITH_MKL)
ExternalProject_Add_StepDependencies(faiss_ep build openblas_ep lapack_ep)
endif ()
endif ()
file(MAKE_DIRECTORY "${FAISS_INCLUDE_DIR}")
add_library(faiss STATIC IMPORTED)
set_target_properties(
faiss
PROPERTIES IMPORTED_LOCATION "${FAISS_STATIC_LIB}"
PROPERTIES
IMPORTED_LOCATION "${FAISS_STATIC_LIB}"
INTERFACE_INCLUDE_DIRECTORIES "${FAISS_INCLUDE_DIR}"
INTERFACE_LINK_LIBRARIES "openblas;lapack" )
)
if (BUILD_FAISS_WITH_MKL)
set_target_properties(
faiss
PROPERTIES
INTERFACE_LINK_LIBRARIES "${MKL_LIBS}")
else ()
set_target_properties(
faiss
PROPERTIES
INTERFACE_LINK_LIBRARIES "openblas;lapack")
endif ()
add_dependencies(faiss faiss_ep)
endmacro()
if(KNOWHERE_WITH_FAISS AND NOT TARGET faiss_ep)
if (KNOWHERE_WITH_FAISS AND NOT TARGET faiss_ep)
resolve_dependency(OpenBLAS)
get_target_property(OPENBLAS_INCLUDE_DIR openblas INTERFACE_INCLUDE_DIRECTORIES)
include_directories(SYSTEM "${OPENBLAS_INCLUDE_DIR}")
link_directories(SYSTEM ${OPENBLAS_PREFIX}/lib)
if (NOT BUILD_FAISS_WITH_MKL)
resolve_dependency(OpenBLAS)
get_target_property(OPENBLAS_INCLUDE_DIR openblas INTERFACE_INCLUDE_DIRECTORIES)
include_directories(SYSTEM "${OPENBLAS_INCLUDE_DIR}")
link_directories(SYSTEM ${OPENBLAS_PREFIX}/lib)
resolve_dependency(LAPACK)
get_target_property(LAPACK_INCLUDE_DIR lapack INTERFACE_INCLUDE_DIRECTORIES)
include_directories(SYSTEM "${LAPACK_INCLUDE_DIR}")
link_directories(SYSTEM "${LAPACK_PREFIX}/lib")
resolve_dependency(LAPACK)
get_target_property(LAPACK_INCLUDE_DIR lapack INTERFACE_INCLUDE_DIRECTORIES)
include_directories(SYSTEM "${LAPACK_INCLUDE_DIR}")
link_directories(SYSTEM "${LAPACK_PREFIX}/lib")
endif ()
resolve_dependency(FAISS)
get_target_property(FAISS_INCLUDE_DIR faiss INTERFACE_INCLUDE_DIRECTORIES)
include_directories(SYSTEM "${FAISS_INCLUDE_DIR}")
link_directories(SYSTEM ${FAISS_PREFIX}/lib/)
endif()
endif ()

57
core/src/index/knowhere/CMakeLists.txt
View File

@ -1,6 +1,3 @@
include_directories(${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64)
include_directories(${INDEX_SOURCE_DIR}/knowhere)
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService)
@ -19,9 +16,9 @@ file(GLOB SRC_FILES
${SPTAG_SOURCE_DIR}/AnnService/src/Core/KDT/*.cpp
${SPTAG_SOURCE_DIR}/AnnService/src/Helper/*.cpp)
if(NOT TARGET SPTAGLibStatic)
if (NOT TARGET SPTAGLibStatic)
add_library(SPTAGLibStatic STATIC ${SRC_FILES} ${HDR_FILES})
endif()
endif ()
set(external_srcs
knowhere/adapter/SptagAdapter.cpp
@ -36,19 +33,13 @@ set(index_srcs
knowhere/index/vector_index/IndexKDT.cpp
knowhere/index/vector_index/IndexIDMAP.cpp
knowhere/index/vector_index/IndexIVF.cpp
knowhere/index/vector_index/IndexGPUIVF.cpp
knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
knowhere/index/vector_index/IndexNSG.cpp
knowhere/index/vector_index/nsg/NSG.cpp
knowhere/index/vector_index/nsg/NSGIO.cpp
knowhere/index/vector_index/nsg/NSGHelper.cpp
knowhere/index/vector_index/helpers/Cloner.cpp
knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
knowhere/index/vector_index/IndexIVFSQ.cpp
knowhere/index/vector_index/IndexGPUIVFSQ.cpp
knowhere/index/vector_index/IndexIVFSQHybrid.cpp
knowhere/index/vector_index/IndexIVFPQ.cpp
knowhere/index/vector_index/IndexGPUIVFPQ.cpp
knowhere/index/vector_index/FaissBaseIndex.cpp
knowhere/index/vector_index/helpers/FaissIO.cpp
knowhere/index/vector_index/helpers/IndexParameter.cpp
@ -57,24 +48,56 @@ set(index_srcs
set(depend_libs
SPTAGLibStatic
faiss
openblas
lapack
arrow
${ARROW_PREFIX}/lib/libjemalloc_pic.a
cudart
cublas
gomp
gfortran
pthread
)
if (BUILD_FAISS_WITH_MKL)
set(depend_libs ${depend_libs}
"-Wl,--start-group \
${MKL_LIB_PATH}/libmkl_intel_ilp64.a \
${MKL_LIB_PATH}/libmkl_gnu_thread.a \
${MKL_LIB_PATH}/libmkl_core.a \
-Wl,--end-group -lgomp -lpthread -lm -ldl"
)
else ()
set(depend_libs ${depend_libs}
lapack
openblas)
endif ()
if(NOT TARGET knowhere)
if (KNOWHERE_GPU_VERSION)
include_directories(${CUDA_INCLUDE_DIRS})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(cuda_lib
cudart
cublas
)
set(depend_libs ${depend_libs}
${cuda_lib}
)
set(index_srcs ${index_srcs}
knowhere/index/vector_index/IndexGPUIVF.cpp
knowhere/index/vector_index/helpers/Cloner.cpp
knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
knowhere/index/vector_index/IndexGPUIVFSQ.cpp
knowhere/index/vector_index/IndexIVFSQHybrid.cpp
knowhere/index/vector_index/IndexGPUIVFPQ.cpp
knowhere/index/vector_index/IndexGPUIDMAP.cpp
)
endif ()
if (NOT TARGET knowhere)
add_library(
knowhere STATIC
${external_srcs}
${index_srcs}
)
endif()
endif ()
target_link_libraries(
knowhere

129
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp Normal file
View File

@ -0,0 +1,129 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include "knowhere/index/vector_index/IndexGPUIDMAP.h"
#include <faiss/AutoTune.h>
#include <faiss/IndexFlat.h>
#include <faiss/MetaIndexes.h>
#include <faiss/index_io.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuCloner.h>
#endif
#include "knowhere/adapter/VectorAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/helpers/FaissIO.h"
namespace knowhere {
VectorIndexPtr
GPUIDMAP::CopyGpuToCpu(const Config& config) {
std::lock_guard<std::mutex> lk(mutex_);
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IDMAP>(new_index);
}
VectorIndexPtr
GPUIDMAP::Clone() {
auto cpu_idx = CopyGpuToCpu(Config());
if (auto idmap = std::dynamic_pointer_cast<IDMAP>(cpu_idx)) {
return idmap->CopyCpuToGpu(gpu_id_, Config());
} else {
KNOWHERE_THROW_MSG("IndexType not Support GpuClone");
}
}
BinarySet
GPUIDMAP::SerializeImpl() {
try {
MemoryIOWriter writer;
{
faiss::Index* index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(index);
faiss::write_index(host_index, &writer);
delete host_index;
}
auto data = std::make_shared<uint8_t>();
data.reset(writer.data_);
BinarySet res_set;
res_set.Append("IVF", data, writer.rp);
return res_set;
} catch (std::exception& e) {
KNOWHERE_THROW_MSG(e.what());
}
}
void
GPUIDMAP::LoadImpl(const BinarySet& index_binary) {
auto binary = index_binary.GetByName("IVF");
MemoryIOReader reader;
{
reader.total = binary->size;
reader.data_ = binary->data.get();
faiss::Index* index = faiss::read_index(&reader);
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) {
ResScope rs(res, gpu_id_, false);
auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index);
index_.reset(device_index);
res_ = res;
} else {
KNOWHERE_THROW_MSG("Load error, can't get gpu resource");
}
delete index;
}
}
VectorIndexPtr
GPUIDMAP::CopyGpuToGpu(const int64_t& device_id, const Config& config) {
auto cpu_index = CopyGpuToCpu(config);
return std::static_pointer_cast<IDMAP>(cpu_index)->CopyCpuToGpu(device_id, config);
}
float*
GPUIDMAP::GetRawVectors() {
KNOWHERE_THROW_MSG("Not support");
}
int64_t*
GPUIDMAP::GetRawIds() {
KNOWHERE_THROW_MSG("Not support");
}
void
GPUIDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) {
ResScope rs(res_, gpu_id_);
index_->search(n, (float*)data, k, distances, labels);
}
} // namespace knowhere

63
core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.h Normal file
View File

@ -0,0 +1,63 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include "IndexGPUIVF.h"
#include "IndexIDMAP.h"
#include "IndexIVF.h"
#include <memory>
#include <utility>
namespace knowhere {
class GPUIDMAP : public IDMAP, public GPUIndex {
public:
explicit GPUIDMAP(std::shared_ptr<faiss::Index> index, const int64_t& device_id, ResPtr& res)
: IDMAP(std::move(index)), GPUIndex(device_id, res) {
}
VectorIndexPtr
CopyGpuToCpu(const Config& config) override;
float*
GetRawVectors() override;
int64_t*
GetRawIds() override;
VectorIndexPtr
Clone() override;
VectorIndexPtr
CopyGpuToGpu(const int64_t& device_id, const Config& config) override;
protected:
void
search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override;
BinarySet
SerializeImpl() override;
void
LoadImpl(const BinarySet& index_binary) override;
};
using GPUIDMAPPtr = std::shared_ptr<GPUIDMAP>;
} // namespace knowhere

113
core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
View File

@ -17,10 +17,18 @@
#include <faiss/IndexFlat.h>
#include <faiss/MetaIndexes.h>
#include <faiss/gpu/GpuCloner.h>
#include <faiss/AutoTune.h>
#include <faiss/clone_index.h>
#include <faiss/index_factory.h>
#include <faiss/index_io.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuCloner.h>
#endif
#include <vector>
#include "knowhere/adapter/VectorAdapter.h"
@ -28,6 +36,13 @@
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/helpers/FaissIO.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIDMAP.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
namespace knowhere {
BinarySet
@ -160,6 +175,8 @@ IDMAP::Clone() {
VectorIndexPtr
IDMAP::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
#ifdef MILVUS_GPU_VERSION
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
ResScope rs(res, device_id, false);
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
@ -170,6 +187,9 @@ IDMAP::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
} else {
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
}
#else
KNOWHERE_THROW_MSG("Calling IDMAP::CopyCpuToGpu when we are using CPU version");
#endif
}
void
@ -177,95 +197,4 @@ IDMAP::Seal() {
// do nothing
}
VectorIndexPtr
GPUIDMAP::CopyGpuToCpu(const Config& config) {
std::lock_guard<std::mutex> lk(mutex_);
faiss::Index* device_index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index);
std::shared_ptr<faiss::Index> new_index;
new_index.reset(host_index);
return std::make_shared<IDMAP>(new_index);
}
VectorIndexPtr
GPUIDMAP::Clone() {
auto cpu_idx = CopyGpuToCpu(Config());
if (auto idmap = std::dynamic_pointer_cast<IDMAP>(cpu_idx)) {
return idmap->CopyCpuToGpu(gpu_id_, Config());
} else {
KNOWHERE_THROW_MSG("IndexType not Support GpuClone");
}
}
BinarySet
GPUIDMAP::SerializeImpl() {
try {
MemoryIOWriter writer;
{
faiss::Index* index = index_.get();
faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(index);
faiss::write_index(host_index, &writer);
delete host_index;
}
auto data = std::make_shared<uint8_t>();
data.reset(writer.data_);
BinarySet res_set;
res_set.Append("IVF", data, writer.rp);
return res_set;
} catch (std::exception& e) {
KNOWHERE_THROW_MSG(e.what());
}
}
void
GPUIDMAP::LoadImpl(const BinarySet& index_binary) {
auto binary = index_binary.GetByName("IVF");
MemoryIOReader reader;
{
reader.total = binary->size;
reader.data_ = binary->data.get();
faiss::Index* index = faiss::read_index(&reader);
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_)) {
ResScope rs(res, gpu_id_, false);
auto device_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), gpu_id_, index);
index_.reset(device_index);
res_ = res;
} else {
KNOWHERE_THROW_MSG("Load error, can't get gpu resource");
}
delete index;
}
}
VectorIndexPtr
GPUIDMAP::CopyGpuToGpu(const int64_t& device_id, const Config& config) {
auto cpu_index = CopyGpuToCpu(config);
return std::static_pointer_cast<IDMAP>(cpu_index)->CopyCpuToGpu(device_id, config);
}
float*
GPUIDMAP::GetRawVectors() {
KNOWHERE_THROW_MSG("Not support");
}
int64_t*
GPUIDMAP::GetRawIds() {
KNOWHERE_THROW_MSG("Not support");
}
void
GPUIDMAP::search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) {
ResScope rs(res_, gpu_id_);
index_->search(n, (float*)data, k, distances, labels);
}
} // namespace knowhere

29
core/src/index/knowhere/knowhere/index/vector_index/IndexIDMAP.h
View File

@ -17,7 +17,6 @@
#pragma once
#include "IndexGPUIVF.h"
#include "IndexIVF.h"
#include <memory>
@ -67,32 +66,4 @@ class IDMAP : public VectorIndex, public FaissBaseIndex {
using IDMAPPtr = std::shared_ptr<IDMAP>;
class GPUIDMAP : public IDMAP, public GPUIndex {
public:
explicit GPUIDMAP(std::shared_ptr<faiss::Index> index, const int64_t& device_id, ResPtr& res)
: IDMAP(std::move(index)), GPUIndex(device_id, res) {
}
VectorIndexPtr
CopyGpuToCpu(const Config& config) override;
float*
GetRawVectors() override;
int64_t*
GetRawIds() override;
VectorIndexPtr
Clone() override;
VectorIndexPtr
CopyGpuToGpu(const int64_t& device_id, const Config& config) override;
protected:
void
search_impl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& cfg) override;
BinarySet
SerializeImpl() override;
void
LoadImpl(const BinarySet& index_binary) override;
};
using GPUIDMAPPtr = std::shared_ptr<GPUIDMAP>;
} // namespace knowhere

23
core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
View File

@ -15,11 +15,19 @@
// specific language governing permissions and limitations
// under the License.
#include <faiss/AutoTune.h>
#include <faiss/IVFlib.h>
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVF.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexIVFPQ.h>
#include <faiss/clone_index.h>
#include <faiss/index_factory.h>
#include <faiss/index_io.h>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuCloner.h>
#endif
#include <chrono>
#include <memory>
@ -29,7 +37,9 @@
#include "knowhere/adapter/VectorAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/common/Log.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#endif
#include "knowhere/index/vector_index/IndexIVF.h"
namespace knowhere {
@ -221,16 +231,17 @@ IVF::search_impl(int64_t n, const float* data, int64_t k, float* distances, int6
faiss::ivflib::search_with_parameters(index_.get(), n, (float*)data, k, distances, labels, params.get());
stdclock::time_point after = stdclock::now();
double search_cost = (std::chrono::duration<double, std::micro>(after - before)).count();
KNOWHERE_LOG_DEBUG << "K=" << k << " NQ=" << n << " NL=" << faiss::indexIVF_stats.nlist
<< " ND=" << faiss::indexIVF_stats.ndis << " NH=" << faiss::indexIVF_stats.nheap_updates
<< " Q=" << faiss::indexIVF_stats.quantization_time
<< " S=" << faiss::indexIVF_stats.search_time;
KNOWHERE_LOG_DEBUG << "IVF search cost: " << search_cost
<< ", quantization cost: " << faiss::indexIVF_stats.quantization_time
<< ", data search cost: " << faiss::indexIVF_stats.search_time;
faiss::indexIVF_stats.quantization_time = 0;
faiss::indexIVF_stats.search_time = 0;
}
VectorIndexPtr
IVF::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
#ifdef MILVUS_GPU_VERSION
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
ResScope rs(res, device_id, false);
auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get());
@ -241,6 +252,10 @@ IVF::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
} else {
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
}
#else
KNOWHERE_THROW_MSG("Calling IVF::CopyCpuToGpu when we are using CPU version");
#endif
}
VectorIndexPtr

15
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp
View File

@ -15,15 +15,22 @@
// specific language governing permissions and limitations
// under the License.
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuCloner.h>
#endif
#include <faiss/index_factory.h>
#include <memory>
#include "knowhere/adapter/VectorAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
namespace knowhere {
@ -54,6 +61,8 @@ IVFSQ::Clone_impl(const std::shared_ptr<faiss::Index>& index) {
VectorIndexPtr
IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
#ifdef MILVUS_GPU_VERSION
if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) {
ResScope rs(res, device_id, false);
@ -65,6 +74,10 @@ IVFSQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) {
} else {
KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource");
}
#else
KNOWHERE_THROW_MSG("Calling IVFSQ::CopyCpuToGpu when we are using CPU version");
#endif
}
} // namespace knowhere

8
core/src/index/knowhere/knowhere/index/vector_index/IndexNSG.cpp
View File

@ -19,8 +19,10 @@
#include "knowhere/adapter/VectorAdapter.h"
#include "knowhere/common/Exception.h"
#include "knowhere/common/Timer.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#endif
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/nsg/NSG.h"
#include "knowhere/index/vector_index/nsg/NSGIO.h"
@ -117,7 +119,11 @@ NSG::Train(const DatasetPtr& dataset, const Config& config) {
}
// TODO(linxj): dev IndexFactory, support more IndexType
#ifdef MILVUS_GPU_VERSION
auto preprocess_index = std::make_shared<GPUIVF>(build_cfg->gpu_id);
#else
auto preprocess_index = std::make_shared<IVF>();
#endif
auto model = preprocess_index->Train(dataset, config);
preprocess_index->set_index_model(model);
preprocess_index->AddWithoutIds(dataset, config);

86
core/src/index/unittest/CMakeLists.txt
View File

@ -2,26 +2,32 @@ include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/thirdparty/SPTAG/AnnService)
include_directories(${INDEX_SOURCE_DIR}/knowhere)
include_directories(${INDEX_SOURCE_DIR})
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
message(STATUS "arrow prefix: ${ARROW_PREFIX}")
message(STATUS "libjemalloc_pic path: ${ARROW_PREFIX}/lib/libjemalloc_pic.a")
set(depend_libs
gtest gmock gtest_main gmock_main
faiss openblas lapack
faiss
arrow "${ARROW_PREFIX}/lib/libjemalloc_pic.a"
)
if (BUILD_FAISS_WITH_MKL)
set(depend_libs ${depend_libs}
"-Wl,--start-group \
${MKL_LIB_PATH}/libmkl_intel_ilp64.a \
${MKL_LIB_PATH}/libmkl_gnu_thread.a \
${MKL_LIB_PATH}/libmkl_core.a \
-Wl,--end-group -lgomp -lpthread -lm -ldl"
)
else ()
set(depend_libs ${depend_libs}
lapack
openblas)
endif ()
set(basic_libs
cudart cublas
gomp gfortran pthread
)
set(util_srcs
${MILVUS_ENGINE_SRC}/external/easyloggingpp/easylogging++.cc
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissIO.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/IndexParameter.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/adapter/Structure.cpp
@ -31,32 +37,49 @@ set(util_srcs
${INDEX_SOURCE_DIR}/unittest/utils.cpp
)
if (KNOWHERE_GPU_VERSION)
include_directories(${CUDA_INCLUDE_DIRS})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(cuda_lib
cudart
cublas
)
set(basic_libs ${basic_libs}
${cuda_lib}
)
set(util_srcs ${util_srcs}
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
)
endif ()
#<IVF-TEST>
set(ivf_srcs
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVF.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/FaissBaseIndex.cpp
)
if(NOT TARGET test_ivf)
if (KNOWHERE_GPU_VERSION)
set(ivf_srcs ${ivf_srcs}
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIDMAP.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFSQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
)
endif ()
if (NOT TARGET test_ivf)
add_executable(test_ivf test_ivf.cpp ${ivf_srcs} ${util_srcs})
endif()
endif ()
target_link_libraries(test_ivf ${depend_libs} ${unittest_libs} ${basic_libs})
#<IDMAP-TEST>
set(idmap_srcs
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexIDMAP.cpp
)
if(NOT TARGET test_idmap)
add_executable(test_idmap test_idmap.cpp ${idmap_srcs} ${ivf_srcs} ${util_srcs})
endif()
if (NOT TARGET test_idmap)
add_executable(test_idmap test_idmap.cpp ${ivf_srcs} ${util_srcs})
endif ()
target_link_libraries(test_idmap ${depend_libs} ${unittest_libs} ${basic_libs})
#<KDT-TEST>
@ -66,25 +89,28 @@ set(kdt_srcs
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/helpers/KDTParameterMgr.cpp
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexKDT.cpp
)
if(NOT TARGET test_kdt)
if (NOT TARGET test_kdt)
add_executable(test_kdt test_kdt.cpp ${kdt_srcs} ${util_srcs})
endif()
endif ()
target_link_libraries(test_kdt
SPTAGLibStatic
${depend_libs} ${unittest_libs} ${basic_libs})
add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})
if (KNOWHERE_GPU_VERSION)
add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})
add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})
add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})
endif ()
install(TARGETS test_ivf DESTINATION unittest)
install(TARGETS test_idmap DESTINATION unittest)
install(TARGETS test_kdt DESTINATION unittest)
install(TARGETS test_gpuresource DESTINATION unittest)
install(TARGETS test_customized_index DESTINATION unittest)
if (KNOWHERE_GPU_VERSION)
install(TARGETS test_gpuresource DESTINATION unittest)
install(TARGETS test_customized_index DESTINATION unittest)
endif ()
#add_subdirectory(faiss_ori)
#add_subdirectory(faiss_benchmark)
add_subdirectory(test_nsg)

19
core/src/index/unittest/Helper.h
View File

@ -18,13 +18,16 @@
#include <memory>
#include <string>
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#endif
int DEVICEID = 0;
constexpr int64_t DIM = 128;
@ -41,16 +44,18 @@ IndexFactory(const std::string& type) {
return std::make_shared<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
#ifdef MILVUS_GPU_VERSION
} else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(DEVICEID);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(DEVICEID);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(DEVICEID);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
#endif
}
}
@ -110,11 +115,15 @@ class TestGpuIndexBase : public ::testing::Test {
protected:
void
SetUp() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
#endif
}
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
};

64
core/src/index/unittest/faiss_benchmark/CMakeLists.txt
View File

@ -1,24 +1,50 @@
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/include)
include_directories(/usr/local/cuda/include)
include_directories(/usr/local/hdf5/include)
if (KNOWHERE_GPU_VERSION)
link_directories(/usr/local/cuda/lib64)
link_directories(/usr/local/hdf5/lib)
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/include)
include_directories(/usr/local/cuda/include)
include_directories(/usr/local/hdf5/include)
set(unittest_libs
gtest gmock gtest_main gmock_main)
link_directories(/usr/local/cuda/lib64)
link_directories(/usr/local/hdf5/lib)
set(depend_libs
faiss openblas lapack hdf5
arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a
)
set(unittest_libs
gtest gmock gtest_main gmock_main)
set(basic_libs
cudart cublas
gomp gfortran pthread
)
set(depend_libs
faiss hdf5
arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a
)
if (BUILD_FAISS_WITH_MKL)
set(depend_libs ${depend_libs}
"-Wl,--start-group \
${MKL_LIB_PATH}/libmkl_intel_ilp64.a \
${MKL_LIB_PATH}/libmkl_gnu_thread.a \
${MKL_LIB_PATH}/libmkl_core.a \
-Wl,--end-group -lgomp -lpthread -lm -ldl"
)
else ()
set(depend_libs ${depend_libs}
lapack
openblas)
endif ()
add_executable(test_faiss_benchmark faiss_benchmark_test.cpp)
target_link_libraries(test_faiss_benchmark ${depend_libs} ${unittest_libs} ${basic_libs})
install(TARGETS test_faiss_benchmark DESTINATION unittest)
set(basic_libs
gomp gfortran pthread
)
include_directories(${CUDA_INCLUDE_DIRS})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(cuda_lib
cudart
cublas
)
set(basic_libs ${basic_libs}
${cuda_lib}
)
add_executable(test_faiss_benchmark faiss_benchmark_test.cpp)
target_link_libraries(test_faiss_benchmark ${depend_libs} ${unittest_libs} ${basic_libs})
install(TARGETS test_faiss_benchmark DESTINATION unittest)
endif ()

63
core/src/index/unittest/faiss_ori/CMakeLists.txt
View File

@ -1,26 +1,49 @@
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/include)
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
if (KNOWHERE_GPU_VERSION)
set(unittest_libs
gtest gmock gtest_main gmock_main)
include_directories(${INDEX_SOURCE_DIR}/thirdparty)
include_directories(${INDEX_SOURCE_DIR}/include)
set(depend_libs
faiss openblas lapack
arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a
)
set(unittest_libs
gtest gmock gtest_main gmock_main)
set(basic_libs
cudart cublas
gomp gfortran pthread
)
set(depend_libs
faiss
arrow ${ARROW_PREFIX}/lib/libjemalloc_pic.a
)
if (BUILD_FAISS_WITH_MKL)
set(depend_libs ${depend_libs}
"-Wl,--start-group \
${MKL_LIB_PATH}/libmkl_intel_ilp64.a \
${MKL_LIB_PATH}/libmkl_gnu_thread.a \
${MKL_LIB_PATH}/libmkl_core.a \
-Wl,--end-group -lgomp -lpthread -lm -ldl"
)
else ()
set(depend_libs ${depend_libs}
lapack
openblas)
endif ()
set(basic_libs
gomp gfortran pthread
)
#<GPU-TEST>
if(NOT TARGET test_gpu)
add_executable(test_gpu gpuresource_test.cpp)
endif()
target_link_libraries(test_gpu ${depend_libs} ${unittest_libs} ${basic_libs})
include_directories(${CUDA_INCLUDE_DIRS})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(cuda_lib
cudart
cublas
)
set(basic_libs ${basic_libs}
${cuda_lib}
)
install(TARGETS test_gpu DESTINATION unittest)
#<GPU-TEST>
if (NOT TARGET test_gpu)
add_executable(test_gpu gpuresource_test.cpp)
endif ()
target_link_libraries(test_gpu ${depend_libs} ${unittest_libs} ${basic_libs})
install(TARGETS test_gpu DESTINATION unittest)
endif ()

6
core/src/index/unittest/test_idmap.cpp
View File

@ -21,8 +21,10 @@
#include "knowhere/adapter/Structure.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIDMAP.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#endif
#include "Helper.h"
#include "unittest/utils.h"
@ -116,6 +118,7 @@ TEST_F(IDMAPTest, idmap_serialize) {
}
}
#ifdef MILVUS_GPU_VERSION
TEST_F(IDMAPTest, copy_test) {
ASSERT_TRUE(!xb.empty());
@ -175,3 +178,4 @@ TEST_F(IDMAPTest, copy_test) {
AssertAnns(device_result, nq, k);
}
}
#endif

39
core/src/index/unittest/test_ivf.cpp
View File

@ -20,13 +20,24 @@
#include <iostream>
#include <thread>
#ifdef MILVUS_GPU_VERSION
#include <faiss/gpu/GpuIndexIVFFlat.h>
#endif
#include "knowhere/common/Exception.h"
#include "knowhere/common/Timer.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#endif
#include "unittest/Helper.h"
#include "unittest/utils.h"
@ -39,8 +50,9 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
protected:
void
SetUp() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
#endif
ParameterType parameter_type;
std::tie(index_type, parameter_type) = GetParam();
// Init_with_default();
@ -54,7 +66,9 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
protected:
@ -64,15 +78,17 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
};
INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
Values(std::make_tuple("IVF", ParameterType::ivf),
std::make_tuple("GPUIVF", ParameterType::ivf),
std::make_tuple("IVFPQ", ParameterType::ivfpq),
std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std::make_tuple("IVFSQ", ParameterType::ivfsq),
Values(
#ifdef MILVUS_GPU_VERSION
std::make_tuple("GPUIVF", ParameterType::ivf),
std::make_tuple("GPUIVFPQ", ParameterType::ivfpq),
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq),
#ifdef CUSTOMIZATION
std::make_tuple("IVFSQHybrid", ParameterType::ivfsq),
std::make_tuple("IVFSQHybrid", ParameterType::ivfsq),
#endif
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)));
#endif
std::make_tuple("IVF", ParameterType::ivf), std::make_tuple("IVFPQ", ParameterType::ivfpq),
std::make_tuple("IVFSQ", ParameterType::ivfsq)));
TEST_P(IVFTest, ivf_basic) {
assert(!xb.empty());
@ -148,6 +164,7 @@ TEST_P(IVFTest, ivf_serialize) {
}
}
#ifdef MILVUS_GPU_VERSION
TEST_P(IVFTest, clone_test) {
assert(!xb.empty());
@ -238,7 +255,9 @@ TEST_P(IVFTest, clone_test) {
}
}
}
#endif
#ifdef MILVUS_GPU_VERSION
#ifdef CUSTOMIZATION
TEST_P(IVFTest, gpu_seal_test) {
std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
@ -271,5 +290,5 @@ TEST_P(IVFTest, gpu_seal_test) {
auto with_seal = tc.RecordSection("With seal");
ASSERT_GE(without_seal, with_seal);
}
#endif
#endif

7
core/src/index/unittest/test_nsg/CMakeLists.txt
View File

@ -4,14 +4,13 @@
add_definitions(-std=c++11 -O3 -lboost -march=native -Wall -DINFO)
find_package(OpenMP)
find_package(OpenMP REQUIRED)
if (OPENMP_FOUND)
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${OpenMP_C_FLAGS}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${OpenMP_CXX_FLAGS}")
else ()
message(FATAL_ERROR "no OpenMP supprot")
endif ()
message(${OpenMP_CXX_FLAGS})
include_directories(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg)
aux_source_directory(${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/nsg nsg_src)
@ -20,9 +19,9 @@ set(interface_src
${INDEX_SOURCE_DIR}/knowhere/knowhere/index/vector_index/IndexNSG.cpp
)
if(NOT TARGET test_nsg)
if (NOT TARGET test_nsg)
add_executable(test_nsg test_nsg.cpp ${interface_src} ${nsg_src} ${util_srcs} ${ivf_srcs})
endif()
endif ()
target_link_libraries(test_nsg ${depend_libs} ${unittest_libs} ${basic_libs})
##############################

6
core/src/index/unittest/test_nsg/test_nsg.cpp
View File

@ -21,7 +21,9 @@
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/FaissBaseIndex.h"
#include "knowhere/index/vector_index/IndexNSG.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
#include "knowhere/index/vector_index/nsg/NSGIO.h"
#include "unittest/utils.h"
@ -37,7 +39,9 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test {
void
SetUp() override {
// Init_with_default();
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
#endif
Generate(256, 1000000 / 100, 1);
index_ = std::make_shared<knowhere::NSG>();
@ -60,7 +64,9 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test {
void
TearDown() override {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
protected:

9
core/src/main.cpp
View File

@ -25,7 +25,7 @@
#include "external/easyloggingpp/easylogging++.h"
#include "metrics/Metrics.h"
#include "server/Server.h"
#include "src/version.h"
#include "src/config.h"
#include "utils/CommonUtil.h"
#include "utils/SignalUtil.h"
@ -52,6 +52,11 @@ print_banner() {
std::cout << std::endl;
std::cout << "Welcome to Milvus!" << std::endl;
std::cout << "Milvus " << BUILD_TYPE << " version: v" << MILVUS_VERSION << ", built at " << BUILD_TIME << std::endl;
#ifdef MILVUS_CPU_VERSION
std::cout << "You are using Milvus CPU version" << std::endl;
#else
std::cout << "You are using Milvus GPU version" << std::endl;
#endif
std::cout << std::endl;
}
@ -95,7 +100,7 @@ main(int argc, char* argv[]) {
char* log_filename_ptr = strdup(optarg);
log_config_file = log_filename_ptr;
free(log_filename_ptr);
std::cout << "Initial log config from: " << log_config_file << std::endl;
std::cout << "Initializing log config from: " << log_config_file << std::endl;
break;
}
case 'p': {

24
core/src/metrics/SystemInfo.cpp
View File

@ -19,7 +19,6 @@
#include "utils/Log.h"
#include <dirent.h>
#include <nvml.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
@ -29,6 +28,10 @@
#include <string>
#include <utility>
#ifdef MILVUS_GPU_VERSION
#include <nvml.h>
#endif
namespace milvus {
namespace server {
@ -60,6 +63,7 @@ SystemInfo::Init() {
total_ram_ = GetPhysicalMemory();
fclose(file);
#ifdef MILVUS_GPU_VERSION
// initialize GPU information
nvmlReturn_t nvmlresult;
nvmlresult = nvmlInit();
@ -72,6 +76,7 @@ SystemInfo::Init() {
SERVER_LOG_ERROR << "Unable to get devidce number";
return;
}
#endif
// initialize network traffic information
std::pair<uint64_t, uint64_t> in_and_out_octets = Octets();
@ -213,6 +218,9 @@ SystemInfo::GPUMemoryTotal() {
if (!initialized_)
Init();
std::vector<uint64_t> result;
#ifdef MILVUS_GPU_VERSION
nvmlMemory_t nvmlMemory;
for (int i = 0; i < num_device_; ++i) {
nvmlDevice_t device;
@ -220,6 +228,8 @@ SystemInfo::GPUMemoryTotal() {
nvmlDeviceGetMemoryInfo(device, &nvmlMemory);
result.push_back(nvmlMemory.total);
}
#endif
return result;
}
@ -228,6 +238,9 @@ SystemInfo::GPUTemperature() {
if (!initialized_)
Init();
std::vector<uint64_t> result;
#ifdef MILVUS_GPU_VERSION
for (int i = 0; i < num_device_; i++) {
nvmlDevice_t device;
nvmlDeviceGetHandleByIndex(i, &device);
@ -235,6 +248,9 @@ SystemInfo::GPUTemperature() {
nvmlDeviceGetTemperature(device, NVML_TEMPERATURE_GPU, &temp);
result.push_back(temp);
}
#endif
return result;
}
@ -283,6 +299,9 @@ SystemInfo::GPUMemoryUsed() {
Init();
std::vector<uint64_t> result;
#ifdef MILVUS_GPU_VERSION
nvmlMemory_t nvmlMemory;
for (int i = 0; i < num_device_; ++i) {
nvmlDevice_t device;
@ -290,6 +309,9 @@ SystemInfo::GPUMemoryUsed() {
nvmlDeviceGetMemoryInfo(device, &nvmlMemory);
result.push_back(nvmlMemory.used);
}
#endif
return result;
}

31
core/src/scheduler/JobMgr.cpp
View File

@ -104,25 +104,20 @@ JobMgr::build_task(const JobPtr& job) {
void
JobMgr::calculate_path(const TaskPtr& task) {
if (task->type_ == TaskType::SearchTask) {
if (task->label()->Type() != TaskLabelType::SPECIFIED_RESOURCE) {
return;
}
std::vector<std::string> path;
auto spec_label = std::static_pointer_cast<SpecResLabel>(task->label());
auto src = res_mgr_->GetDiskResources()[0];
auto dest = spec_label->resource();
ShortestPath(src.lock(), dest.lock(), res_mgr_, path);
task->path() = Path(path, path.size() - 1);
} else if (task->type_ == TaskType::BuildIndexTask) {
auto spec_label = std::static_pointer_cast<SpecResLabel>(task->label());
auto src = res_mgr_->GetDiskResources()[0];
auto dest = spec_label->resource();
std::vector<std::string> path;
ShortestPath(src.lock(), dest.lock(), res_mgr_, path);
task->path() = Path(path, path.size() - 1);
if (task->type_ != TaskType::SearchTask) {
return;
}
if (task->label()->Type() != TaskLabelType::SPECIFIED_RESOURCE) {
return;
}
std::vector<std::string> path;
auto spec_label = std::static_pointer_cast<SpecResLabel>(task->label());
auto src = res_mgr_->GetDiskResources()[0];
auto dest = spec_label->resource();
ShortestPath(src.lock(), dest.lock(), res_mgr_, path);
task->path() = Path(path, path.size() - 1);
}
} // namespace scheduler

13
core/src/scheduler/SchedInst.cpp
View File

@ -18,7 +18,6 @@
#include "scheduler/SchedInst.h"
#include "ResourceFactory.h"
#include "Utils.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "server/Config.h"
#include <set>
@ -55,8 +54,8 @@ load_simple_config() {
// get resources
auto gpu_ids = get_gpu_pool();
int32_t index_build_device_id;
config.GetResourceConfigIndexBuildDevice(index_build_device_id);
int32_t build_gpu_id;
config.GetResourceConfigIndexBuildDevice(build_gpu_id);
// create and connect
ResMgrInst::GetInstance()->Add(ResourceFactory::Create("disk", "DISK", 0, true, false));
@ -70,15 +69,15 @@ load_simple_config() {
for (auto& gpu_id : gpu_ids) {
ResMgrInst::GetInstance()->Add(ResourceFactory::Create(std::to_string(gpu_id), "GPU", gpu_id, true, true));
ResMgrInst::GetInstance()->Connect("cpu", std::to_string(gpu_id), pcie);
if (index_build_device_id == gpu_id) {
if (build_gpu_id == gpu_id) {
find_build_gpu_id = true;
}
}
if (not find_build_gpu_id && index_build_device_id != server::CPU_DEVICE_ID) {
if (not find_build_gpu_id) {
ResMgrInst::GetInstance()->Add(
ResourceFactory::Create(std::to_string(index_build_device_id), "GPU", index_build_device_id, true, true));
ResMgrInst::GetInstance()->Connect("cpu", std::to_string(index_build_device_id), pcie);
ResourceFactory::Create(std::to_string(build_gpu_id), "GPU", build_gpu_id, true, true));
ResMgrInst::GetInstance()->Connect("cpu", std::to_string(build_gpu_id), pcie);
}
}

1
core/src/scheduler/SchedInst.h
View File

@ -106,6 +106,7 @@ class OptimizerInst {
has_cpu = true;
}
}
std::vector<PassPtr> pass_list;
pass_list.push_back(std::make_shared<LargeSQ8HPass>());
pass_list.push_back(std::make_shared<HybridPass>());

11
core/src/scheduler/TaskCreator.cpp
View File

@ -70,15 +70,8 @@ TaskCreator::Create(const DeleteJobPtr& job) {
std::vector<TaskPtr>
TaskCreator::Create(const BuildIndexJobPtr& job) {
std::vector<TaskPtr> tasks;
server::Config& config = server::Config::GetInstance();
int32_t build_index_id;
Status stat = config.GetResourceConfigIndexBuildDevice(build_index_id);
ResourcePtr res_ptr;
if (build_index_id == server::CPU_DEVICE_ID) {
res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
} else {
res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, build_index_id);
}
// TODO(yukun): remove "disk" hardcode here
ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("disk");
for (auto& to_index_file : job->to_index_files()) {
auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));

4
core/src/scheduler/Utils.cpp
View File

@ -19,7 +19,9 @@
#include "server/Config.h"
#include "utils/Log.h"
#ifdef MILVUS_GPU_VERSION
#include <cuda_runtime.h>
#endif
#include <chrono>
#include <set>
#include <string>
@ -38,7 +40,9 @@ get_current_timestamp() {
uint64_t
get_num_gpu() {
int n_devices = 0;
#ifdef MILVUS_GPU_VERSION
cudaGetDeviceCount(&n_devices);
#endif
return n_devices;
}

102
core/src/scheduler/action/PushTaskToNeighbour.cpp
View File

@ -138,41 +138,73 @@ Action::SpecifiedResourceLabelTaskScheduler(const ResourceMgrPtr& res_mgr, Resou
std::shared_ptr<LoadCompletedEvent> event) {
auto task_item = event->task_table_item_;
auto task = event->task_table_item_->task;
// if (resource->type() == ResourceType::DISK) {
// // step 1: calculate shortest path per resource, from disk to compute resource
// auto compute_resources = res_mgr->GetComputeResources();
// std::vector<std::vector<std::string>> paths;
// std::vector<uint64_t> transport_costs;
// for (auto& res : compute_resources) {
// std::vector<std::string> path;
// uint64_t transport_cost = ShortestPath(resource, res, res_mgr, path);
// transport_costs.push_back(transport_cost);
// paths.emplace_back(path);
// }
// if (task->job_.lock()->type() == JobType::BUILD) {
// // step2: Read device id in config
// // get build index gpu resource
// server::Config& config = server::Config::GetInstance();
// int32_t build_index_gpu;
// Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu);
//
// bool find_gpu_res = false;
// if (res_mgr->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) {
// for (uint64_t i = 0; i < compute_resources.size(); ++i) {
// if (compute_resources[i]->name() ==
// res_mgr->GetResource(ResourceType::GPU, build_index_gpu)->name()) {
// find_gpu_res = true;
// Path task_path(paths[i], paths[i].size() - 1);
// task->path() = task_path;
// break;
// }
// }
// }
// if (not find_gpu_res) {
// task->path() = Path(paths[0], paths[0].size() - 1);
// }
// }
// }
if (resource->type() == ResourceType::DISK) {
// step 1: calculate shortest path per resource, from disk to compute resource
auto compute_resources = res_mgr->GetComputeResources();
std::vector<std::vector<std::string>> paths;
std::vector<uint64_t> transport_costs;
for (auto& res : compute_resources) {
std::vector<std::string> path;
uint64_t transport_cost = ShortestPath(resource, res, res_mgr, path);
transport_costs.push_back(transport_cost);
paths.emplace_back(path);
}
// if (task->job_.lock()->type() == JobType::SEARCH) {
// auto label = task->label();
// auto spec_label = std::static_pointer_cast<SpecResLabel>(label);
// if (spec_label->resource().lock()->type() == ResourceType::CPU) {
// std::vector<std::string> spec_path;
// spec_path.push_back(spec_label->resource().lock()->name());
// spec_path.push_back(resource->name());
// task->path() = Path(spec_path, spec_path.size() - 1);
// } else {
// // step 2: select min cost, cost(resource) = avg_cost * task_to_do + transport_cost
// uint64_t min_cost = std::numeric_limits<uint64_t>::max();
// uint64_t min_cost_idx = 0;
// for (uint64_t i = 0; i < compute_resources.size(); ++i) {
// if (compute_resources[i]->TotalTasks() == 0) {
// min_cost_idx = i;
// break;
// }
// uint64_t cost = compute_resources[i]->TaskAvgCost() *
// compute_resources[i]->NumOfTaskToExec() +
// transport_costs[i];
// if (min_cost > cost) {
// min_cost = cost;
// min_cost_idx = i;
// }
// }
//
// // step 3: set path in task
// Path task_path(paths[min_cost_idx], paths[min_cost_idx].size() - 1);
// task->path() = task_path;
// }
//
// } else
if (task->job_.lock()->type() == JobType::BUILD) {
// step2: Read device id in config
// get build index gpu resource
server::Config& config = server::Config::GetInstance();
int32_t build_index_gpu;
Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu);
bool find_gpu_res = false;
if (res_mgr->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) {
for (uint64_t i = 0; i < compute_resources.size(); ++i) {
if (compute_resources[i]->name() ==
res_mgr->GetResource(ResourceType::GPU, build_index_gpu)->name()) {
find_gpu_res = true;
Path task_path(paths[i], paths[i].size() - 1);
task->path() = task_path;
break;
}
}
}
if (not find_gpu_res) {
task->path() = Path(paths[0], paths[0].size() - 1);
}
}
}
if (resource->name() == task->path().Last()) {
resource->WakeupExecutor();

5
core/src/scheduler/job/BuildIndexJob.cpp
View File

@ -50,10 +50,7 @@ void
BuildIndexJob::BuildIndexDone(size_t to_index_id) {
std::unique_lock<std::mutex> lock(mutex_);
to_index_files_.erase(to_index_id);
if (to_index_files_.empty()) {
cv_.notify_all();
}
cv_.notify_all();
SERVER_LOG_DEBUG << "BuildIndexJob " << id() << " finish index file: " << to_index_id;
}

16
core/src/scheduler/job/SearchJob.cpp
View File

@ -49,21 +49,13 @@ void
SearchJob::SearchDone(size_t index_id) {
std::unique_lock<std::mutex> lock(mutex_);
index_files_.erase(index_id);
if (index_files_.empty()) {
cv_.notify_all();
}
cv_.notify_all();
SERVER_LOG_DEBUG << "SearchJob " << id() << " finish index file: " << index_id;
}
ResultIds&
SearchJob::GetResultIds() {
return result_ids_;
}
ResultDistances&
SearchJob::GetResultDistances() {
return result_distances_;
ResultSet&
SearchJob::GetResult() {
return result_;
}
Status&

22
core/src/scheduler/job/SearchJob.h
View File

@ -29,7 +29,6 @@
#include <vector>
#include "Job.h"
#include "db/Types.h"
#include "db/meta/MetaTypes.h"
namespace milvus {
@ -38,9 +37,9 @@ namespace scheduler {
using engine::meta::TableFileSchemaPtr;
using Id2IndexMap = std::unordered_map<size_t, TableFileSchemaPtr>;
using ResultIds = engine::ResultIds;
using ResultDistances = engine::ResultDistances;
using IdDistPair = std::pair<int64_t, double>;
using Id2DistVec = std::vector<IdDistPair>;
using ResultSet = std::vector<Id2DistVec>;
class SearchJob : public Job {
public:
@ -56,11 +55,8 @@ class SearchJob : public Job {
void
SearchDone(size_t index_id);
ResultIds&
GetResultIds();
ResultDistances&
GetResultDistances();
ResultSet&
GetResult();
Status&
GetStatus();
@ -94,11 +90,6 @@ class SearchJob : public Job {
return index_files_;
}
std::mutex&
mutex() {
return mutex_;
}
private:
uint64_t topk_ = 0;
uint64_t nq_ = 0;
@ -108,8 +99,7 @@ class SearchJob : public Job {
Id2IndexMap index_files_;
// TODO: column-base better ?
ResultIds result_ids_;
ResultDistances result_distances_;
ResultSet result_;
Status status_;
std::mutex mutex_;

2
core/src/scheduler/optimizer/OnlyGPUPass.cpp
View File

@ -46,7 +46,7 @@ OnlyGPUPass::Run(const TaskPtr& task) {
auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
task->label() = label;
specified_gpu_id_ = (specified_gpu_id_ + 1) % gpu_id.size();
specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size();
return true;
}

115
core/src/scheduler/task/SearchTask.cpp
View File

@ -219,11 +219,8 @@ XSearchTask::Execute() {
// step 3: pick up topk result
auto spec_k = index_engine_->Count() < topk ? index_engine_->Count() : topk;
{
std::unique_lock<std::mutex> lock(search_job->mutex());
XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, metric_l2,
search_job->GetResultIds(), search_job->GetResultDistances());
}
XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, metric_l2,
search_job->GetResult());
span = rc.RecordSection(hdr + ", reduce topk");
// search_job->AccumReduceCost(span);
@ -243,69 +240,71 @@ XSearchTask::Execute() {
}
void
XSearchTask::MergeTopkToResultSet(const scheduler::ResultIds& src_ids, const scheduler::ResultDistances& src_distances,
size_t src_k, size_t nq, size_t topk, bool ascending, scheduler::ResultIds& tar_ids,
scheduler::ResultDistances& tar_distances) {
if (src_ids.empty()) {
return;
XSearchTask::MergeTopkToResultSet(const std::vector<int64_t>& input_ids, const std::vector<float>& input_distance,
uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending,
scheduler::ResultSet& result) {
if (result.empty()) {
result.resize(nq);
}
size_t tar_k = tar_ids.size() / nq;
size_t buf_k = std::min(topk, src_k + tar_k);
scheduler::ResultIds buf_ids(nq * buf_k, -1);
scheduler::ResultDistances buf_distances(nq * buf_k, 0.0);
for (uint64_t i = 0; i < nq; i++) {
size_t buf_k_j = 0, src_k_j = 0, tar_k_j = 0;
size_t buf_idx, src_idx, tar_idx;
scheduler::Id2DistVec result_buf;
auto& result_i = result[i];
size_t buf_k_multi_i = buf_k * i;
size_t src_k_multi_i = topk * i;
size_t tar_k_multi_i = tar_k * i;
while (buf_k_j < buf_k && src_k_j < src_k && tar_k_j < tar_k) {
src_idx = src_k_multi_i + src_k_j;
tar_idx = tar_k_multi_i + tar_k_j;
buf_idx = buf_k_multi_i + buf_k_j;
if ((ascending && src_distances[src_idx] < tar_distances[tar_idx]) ||
(!ascending && src_distances[src_idx] > tar_distances[tar_idx])) {
buf_ids[buf_idx] = src_ids[src_idx];
buf_distances[buf_idx] = src_distances[src_idx];
src_k_j++;
} else {
buf_ids[buf_idx] = tar_ids[tar_idx];
buf_distances[buf_idx] = tar_distances[tar_idx];
tar_k_j++;
if (result[i].empty()) {
result_buf.resize(input_k, scheduler::IdDistPair(-1, 0.0));
uint64_t input_k_multi_i = topk * i;
for (auto k = 0; k < input_k; ++k) {
uint64_t idx = input_k_multi_i + k;
auto& result_buf_item = result_buf[k];
result_buf_item.first = input_ids[idx];
result_buf_item.second = input_distance[idx];
}
buf_k_j++;
}
if (buf_k_j < buf_k) {
if (src_k_j < src_k) {
while (buf_k_j < buf_k && src_k_j < src_k) {
buf_idx = buf_k_multi_i + buf_k_j;
src_idx = src_k_multi_i + src_k_j;
buf_ids[buf_idx] = src_ids[src_idx];
buf_distances[buf_idx] = src_distances[src_idx];
src_k_j++;
buf_k_j++;
} else {
size_t tar_size = result_i.size();
uint64_t output_k = std::min(topk, input_k + tar_size);
result_buf.resize(output_k, scheduler::IdDistPair(-1, 0.0));
size_t buf_k = 0, src_k = 0, tar_k = 0;
uint64_t src_idx;
uint64_t input_k_multi_i = topk * i;
while (buf_k < output_k && src_k < input_k && tar_k < tar_size) {
src_idx = input_k_multi_i + src_k;
auto& result_buf_item = result_buf[buf_k];
auto& result_item = result_i[tar_k];
if ((ascending && input_distance[src_idx] < result_item.second) ||
(!ascending && input_distance[src_idx] > result_item.second)) {
result_buf_item.first = input_ids[src_idx];
result_buf_item.second = input_distance[src_idx];
src_k++;
} else {
result_buf_item = result_item;
tar_k++;
}
} else {
while (buf_k_j < buf_k && tar_k_j < tar_k) {
buf_idx = buf_k_multi_i + buf_k_j;
tar_idx = tar_k_multi_i + tar_k_j;
buf_ids[buf_idx] = tar_ids[tar_idx];
buf_distances[buf_idx] = tar_distances[tar_idx];
tar_k_j++;
buf_k_j++;
buf_k++;
}
if (buf_k < output_k) {
if (src_k < input_k) {
while (buf_k < output_k && src_k < input_k) {
src_idx = input_k_multi_i + src_k;
auto& result_buf_item = result_buf[buf_k];
result_buf_item.first = input_ids[src_idx];
result_buf_item.second = input_distance[src_idx];
src_k++;
buf_k++;
}
} else {
while (buf_k < output_k && tar_k < tar_size) {
result_buf[buf_k] = result_i[tar_k];
tar_k++;
buf_k++;
}
}
}
}
result_i.swap(result_buf);
}
tar_ids.swap(buf_ids);
tar_distances.swap(buf_distances);
}
// void

5
core/src/scheduler/task/SearchTask.h
View File

@ -39,9 +39,8 @@ class XSearchTask : public Task {
public:
static void
MergeTopkToResultSet(const scheduler::ResultIds& src_ids, const scheduler::ResultDistances& src_distances,
size_t src_k, size_t nq, size_t topk, bool ascending, scheduler::ResultIds& tar_ids,
scheduler::ResultDistances& tar_distances);
MergeTopkToResultSet(const std::vector<int64_t>& input_ids, const std::vector<float>& input_distance,
uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending, scheduler::ResultSet& result);
// static void
// MergeTopkArray(std::vector<int64_t>& tar_ids, std::vector<float>& tar_distance, uint64_t& tar_input_k,

2
core/src/sdk/grpc/ClientProxy.cpp
View File

@ -17,7 +17,7 @@
#include "sdk/grpc/ClientProxy.h"
#include "grpc/gen-milvus/milvus.grpc.pb.h"
#include "src/version.h"
#include "src/config.h"
#include <memory>
#include <string>

76
core/src/server/Config.cpp
View File

@ -162,6 +162,7 @@ Config::ValidateConfig() {
return s;
}
#ifdef MILVUS_GPU_VERSION
int64_t cache_gpu_cache_capacity;
s = GetCacheConfigGpuCacheCapacity(cache_gpu_cache_capacity);
if (!s.ok()) {
@ -173,6 +174,7 @@ Config::ValidateConfig() {
if (!s.ok()) {
return s;
}
#endif
bool cache_insert_data;
s = GetCacheConfigCacheInsertData(cache_insert_data);
@ -401,7 +403,8 @@ Status
Config::CheckServerConfigDeployMode(const std::string& value) {
if (value != "single" && value != "cluster_readonly" && value != "cluster_writable") {
return Status(SERVER_INVALID_ARGUMENT,
"server_config.deploy_mode is not one of single, cluster_readonly, and cluster_writable.");
"server_config.deploy_mode is not one of "
"single, cluster_readonly, and cluster_writable.");
}
return Status::OK();
}
@ -589,18 +592,15 @@ Config::CheckCacheConfigGpuCacheCapacity(const std::string& value) {
return Status(SERVER_INVALID_ARGUMENT, msg);
} else {
uint64_t gpu_cache_capacity = std::stoi(value) * GB;
int device_id;
Status s = GetResourceConfigIndexBuildDevice(device_id);
int gpu_index;
Status s = GetResourceConfigIndexBuildDevice(gpu_index);
if (!s.ok()) {
return s;
}
if (device_id == server::CPU_DEVICE_ID)
return Status::OK();
size_t gpu_memory;
if (!ValidationUtil::GetGpuMemory(device_id, gpu_memory).ok()) {
std::string msg = "Fail to get GPU memory for GPU device: " + std::to_string(device_id);
if (!ValidationUtil::GetGpuMemory(gpu_index, gpu_memory).ok()) {
std::string msg = "Fail to get GPU memory for GPU device: " + std::to_string(gpu_index);
return Status(SERVER_UNEXPECTED_ERROR, msg);
} else if (gpu_cache_capacity >= gpu_memory) {
std::string msg = "Invalid gpu cache capacity: " + value +
@ -689,21 +689,31 @@ Config::CheckResourceConfigMode(const std::string& value) {
}
Status
CheckGpuDevice(const std::string& value) {
const std::regex pat("gpu(\\d+)");
std::cmatch m;
if (!std::regex_match(value.c_str(), m, pat)) {
std::string msg = "Invalid gpu device: " + value +
". Possible reason: resource_config.search_resources does not match your hardware.";
CheckResource(const std::string& value) {
std::string s = value;
std::transform(s.begin(), s.end(), s.begin(), ::tolower);
#ifdef MILVUS_CPU_VERSION
if (s != "cpu") {
return Status(SERVER_INVALID_ARGUMENT, "Invalid CPU resource: " + s);
}
#else
const std::regex pat("cpu|gpu(\\d+)");
std::smatch m;
if (!std::regex_match(s, m, pat)) {
std::string msg = "Invalid search resource: " + value +
". Possible reason: resource_config.search_resources is not in the format of cpux or gpux";
return Status(SERVER_INVALID_ARGUMENT, msg);
}
int32_t gpu_index = std::stoi(value.substr(3));
if (!ValidationUtil::ValidateGpuIndex(gpu_index).ok()) {
std::string msg = "Invalid gpu device: " + value +
". Possible reason: resource_config.search_resources does not match your hardware.";
return Status(SERVER_INVALID_ARGUMENT, msg);
if (s.compare(0, 3, "gpu") == 0) {
int32_t gpu_index = std::stoi(s.substr(3));
if (!ValidationUtil::ValidateGpuIndex(gpu_index).ok()) {
std::string msg = "Invalid search resource: " + value +
". Possible reason: resource_config.search_resources does not match your hardware.";
return Status(SERVER_INVALID_ARGUMENT, msg);
}
}
#endif
return Status::OK();
}
@ -716,14 +726,10 @@ Config::CheckResourceConfigSearchResources(const std::vector<std::string>& value
return Status(SERVER_INVALID_ARGUMENT, msg);
}
for (auto& device : value) {
if (device == "cpu") {
continue;
}
if (!CheckGpuDevice(device).ok()) {
std::string msg = "Invalid search resource: " + device +
". Possible reason: resource_config.search_resources does not match your hardware.";
return Status(SERVER_INVALID_ARGUMENT, msg);
for (auto& resource : value) {
auto status = CheckResource(resource);
if (!status.ok()) {
return Status(SERVER_INVALID_ARGUMENT, status.message());
}
}
return Status::OK();
@ -731,13 +737,9 @@ Config::CheckResourceConfigSearchResources(const std::vector<std::string>& value
Status
Config::CheckResourceConfigIndexBuildDevice(const std::string& value) {
if (value == "cpu") {
return Status::OK();
}
if (!CheckGpuDevice(value).ok()) {
std::string msg = "Invalid index build device: " + value +
". Possible reason: resource_config.index_build_device does not match your hardware.";
return Status(SERVER_INVALID_ARGUMENT, msg);
auto status = CheckResource(value);
if (!status.ok()) {
return Status(SERVER_INVALID_ARGUMENT, status.message());
}
return Status::OK();
}
@ -1015,10 +1017,10 @@ Config::GetResourceConfigIndexBuildDevice(int32_t& value) {
return s;
}
if (str == "cpu") {
value = CPU_DEVICE_ID;
} else {
if (str != "cpu") {
value = std::stoi(str.substr(3));
} else {
value = -1;
}
return Status::OK();

6
core/src/server/Config.h
View File

@ -93,9 +93,11 @@ static const char* CONFIG_RESOURCE_MODE = "mode";
static const char* CONFIG_RESOURCE_MODE_DEFAULT = "simple";
static const char* CONFIG_RESOURCE_SEARCH_RESOURCES = "search_resources";
static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE = "index_build_device";
#ifdef MILVUS_CPU_VERSION
static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE_DEFAULT = "cpu";
#else
static const char* CONFIG_RESOURCE_INDEX_BUILD_DEVICE_DEFAULT = "gpu0";
const int32_t CPU_DEVICE_ID = -1;
#endif
class Config {
public:

2
core/src/server/Server.cpp
View File

@ -25,7 +25,7 @@
#include "server/DBWrapper.h"
#include "server/Server.h"
#include "server/grpc_impl/GrpcServer.h"
#include "src/version.h"
#include "src/config.h"
#include "utils/Log.h"
#include "utils/LogUtil.h"
#include "utils/SignalUtil.h"

28
core/src/server/grpc_impl/GrpcRequestTask.cpp
View File

@ -28,7 +28,7 @@
#include "scheduler/SchedInst.h"
#include "server/DBWrapper.h"
#include "server/Server.h"
#include "src/version.h"
#include "src/config.h"
#include "utils/CommonUtil.h"
#include "utils/Log.h"
#include "utils/TimeRecorder.h"
@ -637,8 +637,7 @@ SearchTask::OnExecute() {
rc.RecordSection("prepare vector data");
// step 6: search vectors
engine::ResultIds result_ids;
engine::ResultDistances result_distances;
engine::QueryResults results;
auto record_count = (uint64_t)search_param_->query_record_array().size();
#ifdef MILVUS_ENABLE_PROFILING
@ -648,11 +647,11 @@ SearchTask::OnExecute() {
#endif
if (file_id_array_.empty()) {
status = DBWrapper::DB()->Query(table_name_, (size_t)top_k, record_count, nprobe, vec_f.data(), dates,
result_ids, result_distances);
status =
DBWrapper::DB()->Query(table_name_, (size_t)top_k, record_count, nprobe, vec_f.data(), dates, results);
} else {
status = DBWrapper::DB()->Query(table_name_, file_id_array_, (size_t)top_k, record_count, nprobe,
vec_f.data(), dates, result_ids, result_distances);
vec_f.data(), dates, results);
}
#ifdef MILVUS_ENABLE_PROFILING
@ -664,20 +663,23 @@ SearchTask::OnExecute() {
return status;
}
if (result_ids.empty()) {
if (results.empty()) {
return Status::OK(); // empty table
}
size_t result_k = result_ids.size() / record_count;
if (results.size() != record_count) {
std::string msg = "Search " + std::to_string(record_count) + " vectors but only return " +
std::to_string(results.size()) + " results";
return Status(SERVER_ILLEGAL_SEARCH_RESULT, msg);
}
// step 7: construct result array
for (size_t i = 0; i < record_count; i++) {
for (auto& result : results) {
::milvus::grpc::TopKQueryResult* topk_query_result = topk_result_list->add_topk_query_result();
for (size_t j = 0; j < result_k; j++) {
for (auto& pair : result) {
::milvus::grpc::QueryResult* grpc_result = topk_query_result->add_query_result_arrays();
size_t idx = i * result_k + j;
grpc_result->set_id(result_ids[idx]);
grpc_result->set_distance(result_distances[idx]);
grpc_result->set_id(pair.first);
grpc_result->set_distance(pair.second);
}
}

40
core/src/utils/BlockingQueue.inl
View File

@ -15,32 +15,26 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
namespace milvus {
namespace server {
template<typename T>
template <typename T>
void
BlockingQueue<T>::Put(const T &task) {
BlockingQueue<T>::Put(const T& task) {
std::unique_lock<std::mutex> lock(mtx);
full_.wait(lock, [this] {
return (queue_.size() < capacity_);
});
full_.wait(lock, [this] { return (queue_.size() < capacity_); });
queue_.push(task);
empty_.notify_all();
}
template<typename T>
template <typename T>
T
BlockingQueue<T>::Take() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] {
return !queue_.empty();
});
empty_.wait(lock, [this] { return !queue_.empty(); });
T front(queue_.front());
queue_.pop();
@ -48,51 +42,45 @@ BlockingQueue<T>::Take() {
return front;
}
template<typename T>
template <typename T>
size_t
BlockingQueue<T>::Size() {
std::lock_guard<std::mutex> lock(mtx);
return queue_.size();
}
template<typename T>
template <typename T>
T
BlockingQueue<T>::Front() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] {
return !queue_.empty();
});
empty_.wait(lock, [this] { return !queue_.empty(); });
T front(queue_.front());
return front;
}
template<typename T>
template <typename T>
T
BlockingQueue<T>::Back() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] {
return !queue_.empty();
});
empty_.wait(lock, [this] { return !queue_.empty(); });
T back(queue_.back());
return back;
}
template<typename T>
template <typename T>
bool
BlockingQueue<T>::Empty() {
std::unique_lock<std::mutex> lock(mtx);
return queue_.empty();
}
template<typename T>
template <typename T>
void
BlockingQueue<T>::SetCapacity(const size_t capacity) {
capacity_ = (capacity > 0 ? capacity : capacity_);
}
} // namespace server
} // namespace milvus
} // namespace server
} // namespace milvus

11
core/src/utils/ValidationUtil.cpp
View File

@ -20,7 +20,9 @@
#include "db/engine/ExecutionEngine.h"
#include <arpa/inet.h>
#ifdef MILVUS_GPU_VERSION
#include <cuda_runtime.h>
#endif
#include <algorithm>
#include <cmath>
#include <regex>
@ -168,6 +170,7 @@ ValidationUtil::ValidateSearchNprobe(int64_t nprobe, const engine::meta::TableSc
Status
ValidationUtil::ValidateGpuIndex(uint32_t gpu_index) {
#ifdef MILVUS_GPU_VERSION
int num_devices = 0;
auto cuda_err = cudaGetDeviceCount(&num_devices);
if (cuda_err != cudaSuccess) {
@ -181,21 +184,27 @@ ValidationUtil::ValidateGpuIndex(uint32_t gpu_index) {
SERVER_LOG_ERROR << msg;
return Status(SERVER_INVALID_ARGUMENT, msg);
}
#endif
return Status::OK();
}
Status
ValidationUtil::GetGpuMemory(uint32_t gpu_index, size_t& memory) {
#ifdef MILVUS_GPU_VERSION
cudaDeviceProp deviceProp;
auto cuda_err = cudaGetDeviceProperties(&deviceProp, gpu_index);
if (cuda_err) {
std::string msg = "Failed to get gpu properties, cuda error:" + std::to_string(cuda_err);
std::string msg = "Failed to get gpu properties for gpu" + std::to_string(gpu_index) +
" , cuda error:" + std::to_string(cuda_err);
SERVER_LOG_ERROR << msg;
return Status(SERVER_UNEXPECTED_ERROR, msg);
}
memory = deviceProp.totalGlobalMem;
#endif
return Status::OK();
}

2
core/src/wrapper/ConfAdapter.cpp
View File

@ -109,7 +109,7 @@ IVFSQConfAdapter::Match(const TempMetaConf& metaconf) {
conf->nlist = MatchNlist(metaconf.size, metaconf.nlist);
conf->d = metaconf.dim;
conf->metric_type = metaconf.metric_type;
conf->gpu_id = metaconf.gpu_id;
conf->gpu_id = conf->gpu_id;
conf->nbits = 8;
MatchBase(conf);
return conf;

8
core/src/wrapper/KnowhereResource.cpp
View File

@ -16,7 +16,9 @@
// under the License.
#include "wrapper/KnowhereResource.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
#include "server/Config.h"
#include <map>
@ -32,6 +34,8 @@ constexpr int64_t M_BYTE = 1024 * 1024;
Status
KnowhereResource::Initialize() {
#ifdef MILVUS_GPU_VERSION
struct GpuResourceSetting {
int64_t pinned_memory = 300 * M_BYTE;
int64_t temp_memory = 300 * M_BYTE;
@ -73,12 +77,16 @@ KnowhereResource::Initialize() {
iter->second.temp_memory, iter->second.resource_num);
}
#endif
return Status::OK();
}
Status
KnowhereResource::Finalize() {
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource.
#endif
return Status::OK();
}

149
core/src/wrapper/VecImpl.cpp
View File

@ -18,11 +18,16 @@
#include "wrapper/VecImpl.h"
#include "DataTransfer.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "utils/Log.h"
#include "wrapper/WrapperException.h"
#ifdef MILVUS_GPU_VERSION
#include <src/index/knowhere/knowhere/index/vector_index/IndexGPUIVF.h>
#include <src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.h>
#endif
/*
* no parameter check in this layer.
@ -146,19 +151,29 @@ VecIndexImpl::GetType() {
VecIndexPtr
VecIndexImpl::CopyToGpu(const int64_t& device_id, const Config& cfg) {
// TODO(linxj): exception handle
#ifdef MILVUS_GPU_VERSION
auto gpu_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, cfg);
auto new_index = std::make_shared<VecIndexImpl>(gpu_index, ConvertToGpuIndexType(type));
new_index->dim = dim;
return new_index;
#else
WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToGpu when we are using CPU version";
throw WrapperException("Calling VecIndexImpl::CopyToGpu when we are using CPU version");
#endif
}
VecIndexPtr
VecIndexImpl::CopyToCpu(const Config& cfg) {
// TODO(linxj): exception handle
#ifdef MILVUS_GPU_VERSION
auto cpu_index = knowhere::cloner::CopyGpuToCpu(index_, cfg);
auto new_index = std::make_shared<VecIndexImpl>(cpu_index, ConvertToCpuIndexType(type));
new_index->dim = dim;
return new_index;
#else
WRAPPER_LOG_ERROR << "Calling VecIndexImpl::CopyToCpu when we are using CPU version";
throw WrapperException("Calling VecIndexImpl::CopyToCpu when we are using CPU version");
#endif
}
VecIndexPtr
@ -171,11 +186,14 @@ VecIndexImpl::Clone() {
int64_t
VecIndexImpl::GetDeviceId() {
#ifdef MILVUS_GPU_VERSION
if (auto device_idx = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
return device_idx->GetGpuDevice();
}
#else
// else
return -1; // -1 == cpu
#endif
}
float*
@ -226,130 +244,5 @@ BFIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const
return Status::OK();
}
// TODO(linxj): add lock here.
Status
IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt,
const float* xt) {
try {
dim = cfg->d;
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
auto preprocessor = index_->BuildPreprocessor(dataset, cfg);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(dataset, cfg);
index_->set_index_model(model);
index_->Add(dataset, cfg);
if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
auto host_index = device_index->CopyGpuToCpu(Config());
index_ = host_index;
type = ConvertToCpuIndexType(type);
} else {
WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed";
return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
Status
IVFMixIndex::Load(const knowhere::BinarySet& index_binary) {
index_->Load(index_binary);
dim = Dimension();
return Status::OK();
}
knowhere::QuantizerPtr
IVFHybridIndex::LoadQuantizer(const Config& conf) {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
return new_idx->LoadQuantizer(conf);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
}
Status
IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
new_idx->SetQuantizer(q);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
return Status(KNOWHERE_ERROR, "not support");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
Status
IVFHybridIndex::UnsetQuantizer() {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
new_idx->UnsetQuantizer();
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
return Status(KNOWHERE_ERROR, "not support");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
VecIndexPtr
IVFHybridIndex::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
return std::make_shared<IVFHybridIndex>(new_idx->LoadData(q, conf), type);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
}
return nullptr;
}
std::pair<VecIndexPtr, knowhere::QuantizerPtr>
IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) {
try {
// TODO(linxj): Hardcode here
if (auto hybrid_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg);
auto new_idx = std::make_shared<IVFHybridIndex>(pair.first, type);
return std::make_pair(new_idx, pair.second);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
}
return std::make_pair(nullptr, nullptr);
}
} // namespace engine
} // namespace milvus

35
core/src/wrapper/VecImpl.h
View File

@ -77,41 +77,6 @@ class VecIndexImpl : public VecIndex {
std::shared_ptr<knowhere::VectorIndex> index_ = nullptr;
};
class IVFMixIndex : public VecIndexImpl {
public:
explicit IVFMixIndex(std::shared_ptr<knowhere::VectorIndex> index, const IndexType& type)
: VecIndexImpl(std::move(index), type) {
}
Status
BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt,
const float* xt) override;
Status
Load(const knowhere::BinarySet& index_binary) override;
};
class IVFHybridIndex : public IVFMixIndex {
public:
explicit IVFHybridIndex(std::shared_ptr<knowhere::VectorIndex> index, const IndexType& type)
: IVFMixIndex(std::move(index), type) {
}
knowhere::QuantizerPtr
LoadQuantizer(const Config& conf) override;
Status
SetQuantizer(const knowhere::QuantizerPtr& q) override;
Status
UnsetQuantizer() override;
std::pair<VecIndexPtr, knowhere::QuantizerPtr>
CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) override;
VecIndexPtr
LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override;
};
class BFIndex : public VecIndexImpl {
public:
explicit BFIndex(std::shared_ptr<knowhere::VectorIndex> index)

47
core/src/wrapper/VecIndex.cpp
View File

@ -18,19 +18,23 @@
#include "wrapper/VecIndex.h"
#include "VecImpl.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/IndexNSG.h"
#include "utils/Log.h"
#ifdef MILVUS_GPU_VERSION
#include <cuda.h>
#include "knowhere/index/vector_index/IndexGPUIDMAP.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "wrapper/gpu/GPUVecImpl.h"
#endif
namespace milvus {
namespace engine {
@ -119,43 +123,46 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) {
index = std::make_shared<knowhere::IVF>();
break;
}
case IndexType::FAISS_IVFFLAT_GPU: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
break;
}
case IndexType::FAISS_IVFFLAT_MIX: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFFLAT_MIX);
}
case IndexType::FAISS_IVFPQ_CPU: {
index = std::make_shared<knowhere::IVFPQ>();
break;
}
case IndexType::FAISS_IVFPQ_GPU: {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
break;
}
case IndexType::SPTAG_KDT_RNT_CPU: {
index = std::make_shared<knowhere::CPUKDTRNG>();
break;
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
}
case IndexType::FAISS_IVFSQ8_CPU: {
index = std::make_shared<knowhere::IVFSQ>();
break;
}
#ifdef MILVUS_GPU_VERSION
case IndexType::FAISS_IVFFLAT_GPU: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
break;
}
case IndexType::FAISS_IVFPQ_GPU: {
index = std::make_shared<knowhere::GPUIVFPQ>(gpu_device);
break;
}
case IndexType::FAISS_IVFSQ8_MIX: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
}
case IndexType::FAISS_IVFSQ8_GPU: {
index = std::make_shared<knowhere::GPUIVFSQ>(gpu_device);
break;
}
case IndexType::FAISS_IVFFLAT_MIX: {
index = std::make_shared<knowhere::GPUIVF>(gpu_device);
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFFLAT_MIX);
}
#ifdef CUSTOMIZATION
case IndexType::FAISS_IVFSQ8_HYBRID: {
index = std::make_shared<knowhere::IVFSQHybrid>(gpu_device);
return std::make_shared<IVFHybridIndex>(index, IndexType::FAISS_IVFSQ8_HYBRID);
}
#endif
#endif
case IndexType::NSG_MIX: { // TODO(linxj): bug.
index = std::make_shared<knowhere::NSG>(gpu_device);

32
core/src/wrapper/WrapperException.cpp Normal file
View File

@ -0,0 +1,32 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include "wrapper/WrapperException.h"
namespace milvus {
namespace engine {
WrapperException::WrapperException(const std::string& msg) : msg(msg) {
}
const char*
WrapperException::what() const noexcept {
return msg.c_str();
}
} // namespace engine
} // namespace milvus

37
core/src/wrapper/WrapperException.h Normal file
View File

@ -0,0 +1,37 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <exception>
#include <string>
namespace milvus {
namespace engine {
class WrapperException : public std::exception {
public:
explicit WrapperException(const std::string& msg);
const char*
what() const noexcept override;
const std::string msg;
};
} // namespace engine
} // namespace milvus

163
core/src/wrapper/gpu/GPUVecImpl.cpp Normal file
View File

@ -0,0 +1,163 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#include "wrapper/gpu/GPUVecImpl.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/IndexGPUIDMAP.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "src/wrapper/DataTransfer.h"
#include "utils/Log.h"
#include "wrapper/VecImpl.h"
/*
* no parameter check in this layer.
* only responible for index combination
*/
namespace milvus {
namespace engine {
// TODO(linxj): add lock here.
Status
IVFMixIndex::BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt,
const float* xt) {
try {
dim = cfg->d;
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
auto preprocessor = index_->BuildPreprocessor(dataset, cfg);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(dataset, cfg);
index_->set_index_model(model);
index_->Add(dataset, cfg);
if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
auto host_index = device_index->CopyGpuToCpu(Config());
index_ = host_index;
type = ConvertToCpuIndexType(type);
} else {
WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed";
return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
Status
IVFMixIndex::Load(const knowhere::BinarySet& index_binary) {
index_->Load(index_binary);
dim = Dimension();
return Status::OK();
}
knowhere::QuantizerPtr
IVFHybridIndex::LoadQuantizer(const Config& conf) {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
return new_idx->LoadQuantizer(conf);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
}
Status
IVFHybridIndex::SetQuantizer(const knowhere::QuantizerPtr& q) {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
new_idx->SetQuantizer(q);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
return Status(KNOWHERE_ERROR, "not support");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
Status
IVFHybridIndex::UnsetQuantizer() {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
new_idx->UnsetQuantizer();
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
return Status(KNOWHERE_ERROR, "not support");
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_ERROR, e.what());
}
return Status::OK();
}
VecIndexPtr
IVFHybridIndex::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) {
try {
// TODO(linxj): Hardcode here
if (auto new_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
return std::make_shared<IVFHybridIndex>(new_idx->LoadData(q, conf), type);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
}
return nullptr;
}
std::pair<VecIndexPtr, knowhere::QuantizerPtr>
IVFHybridIndex::CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) {
try {
// TODO(linxj): Hardcode here
if (auto hybrid_idx = std::dynamic_pointer_cast<knowhere::IVFSQHybrid>(index_)) {
auto pair = hybrid_idx->CopyCpuToGpuWithQuantizer(device_id, cfg);
auto new_idx = std::make_shared<IVFHybridIndex>(pair.first, type);
return std::make_pair(new_idx, pair.second);
} else {
WRAPPER_LOG_ERROR << "Hybrid mode not support for index type: " << int(type);
}
} catch (knowhere::KnowhereException& e) {
WRAPPER_LOG_ERROR << e.what();
} catch (std::exception& e) {
WRAPPER_LOG_ERROR << e.what();
}
return std::make_pair(nullptr, nullptr);
}
} // namespace engine
} // namespace milvus

67
core/src/wrapper/gpu/GPUVecImpl.h Normal file
View File

@ -0,0 +1,67 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <memory>
#include <utility>
#include "knowhere/index/vector_index/VectorIndex.h"
#include "wrapper/VecImpl.h"
#include "wrapper/VecIndex.h"
namespace milvus {
namespace engine {
class IVFMixIndex : public VecIndexImpl {
public:
explicit IVFMixIndex(std::shared_ptr<knowhere::VectorIndex> index, const IndexType& type)
: VecIndexImpl(std::move(index), type) {
}
Status
BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt,
const float* xt) override;
Status
Load(const knowhere::BinarySet& index_binary) override;
};
class IVFHybridIndex : public IVFMixIndex {
public:
explicit IVFHybridIndex(std::shared_ptr<knowhere::VectorIndex> index, const IndexType& type)
: IVFMixIndex(std::move(index), type) {
}
knowhere::QuantizerPtr
LoadQuantizer(const Config& conf) override;
Status
SetQuantizer(const knowhere::QuantizerPtr& q) override;
Status
UnsetQuantizer() override;
std::pair<VecIndexPtr, knowhere::QuantizerPtr>
CopyToGpuWithQuantizer(const int64_t& device_id, const Config& cfg) override;
VecIndexPtr
LoadData(const knowhere::QuantizerPtr& q, const Config& conf) override;
};
} // namespace engine
} // namespace milvus

1
core/thirdparty/versions.txt vendored
View File

@ -9,5 +9,6 @@ LIBUNWIND_VERSION=1.3.1
GPERFTOOLS_VERSION=2.7
GRPC_VERSION=master
ZLIB_VERSION=v1.2.11
MKL_VERSION=2019.5.281
# vim: set filetype=sh:

6
core/ubuntu_build_deps.sh
View File

@ -1,5 +1,11 @@
#!/bin/bash
wget -P /tmp https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB
sudo apt-key add /tmp/GPG-PUB-KEY-INTEL-SW-PRODUCTS-2019.PUB
sudo sh -c 'echo deb https://apt.repos.intel.com/mkl all main > /etc/apt/sources.list.d/intel-mkl.list'
sudo apt-get -y update && sudo apt-get -y install intel-mkl-gnu-2019.5-281 intel-mkl-core-2019.5-281
sudo apt-get install -y gfortran libmysqlclient-dev mysql-client libcurl4-openssl-dev libboost-system-dev \
libboost-filesystem-dev libboost-serialization-dev libboost-regex-dev

24
core/unittest/CMakeLists.txt
View File

@ -16,18 +16,15 @@
# specific language governing permissions and limitations
# under the License.
#-------------------------------------------------------------------------------
include_directories("${CUDA_TOOLKIT_ROOT_DIR}/include")
foreach(dir ${INDEX_INCLUDE_DIRS})
foreach (dir ${INDEX_INCLUDE_DIRS})
include_directories(${dir})
endforeach()
endforeach ()
include_directories(${MILVUS_SOURCE_DIR})
include_directories(${MILVUS_ENGINE_SRC})
include_directories(${CMAKE_CURRENT_SOURCE_DIR})
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
aux_source_directory(${MILVUS_ENGINE_SRC}/cache cache_files)
aux_source_directory(${MILVUS_ENGINE_SRC}/config config_files)
aux_source_directory(${MILVUS_ENGINE_SRC}/metrics metrics_files)
@ -116,11 +113,22 @@ set(unittest_libs
prometheus-cpp-core
dl
z
${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so
cudart
cublas
)
if (MILVUS_GPU_VERSION)
include_directories("${CUDA_INCLUDE_DIRS}")
link_directories("${CUDA_TOOLKIT_ROOT_DIR}/lib64")
set(unittest_libs ${unittest_libs}
${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs/libnvidia-ml.so
cudart
cublas
)
aux_source_directory(${MILVUS_ENGINE_SRC}/wrapper/gpu wrapper_gpu_files)
set(common_files ${common_files}
${wrapper_gpu_files}
)
endif ()
add_subdirectory(db)
add_subdirectory(wrapper)
add_subdirectory(metrics)

2
core/unittest/db/CMakeLists.txt
View File

@ -20,7 +20,7 @@
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} test_files)
cuda_add_executable(test_db
add_executable(test_db
${common_files}
${test_files}
)

122
core/unittest/db/test_db.cpp
View File

@ -15,25 +15,24 @@
// specific language governing permissions and limitations
// under the License.
#include "db/utils.h"
#include "db/DB.h"
#include "db/DBImpl.h"
#include "db/Constants.h"
#include "db/meta/MetaConsts.h"
#include "db/DBFactory.h"
#include "cache/CpuCacheMgr.h"
#include "utils/CommonUtil.h"
#include "db/Constants.h"
#include "db/DB.h"
#include "db/DBFactory.h"
#include "db/DBImpl.h"
#include "db/meta/MetaConsts.h"
#include "db/utils.h"
#include "server/Config.h"
#include "utils/CommonUtil.h"
#include <gtest/gtest.h>
#include <boost/filesystem.hpp>
#include <thread>
#include <random>
#include <thread>
namespace {
static const char *TABLE_NAME = "test_group";
static const char* TABLE_NAME = "test_group";
static constexpr int64_t TABLE_DIM = 256;
static constexpr int64_t VECTOR_COUNT = 25000;
static constexpr int64_t INSERT_LOOP = 1000;
@ -49,10 +48,10 @@ BuildTableSchema() {
}
void
BuildVectors(int64_t n, std::vector<float> &vectors) {
BuildVectors(int64_t n, std::vector<float>& vectors) {
vectors.clear();
vectors.resize(n * TABLE_DIM);
float *data = vectors.data();
float* data = vectors.data();
for (int i = 0; i < n; i++) {
for (int j = 0; j < TABLE_DIM; j++) data[TABLE_DIM * i + j] = drand48();
data[TABLE_DIM * i] += i / 2000.;
@ -68,16 +67,15 @@ CurrentTmDate(int64_t offset_day = 0) {
tm t;
gmtime_r(&tt, &t);
std::string str = std::to_string(t.tm_year + 1900) + "-" + std::to_string(t.tm_mon + 1)
+ "-" + std::to_string(t.tm_mday);
std::string str =
std::to_string(t.tm_year + 1900) + "-" + std::to_string(t.tm_mon + 1) + "-" + std::to_string(t.tm_mday);
return str;
}
void
ConvertTimeRangeToDBDates(const std::string &start_value,
const std::string &end_value,
std::vector<milvus::engine::meta::DateT> &dates) {
ConvertTimeRangeToDBDates(const std::string& start_value, const std::string& end_value,
std::vector<milvus::engine::meta::DateT>& dates) {
dates.clear();
time_t tt_start, tt_end;
@ -90,8 +88,7 @@ ConvertTimeRangeToDBDates(const std::string &start_value,
return;
}
int64_t days = (tt_end > tt_start) ? (tt_end - tt_start) / DAY_SECONDS : (tt_start - tt_end) /
DAY_SECONDS;
int64_t days = (tt_end > tt_start) ? (tt_end - tt_start) / DAY_SECONDS : (tt_start - tt_end) / DAY_SECONDS;
if (days == 0) {
return;
}
@ -101,13 +98,12 @@ ConvertTimeRangeToDBDates(const std::string &start_value,
tm tm_day;
milvus::server::CommonUtil::ConvertTime(tt_day, tm_day);
int64_t date = tm_day.tm_year * 10000 + tm_day.tm_mon * 100 +
tm_day.tm_mday;//according to db logic
int64_t date = tm_day.tm_year * 10000 + tm_day.tm_mon * 100 + tm_day.tm_mday; // according to db logic
dates.push_back(date);
}
}
} // namespace
} // namespace
TEST_F(DBTest, CONFIG_TEST) {
{
@ -175,8 +171,7 @@ TEST_F(DBTest, DB_TEST) {
BuildVectors(qb, qxb);
std::thread search([&]() {
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
milvus::engine::QueryResults results;
int k = 10;
std::this_thread::sleep_for(std::chrono::seconds(2));
@ -191,17 +186,17 @@ TEST_F(DBTest, DB_TEST) {
prev_count = count;
START_TIMER;
stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), result_ids, result_distances);
stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), results);
ss << "Search " << j << " With Size " << count / milvus::engine::M << " M";
STOP_TIMER(ss.str());
ASSERT_TRUE(stat.ok());
for (auto i = 0; i < qb; ++i) {
ASSERT_EQ(result_ids[i*k], target_ids[i]);
for (auto k = 0; k < qb; ++k) {
ASSERT_EQ(results[k][0].first, target_ids[k]);
ss.str("");
ss << "Result [" << i << "]:";
for (auto t = 0; t < k; t++) {
ss << result_ids[i * k + t] << " ";
ss << "Result [" << k << "]:";
for (auto result : results[k]) {
ss << result.first << " ";
}
/* LOG(DEBUG) << ss.str(); */
}
@ -233,7 +228,7 @@ TEST_F(DBTest, DB_TEST) {
TEST_F(DBTest, SEARCH_TEST) {
std::string config_path(CONFIG_PATH);
config_path += CONFIG_FILE;
milvus::server::Config &config = milvus::server::Config::GetInstance();
milvus::server::Config& config = milvus::server::Config::GetInstance();
milvus::Status s = config.LoadConfigFile(config_path);
milvus::engine::meta::TableSchema table_info = BuildTableSchema();
@ -267,68 +262,64 @@ TEST_F(DBTest, SEARCH_TEST) {
}
// result data
//std::vector<long> nns_gt(k*nq);
// std::vector<long> nns_gt(k*nq);
std::vector<int64_t> nns(k * nq); // nns = nearst neg search
//std::vector<float> dis_gt(k*nq);
// std::vector<float> dis_gt(k*nq);
std::vector<float> dis(k * nq);
// insert data
const int batch_size = 100;
for (int j = 0; j < nb / batch_size; ++j) {
stat = db_->InsertVectors(TABLE_NAME, batch_size, xb.data() + batch_size * j * TABLE_DIM, ids);
if (j == 200) { sleep(1); }
if (j == 200) {
sleep(1);
}
ASSERT_TRUE(stat.ok());
}
milvus::engine::TableIndex index;
index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IDMAP;
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IDMAP;
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
{
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results);
ASSERT_TRUE(stat.ok());
}
{//search by specify index file
{ // search by specify index file
milvus::engine::meta::DatesT dates;
std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results);
ASSERT_TRUE(stat.ok());
}
#ifdef CUSTOMIZATION
//test FAISS_IVFSQ8H optimizer
// test FAISS_IVFSQ8H optimizer
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
{
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results);
ASSERT_TRUE(stat.ok());
}
{
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), result_ids, result_distances);
milvus::engine::QueryResults large_nq_results;
stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), large_nq_results);
ASSERT_TRUE(stat.ok());
}
{//search by specify index file
{ // search by specify index file
milvus::engine::meta::DatesT dates;
std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results);
ASSERT_TRUE(stat.ok());
}
#endif
}
@ -354,8 +345,8 @@ TEST_F(DBTest, PRELOADTABLE_TEST) {
}
milvus::engine::TableIndex index;
index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IDMAP;
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IDMAP;
db_->CreateIndex(TABLE_NAME, index); // wait until build index finish
int64_t prev_cache_usage = milvus::cache::CpuCacheMgr::GetInstance()->CacheUsage();
stat = db_->PreloadTable(TABLE_NAME);
@ -397,12 +388,11 @@ TEST_F(DBTest, SHUTDOWN_TEST) {
ASSERT_FALSE(stat.ok());
milvus::engine::meta::DatesT dates;
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(table_info.table_id_, 1, 1, 1, nullptr, dates, result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(table_info.table_id_, 1, 1, 1, nullptr, dates, results);
ASSERT_FALSE(stat.ok());
std::vector<std::string> file_ids;
stat = db_->Query(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, result_ids, result_distances);
stat = db_->Query(table_info.table_id_, file_ids, 1, 1, 1, nullptr, dates, results);
ASSERT_FALSE(stat.ok());
stat = db_->DeleteTable(table_info.table_id_, dates);
@ -422,12 +412,12 @@ TEST_F(DBTest, INDEX_TEST) {
ASSERT_EQ(vector_ids.size(), nb);
milvus::engine::TableIndex index;
index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IVFSQ8;
index.metric_type_ = (int) milvus::engine::MetricType::IP;
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8;
index.metric_type_ = (int)milvus::engine::MetricType::IP;
stat = db_->CreateIndex(table_info.table_id_, index);
ASSERT_TRUE(stat.ok());
index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IVFFLAT;
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFFLAT;
stat = db_->CreateIndex(table_info.table_id_, index);
ASSERT_TRUE(stat.ok());
@ -456,7 +446,7 @@ TEST_F(DBTest2, ARHIVE_DISK_CHECK) {
stat = db_->AllTables(table_schema_array);
ASSERT_TRUE(stat.ok());
bool bfound = false;
for (auto &schema : table_schema_array) {
for (auto& schema : table_schema_array) {
if (schema.table_id_ == TABLE_NAME) {
bfound = true;
break;

87
core/unittest/db/test_db_mysql.cpp
View File

@ -15,21 +15,21 @@
// specific language governing permissions and limitations
// under the License.
#include "db/utils.h"
#include "db/Constants.h"
#include "db/DB.h"
#include "db/DBImpl.h"
#include "db/Constants.h"
#include "db/meta/MetaConsts.h"
#include "db/utils.h"
#include <gtest/gtest.h>
#include <boost/filesystem.hpp>
#include <thread>
#include <random>
#include <thread>
namespace {
static const char *TABLE_NAME = "test_group";
static const char* TABLE_NAME = "test_group";
static constexpr int64_t TABLE_DIM = 256;
static constexpr int64_t VECTOR_COUNT = 25000;
static constexpr int64_t INSERT_LOOP = 1000;
@ -39,22 +39,22 @@ BuildTableSchema() {
milvus::engine::meta::TableSchema table_info;
table_info.dimension_ = TABLE_DIM;
table_info.table_id_ = TABLE_NAME;
table_info.engine_type_ = (int) milvus::engine::EngineType::FAISS_IDMAP;
table_info.engine_type_ = (int)milvus::engine::EngineType::FAISS_IDMAP;
return table_info;
}
void
BuildVectors(int64_t n, std::vector<float> &vectors) {
BuildVectors(int64_t n, std::vector<float>& vectors) {
vectors.clear();
vectors.resize(n * TABLE_DIM);
float *data = vectors.data();
float* data = vectors.data();
for (int i = 0; i < n; i++) {
for (int j = 0; j < TABLE_DIM; j++) data[TABLE_DIM * i + j] = drand48();
data[TABLE_DIM * i] += i / 2000.;
}
}
} // namespace
} // namespace
TEST_F(MySqlDBTest, DB_TEST) {
milvus::engine::meta::TableSchema table_info = BuildTableSchema();
@ -81,8 +81,7 @@ TEST_F(MySqlDBTest, DB_TEST) {
ASSERT_EQ(target_ids.size(), qb);
std::thread search([&]() {
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
milvus::engine::QueryResults results;
int k = 10;
std::this_thread::sleep_for(std::chrono::seconds(5));
@ -97,25 +96,25 @@ TEST_F(MySqlDBTest, DB_TEST) {
prev_count = count;
START_TIMER;
stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), result_ids, result_distances);
stat = db_->Query(TABLE_NAME, k, qb, 10, qxb.data(), results);
ss << "Search " << j << " With Size " << count / milvus::engine::M << " M";
STOP_TIMER(ss.str());
ASSERT_TRUE(stat.ok());
for (auto i = 0; i < qb; ++i) {
// std::cout << results[k][0].first << " " << target_ids[k] << std::endl;
// ASSERT_EQ(results[k][0].first, target_ids[k]);
for (auto k = 0; k < qb; ++k) {
// std::cout << results[k][0].first << " " << target_ids[k] << std::endl;
// ASSERT_EQ(results[k][0].first, target_ids[k]);
bool exists = false;
for (auto t = 0; t < k; t++) {
if (result_ids[i * k + t] == target_ids[i]) {
for (auto& result : results[k]) {
if (result.first == target_ids[k]) {
exists = true;
}
}
ASSERT_TRUE(exists);
ss.str("");
ss << "Result [" << i << "]:";
for (auto t = 0; t < k; t++) {
ss << result_ids[i * k + t] << " ";
ss << "Result [" << k << "]:";
for (auto result : results[k]) {
ss << result.first << " ";
}
/* LOG(DEBUG) << ss.str(); */
}
@ -129,12 +128,12 @@ TEST_F(MySqlDBTest, DB_TEST) {
int loop = INSERT_LOOP;
for (auto i = 0; i < loop; ++i) {
// if (i==10) {
// db_->InsertVectors(TABLE_NAME, qb, qxb.data(), target_ids);
// ASSERT_EQ(target_ids.size(), qb);
// } else {
// db_->InsertVectors(TABLE_NAME, nb, xb.data(), vector_ids);
// }
// if (i==10) {
// db_->InsertVectors(TABLE_NAME, qb, qxb.data(), target_ids);
// ASSERT_EQ(target_ids.size(), qb);
// } else {
// db_->InsertVectors(TABLE_NAME, nb, xb.data(), vector_ids);
// }
db_->InsertVectors(TABLE_NAME, nb, xb.data(), vector_ids);
std::this_thread::sleep_for(std::chrono::microseconds(1));
}
@ -174,24 +173,25 @@ TEST_F(MySqlDBTest, SEARCH_TEST) {
}
// result data
//std::vector<long> nns_gt(k*nq);
// std::vector<long> nns_gt(k*nq);
std::vector<int64_t> nns(k * nq); // nns = nearst neg search
//std::vector<float> dis_gt(k*nq);
// std::vector<float> dis_gt(k*nq);
std::vector<float> dis(k * nq);
// insert data
const int batch_size = 100;
for (int j = 0; j < nb / batch_size; ++j) {
stat = db_->InsertVectors(TABLE_NAME, batch_size, xb.data() + batch_size * j * TABLE_DIM, ids);
if (j == 200) { sleep(1); }
if (j == 200) {
sleep(1);
}
ASSERT_TRUE(stat.ok());
}
sleep(2); // wait until build index finish
sleep(2); // wait until build index finish
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), result_ids, result_distances);
milvus::engine::QueryResults results;
stat = db_->Query(TABLE_NAME, k, nq, 10, xq.data(), results);
ASSERT_TRUE(stat.ok());
}
@ -203,7 +203,7 @@ TEST_F(MySqlDBTest, ARHIVE_DISK_CHECK) {
stat = db_->AllTables(table_schema_array);
ASSERT_TRUE(stat.ok());
bool bfound = false;
for (auto &schema : table_schema_array) {
for (auto& schema : table_schema_array) {
if (schema.table_id_ == TABLE_NAME) {
bfound = true;
break;
@ -243,7 +243,7 @@ TEST_F(MySqlDBTest, ARHIVE_DISK_CHECK) {
TEST_F(MySqlDBTest, DELETE_TEST) {
milvus::engine::meta::TableSchema table_info = BuildTableSchema();
auto stat = db_->CreateTable(table_info);
// std::cout << stat.ToString() << std::endl;
// std::cout << stat.ToString() << std::endl;
milvus::engine::meta::TableSchema table_info_get;
table_info_get.table_id_ = TABLE_NAME;
@ -269,14 +269,13 @@ TEST_F(MySqlDBTest, DELETE_TEST) {
std::this_thread::sleep_for(std::chrono::microseconds(1));
}
// std::vector<engine::meta::DateT> dates;
// stat = db_->DeleteTable(TABLE_NAME, dates);
//// std::cout << "5 sec start" << std::endl;
// std::this_thread::sleep_for(std::chrono::seconds(5));
//// std::cout << "5 sec finish" << std::endl;
// ASSERT_TRUE(stat.ok());
//
// db_->HasTable(TABLE_NAME, has_table);
// ASSERT_FALSE(has_table);
// std::vector<engine::meta::DateT> dates;
// stat = db_->DeleteTable(TABLE_NAME, dates);
//// std::cout << "5 sec start" << std::endl;
// std::this_thread::sleep_for(std::chrono::seconds(5));
//// std::cout << "5 sec finish" << std::endl;
// ASSERT_TRUE(stat.ok());
//
// db_->HasTable(TABLE_NAME, has_table);
// ASSERT_FALSE(has_table);
}

59
core/unittest/db/test_engine.cpp
View File

@ -26,55 +26,36 @@
TEST_F(EngineTest, FACTORY_TEST) {
{
auto engine_ptr = milvus::engine::EngineFactory::Build(
512,
"/tmp/milvus_index_1",
milvus::engine::EngineType::INVALID,
milvus::engine::MetricType::IP,
1024);
512, "/tmp/milvus_index_1", milvus::engine::EngineType::INVALID, milvus::engine::MetricType::IP, 1024);
ASSERT_TRUE(engine_ptr == nullptr);
}
{
auto engine_ptr = milvus::engine::EngineFactory::Build(
512,
"/tmp/milvus_index_1",
milvus::engine::EngineType::FAISS_IDMAP,
milvus::engine::MetricType::IP,
1024);
512, "/tmp/milvus_index_1", milvus::engine::EngineType::FAISS_IDMAP, milvus::engine::MetricType::IP, 1024);
ASSERT_TRUE(engine_ptr != nullptr);
}
{
auto engine_ptr =
milvus::engine::EngineFactory::Build(512, "/tmp/milvus_index_1", milvus::engine::EngineType::FAISS_IVFFLAT,
milvus::engine::MetricType::IP, 1024);
ASSERT_TRUE(engine_ptr != nullptr);
}
{
auto engine_ptr = milvus::engine::EngineFactory::Build(
512,
"/tmp/milvus_index_1",
milvus::engine::EngineType::FAISS_IVFFLAT,
milvus::engine::MetricType::IP,
1024);
512, "/tmp/milvus_index_1", milvus::engine::EngineType::FAISS_IVFSQ8, milvus::engine::MetricType::IP, 1024);
ASSERT_TRUE(engine_ptr != nullptr);
}
{
auto engine_ptr = milvus::engine::EngineFactory::Build(
512,
"/tmp/milvus_index_1",
milvus::engine::EngineType::FAISS_IVFSQ8,
milvus::engine::MetricType::IP,
1024);
ASSERT_TRUE(engine_ptr != nullptr);
}
{
auto engine_ptr = milvus::engine::EngineFactory::Build(
512,
"/tmp/milvus_index_1",
milvus::engine::EngineType::NSG_MIX,
milvus::engine::MetricType::IP,
1024);
512, "/tmp/milvus_index_1", milvus::engine::EngineType::NSG_MIX, milvus::engine::MetricType::IP, 1024);
ASSERT_TRUE(engine_ptr != nullptr);
}
@ -84,21 +65,17 @@ TEST_F(EngineTest, ENGINE_IMPL_TEST) {
uint16_t dimension = 64;
std::string file_path = "/tmp/milvus_index_1";
auto engine_ptr = milvus::engine::EngineFactory::Build(
dimension,
file_path,
milvus::engine::EngineType::FAISS_IVFFLAT,
milvus::engine::MetricType::IP,
1024);
dimension, file_path, milvus::engine::EngineType::FAISS_IVFFLAT, milvus::engine::MetricType::IP, 1024);
std::vector<float> data;
std::vector<int64_t> ids;
const int row_count = 10000;
data.reserve(row_count*dimension);
data.reserve(row_count * dimension);
ids.reserve(row_count);
for (int64_t i = 0; i < row_count; i++) {
ids.push_back(i);
for (uint16_t k = 0; k < dimension; k++) {
data.push_back(i*dimension + k);
data.push_back(i * dimension + k);
}
}
@ -109,14 +86,14 @@ TEST_F(EngineTest, ENGINE_IMPL_TEST) {
ASSERT_EQ(engine_ptr->Count(), ids.size());
status = engine_ptr->CopyToGpu(0, false);
//ASSERT_TRUE(status.ok());
// ASSERT_TRUE(status.ok());
auto new_engine = engine_ptr->Clone();
ASSERT_EQ(new_engine->Dimension(), dimension);
ASSERT_EQ(new_engine->Count(), ids.size());
status = new_engine->CopyToCpu();
//ASSERT_TRUE(status.ok());
// ASSERT_TRUE(status.ok());
auto engine_build = new_engine->BuildIndex("/tmp/milvus_index_2", milvus::engine::EngineType::FAISS_IVFSQ8);
//ASSERT_TRUE(status.ok());
// ASSERT_TRUE(status.ok());
}

78
core/unittest/db/test_mem.cpp
View File

@ -15,25 +15,24 @@
// specific language governing permissions and limitations
// under the License.
#include "gtest/gtest.h"
#include "db/insert/VectorSource.h"
#include "db/insert/MemTableFile.h"
#include "db/insert/MemTable.h"
#include "db/Constants.h"
#include "db/engine/EngineFactory.h"
#include "db/insert/MemTable.h"
#include "db/insert/MemTableFile.h"
#include "db/insert/VectorSource.h"
#include "db/meta/MetaConsts.h"
#include "metrics/Metrics.h"
#include "db/utils.h"
#include "metrics/Metrics.h"
#include <boost/filesystem.hpp>
#include <thread>
#include <chrono>
#include <cmath>
#include <fstream>
#include <iostream>
#include <cmath>
#include <random>
#include <chrono>
#include <thread>
namespace {
@ -42,8 +41,7 @@ static constexpr int64_t TABLE_DIM = 256;
std::string
GetTableName() {
auto now = std::chrono::system_clock::now();
auto micros = std::chrono::duration_cast<std::chrono::microseconds>(
now.time_since_epoch()).count();
auto micros = std::chrono::duration_cast<std::chrono::microseconds>(now.time_since_epoch()).count();
static std::string table_name = std::to_string(micros);
return table_name;
}
@ -63,11 +61,10 @@ BuildVectors(int64_t n, std::vector<float>& vectors) {
vectors.resize(n * TABLE_DIM);
float* data = vectors.data();
for (int i = 0; i < n; i++) {
for (int j = 0; j < TABLE_DIM; j++)
data[TABLE_DIM * i + j] = drand48();
for (int j = 0; j < TABLE_DIM; j++) data[TABLE_DIM * i + j] = drand48();
}
}
} // namespace
} // namespace
TEST_F(MemManagerTest, VECTOR_SOURCE_TEST) {
milvus::engine::meta::TableSchema table_schema = BuildTableSchema();
@ -86,12 +83,10 @@ TEST_F(MemManagerTest, VECTOR_SOURCE_TEST) {
milvus::engine::VectorSource source(n, vectors.data());
size_t num_vectors_added;
milvus::engine::ExecutionEnginePtr execution_engine_ =
milvus::engine::EngineFactory::Build(table_file_schema.dimension_,
table_file_schema.location_,
(milvus::engine::EngineType)table_file_schema.engine_type_,
(milvus::engine::MetricType)table_file_schema.metric_type_,
table_schema.nlist_);
milvus::engine::ExecutionEnginePtr execution_engine_ = milvus::engine::EngineFactory::Build(
table_file_schema.dimension_, table_file_schema.location_,
(milvus::engine::EngineType)table_file_schema.engine_type_,
(milvus::engine::MetricType)table_file_schema.metric_type_, table_schema.nlist_);
milvus::engine::IDNumbers vector_ids;
status = source.Add(execution_engine_, table_file_schema, 50, num_vectors_added, vector_ids);
@ -129,7 +124,7 @@ TEST_F(MemManagerTest, MEM_TABLE_FILE_TEST) {
status = mem_table_file.Add(source, vector_ids);
ASSERT_TRUE(status.ok());
// std::cout << mem_table_file.GetCurrentMem() << " " << mem_table_file.GetMemLeft() << std::endl;
// std::cout << mem_table_file.GetCurrentMem() << " " << mem_table_file.GetMemLeft() << std::endl;
vector_ids = source->GetVectorIds();
ASSERT_EQ(vector_ids.size(), 100);
@ -141,8 +136,8 @@ TEST_F(MemManagerTest, MEM_TABLE_FILE_TEST) {
std::vector<float> vectors_128M;
BuildVectors(n_max, vectors_128M);
milvus::engine::VectorSourcePtr
source_128M = std::make_shared<milvus::engine::VectorSource>(n_max, vectors_128M.data());
milvus::engine::VectorSourcePtr source_128M =
std::make_shared<milvus::engine::VectorSource>(n_max, vectors_128M.data());
vector_ids.clear();
status = mem_table_file.Add(source_128M, vector_ids);
@ -163,8 +158,8 @@ TEST_F(MemManagerTest, MEM_TABLE_TEST) {
std::vector<float> vectors_100;
BuildVectors(n_100, vectors_100);
milvus::engine::VectorSourcePtr
source_100 = std::make_shared<milvus::engine::VectorSource>(n_100, vectors_100.data());
milvus::engine::VectorSourcePtr source_100 =
std::make_shared<milvus::engine::VectorSource>(n_100, vectors_100.data());
milvus::engine::MemTable mem_table(GetTableName(), impl_, options);
@ -184,8 +179,8 @@ TEST_F(MemManagerTest, MEM_TABLE_TEST) {
BuildVectors(n_max, vectors_128M);
vector_ids.clear();
milvus::engine::VectorSourcePtr
source_128M = std::make_shared<milvus::engine::VectorSource>(n_max, vectors_128M.data());
milvus::engine::VectorSourcePtr source_128M =
std::make_shared<milvus::engine::VectorSource>(n_max, vectors_128M.data());
status = mem_table.Add(source_128M, vector_ids);
ASSERT_TRUE(status.ok());
@ -239,7 +234,7 @@ TEST_F(MemManagerTest2, SERIAL_INSERT_SEARCH_TEST) {
stat = db_->InsertVectors(GetTableName(), nb, xb.data(), vector_ids);
ASSERT_TRUE(stat.ok());
std::this_thread::sleep_for(std::chrono::seconds(3));//ensure raw data write to disk
std::this_thread::sleep_for(std::chrono::seconds(3)); // ensure raw data write to disk
std::random_device rd;
std::mt19937 gen(rd());
@ -259,11 +254,10 @@ TEST_F(MemManagerTest2, SERIAL_INSERT_SEARCH_TEST) {
int topk = 10, nprobe = 10;
for (auto& pair : search_vectors) {
auto& search = pair.second;
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
stat = db_->Query(GetTableName(), topk, 1, nprobe, search.data(), result_ids, result_distances);
ASSERT_EQ(result_ids[0], pair.first);
ASSERT_LT(result_distances[0], 1e-4);
milvus::engine::QueryResults results;
stat = db_->Query(GetTableName(), topk, 1, nprobe, search.data(), results);
ASSERT_EQ(results[0][0].first, pair.first);
ASSERT_LT(results[0][0].second, 1e-4);
}
}
@ -315,8 +309,7 @@ TEST_F(MemManagerTest2, CONCURRENT_INSERT_SEARCH_TEST) {
BuildVectors(qb, qxb);
std::thread search([&]() {
milvus::engine::ResultIds result_ids;
milvus::engine::ResultDistances result_distances;
milvus::engine::QueryResults results;
int k = 10;
std::this_thread::sleep_for(std::chrono::seconds(2));
@ -331,17 +324,17 @@ TEST_F(MemManagerTest2, CONCURRENT_INSERT_SEARCH_TEST) {
prev_count = count;
START_TIMER;
stat = db_->Query(GetTableName(), k, qb, 10, qxb.data(), result_ids, result_distances);
stat = db_->Query(GetTableName(), k, qb, 10, qxb.data(), results);
ss << "Search " << j << " With Size " << count / milvus::engine::M << " M";
STOP_TIMER(ss.str());
ASSERT_TRUE(stat.ok());
for (auto i = 0; i < qb; ++i) {
ASSERT_EQ(result_ids[i * k], target_ids[i]);
for (auto k = 0; k < qb; ++k) {
ASSERT_EQ(results[k][0].first, target_ids[k]);
ss.str("");
ss << "Result [" << i << "]:";
for (auto t = 0; t < k; t++) {
ss << result_ids[i * k + t] << " ";
ss << "Result [" << k << "]:";
for (auto result : results[k]) {
ss << result.first << " ";
}
/* LOG(DEBUG) << ss.str(); */
}
@ -402,7 +395,7 @@ TEST_F(MemManagerTest2, VECTOR_IDS_TEST) {
ASSERT_EQ(vector_ids[0], nb);
ASSERT_TRUE(stat.ok());
nb = 262144; //512M
nb = 262144; // 512M
xb.clear();
BuildVectors(nb, xb);
vector_ids.clear();
@ -414,7 +407,7 @@ TEST_F(MemManagerTest2, VECTOR_IDS_TEST) {
ASSERT_EQ(vector_ids[0], nb / 2);
ASSERT_TRUE(stat.ok());
nb = 65536; //128M
nb = 65536; // 128M
xb.clear();
BuildVectors(nb, xb);
vector_ids.clear();
@ -434,4 +427,3 @@ TEST_F(MemManagerTest2, VECTOR_IDS_TEST) {
ASSERT_EQ(vector_ids[i], i + nb);
}
}

35
core/unittest/db/test_meta.cpp
View File

@ -15,16 +15,16 @@
// specific language governing permissions and limitations
// under the License.
#include "db/utils.h"
#include "db/meta/SqliteMetaImpl.h"
#include "db/Utils.h"
#include "db/Constants.h"
#include "db/Utils.h"
#include "db/meta/MetaConsts.h"
#include "db/meta/SqliteMetaImpl.h"
#include "db/utils.h"
#include <gtest/gtest.h>
#include <thread>
#include <stdlib.h>
#include <time.h>
#include <thread>
TEST_F(MetaTest, TABLE_TEST) {
auto table_id = "meta_test_table";
@ -155,7 +155,7 @@ TEST_F(MetaTest, ARCHIVE_TEST_DAYS) {
status = impl.GetTableFiles(table_file.table_id_, ids, files_get);
ASSERT_TRUE(status.ok());
for (auto &file : files_get) {
for (auto& file : files_get) {
if (days[i] < days_num) {
ASSERT_EQ(file.file_type_, milvus::engine::meta::TableFileSchema::NEW);
}
@ -200,7 +200,7 @@ TEST_F(MetaTest, ARCHIVE_TEST_DISK) {
status = impl.GetTableFiles(table_file.table_id_, ids, files_get);
ASSERT_TRUE(status.ok());
for (auto &file : files_get) {
for (auto& file : files_get) {
if (i >= 5) {
ASSERT_EQ(file.file_type_, milvus::engine::meta::TableFileSchema::NEW);
}
@ -293,16 +293,13 @@ TEST_F(MetaTest, TABLE_FILES_TEST) {
milvus::engine::meta::DatesT dates = {table_file.date_};
std::vector<size_t> ids;
status = impl_->FilesToSearch(table_id, ids, dates, dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
status = impl_->FilesToSearch(table_id, ids, milvus::engine::meta::DatesT(), dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
status = impl_->FilesToSearch(table_id, ids, milvus::engine::meta::DatesT(), dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
ids.push_back(size_t(9999999999));
status = impl_->FilesToSearch(table_id, ids, dates, dated_files);
@ -315,19 +312,15 @@ TEST_F(MetaTest, TABLE_FILES_TEST) {
ASSERT_FALSE(status.ok());
file_types = {
milvus::engine::meta::TableFileSchema::NEW,
milvus::engine::meta::TableFileSchema::NEW_MERGE,
milvus::engine::meta::TableFileSchema::NEW_INDEX,
milvus::engine::meta::TableFileSchema::TO_INDEX,
milvus::engine::meta::TableFileSchema::INDEX,
milvus::engine::meta::TableFileSchema::RAW,
milvus::engine::meta::TableFileSchema::NEW, milvus::engine::meta::TableFileSchema::NEW_MERGE,
milvus::engine::meta::TableFileSchema::NEW_INDEX, milvus::engine::meta::TableFileSchema::TO_INDEX,
milvus::engine::meta::TableFileSchema::INDEX, milvus::engine::meta::TableFileSchema::RAW,
milvus::engine::meta::TableFileSchema::BACKUP,
};
status = impl_->FilesByType(table.table_id_, file_types, file_ids);
ASSERT_TRUE(status.ok());
uint64_t total_cnt = new_index_files_cnt + new_merge_files_cnt +
backup_files_cnt + new_files_cnt + raw_files_cnt +
to_index_files_cnt + index_files_cnt;
uint64_t total_cnt = new_index_files_cnt + new_merge_files_cnt + backup_files_cnt + new_files_cnt + raw_files_cnt +
to_index_files_cnt + index_files_cnt;
ASSERT_EQ(file_ids.size(), total_cnt);
status = impl_->DeleteTableFiles(table_id);

47
core/unittest/db/test_meta_mysql.cpp
View File

@ -15,17 +15,17 @@
// specific language governing permissions and limitations
// under the License.
#include "db/utils.h"
#include "db/meta/MySQLMetaImpl.h"
#include "db/Utils.h"
#include "db/meta/MetaConsts.h"
#include "db/meta/MySQLMetaImpl.h"
#include "db/utils.h"
#include <iostream>
#include <thread>
#include <stdlib.h>
#include <time.h>
#include <gtest/gtest.h>
#include <mysql++/mysql++.h>
#include <stdlib.h>
#include <time.h>
#include <iostream>
#include <thread>
TEST_F(MySqlMetaTest, TABLE_TEST) {
auto table_id = "meta_test_table";
@ -52,7 +52,7 @@ TEST_F(MySqlMetaTest, TABLE_TEST) {
table.table_id_ = "";
status = impl_->CreateTable(table);
// ASSERT_TRUE(status.ok());
// ASSERT_TRUE(status.ok());
status = impl_->DropAll();
ASSERT_TRUE(status.ok());
@ -79,8 +79,8 @@ TEST_F(MySqlMetaTest, TABLE_FILE_TEST) {
uint64_t cnt = 0;
status = impl_->Count(table_id, cnt);
// ASSERT_TRUE(status.ok());
// ASSERT_EQ(cnt, 0UL);
// ASSERT_TRUE(status.ok());
// ASSERT_EQ(cnt, 0UL);
auto file_id = table_file.file_id_;
@ -159,7 +159,7 @@ TEST_F(MySqlMetaTest, ARCHIVE_TEST_DAYS) {
status = impl.GetTableFiles(table_file.table_id_, ids, files_get);
ASSERT_TRUE(status.ok());
for (auto &file : files_get) {
for (auto& file : files_get) {
if (days[i] < days_num) {
ASSERT_EQ(file.file_type_, milvus::engine::meta::TableFileSchema::NEW);
}
@ -167,7 +167,7 @@ TEST_F(MySqlMetaTest, ARCHIVE_TEST_DAYS) {
}
std::vector<int> file_types = {
(int) milvus::engine::meta::TableFileSchema::NEW,
(int)milvus::engine::meta::TableFileSchema::NEW,
};
std::vector<std::string> file_ids;
status = impl.FilesByType(table_id, file_types, file_ids);
@ -219,7 +219,7 @@ TEST_F(MySqlMetaTest, ARCHIVE_TEST_DISK) {
status = impl.GetTableFiles(table_file.table_id_, ids, files_get);
ASSERT_TRUE(status.ok());
for (auto &file : files_get) {
for (auto& file : files_get) {
if (i >= 5) {
ASSERT_EQ(file.file_type_, milvus::engine::meta::TableFileSchema::NEW);
}
@ -313,16 +313,13 @@ TEST_F(MySqlMetaTest, TABLE_FILES_TEST) {
milvus::engine::meta::DatesT dates = {table_file.date_};
std::vector<size_t> ids;
status = impl_->FilesToSearch(table_id, ids, dates, dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
status = impl_->FilesToSearch(table_id, ids, milvus::engine::meta::DatesT(), dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
status = impl_->FilesToSearch(table_id, ids, milvus::engine::meta::DatesT(), dated_files);
ASSERT_EQ(dated_files[table_file.date_].size(),
to_index_files_cnt + raw_files_cnt + index_files_cnt);
ASSERT_EQ(dated_files[table_file.date_].size(), to_index_files_cnt + raw_files_cnt + index_files_cnt);
ids.push_back(size_t(9999999999));
status = impl_->FilesToSearch(table_id, ids, dates, dated_files);
@ -335,19 +332,15 @@ TEST_F(MySqlMetaTest, TABLE_FILES_TEST) {
ASSERT_FALSE(status.ok());
file_types = {
milvus::engine::meta::TableFileSchema::NEW,
milvus::engine::meta::TableFileSchema::NEW_MERGE,
milvus::engine::meta::TableFileSchema::NEW_INDEX,
milvus::engine::meta::TableFileSchema::TO_INDEX,
milvus::engine::meta::TableFileSchema::INDEX,
milvus::engine::meta::TableFileSchema::RAW,
milvus::engine::meta::TableFileSchema::NEW, milvus::engine::meta::TableFileSchema::NEW_MERGE,
milvus::engine::meta::TableFileSchema::NEW_INDEX, milvus::engine::meta::TableFileSchema::TO_INDEX,
milvus::engine::meta::TableFileSchema::INDEX, milvus::engine::meta::TableFileSchema::RAW,
milvus::engine::meta::TableFileSchema::BACKUP,
};
status = impl_->FilesByType(table.table_id_, file_types, file_ids);
ASSERT_TRUE(status.ok());
uint64_t total_cnt = new_index_files_cnt + new_merge_files_cnt +
backup_files_cnt + new_files_cnt + raw_files_cnt +
to_index_files_cnt + index_files_cnt;
uint64_t total_cnt = new_index_files_cnt + new_merge_files_cnt + backup_files_cnt + new_files_cnt + raw_files_cnt +
to_index_files_cnt + index_files_cnt;
ASSERT_EQ(file_ids.size(), total_cnt);
status = impl_->DeleteTableFiles(table_id);

31
core/unittest/db/test_misc.cpp
View File

@ -16,15 +16,15 @@
// under the License.
#include "db/Options.h"
#include "db/meta/SqliteMetaImpl.h"
#include "db/engine/EngineFactory.h"
#include "db/Utils.h"
#include "utils/Status.h"
#include "db/engine/EngineFactory.h"
#include "db/meta/SqliteMetaImpl.h"
#include "utils/Exception.h"
#include "utils/Status.h"
#include <gtest/gtest.h>
#include <thread>
#include <boost/filesystem.hpp>
#include <thread>
#include <vector>
TEST(DBMiscTest, EXCEPTION_TEST) {
@ -40,7 +40,7 @@ TEST(DBMiscTest, EXCEPTION_TEST) {
TEST(DBMiscTest, OPTIONS_TEST) {
try {
milvus::engine::ArchiveConf archive("$$##");
} catch (std::exception &ex) {
} catch (std::exception& ex) {
ASSERT_TRUE(true);
}
@ -61,10 +61,7 @@ TEST(DBMiscTest, OPTIONS_TEST) {
{
milvus::engine::ArchiveConf archive("delete");
milvus::engine::ArchiveConf::CriteriaT criterial = {
{"disk", 1024},
{"days", 100}
};
milvus::engine::ArchiveConf::CriteriaT criterial = {{"disk", 1024}, {"days", 100}};
archive.SetCriterias(criterial);
auto crit = archive.GetCriterias();
@ -95,17 +92,17 @@ TEST(DBMiscTest, UTILS_TEST) {
auto status = milvus::engine::utils::CreateTablePath(options, TABLE_NAME);
ASSERT_TRUE(status.ok());
ASSERT_TRUE(boost::filesystem::exists(options.path_));
for (auto &path : options.slave_paths_) {
for (auto& path : options.slave_paths_) {
ASSERT_TRUE(boost::filesystem::exists(path));
}
// options.slave_paths.push_back("/");
// status = engine::utils::CreateTablePath(options, TABLE_NAME);
// ASSERT_FALSE(status.ok());
//
// options.path = "/";
// status = engine::utils::CreateTablePath(options, TABLE_NAME);
// ASSERT_FALSE(status.ok());
// options.slave_paths.push_back("/");
// status = engine::utils::CreateTablePath(options, TABLE_NAME);
// ASSERT_FALSE(status.ok());
//
// options.path = "/";
// status = engine::utils::CreateTablePath(options, TABLE_NAME);
// ASSERT_FALSE(status.ok());
milvus::engine::meta::TableFileSchema file;
file.id_ = 50;

396
core/unittest/db/test_search.cpp
View File

@ -19,97 +19,73 @@
#include <cmath>
#include <vector>
#include "scheduler/job/SearchJob.h"
#include "scheduler/task/SearchTask.h"
#include "utils/TimeRecorder.h"
#include "utils/ThreadPool.h"
#include "utils/TimeRecorder.h"
namespace {
namespace ms = milvus::scheduler;
void
BuildResult(ms::ResultIds& output_ids,
ms::ResultDistances & output_distances,
size_t input_k,
size_t topk,
size_t nq,
bool ascending) {
BuildResult(std::vector<int64_t>& output_ids, std::vector<float>& output_distance, uint64_t input_k, uint64_t topk,
uint64_t nq, bool ascending) {
output_ids.clear();
output_ids.resize(nq * topk);
output_distances.clear();
output_distances.resize(nq * topk);
output_distance.clear();
output_distance.resize(nq * topk);
for (size_t i = 0; i < nq; i++) {
//insert valid items
for (size_t j = 0; j < input_k; j++) {
for (uint64_t i = 0; i < nq; i++) {
// insert valid items
for (uint64_t j = 0; j < input_k; j++) {
output_ids[i * topk + j] = (int64_t)(drand48() * 100000);
output_distances[i * topk + j] = ascending ? (j + drand48()) : ((input_k - j) + drand48());
output_distance[i * topk + j] = ascending ? (j + drand48()) : ((input_k - j) + drand48());
}
//insert invalid items
for (size_t j = input_k; j < topk; j++) {
// insert invalid items
for (uint64_t j = input_k; j < topk; j++) {
output_ids[i * topk + j] = -1;
output_distances[i * topk + j] = -1.0;
output_distance[i * topk + j] = -1.0;
}
}
}
void
CopyResult(ms::ResultIds& output_ids,
ms::ResultDistances& output_distances,
size_t output_topk,
ms::ResultIds& input_ids,
ms::ResultDistances& input_distances,
size_t input_topk,
size_t nq) {
CopyResult(std::vector<int64_t>& output_ids, std::vector<float>& output_distance, uint64_t output_topk,
std::vector<int64_t>& input_ids, std::vector<float>& input_distance, uint64_t input_topk, uint64_t nq) {
ASSERT_TRUE(input_ids.size() >= nq * input_topk);
ASSERT_TRUE(input_distances.size() >= nq * input_topk);
ASSERT_TRUE(input_distance.size() >= nq * input_topk);
ASSERT_TRUE(output_topk <= input_topk);
output_ids.clear();
output_ids.resize(nq * output_topk);
output_distances.clear();
output_distances.resize(nq * output_topk);
output_distance.clear();
output_distance.resize(nq * output_topk);
for (size_t i = 0; i < nq; i++) {
for (size_t j = 0; j < output_topk; j++) {
for (uint64_t i = 0; i < nq; i++) {
for (uint64_t j = 0; j < output_topk; j++) {
output_ids[i * output_topk + j] = input_ids[i * input_topk + j];
output_distances[i * output_topk + j] = input_distances[i * input_topk + j];
output_distance[i * output_topk + j] = input_distance[i * input_topk + j];
}
}
}
void
CheckTopkResult(const ms::ResultIds& input_ids_1,
const ms::ResultDistances& input_distances_1,
size_t input_k_1,
const ms::ResultIds& input_ids_2,
const ms::ResultDistances& input_distances_2,
size_t input_k_2,
size_t topk,
size_t nq,
bool ascending,
const ms::ResultIds& result_ids,
const ms::ResultDistances& result_distances) {
ASSERT_EQ(result_ids.size(), result_distances.size());
ASSERT_EQ(input_ids_1.size(), input_distances_1.size());
ASSERT_EQ(input_ids_2.size(), input_distances_2.size());
CheckTopkResult(const std::vector<int64_t>& input_ids_1, const std::vector<float>& input_distance_1,
const std::vector<int64_t>& input_ids_2, const std::vector<float>& input_distance_2, uint64_t topk,
uint64_t nq, bool ascending, const milvus::scheduler::ResultSet& result) {
ASSERT_EQ(result.size(), nq);
ASSERT_EQ(input_ids_1.size(), input_distance_1.size());
ASSERT_EQ(input_ids_2.size(), input_distance_2.size());
size_t result_k = result_distances.size() / nq;
ASSERT_EQ(result_k, std::min(topk, input_k_1 + input_k_2));
for (size_t i = 0; i < nq; i++) {
std::vector<float>
src_vec(input_distances_1.begin() + i * topk, input_distances_1.begin() + (i + 1) * topk);
src_vec.insert(src_vec.end(),
input_distances_2.begin() + i * topk,
input_distances_2.begin() + (i + 1) * topk);
for (int64_t i = 0; i < nq; i++) {
std::vector<float> src_vec(input_distance_1.begin() + i * topk, input_distance_1.begin() + (i + 1) * topk);
src_vec.insert(src_vec.end(), input_distance_2.begin() + i * topk, input_distance_2.begin() + (i + 1) * topk);
if (ascending) {
std::sort(src_vec.begin(), src_vec.end());
} else {
std::sort(src_vec.begin(), src_vec.end(), std::greater<float>());
}
//erase invalid items
// erase invalid items
std::vector<float>::iterator iter;
for (iter = src_vec.begin(); iter != src_vec.end();) {
if (*iter < 0.0)
@ -118,38 +94,36 @@ CheckTopkResult(const ms::ResultIds& input_ids_1,
++iter;
}
size_t n = std::min(topk, result_ids.size() / nq);
for (size_t j = 0; j < n; j++) {
size_t idx = i * n + j;
if (result_ids[idx] < 0) {
uint64_t n = std::min(topk, result[i].size());
for (uint64_t j = 0; j < n; j++) {
if (result[i][j].first < 0) {
continue;
}
if (src_vec[j] != result_distances[idx]) {
std::cout << src_vec[j] << " " << result_distances[idx] << std::endl;
if (src_vec[j] != result[i][j].second) {
std::cout << src_vec[j] << " " << result[i][j].second << std::endl;
}
ASSERT_TRUE(src_vec[j] == result_distances[idx]);
ASSERT_TRUE(src_vec[j] == result[i][j].second);
}
}
}
} // namespace
} // namespace
void
MergeTopkToResultSetTest(size_t topk_1, size_t topk_2, size_t nq, size_t topk, bool ascending) {
ms::ResultIds ids1, ids2;
ms::ResultDistances dist1, dist2;
ms::ResultIds result_ids;
ms::ResultDistances result_distances;
MergeTopkToResultSetTest(uint64_t topk_1, uint64_t topk_2, uint64_t nq, uint64_t topk, bool ascending) {
std::vector<int64_t> ids1, ids2;
std::vector<float> dist1, dist2;
ms::ResultSet result;
BuildResult(ids1, dist1, topk_1, topk, nq, ascending);
BuildResult(ids2, dist2, topk_2, topk, nq, ascending);
ms::XSearchTask::MergeTopkToResultSet(ids1, dist1, topk_1, nq, topk, ascending, result_ids, result_distances);
ms::XSearchTask::MergeTopkToResultSet(ids2, dist2, topk_2, nq, topk, ascending, result_ids, result_distances);
CheckTopkResult(ids1, dist1, topk_1, ids2, dist2, topk_2, topk, nq, ascending, result_ids, result_distances);
ms::XSearchTask::MergeTopkToResultSet(ids1, dist1, topk_1, nq, topk, ascending, result);
ms::XSearchTask::MergeTopkToResultSet(ids2, dist2, topk_2, nq, topk, ascending, result);
CheckTopkResult(ids1, dist1, ids2, dist2, topk, nq, ascending, result);
}
TEST(DBSearchTest, MERGE_RESULT_SET_TEST) {
size_t NQ = 15;
size_t TOP_K = 64;
uint64_t NQ = 15;
uint64_t TOP_K = 64;
/* test1, id1/dist1 valid, id2/dist2 empty */
MergeTopkToResultSetTest(TOP_K, 0, NQ, TOP_K, true);
@ -168,21 +142,21 @@ TEST(DBSearchTest, MERGE_RESULT_SET_TEST) {
MergeTopkToResultSetTest(TOP_K / 2, TOP_K / 3, NQ, TOP_K, false);
}
//void MergeTopkArrayTest(size_t topk_1, size_t topk_2, size_t nq, size_t topk, bool ascending) {
// void MergeTopkArrayTest(uint64_t topk_1, uint64_t topk_2, uint64_t nq, uint64_t topk, bool ascending) {
// std::vector<int64_t> ids1, ids2;
// std::vector<float> dist1, dist2;
// ms::ResultSet result;
// BuildResult(ids1, dist1, topk_1, topk, nq, ascending);
// BuildResult(ids2, dist2, topk_2, topk, nq, ascending);
// size_t result_topk = std::min(topk, topk_1 + topk_2);
// uint64_t result_topk = std::min(topk, topk_1 + topk_2);
// ms::XSearchTask::MergeTopkArray(ids1, dist1, topk_1, ids2, dist2, topk_2, nq, topk, ascending);
// if (ids1.size() != result_topk * nq) {
// std::cout << ids1.size() << " " << result_topk * nq << std::endl;
// }
// ASSERT_TRUE(ids1.size() == result_topk * nq);
// ASSERT_TRUE(dist1.size() == result_topk * nq);
// for (size_t i = 0; i < nq; i++) {
// for (size_t k = 1; k < result_topk; k++) {
// for (uint64_t i = 0; i < nq; i++) {
// for (uint64_t k = 1; k < result_topk; k++) {
// float f0 = dist1[i * topk + k - 1];
// float f1 = dist1[i * topk + k];
// if (ascending) {
@ -200,9 +174,9 @@ TEST(DBSearchTest, MERGE_RESULT_SET_TEST) {
// }
//}
//TEST(DBSearchTest, MERGE_ARRAY_TEST) {
// size_t NQ = 15;
// size_t TOP_K = 64;
// TEST(DBSearchTest, MERGE_ARRAY_TEST) {
// uint64_t NQ = 15;
// uint64_t TOP_K = 64;
//
// /* test1, id1/dist1 valid, id2/dist2 empty */
// MergeTopkArrayTest(TOP_K, 0, NQ, TOP_K, true);
@ -228,26 +202,26 @@ TEST(DBSearchTest, MERGE_RESULT_SET_TEST) {
//}
TEST(DBSearchTest, REDUCE_PERF_TEST) {
int32_t index_file_num = 478; /* sift1B dataset, index files num */
int32_t index_file_num = 478; /* sift1B dataset, index files num */
bool ascending = true;
std::vector<size_t> thread_vec = {4};
std::vector<size_t> nq_vec = {1000};
std::vector<size_t> topk_vec = {64};
size_t NQ = nq_vec[nq_vec.size() - 1];
size_t TOPK = topk_vec[topk_vec.size() - 1];
std::vector<int32_t> thread_vec = {4, 8};
std::vector<int32_t> nq_vec = {1, 10, 100};
std::vector<int32_t> topk_vec = {1, 4, 16, 64};
int32_t NQ = nq_vec[nq_vec.size() - 1];
int32_t TOPK = topk_vec[topk_vec.size() - 1];
std::vector<ms::ResultIds> id_vec;
std::vector<ms::ResultDistances> dist_vec;
ms::ResultIds input_ids;
ms::ResultDistances input_distances;
std::vector<std::vector<int64_t>> id_vec;
std::vector<std::vector<float>> dist_vec;
std::vector<int64_t> input_ids;
std::vector<float> input_distance;
int32_t i, k, step;
/* generate testing data */
for (i = 0; i < index_file_num; i++) {
BuildResult(input_ids, input_distances, TOPK, TOPK, NQ, ascending);
BuildResult(input_ids, input_distance, TOPK, TOPK, NQ, ascending);
id_vec.push_back(input_ids);
dist_vec.push_back(input_distances);
dist_vec.push_back(input_distance);
}
for (int32_t max_thread_num : thread_vec) {
@ -256,144 +230,136 @@ TEST(DBSearchTest, REDUCE_PERF_TEST) {
for (int32_t nq : nq_vec) {
for (int32_t top_k : topk_vec) {
ms::ResultIds final_result_ids, final_result_ids_2, final_result_ids_3;
ms::ResultDistances final_result_distances, final_result_distances_2, final_result_distances_3;
ms::ResultSet final_result, final_result_2, final_result_3;
std::vector<ms::ResultIds> id_vec_1(index_file_num);
std::vector<ms::ResultDistances> dist_vec_1(index_file_num);
std::vector<std::vector<int64_t>> id_vec_1(index_file_num);
std::vector<std::vector<float>> dist_vec_1(index_file_num);
for (i = 0; i < index_file_num; i++) {
CopyResult(id_vec_1[i], dist_vec_1[i], top_k, id_vec[i], dist_vec[i], TOPK, nq);
}
std::string str1 = "Method-1 " + std::to_string(max_thread_num) + " " +
std::to_string(nq) + " " + std::to_string(top_k);
std::string str1 = "Method-1 " + std::to_string(max_thread_num) + " " + std::to_string(nq) + " " +
std::to_string(top_k);
milvus::TimeRecorder rc1(str1);
///////////////////////////////////////////////////////////////////////////////////////
/* method-1 */
for (i = 0; i < index_file_num; i++) {
ms::XSearchTask::MergeTopkToResultSet(id_vec_1[i],
dist_vec_1[i],
top_k,
nq,
top_k,
ascending,
final_result_ids,
final_result_distances);
ASSERT_EQ(final_result_ids.size(), nq * top_k);
ASSERT_EQ(final_result_distances.size(), nq * top_k);
ms::XSearchTask::MergeTopkToResultSet(id_vec_1[i], dist_vec_1[i], top_k, nq, top_k, ascending,
final_result);
ASSERT_EQ(final_result.size(), nq);
}
rc1.RecordSection("reduce done");
// ///////////////////////////////////////////////////////////////////////////////////////
// /* method-2 */
// std::vector<std::vector<int64_t>> id_vec_2(index_file_num);
// std::vector<std::vector<float>> dist_vec_2(index_file_num);
// std::vector<size_t> k_vec_2(index_file_num);
// for (i = 0; i < index_file_num; i++) {
// CopyResult(id_vec_2[i], dist_vec_2[i], top_k, id_vec[i], dist_vec[i], TOPK, nq);
// k_vec_2[i] = top_k;
// }
//
// std::string str2 = "Method-2 " + std::to_string(max_thread_num) + " " +
// std::to_string(nq) + " " + std::to_string(top_k);
// milvus::TimeRecorder rc2(str2);
//
// for (step = 1; step < index_file_num; step *= 2) {
// for (i = 0; i + step < index_file_num; i += step * 2) {
// ms::XSearchTask::MergeTopkArray(id_vec_2[i], dist_vec_2[i], k_vec_2[i],
// id_vec_2[i + step], dist_vec_2[i + step], k_vec_2[i + step],
// nq, top_k, ascending);
// }
// }
// ms::XSearchTask::MergeTopkToResultSet(id_vec_2[0],
// dist_vec_2[0],
// k_vec_2[0],
// nq,
// top_k,
// ascending,
// final_result_2);
// ASSERT_EQ(final_result_2.size(), nq);
//
// rc2.RecordSection("reduce done");
//
// for (i = 0; i < nq; i++) {
// ASSERT_EQ(final_result[i].size(), final_result_2[i].size());
// for (k = 0; k < final_result[i].size(); k++) {
// if (final_result[i][k].first != final_result_2[i][k].first) {
// std::cout << i << " " << k << std::endl;
// }
// ASSERT_EQ(final_result[i][k].first, final_result_2[i][k].first);
// ASSERT_EQ(final_result[i][k].second, final_result_2[i][k].second);
// }
// }
//
// ///////////////////////////////////////////////////////////////////////////////////////
// /* method-3 parallel */
// std::vector<std::vector<int64_t>> id_vec_3(index_file_num);
// std::vector<std::vector<float>> dist_vec_3(index_file_num);
// std::vector<size_t> k_vec_3(index_file_num);
// for (i = 0; i < index_file_num; i++) {
// CopyResult(id_vec_3[i], dist_vec_3[i], top_k, id_vec[i], dist_vec[i], TOPK, nq);
// k_vec_3[i] = top_k;
// }
//
// std::string str3 = "Method-3 " + std::to_string(max_thread_num) + " " +
// std::to_string(nq) + " " + std::to_string(top_k);
// milvus::TimeRecorder rc3(str3);
//
// for (step = 1; step < index_file_num; step *= 2) {
// for (i = 0; i + step < index_file_num; i += step * 2) {
// threads_list.push_back(
// threadPool.enqueue(ms::XSearchTask::MergeTopkArray,
// std::ref(id_vec_3[i]),
// std::ref(dist_vec_3[i]),
// std::ref(k_vec_3[i]),
// std::ref(id_vec_3[i + step]),
// std::ref(dist_vec_3[i + step]),
// std::ref(k_vec_3[i + step]),
// nq,
// top_k,
// ascending));
// }
//
// while (threads_list.size() > 0) {
// int nready = 0;
// for (auto it = threads_list.begin(); it != threads_list.end(); it = it) {
// auto &p = *it;
// std::chrono::milliseconds span(0);
// if (p.wait_for(span) == std::future_status::ready) {
// threads_list.erase(it++);
// ++nready;
// } else {
// ++it;
// }
// }
//
// if (nready == 0) {
// std::this_thread::yield();
// }
// }
// }
// ms::XSearchTask::MergeTopkToResultSet(id_vec_3[0],
// dist_vec_3[0],
// k_vec_3[0],
// nq,
// top_k,
// ascending,
// final_result_3);
// ASSERT_EQ(final_result_3.size(), nq);
//
// rc3.RecordSection("reduce done");
//
// for (i = 0; i < nq; i++) {
// ASSERT_EQ(final_result[i].size(), final_result_3[i].size());
// for (k = 0; k < final_result[i].size(); k++) {
// ASSERT_EQ(final_result[i][k].first, final_result_3[i][k].first);
// ASSERT_EQ(final_result[i][k].second, final_result_3[i][k].second);
// }
// }
// ///////////////////////////////////////////////////////////////////////////////////////
// /* method-2 */
// std::vector<std::vector<int64_t>> id_vec_2(index_file_num);
// std::vector<std::vector<float>> dist_vec_2(index_file_num);
// std::vector<uint64_t> k_vec_2(index_file_num);
// for (i = 0; i < index_file_num; i++) {
// CopyResult(id_vec_2[i], dist_vec_2[i], top_k, id_vec[i], dist_vec[i], TOPK, nq);
// k_vec_2[i] = top_k;
// }
//
// std::string str2 = "Method-2 " + std::to_string(max_thread_num) + " " +
// std::to_string(nq) + " " + std::to_string(top_k);
// milvus::TimeRecorder rc2(str2);
//
// for (step = 1; step < index_file_num; step *= 2) {
// for (i = 0; i + step < index_file_num; i += step * 2) {
// ms::XSearchTask::MergeTopkArray(id_vec_2[i], dist_vec_2[i], k_vec_2[i],
// id_vec_2[i + step], dist_vec_2[i + step],
// k_vec_2[i + step], nq, top_k, ascending);
// }
// }
// ms::XSearchTask::MergeTopkToResultSet(id_vec_2[0],
// dist_vec_2[0],
// k_vec_2[0],
// nq,
// top_k,
// ascending,
// final_result_2);
// ASSERT_EQ(final_result_2.size(), nq);
//
// rc2.RecordSection("reduce done");
//
// for (i = 0; i < nq; i++) {
// ASSERT_EQ(final_result[i].size(), final_result_2[i].size());
// for (k = 0; k < final_result[i].size(); k++) {
// if (final_result[i][k].first != final_result_2[i][k].first) {
// std::cout << i << " " << k << std::endl;
// }
// ASSERT_EQ(final_result[i][k].first, final_result_2[i][k].first);
// ASSERT_EQ(final_result[i][k].second, final_result_2[i][k].second);
// }
// }
//
// ///////////////////////////////////////////////////////////////////////////////////////
// /* method-3 parallel */
// std::vector<std::vector<int64_t>> id_vec_3(index_file_num);
// std::vector<std::vector<float>> dist_vec_3(index_file_num);
// std::vector<uint64_t> k_vec_3(index_file_num);
// for (i = 0; i < index_file_num; i++) {
// CopyResult(id_vec_3[i], dist_vec_3[i], top_k, id_vec[i], dist_vec[i], TOPK, nq);
// k_vec_3[i] = top_k;
// }
//
// std::string str3 = "Method-3 " + std::to_string(max_thread_num) + " " +
// std::to_string(nq) + " " + std::to_string(top_k);
// milvus::TimeRecorder rc3(str3);
//
// for (step = 1; step < index_file_num; step *= 2) {
// for (i = 0; i + step < index_file_num; i += step * 2) {
// threads_list.push_back(
// threadPool.enqueue(ms::XSearchTask::MergeTopkArray,
// std::ref(id_vec_3[i]),
// std::ref(dist_vec_3[i]),
// std::ref(k_vec_3[i]),
// std::ref(id_vec_3[i + step]),
// std::ref(dist_vec_3[i + step]),
// std::ref(k_vec_3[i + step]),
// nq,
// top_k,
// ascending));
// }
//
// while (threads_list.size() > 0) {
// int nready = 0;
// for (auto it = threads_list.begin(); it != threads_list.end(); it = it) {
// auto &p = *it;
// std::chrono::milliseconds span(0);
// if (p.wait_for(span) == std::future_status::ready) {
// threads_list.erase(it++);
// ++nready;
// } else {
// ++it;
// }
// }
//
// if (nready == 0) {
// std::this_thread::yield();
// }
// }
// }
// ms::XSearchTask::MergeTopkToResultSet(id_vec_3[0],
// dist_vec_3[0],
// k_vec_3[0],
// nq,
// top_k,
// ascending,
// final_result_3);
// ASSERT_EQ(final_result_3.size(), nq);
//
// rc3.RecordSection("reduce done");
//
// for (i = 0; i < nq; i++) {
// ASSERT_EQ(final_result[i].size(), final_result_3[i].size());
// for (k = 0; k < final_result[i].size(); k++) {
// ASSERT_EQ(final_result[i][k].first, final_result_3[i][k].first);
// ASSERT_EQ(final_result[i][k].second, final_result_3[i][k].second);
// }
// }
}
}
}

117
core/unittest/db/utils.cpp
View File

@ -15,83 +15,92 @@
// specific language governing permissions and limitations
// under the License.
#include <boost/filesystem.hpp>
#include <iostream>
#include <thread>
#include <memory>
#include <string>
#include <boost/filesystem.hpp>
#include <thread>
#include "db/utils.h"
#include "cache/GpuCacheMgr.h"
#include "cache/CpuCacheMgr.h"
#include "cache/GpuCacheMgr.h"
#include "db/DBFactory.h"
#include "db/Options.h"
#include "db/utils.h"
#ifdef MILVUS_GPU_VERSION
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#endif
#include "utils/CommonUtil.h"
INITIALIZE_EASYLOGGINGPP
namespace {
static const char
* CONFIG_STR = "# All the following configurations are default values.\n"
"\n"
"server_config:\n"
" address: 0.0.0.0 # milvus server ip address (IPv4)\n"
" port: 19530 # port range: 1025 ~ 65534\n"
" deploy_mode: single \n"
" time_zone: UTC+8\n"
"\n"
"db_config:\n"
" primary_path: /tmp/milvus # path used to store data and meta\n"
" secondary_path: # path used to store data only, split by semicolon\n"
"\n"
" backend_url: sqlite://:@:/ \n"
" \n"
" # Replace 'dialect' with 'mysql' or 'sqlite'\n"
"\n"
" insert_buffer_size: 4 # GB, maximum insert buffer size allowed\n"
"\n"
"metric_config:\n"
" enable_monitor: false # enable monitoring or not\n"
" collector: prometheus # prometheus\n"
" prometheus_config:\n"
" port: 8080 # port prometheus used to fetch metrics\n"
"\n"
"cache_config:\n"
" cpu_mem_capacity: 16 # GB, CPU memory used for cache\n"
" cpu_mem_threshold: 0.85 # percentage of data kept when cache cleanup triggered\n"
" cache_insert_data: false # whether load inserted data into cache\n"
"\n"
"engine_config:\n"
" use_blas_threshold: 20\n"
"\n"
"resource_config:\n"
" search_resources:\n"
" - gpu0\n"
" index_build_device: gpu0 # GPU used for building index";
static const char* CONFIG_STR =
"# All the following configurations are default values.\n"
"\n"
"server_config:\n"
" address: 0.0.0.0 # milvus server ip address (IPv4)\n"
" port: 19530 # port range: 1025 ~ 65534\n"
" deploy_mode: single \n"
" time_zone: UTC+8\n"
"\n"
"db_config:\n"
" primary_path: /tmp/milvus # path used to store data and meta\n"
" secondary_path: # path used to store data only, split by semicolon\n"
"\n"
" backend_url: sqlite://:@:/ \n"
" \n"
" # Replace 'dialect' with 'mysql' or 'sqlite'\n"
"\n"
" insert_buffer_size: 4 # GB, maximum insert buffer size allowed\n"
"\n"
"metric_config:\n"
" enable_monitor: false # enable monitoring or not\n"
" collector: prometheus # prometheus\n"
" prometheus_config:\n"
" port: 8080 # port prometheus used to fetch metrics\n"
"\n"
"cache_config:\n"
" cpu_mem_capacity: 16 # GB, CPU memory used for cache\n"
" cpu_mem_threshold: 0.85 # percentage of data kept when cache cleanup triggered\n"
" cache_insert_data: false # whether load inserted data into cache\n"
"\n"
"engine_config:\n"
" use_blas_threshold: 20\n"
"\n"
"resource_config:\n"
#ifdef MILVUS_CPU_VERSION
" search_resources:\n"
" - cpu\n"
" index_build_device: cpu # CPU used for building index";
#else
" search_resources:\n"
" - cpu\n"
" - gpu0\n"
" index_build_device: gpu0 # GPU used for building index";
#endif
void
WriteToFile(const std::string& file_path, const char* content) {
std::fstream fs(file_path.c_str(), std::ios_base::out);
//write data to file
// write data to file
fs << content;
fs.close();
}
class DBTestEnvironment : public ::testing::Environment {
public:
explicit DBTestEnvironment(const std::string& uri)
: uri_(uri) {
explicit DBTestEnvironment(const std::string& uri) : uri_(uri) {
}
std::string getURI() const {
std::string
getURI() const {
return uri_;
}
void SetUp() override {
void
SetUp() override {
getURI();
}
@ -101,32 +110,32 @@ class DBTestEnvironment : public ::testing::Environment {
DBTestEnvironment* test_env = nullptr;
} // namespace
} // namespace
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
void
BaseTest::InitLog() {
el::Configurations defaultConf;
defaultConf.setToDefault();
defaultConf.set(el::Level::Debug,
el::ConfigurationType::Format, "[%thread-%datetime-%level]: %msg (%fbase:%line)");
defaultConf.set(el::Level::Debug, el::ConfigurationType::Format, "[%thread-%datetime-%level]: %msg (%fbase:%line)");
el::Loggers::reconfigureLogger("default", defaultConf);
}
void
BaseTest::SetUp() {
InitLog();
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0, 1024 * 1024 * 200, 1024 * 1024 * 300, 2);
#endif
}
void
BaseTest::TearDown() {
milvus::cache::CpuCacheMgr::GetInstance()->ClearCache();
milvus::cache::GpuCacheMgr::GetInstance(0)->ClearCache();
#ifdef MILVUS_GPU_VERSION
knowhere::FaissGpuResourceMgr::GetInstance().Free();
#endif
}
milvus::engine::DBOptions
@ -249,7 +258,7 @@ MySqlMetaTest::GetOptions() {
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
int
main(int argc, char **argv) {
main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
std::string uri;

68
core/unittest/db/utils.h
View File

@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <gtest/gtest.h>
@ -23,35 +22,42 @@
#include <memory>
#include "db/DB.h"
#include "db/meta/SqliteMetaImpl.h"
#include "db/meta/MySQLMetaImpl.h"
#include "scheduler/SchedInst.h"
#include "db/meta/SqliteMetaImpl.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/SchedInst.h"
#define TIMING
#ifdef TIMING
#define INIT_TIMER auto start = std::chrono::high_resolution_clock::now();
#define START_TIMER start = std::chrono::high_resolution_clock::now();
#define STOP_TIMER(name) LOG(DEBUG) << "RUNTIME of " << name << ": " << \
std::chrono::duration_cast<std::chrono::milliseconds>( \
std::chrono::high_resolution_clock::now()-start).count() << " ms ";
#define START_TIMER start = std::chrono::high_resolution_clock::now();
#define STOP_TIMER(name) \
LOG(DEBUG) << "RUNTIME of " << name << ": " \
<< std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - \
start) \
.count() \
<< " ms ";
#else
#define INIT_TIMER
#define START_TIMER
#define STOP_TIMER(name)
#endif
static const char *CONFIG_PATH = "/tmp/milvus_test";
static const char *CONFIG_FILE = "/server_config.yaml";
static const char* CONFIG_PATH = "/tmp/milvus_test";
static const char* CONFIG_FILE = "/server_config.yaml";
class BaseTest : public ::testing::Test {
protected:
void InitLog();
void
InitLog();
void SetUp() override;
void TearDown() override;
virtual milvus::engine::DBOptions GetOptions();
void
SetUp() override;
void
TearDown() override;
virtual milvus::engine::DBOptions
GetOptions();
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@ -59,33 +65,38 @@ class DBTest : public BaseTest {
protected:
milvus::engine::DBPtr db_;
void SetUp() override;
void TearDown() override;
void
SetUp() override;
void
TearDown() override;
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class DBTest2 : public DBTest {
protected:
milvus::engine::DBOptions GetOptions() override;
milvus::engine::DBOptions
GetOptions() override;
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class EngineTest : public DBTest {
};
class EngineTest : public DBTest {};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class MetaTest : public BaseTest {
protected:
std::shared_ptr<milvus::engine::meta::SqliteMetaImpl> impl_;
void SetUp() override;
void TearDown() override;
void
SetUp() override;
void
TearDown() override;
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class MySqlDBTest : public DBTest {
protected:
milvus::engine::DBOptions GetOptions() override;
milvus::engine::DBOptions
GetOptions() override;
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
@ -93,15 +104,16 @@ class MySqlMetaTest : public BaseTest {
protected:
std::shared_ptr<milvus::engine::meta::MySQLMetaImpl> impl_;
void SetUp() override;
void TearDown() override;
milvus::engine::DBOptions GetOptions() override;
void
SetUp() override;
void
TearDown() override;
milvus::engine::DBOptions
GetOptions() override;
};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class MemManagerTest : public MetaTest {
};
class MemManagerTest : public MetaTest {};
/////////////////////////////////////////////////////////////////////////////////////////////////////////////////
class MemManagerTest2 : public DBTest {
};
class MemManagerTest2 : public DBTest {};

4
core/unittest/main.cpp
View File

@ -15,15 +15,15 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "external/easyloggingpp/easylogging++.h"
INITIALIZE_EASYLOGGINGPP
int
main(int argc, char **argv) {
main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
return RUN_ALL_TESTS();
}

1
core/unittest/metrics/test_metricbase.cpp
View File

@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#include "metrics/Metrics.h"
#include <gtest/gtest.h>

38
core/unittest/metrics/test_metrics.cpp
View File

@ -15,19 +15,19 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include <chrono>
#include <map>
#include <memory>
#include <string>
#include <thread>
#include <gtest/gtest.h>
#include "cache/CpuCacheMgr.h"
#include "server/Config.h"
#include "metrics/Metrics.h"
#include "metrics/utils.h"
#include "db/DB.h"
#include "db/meta/SqliteMetaImpl.h"
#include "metrics/Metrics.h"
#include "metrics/utils.h"
#include "server/Config.h"
TEST_F(MetricTest, METRIC_TEST) {
milvus::server::Config::GetInstance().SetMetricConfigCollector("zabbix");
@ -36,15 +36,15 @@ TEST_F(MetricTest, METRIC_TEST) {
milvus::server::Metrics::GetInstance();
milvus::server::SystemInfo::GetInstance().Init();
// server::Metrics::GetInstance().Init();
// server::Metrics::GetInstance().exposer_ptr()->RegisterCollectable(server::Metrics::GetInstance().registry_ptr());
// server::Metrics::GetInstance().Init();
// server::Metrics::GetInstance().exposer_ptr()->RegisterCollectable(server::Metrics::GetInstance().registry_ptr());
milvus::server::Metrics::GetInstance().Init();
// server::PrometheusMetrics::GetInstance().exposer_ptr()->RegisterCollectable(server::PrometheusMetrics::GetInstance().registry_ptr());
// server::PrometheusMetrics::GetInstance().exposer_ptr()->RegisterCollectable(server::PrometheusMetrics::GetInstance().registry_ptr());
milvus::cache::CpuCacheMgr::GetInstance()->SetCapacity(1UL * 1024 * 1024 * 1024);
std::cout << milvus::cache::CpuCacheMgr::GetInstance()->CacheCapacity() << std::endl;
static const char *group_name = "test_group";
static const char* group_name = "test_group";
static const int group_dim = 256;
milvus::engine::meta::TableSchema group_info;
@ -61,22 +61,21 @@ TEST_F(MetricTest, METRIC_TEST) {
int d = 256;
int nb = 50;
float *xb = new float[d * nb];
float* xb = new float[d * nb];
for (int i = 0; i < nb; i++) {
for (int j = 0; j < d; j++) xb[d * i + j] = drand48();
xb[d * i] += i / 2000.;
}
int qb = 5;
float *qxb = new float[d * qb];
float* qxb = new float[d * qb];
for (int i = 0; i < qb; i++) {
for (int j = 0; j < d; j++) qxb[d * i + j] = drand48();
qxb[d * i] += i / 2000.;
}
std::thread search([&]() {
// milvus::engine::ResultIds result_ids;
// milvus::engine::ResultDistances result_distances;
milvus::engine::QueryResults results;
int k = 10;
std::this_thread::sleep_for(std::chrono::seconds(2));
@ -91,17 +90,16 @@ TEST_F(MetricTest, METRIC_TEST) {
prev_count = count;
START_TIMER;
// stat = db_->Query(group_name, k, qb, qxb, result_ids, result_distances);
ss << "Search " << j << " With Size " << (float) (count * group_dim * sizeof(float)) / (1024 * 1024)
<< " M";
// stat = db_->Query(group_name, k, qb, qxb, results);
ss << "Search " << j << " With Size " << (float)(count * group_dim * sizeof(float)) / (1024 * 1024) << " M";
for (auto k = 0; k < qb; ++k) {
// ASSERT_EQ(results[k][0].first, target_ids[k]);
// ASSERT_EQ(results[k][0].first, target_ids[k]);
ss.str("");
ss << "Result [" << k << "]:";
// for (auto result : results[k]) {
// ss << result.first << " ";
// }
// for (auto result : results[k]) {
// ss << result.first << " ";
// }
}
ASSERT_TRUE(count >= prev_count);
std::this_thread::sleep_for(std::chrono::seconds(1));
@ -154,5 +152,3 @@ TEST_F(MetricTest, COLLECTOR_METRICS_TEST) {
milvus::server::MetricCollector metric_collector();
}

1
core/unittest/metrics/test_prometheus.cpp
View File

@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#include "metrics/PrometheusMetrics.h"
#include "server/Config.h"

38
core/unittest/metrics/utils.cpp
View File

@ -15,14 +15,13 @@
// specific language governing permissions and limitations
// under the License.
#include <iostream>
#include <thread>
#include <string>
#include <boost/filesystem.hpp>
#include <iostream>
#include <string>
#include <thread>
#include "metrics/utils.h"
#include "db/DBFactory.h"
#include "metrics/utils.h"
INITIALIZE_EASYLOGGINGPP
@ -30,13 +29,16 @@ namespace {
class DBTestEnvironment : public ::testing::Environment {
public:
explicit DBTestEnvironment(const std::string& uri) : uri_(uri) {}
explicit DBTestEnvironment(const std::string& uri) : uri_(uri) {
}
std::string getURI() const {
std::string
getURI() const {
return uri_;
}
void SetUp() override {
void
SetUp() override {
getURI();
}
@ -46,36 +48,40 @@ class DBTestEnvironment : public ::testing::Environment {
DBTestEnvironment* test_env = nullptr;
} // namespace
} // namespace
void MetricTest::InitLog() {
void
MetricTest::InitLog() {
el::Configurations defaultConf;
defaultConf.setToDefault();
defaultConf.set(el::Level::Debug,
el::ConfigurationType::Format, "[%thread-%datetime-%level]: %msg (%fbase:%line)");
defaultConf.set(el::Level::Debug, el::ConfigurationType::Format, "[%thread-%datetime-%level]: %msg (%fbase:%line)");
el::Loggers::reconfigureLogger("default", defaultConf);
}
milvus::engine::DBOptions MetricTest::GetOptions() {
milvus::engine::DBOptions
MetricTest::GetOptions() {
auto options = milvus::engine::DBFactory::BuildOption();
options.meta_.path_ = "/tmp/milvus_test";
options.meta_.backend_uri_ = "sqlite://:@:/";
return options;
}
void MetricTest::SetUp() {
void
MetricTest::SetUp() {
boost::filesystem::remove_all("/tmp/milvus_test");
InitLog();
auto options = GetOptions();
db_ = milvus::engine::DBFactory::Build(options);
}
void MetricTest::TearDown() {
void
MetricTest::TearDown() {
db_->Stop();
boost::filesystem::remove_all("/tmp/milvus_test");
}
int main(int argc, char **argv) {
int
main(int argc, char** argv) {
::testing::InitGoogleTest(&argc, argv);
std::string uri;

32
core/unittest/metrics/utils.h
View File

@ -15,7 +15,6 @@
// specific language governing permissions and limitations
// under the License.
#pragma once
#include <gtest/gtest.h>
@ -23,17 +22,20 @@
//#include <src/db/MySQLMetaImpl.h>
#include "db/DB.h"
#include "db/meta/SqliteMetaImpl.h"
#include "db/meta/MySQLMetaImpl.h"
#include "db/meta/SqliteMetaImpl.h"
#define TIMING
#ifdef TIMING
#define INIT_TIMER auto start = std::chrono::high_resolution_clock::now();
#define START_TIMER start = std::chrono::high_resolution_clock::now();
#define STOP_TIMER(name) LOG(DEBUG) << "RUNTIME of " << name << ": " << \
std::chrono::duration_cast<std::chrono::milliseconds>( \
std::chrono::high_resolution_clock::now()-start).count() << " ms ";
#define START_TIMER start = std::chrono::high_resolution_clock::now();
#define STOP_TIMER(name) \
LOG(DEBUG) << "RUNTIME of " << name << ": " \
<< std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::high_resolution_clock::now() - \
start) \
.count() \
<< " ms ";
#else
#define INIT_TIMER
#define START_TIMER
@ -41,10 +43,10 @@
#endif
void
ASSERT_STATS(milvus::Status &stat);
ASSERT_STATS(milvus::Status& stat);
//class TestEnv : public ::testing::Environment {
//public:
// class TestEnv : public ::testing::Environment {
// public:
//
// static std::string getURI() {
// if (const char* uri = std::getenv("MILVUS_DBMETA_URI")) {
@ -68,8 +70,12 @@ class MetricTest : public ::testing::Test {
protected:
milvus::engine::DBPtr db_;
void InitLog();
void SetUp() override;
void TearDown() override;
virtual milvus::engine::DBOptions GetOptions();
void
InitLog();
void
SetUp() override;
void
TearDown() override;
virtual milvus::engine::DBOptions
GetOptions();
};

2
core/unittest/scheduler/CMakeLists.txt
View File

@ -19,7 +19,7 @@
aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR} test_files)
cuda_add_executable(test_scheduler
add_executable(test_scheduler
${common_files}
${entry_file}
${test_files}

12
core/unittest/scheduler/task_test.cpp
View File

@ -15,11 +15,9 @@
// specific language governing permissions and limitations
// under the License.
#include "scheduler/task/SearchTask.h"
#include "scheduler/task/BuildIndexTask.h"
#include <gtest/gtest.h>
#include "scheduler/task/BuildIndexTask.h"
#include "scheduler/task/SearchTask.h"
namespace milvus {
namespace scheduler {
@ -34,7 +32,5 @@ TEST(TaskTest, INVALID_INDEX) {
build_task->Execute();
}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

13
core/unittest/scheduler/test_algorithm.cpp
View File

@ -15,15 +15,13 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include "scheduler/resource/Resource.h"
#include "scheduler/Algorithm.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/ResourceMgr.h"
#include "scheduler/resource/CpuResource.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/Algorithm.h"
#include "scheduler/resource/Resource.h"
namespace milvus {
namespace scheduler {
@ -103,6 +101,5 @@ TEST_F(AlgorithmTest, SHORTESTPATH_TEST) {
}
}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

13
core/unittest/scheduler/test_event.cpp
View File

@ -15,14 +15,10 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include "scheduler/resource/Resource.h"
#include "scheduler/event/Event.h"
#include "scheduler/event/StartUpEvent.h"
#include "scheduler/resource/Resource.h"
namespace milvus {
namespace scheduler {
@ -51,7 +47,6 @@ TEST(EventTest, FINISH_TASK_EVENT) {
std::cout << *EventPtr(event);
}
TEST(EventTest, TASKTABLE_UPDATED_EVENT) {
ResourcePtr res(nullptr);
auto event = std::make_shared<TaskTableUpdatedEvent>(res);
@ -60,7 +55,5 @@ TEST(EventTest, TASKTABLE_UPDATED_EVENT) {
std::cout << *EventPtr(event);
}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

9
core/unittest/scheduler/test_normal.cpp
View File

@ -15,21 +15,20 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include "scheduler/ResourceFactory.h"
#include "scheduler/ResourceMgr.h"
#include "scheduler/SchedInst.h"
#include "scheduler/Scheduler.h"
#include "scheduler/task/TestTask.h"
#include "scheduler/tasklabel/DefaultLabel.h"
#include "scheduler/SchedInst.h"
#include "utils/Log.h"
#include <gtest/gtest.h>
namespace {
namespace ms = milvus::scheduler;
} // namespace
} // namespace
TEST(NormalTest, INST_TEST) {
// ResourceMgr only compose resources, provide unified event
@ -62,7 +61,7 @@ TEST(NormalTest, INST_TEST) {
}
}
for (auto &task : tasks) {
for (auto& task : tasks) {
task->Wait();
ASSERT_EQ(task->load_count_, 1);
ASSERT_EQ(task->exec_count_, 1);

23
core/unittest/scheduler/test_resource.cpp
View File

@ -15,18 +15,16 @@
// specific language governing permissions and limitations
// under the License.
#include "scheduler/resource/Resource.h"
#include "scheduler/resource/DiskResource.h"
#include <gtest/gtest.h>
#include "scheduler/ResourceFactory.h"
#include "scheduler/resource/CpuResource.h"
#include "scheduler/resource/DiskResource.h"
#include "scheduler/resource/GpuResource.h"
#include "scheduler/resource/Resource.h"
#include "scheduler/resource/TestResource.h"
#include "scheduler/task/Task.h"
#include "scheduler/task/TestTask.h"
#include "scheduler/tasklabel/DefaultLabel.h"
#include "scheduler/ResourceFactory.h"
#include <gtest/gtest.h>
namespace milvus {
namespace scheduler {
@ -158,17 +156,13 @@ class ResourceAdvanceTest : public testing::Test {
void
WaitLoader(uint64_t count) {
std::unique_lock<std::mutex> lock(load_mutex_);
cv_.wait(lock, [&] {
return load_count_ == count;
});
cv_.wait(lock, [&] { return load_count_ == count; });
}
void
WaitExecutor(uint64_t count) {
std::unique_lock<std::mutex> lock(exec_mutex_);
cv_.wait(lock, [&] {
return exec_count_ == count;
});
cv_.wait(lock, [&] { return exec_count_ == count; });
}
ResourcePtr disk_resource_;
@ -287,6 +281,5 @@ TEST_F(ResourceAdvanceTest, TEST_RESOURCE_TEST) {
}
}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

5
core/unittest/scheduler/test_resource_factory.cpp
View File

@ -15,15 +15,14 @@
// specific language governing permissions and limitations
// under the License.
#include "scheduler/ResourceFactory.h"
#include <gtest/gtest.h>
#include "scheduler/ResourceFactory.h"
namespace {
namespace ms = milvus::scheduler;
} // namespace
} // namespace
TEST(ResourceFactoryTest, CREATE) {
auto disk = ms::ResourceFactory::Create("ssd", "DISK", 0);

28
core/unittest/scheduler/test_resource_mgr.cpp
View File

@ -15,16 +15,14 @@
// specific language governing permissions and limitations
// under the License.
#include <gtest/gtest.h>
#include "scheduler/ResourceMgr.h"
#include "scheduler/resource/CpuResource.h"
#include "scheduler/resource/GpuResource.h"
#include "scheduler/resource/DiskResource.h"
#include "scheduler/resource/GpuResource.h"
#include "scheduler/resource/TestResource.h"
#include "scheduler/task/TestTask.h"
#include "scheduler/tasklabel/DefaultLabel.h"
#include "scheduler/ResourceMgr.h"
#include <gtest/gtest.h>
namespace milvus {
namespace scheduler {
@ -101,10 +99,13 @@ TEST_F(ResourceMgrBaseTest, GET_ALL_RESOURCES) {
bool disk = false, cpu = false, gpu = false;
auto resources = mgr1_->GetAllResources();
ASSERT_EQ(resources.size(), 3);
for (auto &res : resources) {
if (res->type() == ResourceType::DISK) disk = true;
if (res->type() == ResourceType::CPU) cpu = true;
if (res->type() == ResourceType::GPU) gpu = true;
for (auto& res : resources) {
if (res->type() == ResourceType::DISK)
disk = true;
if (res->type() == ResourceType::CPU)
cpu = true;
if (res->type() == ResourceType::GPU)
gpu = true;
}
ASSERT_TRUE(disk);
@ -183,9 +184,7 @@ class ResourceMgrAdvanceTest : public testing::Test {
TEST_F(ResourceMgrAdvanceTest, REGISTER_SUBSCRIBER) {
bool flag = false;
auto callback = [&](EventPtr event) {
flag = true;
};
auto callback = [&](EventPtr event) { flag = true; };
mgr1_->RegisterSubscriber(callback);
TableFileSchemaPtr dummy = nullptr;
auto label = std::make_shared<DefaultLabel>();
@ -194,6 +193,5 @@ TEST_F(ResourceMgrAdvanceTest, REGISTER_SUBSCRIBER) {
ASSERT_TRUE(flag);
}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

71
core/unittest/scheduler/test_scheduler.cpp
View File

@ -17,14 +17,14 @@
#include <gtest/gtest.h>
#include "scheduler/Scheduler.h"
#include "scheduler/tasklabel/DefaultLabel.h"
#include "scheduler/tasklabel/SpecResLabel.h"
#include "scheduler/task/TestTask.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/resource/Resource.h"
#include "cache/DataObj.h"
#include "cache/GpuCacheMgr.h"
#include "scheduler/ResourceFactory.h"
#include "scheduler/Scheduler.h"
#include "scheduler/resource/Resource.h"
#include "scheduler/task/TestTask.h"
#include "scheduler/tasklabel/DefaultLabel.h"
#include "scheduler/tasklabel/SpecResLabel.h"
#include "utils/Error.h"
#include "wrapper/VecIndex.h"
@ -33,59 +33,61 @@ namespace scheduler {
class MockVecIndex : public engine::VecIndex {
public:
virtual Status BuildAll(const int64_t& nb,
const float* xb,
const int64_t* ids,
const engine::Config& cfg,
const int64_t& nt = 0,
const float* xt = nullptr) {
virtual Status
BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg, const int64_t& nt = 0,
const float* xt = nullptr) {
}
engine::VecIndexPtr Clone() override {
engine::VecIndexPtr
Clone() override {
return milvus::engine::VecIndexPtr();
}
int64_t GetDeviceId() override {
int64_t
GetDeviceId() override {
return 0;
}
engine::IndexType GetType() override {
engine::IndexType
GetType() override {
return engine::IndexType::INVALID;
}
virtual Status Add(const int64_t& nb,
const float* xb,
const int64_t* ids,
const engine::Config& cfg = engine::Config()) {
virtual Status
Add(const int64_t& nb, const float* xb, const int64_t* ids, const engine::Config& cfg = engine::Config()) {
}
virtual Status Search(const int64_t& nq,
const float* xq,
float* dist,
int64_t* ids,
const engine::Config& cfg = engine::Config()) {
virtual Status
Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids,
const engine::Config& cfg = engine::Config()) {
}
engine::VecIndexPtr CopyToGpu(const int64_t& device_id, const engine::Config& cfg) override {
engine::VecIndexPtr
CopyToGpu(const int64_t& device_id, const engine::Config& cfg) override {
}
engine::VecIndexPtr CopyToCpu(const engine::Config& cfg) override {
engine::VecIndexPtr
CopyToCpu(const engine::Config& cfg) override {
}
virtual int64_t Dimension() {
virtual int64_t
Dimension() {
return dimension_;
}
virtual int64_t Count() {
virtual int64_t
Count() {
return ntotal_;
}
virtual knowhere::BinarySet Serialize() {
virtual knowhere::BinarySet
Serialize() {
knowhere::BinarySet binset;
return binset;
}
virtual Status Load(const knowhere::BinarySet& index_binary) {
virtual Status
Load(const knowhere::BinarySet& index_binary) {
}
public:
@ -185,7 +187,7 @@ TEST_F(SchedulerTest, PUSH_TASK_TO_NEIGHBOUR_RANDOMLY_TEST) {
}
sleep(3);
// ASSERT_EQ(res_mgr_->GetResource(ResourceType::GPU, 1)->task_table().Size(), NUM);
// ASSERT_EQ(res_mgr_->GetResource(ResourceType::GPU, 1)->task_table().Size(), NUM);
}
class SchedulerTest2 : public testing::Test {
@ -240,7 +242,7 @@ class SchedulerTest2 : public testing::Test {
std::shared_ptr<Scheduler> scheduler_;
};
//TEST_F(SchedulerTest2, SPECIFIED_RESOURCE_TEST) {
// TEST_F(SchedulerTest2, SPECIFIED_RESOURCE_TEST) {
// const uint64_t NUM = 2;
// std::vector<std::shared_ptr<TestTask>> tasks;
// TableFileSchemaPtr dummy = std::make_shared<TableFileSchema>();
@ -257,6 +259,5 @@ class SchedulerTest2 : public testing::Test {
// ASSERT_EQ(res_mgr_->GetResource(ResourceType::GPU, 1)->task_table().Size(), NUM);
//}
} // namespace scheduler
} // namespace milvus
} // namespace scheduler
} // namespace milvus

2
core/unittest/server/CMakeLists.txt
View File

@ -47,7 +47,7 @@ set(server_test_files
${test_files}
)
cuda_add_executable(test_server ${server_test_files})
add_executable(test_server ${server_test_files})
set(client_grpc_lib
grpcpp_channelz

68
core/unittest/server/test_cache.cpp
View File

@ -38,70 +38,70 @@ class LessItemCacheMgr : public milvus::cache::CacheMgr<milvus::cache::DataObjPt
class MockVecIndex : public milvus::engine::VecIndex {
public:
MockVecIndex(int64_t dim, int64_t total)
: dimension_(dim),
ntotal_(total) {
MockVecIndex(int64_t dim, int64_t total) : dimension_(dim), ntotal_(total) {
}
virtual milvus::Status BuildAll(const int64_t &nb,
const float *xb,
const int64_t *ids,
const milvus::engine::Config &cfg,
const int64_t &nt = 0,
const float *xt = nullptr) {
virtual milvus::Status
BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const milvus::engine::Config& cfg,
const int64_t& nt = 0, const float* xt = nullptr) {
return milvus::Status();
}
milvus::engine::VecIndexPtr Clone() override {
milvus::engine::VecIndexPtr
Clone() override {
return milvus::engine::VecIndexPtr();
}
int64_t GetDeviceId() override {
int64_t
GetDeviceId() override {
return 0;
}
milvus::engine::IndexType GetType() override {
milvus::engine::IndexType
GetType() override {
return milvus::engine::IndexType::INVALID;
}
virtual milvus::Status Add(const int64_t &nb,
const float *xb,
const int64_t *ids,
const milvus::engine::Config &cfg = milvus::engine::Config()) {
virtual milvus::Status
Add(const int64_t& nb, const float* xb, const int64_t* ids,
const milvus::engine::Config& cfg = milvus::engine::Config()) {
return milvus::Status();
}
virtual milvus::Status Search(const int64_t &nq,
const float *xq,
float *dist,
int64_t *ids,
const milvus::engine::Config &cfg = milvus::engine::Config()) {
virtual milvus::Status
Search(const int64_t& nq, const float* xq, float* dist, int64_t* ids,
const milvus::engine::Config& cfg = milvus::engine::Config()) {
return milvus::Status();
}
milvus::engine::VecIndexPtr CopyToGpu(const int64_t &device_id,
const milvus::engine::Config &cfg) override {
milvus::engine::VecIndexPtr
CopyToGpu(const int64_t& device_id, const milvus::engine::Config& cfg) override {
return nullptr;
}
milvus::engine::VecIndexPtr CopyToCpu(const milvus::engine::Config &cfg) override {
milvus::engine::VecIndexPtr
CopyToCpu(const milvus::engine::Config& cfg) override {
return nullptr;
}
virtual int64_t Dimension() {
virtual int64_t
Dimension() {
return dimension_;
}
virtual int64_t Count() {
virtual int64_t
Count() {
return ntotal_;
}
virtual knowhere::BinarySet Serialize() {
virtual knowhere::BinarySet
Serialize() {
knowhere::BinarySet binset;
return binset;
}
virtual milvus::Status Load(const knowhere::BinarySet &index_binary) {
virtual milvus::Status
Load(const knowhere::BinarySet& index_binary) {
return milvus::Status();
}
@ -110,7 +110,7 @@ class MockVecIndex : public milvus::engine::VecIndex {
int64_t ntotal_ = 0;
};
} // namespace
} // namespace
TEST(CacheTest, DUMMY_TEST) {
milvus::engine::Config cfg;
@ -141,7 +141,7 @@ TEST(CacheTest, CPU_CACHE_TEST) {
uint64_t item_count = 20;
for (uint64_t i = 0; i < item_count; i++) {
//each vector is 1k byte, total size less than 1G
// each vector is 1k byte, total size less than 1G
milvus::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 1000000);
milvus::cache::DataObjPtr data_obj = std::static_pointer_cast<milvus::cache::DataObj>(mock_index);
cpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
@ -165,7 +165,7 @@ TEST(CacheTest, CPU_CACHE_TEST) {
g_num = 5;
cpu_mgr->SetCapacity(g_num * gbyte);
//each vector is 1k byte, total size less than 6G
// each vector is 1k byte, total size less than 6G
milvus::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 6000000);
milvus::cache::DataObjPtr data_obj = std::static_pointer_cast<milvus::cache::DataObj>(mock_index);
cpu_mgr->InsertItem("index_6g", data_obj);
@ -179,7 +179,7 @@ TEST(CacheTest, GPU_CACHE_TEST) {
auto gpu_mgr = milvus::cache::GpuCacheMgr::GetInstance(0);
for (int i = 0; i < 20; i++) {
//each vector is 1k byte
// each vector is 1k byte
milvus::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 1000);
milvus::cache::DataObjPtr data_obj = std::static_pointer_cast<milvus::cache::DataObj>(mock_index);
gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
@ -192,7 +192,7 @@ TEST(CacheTest, GPU_CACHE_TEST) {
for (auto i = 0; i < 3; i++) {
// TODO(myh): use gpu index to mock
//each vector is 1k byte, total size less than 2G
// each vector is 1k byte, total size less than 2G
milvus::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 2000000);
milvus::cache::DataObjPtr data_obj = std::static_pointer_cast<milvus::cache::DataObj>(mock_index);
std::cout << data_obj->Size() << std::endl;
@ -223,7 +223,7 @@ TEST(CacheTest, INVALID_TEST) {
{
LessItemCacheMgr mgr;
for (int i = 0; i < 20; i++) {
//each vector is 1k byte
// each vector is 1k byte
milvus::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 2);
milvus::cache::DataObjPtr data_obj = std::static_pointer_cast<milvus::cache::DataObj>(mock_index);
mgr.InsertItem("index_" + std::to_string(i), data_obj);

Some files were not shown because too many files have changed in this diff Show More