AnswerDotAI / AnswerDotAI/RAGatouille

`RAG.search` is not thread-safe

Open
#262 1 comment 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
4k
Forks
276
PR merge metrics
No merged PRs in 30d

Description

From the digging that I've done, it appears that the following code is _not thread safe_ (using Langchain):
```py
RAG = RAGPretrainedModel.from_index(index_path)
retriever = RAG.as_langchain_retriever(k=k)

chain = RunnablePassthrough.assign(
passages=itemgetter("query") | retriever
) | prompt | llm

await chain.abatch([{"query": ...}, {"query": ...}, ...])
```

This results in the following error (truncated for security, should only contain relevant files):
```
Traceback (most recent call last):
...
File ".../lib/python3.11/site-packages/ragatouille/RAGPretrainedModel.py", line 315, in search
return self.model.search(
^^^^^^^^^^^^^^^^^^
File ".../lib/python3.11/site-packages/ragatouille/models/colbert.py", line 394, in search
results = self.model_index.search(
^^^^^^^^^^^^^^^^^^^^^^^^
File ".../lib/python3.11/site-packages/ragatouille/models/index.py", line 343, in search
if k > (32 * self.searcher.config.ncells):
~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~
TypeError: unsupported operand type(s) for *: 'int' and 'NoneType'
```

The culprit, I believe, is in the `ModelIndex.search` function:
```py
class PLAIDModelIndex(ModelIndex):
...
def search(self, ...):
# This will result in a race condition when run in parallel on multiple threads!!!!
if self.searcher is None or force_reload:
self._load_searcher(
checkpoint,
collection,
index_name,
force_fast,
)
assert self.searcher is not None
...
```

My intuition is that this would be a pretty common use-case. For reference, the official ColBERT implementation of `server.py` [initializes the Searcher at the beginning before the API calls](https://github.com/stanford-futuredata/ColBERT/blob/main/server.py). I think this searcher should be initialized before a call to `RAG.search` is ever made to prevent this race condition, OR there should be a batch function on `RAG.as_langchain_retriever`.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in ragatouille/models/index.py at PLAIDModelIndex.search and _load_searcher, then trace the call from ragatouille/RAGPretrainedModel.py through models/colbert.py. Reproduce the concurrent LangChain abatch scenario and verify that parallel searches initialize or use the searcher safely without the NoneType ncells error.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning, search
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.