Quellcode durchsuchen

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

### What problem does this PR solve?

Added current task into task executor's hearbeat

### Type of change

- [x] Refactoring
tags/v0.14.0
Zhichang Yu vor 11 Monaten
Ursprung
Commit
77bdeb32bd
Es ist kein Account mit der E-Mail-Adresse des Committers verbunden
1 geänderte Dateien mit 32 neuen und 18 gelöschten Zeilen
  1. 32
    18
      rag/svr/task_executor.py

+ 32
- 18
rag/svr/task_executor.py Datei anzeigen

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
return None return None


if TaskService.do_cancel(msg["id"]): if TaskService.do_cancel(msg["id"]):
DONE_TASKS += 1
with mt_lock:
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:
DONE_TASKS += 1
with mt_lock:
DONE_TASKS += 1
logging.warning("{} empty task!".format(msg["id"])) logging.warning("{} empty task!".format(msg["id"]))
return None return None






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)
DONE_TASKS += 1
logging.exception(f"handle_task done for task {json.dumps(task)}")
with mt_lock:
DONE_TASKS += 1
CURRENT_TASK = None
logging.info(f"handle_task done for task {json.dumps(task)}")
except Exception: except Exception:
FAILED_TASKS += 1
with mt_lock:
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()




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:
PENDING_TASKS = int(group_info["pending"]) PENDING_TASKS = int(group_info["pending"])
LAG_TASKS = int(group_info["lag"]) LAG_TASKS = int(group_info["lag"])


heartbeat = json.dumps({
"name": CONSUMER_NAME,
"now": now.isoformat(),
"boot_at": BOOT_AT,
"done": DONE_TASKS,
"failed": FAILED_TASKS,
"pending": PENDING_TASKS,
"lag": LAG_TASKS,
})
with mt_lock:
heartbeat = json.dumps({
"name": CONSUMER_NAME,
"now": now.isoformat(),
"boot_at": BOOT_AT,
"pending": PENDING_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}")


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()

Laden…
Abbrechen
Speichern