如何利用Celery按GPU内存权重调度不同任务？

核心结论：给Task B设置权重为4是实现你要的调度顺序的关键一步，但单纯依赖Celery默认权重机制还不够，必须配合GPU资源占用跟踪与分配逻辑，才能精准控制任务执行顺序和资源利用。

具体实现思路 #

1. 给任务标记权重与资源占用量 #

首先为两类任务分别设置权重和GPU内存占用属性：

• Task A：权重设为1，对应2G GPU内存占用
• Task B：权重设为4，对应8G GPU内存占用

Celery的weight参数会让调度器优先选择权重更高的任务，但默认调度器不感知资源占用，所以需要额外逻辑管控GPU资源。

2. 自定义GPU资源跟踪与调度控制 #

由于Celery默认不支持硬件资源感知调度，需要通过信号钩子+**分布式存储（如Redis）**来实时跟踪GPU内存使用情况，确保只有资源足够时才执行任务：

代码示例

首先定义Celery任务并标记属性：

from celery import Celery

app = Celery('video_retrieval', broker='redis://localhost:6379/0')

# Task A：权重1，占用2G GPU内存
@app.task(weight=1, gpu_memory=2)
def task_a(video_path):
    # 视频目标检索逻辑（占用2G GPU）
    pass

# Task B：权重4，占用8G GPU内存
@app.task(weight=4, gpu_memory=8)
def task_b(video_path):
    # 视频目标检索逻辑（占用8G GPU）
    pass

然后通过Celery信号实现资源管控：

from celery.signals import task_prerun, task_postrun
import redis

# 用Redis存储全局GPU内存使用量（适配多进程/多用户分布式场景）
redis_client = redis.Redis(host='localhost', port=6379, db=0)
GPU_TOTAL_MEM = 8  # 总GPU内存8G

@task_prerun.connect
def check_and_allocate_gpu(sender, **kwargs):
    # 获取当前任务所需GPU内存
    required_mem = getattr(sender, 'gpu_memory', 0)
    if required_mem <= 0:
        return
    
    # 获取当前已使用GPU内存
    current_used = int(redis_client.get('gpu_used') or 0)
    
    # 检查资源是否足够，不足则重试任务
    if current_used + required_mem > GPU_TOTAL_MEM:
        sender.retry(countdown=3, max_retries=None)
    
    # 分配资源
    redis_client.set('gpu_used', current_used + required_mem)

@task_postrun.connect
def release_gpu(sender, **kwargs):
    # 释放当前任务占用的GPU内存
    required_mem = getattr(sender, 'gpu_memory', 0)
    if required_mem <= 0:
        return
    
    current_used = int(redis_client.get('gpu_used') or 0)
    new_used = max(current_used - required_mem, 0)
    redis_client.set('gpu_used', new_used)

3. 实现你期望的调度顺序 #

结合上述配置，就能实现你给出的示例流程：

1. 初始GPU空闲（0G使用），调度器会优先选择总权重最高的可执行任务组合：4个Task A（总权重4）刚好占满8G，因此先执行A1-A4；此时B1、B2、A5、A6因资源不足进入等待。
2. A1-A4执行完成后，GPU释放8G内存，此时等待队列中权重最高的是B1（权重4），因此执行B1；剩余任务继续等待。
3. B1执行完成后，GPU再次释放8G，执行下一个高权重任务B2；A5、A6继续等待。
4. B2执行完成后，GPU释放8G，调度器选择A5、A6（总权重2，且占用4G内存）执行。

注意事项 #

• 若需要更精细化的调度（比如优先选择填满GPU的任务组合），可以自定义Celery调度器（继承celery.schedulers.base.BaseScheduler），在调度逻辑中加入资源利用率优先的判断。
• Redis用于存储GPU使用量是为了适配多Worker进程/多机器的分布式场景，若单Worker单进程部署，也可以用全局变量替代，但可靠性稍差。

内容的提问来源于stack exchange，提问作者UnPapeur