Fix session locks to handle ingestion race conditions

src/exceptions.py

@@ -13,3 +13,11 @@ class NoVectorStoresFoundError(Exception):
     """
 
     pass
+
+
+class IngestionPipelineError(Exception):
+    """
+    Raised when all ingestion pipelines fail.
+    """
+
+    pass

src/ingest.py

@@ -28,6 +28,7 @@
     DEFAULT_LLAMA_STACK_RETRY_DELAY,
     DEFAULT_LLAMA_STACK_WAITING_RETRIES,
 )
+from src.exceptions import IngestionPipelineError
 from src.types import (
     RAW_PIPELINES_TYPE,
     Pipeline,
@@ -585,6 +586,9 @@ async def create_vector_db(
         """
         creates vector database and inserts documents.
         uses async (concurrent) or sync (sequential) mode based on config.
+
+        all (LlamaStack) client calls are wrapped in thread (asyncio.to_thread)
+        to avoid blocking the concurrent execution.
         """
         if not documents:
             logger.warning(f"No documents to insert for {vector_store_name}")
@@ -594,6 +598,9 @@ async def create_vector_db(
 
         vector_store = None
         try:
+            # asyncio.to_thread moves the synchronous llama-stack-client
+            # call to a separate thread so the asyncio (event) loop keeps
+            # running.
             vector_store = await asyncio.to_thread(
                 self.client.vector_stores.create, name=vector_store_name
             )
@@ -605,6 +612,8 @@ async def create_vector_db(
                     f"Vector DB '{vector_store_name}' already exists, continuing..."
                 )
 
+                # vector store exists, find it by name to get its ID
+                # to continue ingestion
                 vector_stores = await asyncio.to_thread(self.client.vector_stores.list)
                 for vs in vector_stores or []:
                     if vs.name == vector_store_name:
@@ -617,8 +626,10 @@ async def create_vector_db(
                     )
                     return False
 
-                # check if the store already has files — skip insertion to
-                # avoid duplicate documents from parallel sessions
+                # if the existing store already contains files, we skip
+                # ingestion entirely. Simple rule to prevent duplicate documents
+                # when multiple parallel sessions race to populate the same store.
+                # TODO: investigate more robust approaches
                 try:
                     existing_files = await asyncio.to_thread(
                         self.client.vector_stores.files.list,
@@ -638,6 +649,10 @@ async def create_vector_db(
                 logger.error(f"Failed to register vector DB '{vector_store_name}': {e}")
                 return False
 
+        # ingest documents into the vector store.
+        # - async mode: all documents are inserted concurrently using
+        # asyncio.gather (each ingestion action runs in its own thread).
+        # - sync mode: documents are inserted one at a time sequentially.
         try:
             if self.ingestion_mode == "async":
                 logger.info(
@@ -833,7 +848,9 @@ async def process_with_error_handling(pipeline: Pipeline) -> bool:
 
         if successful == 0:
             logger.warning("all pipeline(s) failed. Check logs for details.")
-            sys.exit(1)
+            raise IngestionPipelineError(
+                "All ingestion pipelines failed. Check logs for details."
+            )
         elif failed > 0:
             logger.warning(f"{failed} pipeline(s) failed. Check logs for details.")
         else:

src/utils.py

@@ -31,6 +31,23 @@ def __setitem__(self, key, value):
 
 submission_states: "ObservableDict[str, WorkflowState]" = ObservableDict()
 
+# lock/mutex shared across all Streamlit sessions in this process.
+#
+# Fix: commit 61398cf565654f2e2d2da5247af488ec098b2d2e
+#
+# the lock was moved in utils.py (from streamlit_app.py) in order to
+# be cached similar to all imported modules (sys.modules) of Python:
+# - https://docs.python.org/3/reference/import.html#the-module-cache
+#
+# this allows us to keep the same lock between different streamlit
+# sessions (of the same streamlit process). With the previous approach,
+# (lock object was located in streamlit_app.py), each streamlit session
+# was generating again the lock object (since streamlit_app.py was
+# re-run for each session)
+
+# see: https://docs.streamlit.io/develop/concepts/design/multithreading
+ingestion_lock = threading.Lock()
+
 
 def clean_text(text: "str") -> "str":
     """

src/workflow.py

@@ -71,7 +71,8 @@ def _call_openai_llm(self, state: "WorkflowState") -> "str":
 
         messages = self._convert_messages_to_openai_format(state)
         completion = self.rag_service.openai_client.chat.completions.create(
-            model=INFERENCE_MODEL, messages=messages
+            model=INFERENCE_MODEL,
+            messages=messages,  # type: ignore[invalid-argument-type]
         )
         return completion.choices[0].message.content or ""
 
@@ -205,7 +206,7 @@ def llm_node(state: "WorkflowState") -> "WorkflowState":
                     rag_response = client_to_use.responses.create(
                         model=INFERENCE_MODEL,
                         input=rag_prompt,
-                        tools=[file_search_tool],
+                        tools=[file_search_tool],  # type: ignore[invalid-argument-type]
                     )
                     rag_end_time = time.time()
                     state["rag_query_time"] = rag_end_time - rag_start_time

streamlit_app.py

@@ -1,7 +1,6 @@
 import asyncio
 import json
 import os
-import threading
 import time
 import uuid
 from concurrent.futures import Future, ThreadPoolExecutor
@@ -24,17 +23,12 @@
 from src.types import Pipeline, WorkflowState
 from src.utils import (
     check_llama_stack_availability,
+    ingestion_lock,
     logger,
     submission_states,
 )
 from src.workflow import Workflow
 
-# lock shared across all Streamlit sessions in this process. Streamlit runs
-# each session in a separate thread but within the same process, so this
-# lock will handle parallel ingestion attempts.
-# see: https://docs.streamlit.io/develop/concepts/design/multithreading
-_ingestion_lock = threading.Lock()
-
 # API_KEY: OpenAI API key (not used directly but may be needed
 API_KEY = os.getenv("OPENAI_API_KEY", "not applicable")
 
@@ -930,7 +924,7 @@ def main() -> "None":
             # vector stores. All other sessions will get the lock only after
             # that's completed, check again (stores now exist), and finally
             # will skip ingestion.
-            with _ingestion_lock:
+            with ingestion_lock:
                 loop = get_or_create_event_loop()
                 loop.run_until_complete(check_and_run_ingestion_if_needed())
             st.rerun()

tests/test_ingest.py

@@ -5,6 +5,7 @@
 import pytest
 import yaml
 
+from src.exceptions import IngestionPipelineError
 from src.ingest import IngestionService
 from src.types import Pipeline, SourceConfig, SourceTypes
 
@@ -591,7 +592,7 @@ async def test_run_with_failed_pipelines(
 
         with patch.object(service, "process_pipeline", return_value=False):
             with patch.object(service, "save_file_metadata"):
-                with pytest.raises(SystemExit):
+                with pytest.raises(IngestionPipelineError):
                     await service.run()
 
     @pytest.mark.anyio