AnswerDotAI / AnswerDotAI/RAGatouille
`RAG.search` is not thread-safe
- 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