hnsw_exception.py

import numpy as np
import threading
import hnswlib

def func1():
	global index, count, lock
	while True:
		if count > 5000:
			print(f"Max elem reached {count}")
			return
		num_elements = 1
		# Generating sample data
		data = np.float32(np.random.random((num_elements, 512)))
		data_labels = np.arange(num_elements)
		try:
			index.add_items(data, data_labels)
		except Exception as e:
			print(repr(e))
		try:
			with lock:
				count += 1
		except:
			pass

def func2():
	global index,count
	while True:
		nsearch = count if count < 4 else 4
		#print(f"Searching for {kn}")
		data = np.float32(np.random.random((1, 512)))
		try:
			#print("Searching")
			index.set_ef(50) # ef should always be > k
			labels, distances = index.knn_query(data, k = nsearch)
		except Exception as e:
			print(f"Search: {nsearch} {count} {repr(e)}")

print('Preparing index with a single entry')
index = hnswlib.Index(space='l2', dim=512)
index.init_index(max_elements=5000, ef_construction=200, M=64)
count = 1
lock = threading.Lock()
num_elements = 1
# Generating sample data
data = np.float32(np.random.random((num_elements, 512)))
data_labels = np.arange(num_elements)
index.add_items(data, data_labels)

print('Starting threads')
thread2 = threading.Thread(target=func1)
thread1 = threading.Thread(target=func2)
thread3 = threading.Thread(target=func1)
thread1.start()
thread2.start()
thread3.start()

print('Joining threads')
thread1.join()
thread2.join()
thread3.join()