Celery 是一个Python的分布式异步任务框架，在Python后端开发中经常使用，但在某些场景下需要对任务进行编排，同一优先级的任务可以同时执行，低优先级的任务需要在所有高优先级任务完成之后再执行。本文记录一下利用Celery进行任务优先级编排的方法。
整个调度流程由调度入口，任务实现，其他组件三部分组成。
1. 调度入口
调度入口是整个调度流程的开始，并由入口调用next_task_process，next_task_process根据当前优先级任务状态是否调用下一优先级的任务。
def start_task_process(data: dict):
"""
执行任务调度
:param data: 任务信息, 包含sub_task 和 params 两个字段
sub_tasks: 子任务优先级优先级列表，优先级高的先执行，每一项为["task_name", task_priority]
如: sub_tasks = [["task1", 1], ["task2", 3], ["task3", 2]]
params: 其他参数，会始终贯穿于整个任务流程，可将下一个任务需要的参数通过params传递过去
data = {
"sub_task:" [["task_name", task_priority], ["task_name2", task_priority2]]
}
:return: None
"""
task_id = str(uuid.uuid4())
data.update({"task_id": task_id, "queue": [], "status": Status.WAITING})
task_group = {}
# 相同优先级的任务放在同一个list中
for item in data["sub_tasks"]:
# 每一个优先级生成一个list, 相同优先级的任务放到同一个list
task_group.setdefault(item[1], []).append(item[0])
# 按照优先级从高到低排序
for i in sorted(task_group.keys(), reverse=True):
data["queue"].append(task_group[i])
# 调度优先级高的任务
if len(data["queue"].queue) < 1:
return
temp_res = TaskTempResult(task_id=task_id)
temp_res.set(temp_res)
for name in data["queue"][0]:
async_result = TASK_LIST[name].apply_async((data, name), link=next_task_process.s())
temp_res.set_celery_task(name, {"id": async_result.task_id, "status": Status.WAITING})
@shared_task(bind=True, retry_kwargs={"max_retries": 5}, soft_time_limit=300)
def next_task_process(self, sub_res):
"""
调度到下一优先级的任务中
:param self: 当前celery任务
:param sub_res: 上一阶段的返回值，其中包括以下字段：
sub_res = {
"id": "task_id",
"name": "name",
"result": True,
"params": {}
}
:return:
"""
try:
if sub_res["id"] is None:
return
temp_res = TaskTempResult(sub_res["id"])
flag = False # 执行下一流程的标志
if temp_res.lock(): # 获取琐，防止同一优先级的多个任务同时调起下一优先级
obj = temp_res.get()
if obj["status"] in Status.STOPPED:
temp_res.unlock()
return
obj["params"].update(sub_res.get("params", {}))
obj["queue"][0].remove(sub_res["name"])
if len(obj["queue"][0]) == 0: # 当前优先级的任务都完成
del obj["queue"][0]
flag = True
# 上一任务状态是成功，更新状态
if sub_res["result"]:
obj["status"] = Status.RUNNING if len(obj["queue"]) > 0 else Status.SUCCEED
else: # 上一任务状态失败，清空队列，置为失败
obj["queue"].clear()
obj["error_step"] = sub_res["name"]
obj["status"] = Status.FAILED
temp_res.set(obj)
temp_res.unlock()
if flag:
for name in obj["queue"][0]:
async_result = TASK_LIST[name].apply_async((obj, name), link=next_task_process.s())
temp_res.set_celery_task(name, {"id": async_result.task_id, "status": Status.WAITING})
else:
raise self.retry(countdown=1)
except SoftTimeLimitExceeded:
print('next_task_process time out, will retry in 1s')
raise self.retry(countdown=1)
2. 任务实现
任务实现有两种方式，一种是直接根据入参返回值实现，另一种是使用装饰器对任务进行封装。
# 任务名称对应的任务实现映射
TASK_LIST = {
"task1": t,
"task2": t2
}
def task_process(func):
"""
任务装饰器，用于封装任务的入参与返回值
:param func:
:return:
"""
def wrap(self, obj, name):
result_cache = TaskTempResult(obj['task_id'])
celery_task_info = {"id": self.request.id, "status": Status.RUNNING}
data = {"id": obj["task_id"], "name": name}
try:
print(f">>>>>>>>>>>id: {obj['task_id'], name}")
obj["status"] = Status.RUNNING
result_cache.set_celery_task(name, celery_task_info)
success, params = func(obj)
celery_task_info['status'] = Status.SUCCEED if success else Status.FAILED
data["result"] = success
if isinstance(params, dict):
data.update(params)
result_cache.set_celery_task(name, celery_task_info)
return data
except Exception as e:
print(e)
celery_task_info["error_info"] = "%s:%s" % (type(e).__name__, str(e))
celery_task_info["status"] = Status.ABORTED
data["error_info"] = "%s:%s" % (type(e).__name__, str(e))
data["result"] = False
result_cache.set_celery_task(name, celery_task_info)
return data
setattr(wrap, "__module__", func.__module__)
setattr(wrap, "__name__", func.__name__)
return wrap
@shared_task(bind=True)
def t(self, obj, name):
"""
具体任务实现，方式1
:param self: celery 任务对象
:param obj: 任务所需参数
:param name: 当前任务名称
:return: res
res = {
"id": "task_id",
"name": "name",
"result": True,
"params": {}
}
"""
# do something ....
res = {
"id": "task_id",
"name": "name",
"result": True,
"params": {}
}
temp_res = TaskTempResult(obj["task_id"])
celery_task_info = {"id": self.request.id, "status": Status.SUCCEED if res["result"] else Status.FAILED}
temp_res.set_celery_task(name, celery_task_info)
return res
@shared_task(bind=True)
@task_process
def t2(obj):
"""
具体任务实现，方式2
通过task_process 装饰器实现
:param obj: 任务所需参数
:return:
result: bool
params: dict|None 其他参数，传到下一流程中
"""
# do something ....
result = True
params = {"key": "value"}
return result, params
3. 其他组件
class Status:
WAITING = 'waiting'
FAILED = 'failed'
RUNNING = 'running'
SUCCEED = 'succeed'
ABORTED = 'aborted'
STOPPED = [FAILED, SUCCEED, ABORTED]
class TaskTempResult:
"""
任务临时结果，缓存在redis
"""
K_db_name = "task_temp_result"
expire_time = 60*60*24*7
def __init__(self, task_id):
self.redis = redis.Redis(host='localhost', port=6379, decode_responses=True)
self._task_id = task_id
def lock(self, time=30):
key = "%s:%s:lock" %(self.K_db_name, self._task_id)
if self.redis.setnx(key, 1) == 0:
return False
if time > 0:
self.redis.expire(key, time)
return True
def unlock(self):
key = "%s:%s:lock" %(self.K_db_name, self._task_id)
self.redis.delete(key)
def set(self, data):
key = "%s:%s" %(self.K_db_name, self._task_id)
self.redis.set(key, json.dumps(data), ex=self.expire_time)
def get(self):
key = "%s:%s" %(self.K_db_name, self._task_id)
data = self.redis.get(key)
return None if data is None else json.loads(data)
def set_celery_task(self, name, data, nx=False):
key = "%s:%s:celery:%s" %(self.K_db_name, self._task_id, name)
self.redis.set(key, json.dumps(data), ex=self.expire_time, nx=nx)
def get_celery_task(self, name):
key = "%s:%s:celery:%s" %(self.K_db_name, self._task_id, name)
data = self.redis.get(key)
return None if data is None else json.loads(data)
使用方法
任选一种方式将任务实现函数完成后，加到TASK_LIST的字典映射中，调用将任务参数及优先级信息传入到start_task_process函数中即可。具体参数如下：
"""
sub_tasks: 子任务优先级优先级列表，优先级高的先执行，每一项为["task_name", task_priority]
如: sub_tasks = [["task1", 1], ["task2", 3], ["task3", 2]]
params: 其他参数，会始终贯穿于整个任务流程，可将下一个任务需要的参数通过params传递过去
"""
data = {
"sub_task:" [["task_name", task_priority], ["task_name2", task_priority2]]
"params": {"a": "a", "b": "b"}
}
start_task_process(data)
