Merge pull request #1447 from danielaskdd/improve-pipeline-file-batch

Improve parallel handling for documents processing

lightrag/api/routers/document_routes.py

@@ -674,27 +674,9 @@ async def run_scanning_process(rag: LightRAG, doc_manager: DocumentManager):
         if not new_files:
             return
 
-        # Get MAX_PARALLEL_INSERT from global_args
+        # Process all files at once
-        max_parallel = global_args.max_parallel_insert
+        await pipeline_index_files(rag, new_files)
-        # Calculate batch size as 2 * MAX_PARALLEL_INSERT
+        logger.info(f"Scanning process completed: {total_files} files Processed.")
-        batch_size = 2 * max_parallel
-
-        # Process files in batches
-        for i in range(0, total_files, batch_size):
-            batch_files = new_files[i : i + batch_size]
-            batch_num = i // batch_size + 1
-            total_batches = (total_files + batch_size - 1) // batch_size
-
-            logger.info(
-                f"Processing batch {batch_num}/{total_batches} with {len(batch_files)} files"
-            )
-            await pipeline_index_files(rag, batch_files)
-
-            # Log progress
-            processed = min(i + batch_size, total_files)
-            logger.info(
-                f"Processed {processed}/{total_files} files ({processed/total_files*100:.1f}%)"
-            )
 
     except Exception as e:
         logger.error(f"Error during scanning process: {str(e)}")

lightrag/lightrag.py

@@ -841,8 +841,8 @@ class LightRAG:
                         "job_name": "Default Job",
                         "job_start": datetime.now().isoformat(),
                         "docs": 0,
-                        "batchs": 0,
+                        "batchs": 0,  # Total number of files to be processed
-                        "cur_batch": 0,
+                        "cur_batch": 0,  # Number of files already processed
                         "request_pending": False,  # Clear any previous request
                         "latest_message": "",
                     }
@@ -867,18 +867,13 @@ class LightRAG:
                     pipeline_status["history_messages"].append(log_message)
                     break
 
-                # 2. split docs into chunks, insert chunks, update doc status
+                log_message = f"Processing {len(to_process_docs)} document(s)"
-                docs_batches = [
-                    list(to_process_docs.items())[i : i + self.max_parallel_insert]
-                    for i in range(0, len(to_process_docs), self.max_parallel_insert)
-                ]
-
-                log_message = f"Processing {len(to_process_docs)} document(s) in {len(docs_batches)} batches"
                 logger.info(log_message)
 
-                # Update pipeline status with current batch information
+                # Update pipeline_status, batchs now represents the total number of files to be processed
                 pipeline_status["docs"] = len(to_process_docs)
-                pipeline_status["batchs"] = len(docs_batches)
+                pipeline_status["batchs"] = len(to_process_docs)
+                pipeline_status["cur_batch"] = 0
                 pipeline_status["latest_message"] = log_message
                 pipeline_status["history_messages"].append(log_message)
 
@@ -892,6 +887,11 @@ class LightRAG:
                 job_name = f"{path_prefix}[{total_files} files]"
                 pipeline_status["job_name"] = job_name
 
+                # Create a counter to track the number of processed files
+                processed_count = 0
+                # Create a semaphore to limit the number of concurrent file processing
+                semaphore = asyncio.Semaphore(self.max_parallel_insert)
+
                 async def process_document(
                     doc_id: str,
                     status_doc: DocProcessingStatus,
@@ -899,14 +899,27 @@ class LightRAG:
                     split_by_character_only: bool,
                     pipeline_status: dict,
                     pipeline_status_lock: asyncio.Lock,
+                    semaphore: asyncio.Semaphore,
                 ) -> None:
                     """Process single document"""
+                    async with semaphore:
+                        nonlocal processed_count
+                        current_file_number = 0
                         try:
                             # Get file path from status document
-                        file_path = getattr(status_doc, "file_path", "unknown_source")
+                            file_path = getattr(
+                                status_doc, "file_path", "unknown_source"
+                            )
 
                             async with pipeline_status_lock:
-                            log_message = f"Processing file: {file_path}"
+                                # Update processed file count and save current file number
+                                processed_count += 1
+                                current_file_number = (
+                                    processed_count  # Save the current file number
+                                )
+                                pipeline_status["cur_batch"] = processed_count
+
+                                log_message = f"Processing file ({current_file_number}/{total_files}): {file_path}"
                                 logger.info(log_message)
                                 pipeline_status["history_messages"].append(log_message)
                                 log_message = f"Processing d-id: {doc_id}"
@@ -987,6 +1000,16 @@ class LightRAG:
                                     }
                                 }
                             )
+
+                            # Call _insert_done after processing each file
+                            await self._insert_done()
+
+                            async with pipeline_status_lock:
+                                log_message = f"Completed processing file {current_file_number}/{total_files}: {file_path}"
+                                logger.info(log_message)
+                                pipeline_status["latest_message"] = log_message
+                                pipeline_status["history_messages"].append(log_message)
+
                         except Exception as e:
                             # Log error and update pipeline status
                             error_msg = f"Failed to process document {doc_id}: {traceback.format_exc()}"
@@ -1021,20 +1044,9 @@ class LightRAG:
                                 }
                             )
 
-                # 3. iterate over batches
+                # Create processing tasks for all documents
-                total_batches = len(docs_batches)
-                for batch_idx, docs_batch in enumerate(docs_batches):
-                    current_batch = batch_idx + 1
-                    log_message = (
-                        f"Start processing batch {current_batch} of {total_batches}."
-                    )
-                    logger.info(log_message)
-                    pipeline_status["cur_batch"] = current_batch
-                    pipeline_status["latest_message"] = log_message
-                    pipeline_status["history_messages"].append(log_message)
-
                 doc_tasks = []
-                    for doc_id, status_doc in docs_batch:
+                for doc_id, status_doc in to_process_docs.items():
                     doc_tasks.append(
                         process_document(
                             doc_id,
@@ -1043,17 +1055,12 @@ class LightRAG:
                             split_by_character_only,
                             pipeline_status,
                             pipeline_status_lock,
+                            semaphore,
                         )
                     )
 
-                    # Process documents in one batch parallelly
+                # Wait for all document processing to complete
                 await asyncio.gather(*doc_tasks)
-                    await self._insert_done()
-
-                    log_message = f"Completed batch {current_batch} of {total_batches}."
-                    logger.info(log_message)
-                    pipeline_status["latest_message"] = log_message
-                    pipeline_status["history_messages"].append(log_message)
 
                 # Check if there's a pending request to process more documents (with lock)
                 has_pending_request = False
@@ -1107,9 +1114,11 @@ class LightRAG:
                 llm_response_cache=self.llm_response_cache,
             )
         except Exception as e:
-            logger.error(
+            error_msg = f"Failed to extract entities and relationships: {str(e)}"
-                f"Failed to extract entities and relationships : {traceback.format_exc()} 。"
+            logger.error(error_msg)
-            )
+            async with pipeline_status_lock:
+                pipeline_status["latest_message"] = error_msg
+                pipeline_status["history_messages"].append(error_msg)
             raise e
 
     async def _insert_done(