Celery动态子任务的同步等待机制：突破编排限制

import time
from celery import Celery, Task
from celery.result import AsyncResult
from typing import List

# 假设 app 已经初始化，并且配置了 Redis 作为 broker 和 result backend
app = Celery('my_app', broker='redis://localhost:6379/0', backend='redis://localhost:6379/1')

# 假设 JobMaster 和 consts 是用于自定义日志和状态管理的模块
# 在实际应用中，您可以替换为自己的日志系统或直接使用 print
class JobMaster:
    @staticmethod
    def get_job(job_id, job_title):
        # 模拟获取一个任务对象，用于记录日志
        print(f"[{job_title}] Getting job {job_id if job_id else 'new'}")
        return type('Job', (object,), {'log_message': lambda self, log_message, **kwargs: print(f"[{job_title}] {log_message}")})(), job_id if job_id else 1 # 模拟返回一个job对象和job_id

class consts:
    IN_PROGRESS = "IN_PROGRESS"
    COMPLETED = "COMPLETED"
    ERRORS_FOUND = "ERRORS_FOUND"

@app.task(bind=True)
def task_dummy_task1(self: Task, part_number: int, job_id: int = None):
    job, job_id = JobMaster.get_job(job_id, job_title="dummy task")
    sleeping_duration = 1 # 缩短等待时间以便测试
    subtask_ids = []
    job.log_message(log_message=f"Entered dummy task 1 with sleeping duration of {sleeping_duration}")

    job.log_message(log_message="In dummy task1, creating subtask a")
    subtask = task_dummy_subtask.apply_async(kwargs={"parent_task_name": "task1_a", "job_id": job_id},
                                             add_to_parent=True)
    subtask_ids.append(subtask.id)

    job.log_message(log_message="In dummy task1, creating subtask b")
    subtask = task_dummy_subtask.apply_async(kwargs={"parent_task_name": "task1_b", "job_id": job_id},
                                             add_to_parent=True)
    subtask_ids.append(subtask.id)

    job.log_message(log_message="In dummy task1, creating subtask c")
    subtask = task_dummy_subtask.apply_async(kwargs={"parent_task_name": "task1_c", "job_id": job_id},
                                             add_to_parent=True)
    subtask_ids.append(subtask.id)

    job.log_message(log_message="In dummy task1, creating intermediary subtask d")
    subtask = intermediary_dummy_subtask_function(parent_task_name="task1_d", job_id=job_id)
    subtask_ids.append(subtask.id)

    job.log_message(log_message="In dummy task1, creating intermediary subtask e")
    subtask = intermediary_dummy_subtask_function(parent_task_name="task1_e", job_id=job_id)
    subtask_ids.append(subtask.id)

    time.sleep(sleeping_duration) # 主任务执行一些自己的逻辑

    # 等待所有动态子任务完成
    wait_for_tasks_to_complete(async_ids=subtask_ids, job_id=job_id,
                                    msg="Waiting in dummy task1 for subtasks to complete")

    job.log_message(log_message="Finished dummy task1 main body")

    return part_number

@app.task
def task_dummy_subtask(parent_task_name: str, job_id: int):
    job, _ = JobMaster.get_job(job_id, job_title=f"subtask-{parent_task_name}")
    sleep_time = 2 # 模拟子任务耗时
    job.log_message(log_message=f"Subtask {parent_task_name} started, will sleep for {sleep_time}s")
    time.sleep(sleep_time)
    job.log_message(log_message=f"Subtask {parent_task_name} finished")
    return f"Result from {parent_task_name}"

def intermediary_dummy_subtask_function(parent_task_name, job_id) -> AsyncResult:
    job, _ = JobMaster.get_job(job_id, job_title="dummy task")
    job.log_message(
        log_message=f"Intermediary function for {parent_task_name} has been reached, will now make a task")
    r = task_dummy_subtask.apply_async(kwargs={"parent_task_name": parent_task_name, "job_id": job_id},
                                       add_to_parent=True)
    return r

def wait_for_tasks_to_complete(async_ids: List[str], job_id: int = None, msg: str = None, timeout: int = 300):
    job, _ = JobMaster.get_job(job_id, job_title="waiting for refresh data")
    job.log_message(log_message=f"Waiting for {len(async_ids)} tasks to complete, {msg}", status=consts.IN_PROGRESS,
                    job_score=0)
    job.log_message(log_message=f"tasks: {async_ids}", status=consts.IN_PROGRESS, job_score=0)

    # 创建一个可变的列表用于跟踪未完成的任务ID
    remaining_async_ids = list(async_ids) 

    count_down = timeout
    while count_down > 0:
        # 遍历 remaining_async_ids 的副本，因为我们可能在循环中修改它
        for async_id in list(remaining_async_ids): 
            result = app.AsyncResult(async_id)  # 获取任务结果对象
            status = result.status

            if status == "SUCCESS":
                # 任务成功完成
                returned_value = result.result
                job.log_message(log_message=f"Confirmed status SUCCESS for task {async_id} with {returned_value=}")
                remaining_async_ids.remove(async_id) # 从待处理列表中移除
            elif status in ["PENDING", "STARTED", "RETRY"]:
                # 任务仍在进行中或等待执行
                pass
            elif status in ["FAILURE", "REVOKED"]:
                # 任务失败或被撤销，需要根据业务逻辑处理
                job.log_message(log_message=f"Task {async_id} failed or revoked with status {status}. Error: {result.info}",
                                status=consts.ERRORS_FOUND)
                # 可以在这里选择抛出异常，或将失败任务从列表中移除并继续等待其他任务
                remaining_async_ids.remove(async_id) 
                # 示例：如果一个失败就认为整体失败，可以立即返回或抛出异常
                # raise Exception(f"Subtask {async_id} failed!")

        if not remaining_async_ids: # 所有任务都已完成或处理完毕
            job.log_message(log_message="Finished waiting for refresh data, all tasks succeeded or handled",
                            status=consts.COMPLETED, job_score=100)
            return

        count_down -= 1
        if count_down % 10 == 0 or count_down == timeout -1: # 每隔一段时间或首次轮询时打印进度
            job.log_message(log_message=f"There are {len(remaining_async_ids)} tasks remaining. Timeout in {count_down}s")
        time.sleep(1) # 每秒轮询一次，避免CPU空转

    # 超时处理
    job.log_message(log_message=f"After waiting for {timeout=}s, some tasks did not complete on time. Remaining tasks: {remaining_async_ids}",
                    status=consts.ERRORS_FOUND, job_score=100)
    # 可以在这里抛出异常或返回特定状态