gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: resolve "Task exception was never retrieved" warnings in async task handling

Browse source 
- Handle multiple simultaneous exceptions correctly
- Maintain fast-fail behavior while ensuring proper exception cleanup to
  prevent asyncio warnings
This commit is contained in:
yangdx 2025-09-04 12:40:41 +08:00
parent c903b14849
commit 83b54975a2
1 changed files with 101 additions and 43 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

144
lightrag/operate.py
View file

@ -679,19 +679,34 @@ async def _rebuild_knowledge_from_chunks(
# Execute all tasks in parallel with semaphore control and early failure detection # Execute all tasks in parallel with semaphore control and early failure detection
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
# Check if any task raised an exception # Check if any task raised an exception and ensure all exceptions are retrieved
first_exception = None
for task in done: for task in done:
if task.exception(): try:
# If a task failed, cancel all pending tasks exception = task.exception()
for pending_task in pending: if exception is not None:
pending_task.cancel() if first_exception is None:
first_exception = exception
else:
# Task completed successfully, retrieve result to mark as processed
task.result()
except Exception as e:
if first_exception is None:
first_exception = e
# Wait for cancellation to complete # If any task failed, cancel all pending tasks and raise the first exception
if pending: if first_exception is not None:
await asyncio.wait(pending) # Cancel all pending tasks
for pending_task in pending:
pending_task.cancel()
# Re-raise the exception to notify the caller # Wait for cancellation to complete
raise task.exception() if pending:
await asyncio.wait(pending)
# Re-raise the first exception to notify the caller
raise first_exception
# Final status report # Final status report
status_message = f"KG rebuild completed: {rebuilt_entities_count} entities and {rebuilt_relationships_count} relationships rebuilt successfully." status_message = f"KG rebuild completed: {rebuilt_entities_count} entities and {rebuilt_relationships_count} relationships rebuilt successfully."
@ -1619,17 +1634,32 @@ async def merge_nodes_and_edges(
entity_tasks, return_when=asyncio.FIRST_EXCEPTION entity_tasks, return_when=asyncio.FIRST_EXCEPTION
) )
# Check if any task raised an exception # Check if any task raised an exception and ensure all exceptions are retrieved
first_exception = None
successful_results = []
for task in done: for task in done:
if task.exception(): try:
# If a task failed, cancel all pending tasks exception = task.exception()
for pending_task in pending: if exception is not None:
pending_task.cancel() if first_exception is None:
# Wait for cancellation to complete first_exception = exception
if pending: else:
await asyncio.wait(pending) successful_results.append(task.result())
# Re-raise the exception to notify the caller except Exception as e:
raise task.exception() if first_exception is None:
first_exception = e
# If any task failed, cancel all pending tasks and raise the first exception
if first_exception is not None:
# Cancel all pending tasks
for pending_task in pending:
pending_task.cancel()
# Wait for cancellation to complete
if pending:
await asyncio.wait(pending)
# Re-raise the first exception to notify the caller
raise first_exception
# If all tasks completed successfully, collect results # If all tasks completed successfully, collect results
processed_entities = [task.result() for task in entity_tasks] processed_entities = [task.result() for task in entity_tasks]
@ -1737,17 +1767,32 @@ async def merge_nodes_and_edges(
edge_tasks, return_when=asyncio.FIRST_EXCEPTION edge_tasks, return_when=asyncio.FIRST_EXCEPTION
) )
# Check if any task raised an exception # Check if any task raised an exception and ensure all exceptions are retrieved
first_exception = None
successful_results = []
for task in done: for task in done:
if task.exception(): try:
# If a task failed, cancel all pending tasks exception = task.exception()
for pending_task in pending: if exception is not None:
pending_task.cancel() if first_exception is None:
# Wait for cancellation to complete first_exception = exception
if pending: else:
await asyncio.wait(pending) successful_results.append(task.result())
# Re-raise the exception to notify the caller except Exception as e:
raise task.exception() if first_exception is None:
first_exception = e
# If any task failed, cancel all pending tasks and raise the first exception
if first_exception is not None:
# Cancel all pending tasks
for pending_task in pending:
pending_task.cancel()
# Wait for cancellation to complete
if pending:
await asyncio.wait(pending)
# Re-raise the first exception to notify the caller
raise first_exception
# If all tasks completed successfully, collect results # If all tasks completed successfully, collect results
for task in edge_tasks: for task in edge_tasks:
@ -1995,23 +2040,36 @@ async def extract_entities(
# This allows us to cancel remaining tasks if any task fails # This allows us to cancel remaining tasks if any task fails
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION) done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
# Check if any task raised an exception # Check if any task raised an exception and ensure all exceptions are retrieved
first_exception = None
chunk_results = []
for task in done: for task in done:
if task.exception(): try:
# If a task failed, cancel all pending tasks exception = task.exception()
# This prevents unnecessary processing since the parent function will abort anyway if exception is not None:
for pending_task in pending: if first_exception is None:
pending_task.cancel() first_exception = exception
else:
chunk_results.append(task.result())
except Exception as e:
if first_exception is None:
first_exception = e
# Wait for cancellation to complete # If any task failed, cancel all pending tasks and raise the first exception
if pending: if first_exception is not None:
await asyncio.wait(pending) # Cancel all pending tasks
for pending_task in pending:
pending_task.cancel()
# Re-raise the exception to notify the caller # Wait for cancellation to complete
raise task.exception() if pending:
await asyncio.wait(pending)
# If all tasks completed successfully, collect results # Re-raise the first exception to notify the caller
chunk_results = [task.result() for task in tasks] raise first_exception
# If all tasks completed successfully, chunk_results already contains the results
# Return the chunk_results for later processing in merge_nodes_and_edges # Return the chunk_results for later processing in merge_nodes_and_edges
return chunk_results return chunk_results