浏览代码

fix: single task executor getting all tasks from Redis queue (#7330)

### What problem does this PR solve?

Currently, as long as there are tasks in Redis, this loop will keep
getting the tasks. This will lead to a single task executor with many
tasks in the pending state. Then we need to wait for the pending tasks
to get them back in the queue.

In first place, if we set the `MAX_CONCURRENT_TASKS` to X, then only X
tasks should be picked from the queue, and others should be left in the
queue for other `task_executors` or be picked after 1 of the spots in
the current executor gets free. This PR ensures this behavior.

The additional changes were due to the Ruff linting in pre-commit. But I
believe these are expected to keep the coding style.

### Type of change

- [X] Bug Fix (non-breaking change which fixes an issue)
- [ ] New Feature (non-breaking change which adds functionality)
- [ ] Documentation Update
- [ ] Refactoring
- [ ] Performance Improvement
- [ ] Other (please describe):

Co-authored-by: Zhichang Yu <yuzhichang@gmail.com>
tags/v0.19.1
Wanderson Pinto dos Santos 4 个月前
父节点
当前提交
0e03542db5
没有帐户链接到提交者的电子邮件
共有 1 个文件被更改,包括 5 次插入4 次删除
  1. 5
    4
      rag/svr/task_executor.py

+ 5
- 4
rag/svr/task_executor.py 查看文件

@@ -100,7 +100,7 @@ CURRENT_TASKS = {}
MAX_CONCURRENT_TASKS = int(os.environ.get('MAX_CONCURRENT_TASKS', "5"))
MAX_CONCURRENT_CHUNK_BUILDERS = int(os.environ.get('MAX_CONCURRENT_CHUNK_BUILDERS', "1"))
MAX_CONCURRENT_MINIO = int(os.environ.get('MAX_CONCURRENT_MINIO', '10'))
task_limiter = trio.CapacityLimiter(MAX_CONCURRENT_TASKS)
task_limiter = trio.Semaphore(MAX_CONCURRENT_TASKS)
chunk_limiter = trio.CapacityLimiter(MAX_CONCURRENT_CHUNK_BUILDERS)
minio_limiter = trio.CapacityLimiter(MAX_CONCURRENT_MINIO)
kg_limiter = trio.CapacityLimiter(2)
@@ -736,9 +736,10 @@ def recover_pending_tasks():
stop_event.wait(60)
async def task_manager():
global task_limiter
async with task_limiter:
try:
await handle_task()
finally:
task_limiter.release()


async def main():
@@ -767,8 +768,8 @@ async def main():
async with trio.open_nursery() as nursery:
nursery.start_soon(report_status)
while not stop_event.is_set():
await task_limiter.acquire()
nursery.start_soon(task_manager)
await trio.sleep(0.1)
logging.error("BUG!!! You should not reach here!!!")

if __name__ == "__main__":

正在加载...
取消
保存