TBNilles/LEANN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(embedding-server): ensure shutdown-capable ZMQ threads create/bind their own REP sockets and poll with timeouts; fix undefined socket causing startup crash and CI hangs on Ubuntu 22.04

Browse Source
This commit is contained in:
Andy Lee
2025-08-13 13:53:08 -07:00
parent b381278c3e
commit 4b714f3b44
2 changed files with 211 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
packages/leann-backend-diskann/leann_backend_diskann/diskann_embedding_server.py
View File

@@ -222,22 +222,33 @@ def create_diskann_embedding_server(
raise
def zmq_server_thread_with_shutdown(shutdown_event):
"""ZMQ server thread that respects shutdown signal."""
"""ZMQ server thread that respects shutdown signal.
This creates its own REP socket, binds to zmq_port, and periodically
checks shutdown_event using recv timeouts to exit cleanly.
"""
logger.info("DiskANN ZMQ server thread started with shutdown support")
# Set receive timeout so we can check shutdown_event periodically
socket.setsockopt(zmq.RCVTIMEO, 1000) # 1 second timeout
context = zmq.Context()
rep_socket = context.socket(zmq.REP)
rep_socket.bind(f"tcp://*:{zmq_port}")
logger.info(f"DiskANN ZMQ REP server listening on port {zmq_port}")
# Set receive timeout so we can check shutdown_event periodically
rep_socket.setsockopt(zmq.RCVTIMEO, 1000) # 1 second timeout
rep_socket.setsockopt(zmq.SNDTIMEO, 300000)
try:
while not shutdown_event.is_set():
try:
e2e_start = time.time()
# REP socket receives single-part messages
message = socket.recv(zmq.NOBLOCK)
message = rep_socket.recv()
# Check for empty messages - REP socket requires response to every request
if not message:
logger.warning("Received empty message, sending empty response")
socket.send(b"")
rep_socket.send(b"")
continue
# Try protobuf first (same logic as original)
@@ -252,7 +263,8 @@ def create_diskann_embedding_server(
# Look up texts by node IDs
for nid in node_ids:
try:
txt = passages.get_text(nid)
passage_data = passages.get_passage(str(nid))
txt = passage_data["text"]
if not txt:
raise RuntimeError(f"FATAL: Empty text for passage ID {nid}")
texts.append(txt)
@@ -271,14 +283,16 @@ def create_diskann_embedding_server(
):
texts = request
is_text_request = True
logger.info(f"ZMQ received msgpack text request for {len(texts)} texts")
logger.info(
f"ZMQ received msgpack text request for {len(texts)} texts"
)
else:
raise ValueError("Not a valid msgpack text request")
except Exception:
logger.error("Both protobuf and msgpack parsing failed!")
# Send error response
resp_proto = embedding_pb2.NodeEmbeddingResponse()
socket.send(resp_proto.SerializeToString())
rep_socket.send(resp_proto.SerializeToString())
continue
# Process the request
@@ -296,7 +310,7 @@ def create_diskann_embedding_server(
else:
resp_proto = embedding_pb2.NodeEmbeddingResponse()
response_data = resp_proto.SerializeToString()
socket.send(response_data)
rep_socket.send(response_data)
continue
# Prepare response based on request type
@@ -317,7 +331,7 @@ def create_diskann_embedding_server(
response_data = resp_proto.SerializeToString()
# Send response back to the client
socket.send(response_data)
rep_socket.send(response_data)
e2e_end = time.time()
logger.info(f"⏱️ ZMQ E2E time: {e2e_end - e2e_start:.6f}s")
@@ -331,12 +345,21 @@ def create_diskann_embedding_server(
try:
# Send error response for REP socket
resp_proto = embedding_pb2.NodeEmbeddingResponse()
socket.send(resp_proto.SerializeToString())
rep_socket.send(resp_proto.SerializeToString())
except Exception:
pass
else:
logger.info("Shutdown in progress, ignoring ZMQ error")
break
finally:
try:
rep_socket.close(0)
except Exception:
pass
try:
context.term()
except Exception:
pass
logger.info("DiskANN ZMQ server thread exiting gracefully")

49
packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_embedding_server.py
View File

@@ -241,33 +241,55 @@ def create_hnsw_embedding_server(
socket.send(msgpack.packb([[], []]))
def zmq_server_thread_with_shutdown(shutdown_event):
"""ZMQ server thread that respects shutdown signal."""
"""ZMQ server thread that respects shutdown signal.
Creates its own REP socket bound to zmq_port and polls with timeouts
to allow graceful shutdown.
"""
logger.info("ZMQ server thread started with shutdown support")
# Set receive timeout so we can check shutdown_event periodically
socket.setsockopt(zmq.RCVTIMEO, 1000) # 1 second timeout
context = zmq.Context()
rep_socket = context.socket(zmq.REP)
rep_socket.bind(f"tcp://*:{zmq_port}")
logger.info(f"HNSW ZMQ REP server listening on port {zmq_port}")
rep_socket.setsockopt(zmq.RCVTIMEO, 1000)
rep_socket.setsockopt(zmq.SNDTIMEO, 300000)
try:
while not shutdown_event.is_set():
try:
e2e_start = time.time()
logger.debug("🔍 Waiting for ZMQ message...")
request_bytes = socket.recv(zmq.NOBLOCK)
request_bytes = rep_socket.recv()
# Rest of the processing logic (same as original)
request = msgpack.unpackb(request_bytes)
if len(request) == 1 and request[0] == "__QUERY_MODEL__":
response_bytes = msgpack.packb([model_name])
socket.send(response_bytes)
rep_socket.send(response_bytes)
continue
node_ids = request
# Handle direct text embedding request
if (
isinstance(request, list)
and request
and all(isinstance(item, str) for item in request)
):
embeddings = compute_embeddings(request, model_name, mode=embedding_mode)
rep_socket.send(msgpack.packb(embeddings.tolist()))
e2e_end = time.time()
logger.info(f"⏱️ Text embedding E2E time: {e2e_end - e2e_start:.6f}s")
continue
node_ids = request if isinstance(request, list) else []
logger.info(f"ZMQ received {len(node_ids)} node IDs")
texts = []
for nid in node_ids:
try:
txt = passages.get_text(nid)
passage_data = passages.get_passage(str(nid))
txt = passage_data["text"]
if not txt:
raise RuntimeError(f"FATAL: Empty text for passage ID {nid}")
texts.append(txt)
@@ -295,7 +317,7 @@ def create_hnsw_embedding_server(
]
response_bytes = msgpack.packb(response_payload, use_single_float=True)
socket.send(response_bytes)
rep_socket.send(response_bytes)
e2e_end = time.time()
logger.info(f"⏱️ ZMQ E2E time: {e2e_end - e2e_start:.6f}s")
@@ -306,12 +328,21 @@ def create_hnsw_embedding_server(
if not shutdown_event.is_set():
logger.error(f"Error in ZMQ server loop: {e}")
try:
socket.send(msgpack.packb([[], []]))
rep_socket.send(msgpack.packb([[], []]))
except Exception:
pass
else:
logger.info("Shutdown in progress, ignoring ZMQ error")
break
finally:
try:
rep_socket.close(0)
except Exception:
pass
try:
context.term()
except Exception:
pass
logger.info("ZMQ server thread exiting gracefully")