def test_not_connect(self): client = Milvus() with pytest.raises(NotConnectError): client.create_table({}) with pytest.raises(NotConnectError): client.has_table("a") with pytest.raises(NotConnectError): client.describe_table("a") with pytest.raises(NotConnectError): client.drop_table("a") with pytest.raises(NotConnectError): client.create_index("a") with pytest.raises(NotConnectError): client.insert("a", [], None) with pytest.raises(NotConnectError): client.count_table("a") with pytest.raises(NotConnectError): client.show_tables() with pytest.raises(NotConnectError): client.search("a", 1, 2, []) with pytest.raises(NotConnectError): client.search_in_files("a", [], [], 2, 1) with pytest.raises(NotConnectError): client._cmd("") with pytest.raises(NotConnectError): client.preload_table("a") with pytest.raises(NotConnectError): client.describe_index("a") with pytest.raises(NotConnectError): client.drop_index("")
class MilvusClient(object): def __init__(self, table_name=None, ip=None, port=None): self._milvus = Milvus() self._table_name = table_name try: if not ip: self._milvus.connect( host = SERVER_HOST_DEFAULT, port = SERVER_PORT_DEFAULT) else: self._milvus.connect( host = ip, port = port) except Exception as e: raise e def __str__(self): return 'Milvus table %s' % self._table_name def check_status(self, status): if not status.OK(): logger.error(status.message) raise Exception("Status not ok") def create_table(self, table_name, dimension, index_file_size, metric_type): if not self._table_name: self._table_name = table_name if metric_type == "l2": metric_type = MetricType.L2 elif metric_type == "ip": metric_type = MetricType.IP else: logger.error("Not supported metric_type: %s" % metric_type) create_param = {'table_name': table_name, 'dimension': dimension, 'index_file_size': index_file_size, "metric_type": metric_type} status = self._milvus.create_table(create_param) self.check_status(status) @time_wrapper def insert(self, X, ids=None): status, result = self._milvus.add_vectors(self._table_name, X, ids) self.check_status(status) return status, result @time_wrapper def create_index(self, index_type, nlist): if index_type == "flat": index_type = IndexType.FLAT elif index_type == "ivf_flat": index_type = IndexType.IVFLAT elif index_type == "ivf_sq8": index_type = IndexType.IVF_SQ8 elif index_type == "nsg": index_type = IndexType.NSG elif index_type == "ivf_sq8h": index_type = IndexType.IVF_SQ8H elif index_type == "ivf_pq": index_type = IndexType.IVF_PQ index_params = { "index_type": index_type, "nlist": nlist, } logger.info("Building index start, table_name: %s, index_params: %s" % (self._table_name, json.dumps(index_params))) status = self._milvus.create_index(self._table_name, index=index_params) self.check_status(status) def describe_index(self): return self._milvus.describe_index(self._table_name) def drop_index(self): logger.info("Drop index: %s" % self._table_name) return self._milvus.drop_index(self._table_name) @time_wrapper def query(self, X, top_k, nprobe): status, result = self._milvus.search_vectors(self._table_name, top_k, nprobe, X) self.check_status(status) return status, result def count(self): return self._milvus.get_table_row_count(self._table_name)[1] def delete(self, timeout=60): logger.info("Start delete table: %s" % self._table_name) self._milvus.delete_table(self._table_name) i = 0 while i < timeout: if self.count(): time.sleep(1) i = i + 1 continue else: break if i < timeout: logger.error("Delete table timeout") def describe(self): return self._milvus.describe_table(self._table_name) def show_tables(self): return self._milvus.show_tables() def exists_table(self): status, res = self._milvus.has_table(self._table_name) self.check_status(status) return res @time_wrapper def preload_table(self): return self._milvus.preload_table(self._table_name, timeout=3000)
class MilvusClient(object): def __init__(self, table_name=None, ip=None, port=None, timeout=60): self._milvus = Milvus() self._table_name = table_name try: i = 1 start_time = time.time() if not ip: self._milvus.connect(host=SERVER_HOST_DEFAULT, port=SERVER_PORT_DEFAULT) else: # retry connect for remote server while time.time() < start_time + timeout: try: self._milvus.connect(host=ip, port=port) if self._milvus.connected() is True: logger.debug( "Try connect times: %d, %s" % (i, round(time.time() - start_time, 2))) break except Exception as e: logger.debug("Milvus connect failed") i = i + 1 except Exception as e: raise e def __str__(self): return 'Milvus table %s' % self._table_name def check_status(self, status): if not status.OK(): logger.error(status.message) # raise Exception("Status not ok") def create_table(self, table_name, dimension, index_file_size, metric_type): if not self._table_name: self._table_name = table_name if metric_type == "l2": metric_type = MetricType.L2 elif metric_type == "ip": metric_type = MetricType.IP elif metric_type == "jaccard": metric_type = MetricType.JACCARD elif metric_type == "hamming": metric_type = MetricType.HAMMING else: logger.error("Not supported metric_type: %s" % metric_type) create_param = { 'table_name': table_name, 'dimension': dimension, 'index_file_size': index_file_size, "metric_type": metric_type } status = self._milvus.create_table(create_param) self.check_status(status) @time_wrapper def insert(self, X, ids=None): status, result = self._milvus.add_vectors(self._table_name, X, ids) self.check_status(status) return status, result @time_wrapper def create_index(self, index_type, nlist): index_params = { "index_type": INDEX_MAP[index_type], "nlist": nlist, } logger.info("Building index start, table_name: %s, index_params: %s" % (self._table_name, json.dumps(index_params))) status = self._milvus.create_index(self._table_name, index=index_params) self.check_status(status) def describe_index(self): status, result = self._milvus.describe_index(self._table_name) index_type = None for k, v in INDEX_MAP.items(): if result._index_type == v: index_type = k break nlist = result._nlist res = {"index_type": index_type, "nlist": nlist} return res def drop_index(self): logger.info("Drop index: %s" % self._table_name) return self._milvus.drop_index(self._table_name) @time_wrapper def query(self, X, top_k, nprobe): status, result = self._milvus.search_vectors(self._table_name, top_k, nprobe, X) self.check_status(status) return result def count(self): return self._milvus.get_table_row_count(self._table_name)[1] def delete(self, timeout=60): logger.info("Start delete table: %s" % self._table_name) self._milvus.delete_table(self._table_name) i = 0 while i < timeout: if self.count(): time.sleep(1) i = i + 1 continue else: break if i >= timeout: logger.error("Delete table timeout") def describe(self): return self._milvus.describe_table(self._table_name) def show_tables(self): return self._milvus.show_tables() def exists_table(self, table_name=None): if table_name is None: table_name = self._table_name status, res = self._milvus.has_table(table_name) self.check_status(status) return res @time_wrapper def preload_table(self): return self._milvus.preload_table(self._table_name, timeout=3000) def get_server_version(self): status, res = self._milvus.server_version() return res 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_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): status, res = self._milvus._cmd(command) logger.info("Server command: %s, result: %s" % (command, res)) self.check_status(status) return res
class MilvusClient(object): def __init__(self, table_name=None, host=None, port=None): self._milvus = Milvus() self._table_name = table_name try: if not host: self._milvus.connect(host=SERVER_HOST_DEFAULT, port=SERVER_PORT_DEFAULT) else: self._milvus.connect(host=host, port=port) except Exception as e: raise e def __str__(self): return 'Milvus table %s' % self._table_name def check_status(self, status): if not status.OK(): logger.error(status.message) raise Exception("Status not ok") def create_table(self, table_name, dimension, index_file_size, metric_type): if not self._table_name: self._table_name = table_name if metric_type == "l2": metric_type = MetricType.L2 elif metric_type == "ip": metric_type = MetricType.IP else: logger.error("Not supported metric_type: %s" % metric_type) self._metric_type = metric_type create_param = { 'table_name': table_name, 'dimension': dimension, 'index_file_size': index_file_size, "metric_type": metric_type } status = self._milvus.create_table(create_param) self.check_status(status) @time_wrapper def insert(self, X, ids): if self._metric_type == MetricType.IP: logger.info("Set normalize for metric_type: Inner Product") X = sklearn.preprocessing.normalize(X, axis=1, norm='l2') X = X.astype(numpy.float32) status, result = self._milvus.add_vectors(self._table_name, X.tolist(), ids=ids) self.check_status(status) return status, result @time_wrapper def create_index(self, index_type, nlist): if index_type == "flat": index_type = IndexType.FLAT elif index_type == "ivf_flat": index_type = IndexType.IVFLAT elif index_type == "ivf_sq8": index_type = IndexType.IVF_SQ8 elif index_type == "ivf_sq8h": index_type = IndexType.IVF_SQ8H elif index_type == "nsg": index_type = IndexType.NSG elif index_type == "ivf_pq": index_type = IndexType.IVF_PQ index_params = { "index_type": index_type, "nlist": nlist, } logger.info("Building index start, table_name: %s, index_params: %s" % (self._table_name, json.dumps(index_params))) status = self._milvus.create_index(self._table_name, index=index_params, timeout=6 * 3600) self.check_status(status) def describe_index(self): return self._milvus.describe_index(self._table_name) def drop_index(self): logger.info("Drop index: %s" % self._table_name) return self._milvus.drop_index(self._table_name) @time_wrapper def query(self, X, top_k, nprobe): if self._metric_type == MetricType.IP: logger.info("Set normalize for metric_type: Inner Product") X = sklearn.preprocessing.normalize(X, axis=1, norm='l2') X = X.astype(numpy.float32) status, results = self._milvus.search_vectors(self._table_name, top_k, nprobe, X.tolist()) self.check_status(status) ids = [] for result in results: tmp_ids = [] for item in result: tmp_ids.append(item.id) ids.append(tmp_ids) return ids def count(self): return self._milvus.get_table_row_count(self._table_name)[1] def delete(self, table_name): logger.info("Start delete table: %s" % table_name) return self._milvus.delete_table(table_name) def describe(self): return self._milvus.describe_table(self._table_name) def exists_table(self, table_name): return self._milvus.has_table(table_name) def get_server_version(self): status, res = self._milvus.server_version() self.check_status(status) return res @time_wrapper def preload_table(self): return self._milvus.preload_table(self._table_name)