milvus/tests/benchmark/milvus_benchmark/client.py

import random
import logging
import time
# import traceback
from multiprocessing import Process
from pymilvus import Milvus, DataType
import utils as util
import config
from milvus_benchmark.runners import utils

logger = logging.getLogger("milvus_benchmark.client")


# yaml file and code file comparison table of Index parameters
INDEX_MAP = {
    "flat": "FLAT",
    "ivf_flat": "IVF_FLAT",
    "ivf_sq8": "IVF_SQ8",
    "ivf_pq": "IVF_PQ",
    "hnsw": "HNSW",
    "annoy": "ANNOY",
    "bin_flat": "BIN_FLAT",
    "bin_ivf_flat": "BIN_IVF_FLAT",
}
epsilon = 0.1
DEFAULT_WARM_QUERY_TOPK = 1
DEFAULT_WARM_QUERY_NQ = 1


def time_wrapper(func):
    """
    This decorator prints the execution time for the decorated function.
    """

    def wrapper(*args, **kwargs):
        """
        log: Specify output log
        rps: Specify the rps of the return interface
        """
        start = time.time()
        log = kwargs.get("log", True)
        kwargs.pop("log", None)
        rps = kwargs.get("rps", False)
        kwargs.pop("rps", None)
        result = func(*args, **kwargs)
        end = time.time()
        if log:
            _rps = round(end - start, 2)
            logger.debug("Milvus {} run in {}s".format(func.__name__, _rps))
            if rps is not False:
                return result, _rps
        return result
    return wrapper


