gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix linting

Browse Source
This commit is contained in:
yangdx
2025-04-28 23:21:34 +08:00
parent a4a0643ad0
commit 1fc26127d5
1 changed files with 48 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
lightrag/utils.py
View File

@@ -271,8 +271,10 @@ def compute_mdhash_id(content: str, prefix: str = "") -> str:
# Custom exception class # Custom exception class
class QueueFullError(Exception): class QueueFullError(Exception):
"""Raised when the queue is full and the wait times out""" """Raised when the queue is full and the wait times out"""
pass pass
def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000): def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
""" """
Enhanced priority-limited asynchronous function call decorator Enhanced priority-limited asynchronous function call decorator
@@ -283,6 +285,7 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
Returns: Returns:
Decorator function Decorator function
""" """
def final_decro(func): def final_decro(func):
queue = asyncio.PriorityQueue(maxsize=max_queue_size) queue = asyncio.PriorityQueue(maxsize=max_queue_size)
tasks = set() tasks = set()
@@ -302,9 +305,13 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
try: try:
# Use timeout to get tasks, allowing periodic checking of shutdown signal # Use timeout to get tasks, allowing periodic checking of shutdown signal
try: try:
priority, count, future, args, kwargs = await asyncio.wait_for( (
queue.get(), timeout=1.0 priority,
) count,
future,
args,
kwargs,
) = await asyncio.wait_for(queue.get(), timeout=1.0)
except asyncio.TimeoutError: except asyncio.TimeoutError:
# Timeout is just to check shutdown signal, continue to next iteration # Timeout is just to check shutdown signal, continue to next iteration
continue continue
@@ -325,7 +332,9 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
future.cancel() future.cancel()
logger.debug("limit_async: Task cancelled during execution") logger.debug("limit_async: Task cancelled during execution")
except Exception as e: except Exception as e:
logger.error(f"limit_async: Error in decorated function: {str(e)}") logger.error(
f"limit_async: Error in decorated function: {str(e)}"
)
if not future.done(): if not future.done():
future.set_exception(e) future.set_exception(e)
finally: finally:
@@ -351,7 +360,9 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
active_tasks_count = len(tasks) active_tasks_count = len(tasks)
workers_needed = max_size - active_tasks_count workers_needed = max_size - active_tasks_count
if workers_needed > 0: if workers_needed > 0:
logger.info(f"limit_async: Creating {workers_needed} new workers") logger.info(
f"limit_async: Creating {workers_needed} new workers"
)
for _ in range(workers_needed): for _ in range(workers_needed):
task = asyncio.create_task(worker()) task = asyncio.create_task(worker())
tasks.add(task) tasks.add(task)
@@ -373,12 +384,17 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
# Try to acquire the lock, wait up to 5 seconds # Try to acquire the lock, wait up to 5 seconds
lock_acquired = await asyncio.wait_for(lock.acquire(), timeout=5.0) lock_acquired = await asyncio.wait_for(lock.acquire(), timeout=5.0)
except asyncio.TimeoutError: except asyncio.TimeoutError:
logger.error("limit_async: Timeout acquiring lock in ensure_workers") logger.error(
"limit_async: Timeout acquiring lock in ensure_workers"
)
# Even if acquiring the lock times out, continue trying to create workers # Even if acquiring the lock times out, continue trying to create workers
try: try:
# Start the health check task (if not already started) # Start the health check task (if not already started)
if worker_health_check_task is None or worker_health_check_task.done(): if (
worker_health_check_task is None
or worker_health_check_task.done()
):
worker_health_check_task = asyncio.create_task(health_check()) worker_health_check_task = asyncio.create_task(health_check())
# Directly remove completed tasks from the tasks set # Directly remove completed tasks from the tasks set
@@ -422,7 +438,9 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
try: try:
await asyncio.wait_for(queue.join(), timeout=5.0) await asyncio.wait_for(queue.join(), timeout=5.0)
except asyncio.TimeoutError: except asyncio.TimeoutError:
logger.warning("limit_async: Timeout waiting for queue to empty during shutdown") logger.warning(
"limit_async: Timeout waiting for queue to empty during shutdown"
)
# Cancel all worker tasks # Cancel all worker tasks
for task in list(tasks): for task in list(tasks):
@@ -444,7 +462,9 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
logger.info("limit_async: Priority queue workers shutdown complete") logger.info("limit_async: Priority queue workers shutdown complete")
@wraps(func) @wraps(func)
async def wait_func(*args, _priority=10, _timeout=None, _queue_timeout=None, **kwargs): async def wait_func(
*args, _priority=10, _timeout=None, _queue_timeout=None, **kwargs
):
""" """
Execute the function with priority-based concurrency control Execute the function with priority-based concurrency control
Args: Args:
@@ -479,10 +499,12 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
try: try:
await asyncio.wait_for( await asyncio.wait_for(
queue.put((_priority, current_count, future, args, kwargs)), queue.put((_priority, current_count, future, args, kwargs)),
timeout=_queue_timeout timeout=_queue_timeout,
) )
except asyncio.TimeoutError: except asyncio.TimeoutError:
raise QueueFullError(f"Queue full, timeout after {_queue_timeout} seconds") raise QueueFullError(
f"Queue full, timeout after {_queue_timeout} seconds"
)
else: else:
# No timeout, may wait indefinitely # No timeout, may wait indefinitely
await queue.put((_priority, current_count, future, args, kwargs)) await queue.put((_priority, current_count, future, args, kwargs))
@@ -502,7 +524,9 @@ def priority_limit_async_func_call(max_size: int, max_queue_size: int = 1000):
# Cancel the future # Cancel the future
if not future.done(): if not future.done():
future.cancel() future.cancel()
raise TimeoutError(f"limit_async: Task timed out after {_timeout} seconds") raise TimeoutError(
f"limit_async: Task timed out after {_timeout} seconds"
)
else: else:
# Wait for the result without timeout # Wait for the result without timeout
return await future return await future