herocode/herodb
12
0
Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases Wiki Activity

WIP4 implementation lanceDB: removed blocking Tokio runtime usage during embeddings and isolated all embedding work off the async runtime

Browse Source
This commit is contained in:
Maxime Van Hees
2025-09-29 15:54:12 +02:00
parent 644946f1ca
commit 4aa49e0d5c
5 changed files with 132 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
src/rpc.rs
View File

@@ -1057,10 +1057,22 @@ impl RpcServer for RpcServerImpl {
if !server.has_write_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "write permission denied", None::<()>));
}
// Resolve embedder and run blocking embedding off the async runtime
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
let embedder = server.get_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let vector = embedder.embed(&text)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let vector = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
};
server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.store_vector(&name, &id, vector, meta.unwrap_or_default(), Some(text)).await
@@ -1087,10 +1099,21 @@ impl RpcServer for RpcServerImpl {
if !server.has_read_permission() {
return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, "read permission denied", None::<()>));
}
// Resolve embedder and run embedding on a plain OS thread (avoid dropping any runtime in async context)
let embedder = server.get_embedder_for(&name)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let qv = embedder.embed(&text)
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?;
let (tx, rx) = tokio::sync::oneshot::channel();
let emb_arc = embedder.clone();
let text_cl = text.clone();
std::thread::spawn(move || {
let res = emb_arc.embed(&text_cl);
let _ = tx.send(res);
});
let qv = match rx.await {
Ok(Ok(v)) => v,
Ok(Err(e)) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>)),
Err(recv_err) => return Err(jsonrpsee::types::ErrorObjectOwned::owned(-32000, format!("embedding thread error: {}", recv_err), None::<()>)),
};
let results = server.lance_store()
.map_err(|e| jsonrpsee::types::ErrorObjectOwned::owned(-32000, e.0, None::<()>))?
.search_vectors(&name, qv, k, filter, return_fields).await