Run populate_vector_db in thread (#143)

Author: yoomlam

app/gunicorn.conf.py

@@ -10,11 +10,12 @@
 """
 
 import os
+import threading
 
 from src.app_config import config as app_config
 from src.ingestion import rag_utils
 
-bind = app_config.host + ':' + str(app_config.port)
+bind = app_config.host + ":" + str(app_config.port)
 # Calculates the number of usable cores and doubles it. Recommended number of workers per core is two.
 # https://docs.gunicorn.org/en/latest/design.html#how-many-workers
 # We use 'os.sched_getaffinity(pid)' not 'os.cpu_count()' because it returns only allowable CPUs.
@@ -28,4 +29,5 @@
 # https://stackoverflow.com/questions/24101724/gunicorn-with-multiple-workers-is-there-an-easy-way-to-execute-certain-code-onl
 def when_ready(server: object) -> None:
     print("when_ready()", server)
-    rag_utils.populate_vector_db()
+    populate_rag_db_thread = threading.Thread(target=rag_utils.populate_vector_db)
+    populate_rag_db_thread.start()

app/src/ingestion/rag_utils.py

@@ -39,10 +39,14 @@ def populate_vector_db() -> None:
     collection_name = doc_store._collection_name
 
     # Clear existing collection if any
-    if doc_store.count_documents() > 0:
+    if (doc_count := doc_store.count_documents()) > 0:
         try:
             # Don't delete collection since it's referenced by existing pipelines upon their startup
-            logger.info("Clearing out existing vector DB collection=%r", collection_name)
+            logger.info(
+                "Clearing out existing vector DB collection=%r with %d docs",
+                collection_name,
+                doc_count,
+            )
             # recreate_index=True results in a new id for the collection, which breaks existing pipelines
             doc_store.delete_all_documents(recreate_index=False)
         except DocumentStoreError as e:
@@ -65,7 +69,6 @@ def populate_vector_db() -> None:
     pipeline = _create_ingest_pipeline(doc_store)
     pipeline.run({"converter": {"sources": files_to_ingest}})
     logger.info("Ingested documents doc_count=%d", doc_store.count_documents())
-
     logger.info("ChromaDB collections: %s", chroma_client.list_collections())
 
 
@@ -78,16 +81,18 @@ def download_s3_folder_to_local(s3_folder: str = "files_to_ingest_into_vector_db
     except PermissionError as e:
         logger.error("Error creating directories for %s: %s", s3_folder, e)
         local_folder = f"/tmp/{s3_folder}"  # nosec B108
-    logger.info("Downloading s3://%s/%s to local folder %s", bucket, s3_folder, local_folder)
 
     if config.environment == "local":
         assert os.path.exists(
             local_folder
         ), f"Local folder {local_folder} should exist with manually downloaded files from S3"
         return local_folder
 
+    logger.info("Downloading s3://%s/%s to local folder %s", bucket, s3_folder, local_folder)
     s3 = file_util.get_s3_client()
     paginator = s3.get_paginator("list_objects_v2")
+    if not s3_folder.endswith("/"):
+        s3_folder += "/"
     try:
         for result in paginator.paginate(Bucket=bucket, Prefix=s3_folder):
             for obj in result.get("Contents", []):