class MilvusClient(object):
    def __init__(self, collection_name=None, host=None, port=None, timeout=300):
        self._collection_name = collection_name
        self._collection_info = None
        self._dimension = None
        start_time = time.time()
        if not host:
            host = config.SERVER_HOST_DEFAULT
        if not port:
            port = config.SERVER_PORT_DEFAULT
        # retry connect remote server
        i = 0
        while time.time() < start_time + timeout:
            try:
                self._milvus = Milvus(
                    host=host,
                    port=port,
                    try_connect=False,
                    pre_ping=False)
                break
            except Exception as e:
                logger.error(str(e))
                logger.error("Milvus connect failed: %d times" % i)
                i = i + 1
                time.sleep(30)

        if time.time() > start_time + timeout:
            raise Exception("Server connect timeout")
        # self._metric_type = None

    def __str__(self):
        return 'Milvus collection %s' % self._collection_name

    def set_collection(self, collection_name):
        """ Setting collection name """
        self._collection_name = collection_name

    # TODO: server not support
    # def check_status(self, status):
    #     if not status.OK():
    #         logger.error(status.message)
    #         logger.error(self._milvus.server_status())
    #         logger.error(self.count())
    #         raise Exception("Status not ok")

    def check_result_ids(self, result):
        for index, item in enumerate(result):
            if item[0].distance >= epsilon:
                logger.error(index)
                logger.error(item[0].distance)
                raise Exception("Distance wrong")

    @property
    def collection_name(self):
        return self._collection_name

    # only support the given field name
    def create_collection(self, dimension, data_type=DataType.FLOAT_VECTOR, auto_id=False,
                          collection_name=None, other_fields=None):
        self._dimension = dimension
        if not collection_name:
            collection_name = self._collection_name
        vec_field_name = utils.get_default_field_name(data_type)
        fields = [
            {"name": vec_field_name, "type": data_type, "params": {"dim": dimension}},
            {"name": "id", "type": DataType.INT64, "is_primary": True}
        ]
        if other_fields:
            other_fields = other_fields.split(",")
            for other_field_name in other_fields:
                if other_field_name.startswith("int"):
                    field_type = DataType.INT64
                elif other_field_name.startswith("float"):
                    field_type = DataType.FLOAT
                elif other_field_name.startswith("double"):
                    field_type = DataType.DOUBLE
                else:
                    raise Exception("Create collection: Field name not supported")
                fields.append({"name": other_field_name, "type": field_type})
        create_param = {
            "fields": fields,
            "auto_id": auto_id}
        try:
            self._milvus.create_collection(collection_name, create_param)
            logger.info("Create collection: <%s> successfully" % collection_name)
        except Exception as e:
            logger.error(str(e))
            raise

    def create_partition(self, tag, collection_name=None):
        if not collection_name:
            collection_name = self._collection_name
        self._milvus.create_partition(collection_name, tag)

    @time_wrapper
    def insert(self, entities, collection_name=None, timeout=None):
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        try:
            insert_res = self._milvus.insert(tmp_collection_name, entities, timeout=timeout)
            return insert_res.primary_keys
        except Exception as e:
            logger.error(str(e))

    @time_wrapper
    def insert_flush(self, entities, _async=False, collection_name=None):
        # the method that included insert and flush
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        try:
            insert_res = self._milvus.insert(tmp_collection_name, entities)
            return insert_res.primary_keys
        except Exception as e:
            logger.error(str(e))
        self._milvus.flush([tmp_collection_name], _async=_async)

    def get_dimension(self):
        info = self.get_info()
        for field in info["fields"]:
            if field["type"] in [DataType.FLOAT_VECTOR, DataType.BINARY_VECTOR]:
                return field["params"]["dim"]

    def get_rand_ids(self, length):
        segment_ids = []
        while True:
            stats = self.get_stats()
            segments = stats["partitions"][0]["segments"]
            # random choice one segment
            segment = random.choice(segments)
            try:
                segment_ids = self._milvus.list_id_in_segment(self._collection_name, segment["id"])
            except Exception as e:
                logger.error(str(e))
            if not len(segment_ids):
                continue
            elif len(segment_ids) > length:
                return random.sample(segment_ids, length)
            else:
                logger.debug("Reset length: %d" % len(segment_ids))
                return segment_ids

    # def get_rand_ids_each_segment(self, length):
    #     res = []
    #     status, stats = self._milvus.get_collection_stats(self._collection_name)
    #     self.check_status(status)
    #     segments = stats["partitions"][0]["segments"]
    #     segments_num = len(segments)
    #     # random choice from each segment
    #     for segment in segments:
    #         status, segment_ids = self._milvus.list_id_in_segment(self._collection_name, segment["name"])
    #         self.check_status(status)
    #         res.extend(segment_ids[:length])
    #     return segments_num, res

    # def get_rand_entities(self, length):
    #     ids = self.get_rand_ids(length)
    #     status, get_res = self._milvus.get_entity_by_id(self._collection_name, ids)
    #     self.check_status(status)
    #     return ids, get_res


    @time_wrapper
    def get_entities(self, get_ids):
        get_res = self._milvus.get_entity_by_id(self._collection_name, get_ids)
        return get_res

    @time_wrapper
    def delete(self, ids, collection_name=None):
        # delete entity by id
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        self._milvus.delete_entity_by_id(tmp_collection_name, ids)

    def delete_rand(self):
        delete_id_length = random.randint(1, 100)
        # count_before = self.count()
        logger.debug("%s: length to delete: %d" % (self._collection_name, delete_id_length))
        delete_ids = self.get_rand_ids(delete_id_length)
        self.delete(delete_ids)
        self.flush()
        logger.info("%s: count after delete: %d" % (self._collection_name, self.count()))
        get_res = self._milvus.get_entity_by_id(self._collection_name, delete_ids)
        for item in get_res:
            assert not item
        # if count_before - len(delete_ids) < self.count():
        #     logger.error(delete_ids)
        #     raise Exception("Error occured")

    @time_wrapper
    def flush(self, _async=False, collection_name=None, timeout=None):
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        self._milvus.flush([tmp_collection_name], _async=_async, timeout=timeout)

    @time_wrapper
    def compact(self, collection_name=None):
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        status = self._milvus.compact(tmp_collection_name)
        self.check_status(status)

    @time_wrapper
    def get(self, ids, collection_name=None, timeout=None):
        """ only support "in" in expr """
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        # res = self._milvus.get(tmp_collection_name, ids, output_fields=None, partition_names=None)
        ids_expr = "id in %s" % (str(ids))
        res = self._milvus.query(tmp_collection_name, ids_expr, output_fields=None, partition_names=None, timeout=timeout)
        return res

    @time_wrapper
    def create_index(self, field_name, index_type, metric_type, _async=False, index_param=None):
        index_type = INDEX_MAP[index_type]
        metric_type = utils.metric_type_trans(metric_type)
        logger.info("Building index start, collection_name: %s, index_type: %s, metric_type: %s" % (
            self._collection_name, index_type, metric_type))
        if index_param:
            logger.info(index_param)
        index_params = {
            "index_type": index_type,
            "metric_type": metric_type,
            "params": index_param
        }
        self._milvus.create_index(self._collection_name, field_name, index_params, _async=_async)

    # TODO: need to check
    def describe_index(self, field_name, collection_name=None):
        # stats = self.get_stats()
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        info = self._milvus.describe_index(tmp_collection_name, field_name)
        logger.info(info)
        index_info = {"index_type": "flat", "metric_type": None, "index_param": None}
        if info:
            index_info = {"index_type": info["index_type"], "metric_type": info["metric_type"], "index_param": info["params"]}
            # transfer index type name
            for k, v in INDEX_MAP.items():
                if index_info['index_type'] == v:
                    index_info['index_type'] = k
        return index_info

    def drop_index(self, field_name):
        logger.info("Drop index: %s" % self._collection_name)
        return self._milvus.drop_index(self._collection_name, field_name)

    @time_wrapper
    def query(self, vector_query, filter_query=None, collection_name=None, guarantee_timestamp=None, timeout=300):
        """ This method corresponds to the search method of milvus """
        tmp_collection_name = self._collection_name if collection_name is None else collection_name

        params = util.search_param_analysis(vector_query, filter_query)
        params.update({"timeout": timeout})

        if guarantee_timestamp is not None:
            params.update({"guarantee_timestamp": guarantee_timestamp})

        result = self._milvus.search(tmp_collection_name, **params)

        # must_params = [vector_query]
        # if filter_query:
        #     must_params.extend(filter_query)
        # query = {
        #     "bool": {"must": must_params}
        # }
        # result = self._milvus.search(tmp_collection_name, query, timeout=timeout)
        return result

    @time_wrapper
    def warm_query(self, index_field_name, search_param, metric_type, times=2):
        query_vectors = [[random.random() for _ in range(self._dimension)] for _ in range(DEFAULT_WARM_QUERY_NQ)]
        # index_info = self.describe_index(index_field_name)
        vector_query = {"vector": {index_field_name: {
            "topk": DEFAULT_WARM_QUERY_TOPK, 
            "query": query_vectors, 
            "metric_type": metric_type,
            "params": search_param}
        }}
        must_params = [vector_query]
        # query = {
        #     "bool": {"must": must_params}
        # }
        logger.debug("Start warm up query")
        for i in range(times):
            params = util.search_param_analysis(vector_query, None)
            self._milvus.search(self._collection_name, **params)

            # self._milvus.search(self._collection_name, query)
        logger.debug("End warm up query")

    @time_wrapper
    def load_and_query(self, vector_query, filter_query=None, collection_name=None, timeout=120):
        tmp_collection_name = self._collection_name if collection_name is None else collection_name
        must_params = [vector_query]
        if filter_query:
            must_params.extend(filter_query)

        self.load_collection(tmp_collection_name)

        params = util.search_param_analysis(vector_query, filter_query)
        params.update({"timeout": timeout})
        result = self._milvus.search(tmp_collection_name, **params)

        # result = self._milvus.search(tmp_collection_name, query, timeout=timeout)
        return result

    def get_ids(self, result):
        # idss = result._entities.ids
        ids = []
        # len_idss = len(idss)
        # len_r = len(result)
        # top_k = len_idss // len_r
        # for offset in range(0, len_idss, top_k):
        #     ids.append(idss[offset: min(offset + top_k, len_idss)])
        for res in result:
            ids.append(res.ids)
        return ids

    def query_rand(self, nq_max=100, timeout=None):
        # for ivf search
        dimension = 128
        top_k = random.randint(1, 100)
        nq = random.randint(1, nq_max)
        nprobe = random.randint(1, 100)
        search_param = {"nprobe": nprobe}
        query_vectors = [[random.random() for _ in range(dimension)] for _ in range(nq)]
        metric_type = random.choice(["l2", "ip"])
        logger.info("%s, Search nq: %d, top_k: %d, nprobe: %d" % (self._collection_name, nq, top_k, nprobe))
        vec_field_name = utils.get_default_field_name()
        vector_query = {"vector": {vec_field_name: {
            "topk": top_k,
            "query": query_vectors,
            "metric_type": utils.metric_type_trans(metric_type),
            "params": search_param}
        }}
        self.query(vector_query, timeout=timeout)

    def load_query_rand(self, nq_max=100, timeout=None):
        # for ivf search
        dimension = 128
        top_k = random.randint(1, 100)
        nq = random.randint(1, nq_max)
        nprobe = random.randint(1, 100)
        search_param = {"nprobe": nprobe}
        query_vectors = [[random.random() for _ in range(dimension)] for _ in range(nq)]
        metric_type = random.choice(["l2", "ip"])
        logger.info("%s, Search nq: %d, top_k: %d, nprobe: %d" % (self._collection_name, nq, top_k, nprobe))
        vec_field_name = utils.get_default_field_name()
        vector_query = {"vector": {vec_field_name: {
            "topk": top_k,
            "query": query_vectors,
            "metric_type": utils.metric_type_trans(metric_type),
            "params": search_param}
        }}
        self.load_and_query(vector_query, timeout=timeout)

    # TODO: need to check
    def count(self, collection_name=None):
        if collection_name is None:
            collection_name = self._collection_name
        row_count = self._milvus.get_collection_stats(collection_name)["row_count"]
        logger.debug("Row count: %d in collection: <%s>" % (row_count, collection_name))
        return row_count

    def drop(self, loop_cunt=120, collection_name=None):
        """
        drop steps:
        1.drop collection
        2.check collection exist
        3.Set timeout to exit
        """
        loop_cunt = int(loop_cunt)
        if collection_name is None:
            collection_name = self._collection_name
        logger.info("Start delete collection: %s" % collection_name)
        self._milvus.drop_collection(collection_name)
        i = 0
        while i < loop_cunt:
            try:
                res = self._milvus.has_collection(collection_name)
                if res:
                    time.sleep(1)
                    i = i + 1
                    continue
                else:
                    break
            except Exception as e:
                logger.warning("Collection count failed: {}".format(str(e)))
                break
        if i >= loop_cunt:
            logger.error("Delete collection timeout")

    def get_stats(self):
        return self._milvus.get_collection_stats(self._collection_name)

    def get_info(self, collection_name=None):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.describe_collection(collection_name)

    def show_collections(self):
        return self._milvus.list_collections()

    def exists_collection(self, collection_name=None):
        if collection_name is None:
            collection_name = self._collection_name
        res = self._milvus.has_collection(collection_name)
        return res

    def clean_db(self):
        collection_names = self.show_collections()
        for name in collection_names:
            self.drop(collection_name=name)

    @time_wrapper
    def load_collection(self, collection_name=None, timeout=3000):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.load_collection(collection_name, timeout=timeout)

    @time_wrapper
    def release_collection(self, collection_name=None, timeout=3000):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.release_collection(collection_name, timeout=timeout)

    @time_wrapper
    def load_partitions(self, tag_names, collection_name=None, timeout=3000):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.load_partitions(collection_name, tag_names, timeout=timeout)

    @time_wrapper
    def release_partitions(self, tag_names, collection_name=None, timeout=3000):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.release_partitions(collection_name, tag_names, timeout=timeout)

    @time_wrapper
    def get_query_segment_info(self, collection_name=None, timeout=300, **kwargs):
        if collection_name is None:
            collection_name = self._collection_name
        return self._milvus.get_query_segment_info(collection_name, timeout=timeout, **kwargs)

    @time_wrapper
    def scene_test(self, collection_name=None, vectors=None, ids=None):
        """
        Scene test steps：
        1.create collection with the specified collection name
        2.insert data
        3.flush data
        4.create index
        5.drop collection
        """

        logger.debug("[scene_test] Start scene test : %s" % collection_name)
        self.create_collection(dimension=128, collection_name=collection_name)
        time.sleep(1)

        collection_info = self.get_info(collection_name)

        entities = utils.generate_entities(collection_info, vectors, ids)
        logger.debug("[scene_test] Start insert : %s" % collection_name)
        self.insert(entities, collection_name=collection_name)
        logger.debug("[scene_test] Start flush : %s" % collection_name)
        self.flush(collection_name=collection_name)

        logger.debug("[scene_test] Start create index : %s" % collection_name)
        self.create_index(field_name='float_vector', index_type="ivf_sq8", metric_type='l2',
                          collection_name=collection_name, index_param={'nlist': 2048})
        time.sleep(59)

        logger.debug("[scene_test] Start drop : %s" % collection_name)
        self.drop(collection_name=collection_name)
        logger.debug("[scene_test]Scene test close : %s" % collection_name)

    # TODO: remove
    # def get_server_version(self):
    #     return self._milvus.server_version()

    # def get_server_mode(self):
    #     return self.cmd("mode")

    # def get_server_commit(self):
    #     return self.cmd("build_commit_id")

    # def get_server_config(self):
    #     return json.loads(self.cmd("get_milvus_config"))

    # def get_mem_info(self):
    #     result = json.loads(self.cmd("get_system_info"))
    #     result_human = {
    #         # unit: Gb
    #         "memory_used": round(int(result["memory_used"]) / (1024 * 1024 * 1024), 2)
    #     }
    #     return result_human

    # def cmd(self, command):
    #     res = self._milvus._cmd(command)
    #     logger.info("Server command: %s, result: %s" % (command, res))
    #     return res

    # @time_wrapper
    # def set_config(self, parent_key, child_key, value):
    #     self._milvus.set_config(parent_key, child_key, value)

    # def get_config(self, key):
    #     return self._milvus.get_config(key)