AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-12 18:49:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Added current task into task executor's hearbeat (#3444)

Browse Source 
### What problem does this PR solve?

Added current task into task executor's hearbeat

### Type of change

- [x] Refactoring
This commit is contained in:
Zhichang Yu 2024-11-15 22:55:41 +08:00 committed by GitHub
parent af18217d78
commit 77bdeb32bd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 32 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
rag/svr/task_executor.py
View File

@ -83,11 +83,14 @@ FACTORY = {
CONSUMER_NAME = "task_consumer_" + CONSUMER_NO CONSUMER_NAME = "task_consumer_" + CONSUMER_NO
PAYLOAD: Payload | None = None PAYLOAD: Payload | None = None
BOOT_AT = datetime.now().isoformat() BOOT_AT = datetime.now().isoformat()
DONE_TASKS = 0
FAILED_TASKS = 0
PENDING_TASKS = 0 PENDING_TASKS = 0
LAG_TASKS = 0 LAG_TASKS = 0
mt_lock = threading.Lock()
DONE_TASKS = 0
FAILED_TASKS = 0
CURRENT_TASK = None
def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing..."): def set_progress(task_id, from_page=0, to_page=-1, prog=None, msg="Processing..."):
global PAYLOAD global PAYLOAD
@ -135,11 +138,13 @@ def collect():
return None return None
if TaskService.do_cancel(msg["id"]): if TaskService.do_cancel(msg["id"]):
with mt_lock:
DONE_TASKS += 1 DONE_TASKS += 1
logging.info("Task {} has been canceled.".format(msg["id"])) logging.info("Task {} has been canceled.".format(msg["id"]))
return None return None
task = TaskService.get_task(msg["id"]) task = TaskService.get_task(msg["id"])
if not task: if not task:
with mt_lock:
DONE_TASKS += 1 DONE_TASKS += 1
logging.warning("{} empty task!".format(msg["id"])) logging.warning("{} empty task!".format(msg["id"]))
return None return None
@ -427,16 +432,22 @@ def do_handle_task(r):
def handle_task(): def handle_task():
global PAYLOAD, DONE_TASKS, FAILED_TASKS global PAYLOAD, mt_lock, DONE_TASKS, FAILED_TASKS, CURRENT_TASK
task = collect() task = collect()
if task: if task:
try: try:
logging.info(f"handle_task begin for task {json.dumps(task)}") logging.info(f"handle_task begin for task {json.dumps(task)}")
with mt_lock:
CURRENT_TASK = copy.deepcopy(task)
do_handle_task(task) do_handle_task(task)
with mt_lock:
DONE_TASKS += 1 DONE_TASKS += 1
logging.exception(f"handle_task done for task {json.dumps(task)}") CURRENT_TASK = None
logging.info(f"handle_task done for task {json.dumps(task)}")
except Exception: except Exception:
with mt_lock:
FAILED_TASKS += 1 FAILED_TASKS += 1
CURRENT_TASK = None
logging.exception(f"handle_task got exception for task {json.dumps(task)}") logging.exception(f"handle_task got exception for task {json.dumps(task)}")
if PAYLOAD: if PAYLOAD:
PAYLOAD.ack() PAYLOAD.ack()
@ -444,7 +455,7 @@ def handle_task():
def report_status(): def report_status():
global CONSUMER_NAME, BOOT_AT, DONE_TASKS, FAILED_TASKS, PENDING_TASKS, LAG_TASKS global CONSUMER_NAME, BOOT_AT, PENDING_TASKS, LAG_TASKS, mt_lock, DONE_TASKS, FAILED_TASKS, CURRENT_TASK
REDIS_CONN.sadd("TASKEXE", CONSUMER_NAME) REDIS_CONN.sadd("TASKEXE", CONSUMER_NAME)
while True: while True:
try: try:
@ -454,14 +465,16 @@ def report_status():
PENDING_TASKS = int(group_info["pending"]) PENDING_TASKS = int(group_info["pending"])
LAG_TASKS = int(group_info["lag"]) LAG_TASKS = int(group_info["lag"])
with mt_lock:
heartbeat = json.dumps({ heartbeat = json.dumps({
"name": CONSUMER_NAME, "name": CONSUMER_NAME,
"now": now.isoformat(), "now": now.isoformat(),
"boot_at": BOOT_AT, "boot_at": BOOT_AT,
"done": DONE_TASKS,
"failed": FAILED_TASKS,
"pending": PENDING_TASKS, "pending": PENDING_TASKS,
"lag": LAG_TASKS, "lag": LAG_TASKS,
"done": DONE_TASKS,
"failed": FAILED_TASKS,
"current": CURRENT_TASK,
}) })
REDIS_CONN.zadd(CONSUMER_NAME, heartbeat, now.timestamp()) REDIS_CONN.zadd(CONSUMER_NAME, heartbeat, now.timestamp())
logging.info(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}") logging.info(f"{CONSUMER_NAME} reported heartbeat: {heartbeat}")
@ -474,6 +487,7 @@ def report_status():
time.sleep(30) time.sleep(30)
def main(): def main():
settings.init_settings()
background_thread = threading.Thread(target=report_status) background_thread = threading.Thread(target=report_status)
background_thread.daemon = True background_thread.daemon = True
background_thread.start() background_thread.start()