如何为Celery任务添加权重以适配GPU显存资源调度？

完全可以实现你描述的调度逻辑，核心是通过自定义任务属性、调度器和资源跟踪机制，让Celery能够感知任务的GPU显存占用并动态分配worker槽位。以下是具体实现思路：

1. 标记任务的显存占用属性 #

给每个任务添加自定义属性，明确其显存消耗，后续用于计算所需的worker槽位数：

from celery import Celery

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

@app.task(bind=True, gpu_memory=2)
def task1(self):
    # 任务逻辑（占用2GB显存）
    pass

@app.task(bind=True, gpu_memory=4)
def task2(self):
    # 任务逻辑（占用4GB显存）
    pass

@app.task(bind=True, gpu_memory=8)
def task3(self):
    # 任务逻辑（占用8GB显存）
    pass

2. 自定义GPU感知调度器 #

Celery默认的FIFO调度器不支持权重分配，需要继承基础调度器，实现基于显存占用的槽位计算与调度逻辑：

from celery.schedulers import BaseScheduler
from celery import current_app

class GPUAwareScheduler(BaseScheduler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.total_slots = 4  # 对应worker并发数4，每个槽位等价2GB显存
        # 分布式场景下需改用Redis/Memcached共享存储，避免单进程内存变量失效
        self.used_slots = 0

    def apply_async(self, task, args=None, kwargs=None, **options):
        # 计算任务所需槽位数
        required_slots = task.gpu_memory // 2
        # 检查剩余槽位是否足够
        if self.used_slots + required_slots <= self.total_slots:
            self.used_slots += required_slots
            # 绑定任务完成后的槽位释放回调
            options['link'] = current_app.signature('release_slots', args=[required_slots])
            return super().apply_async(task, args, kwargs, **options)
        else:
            # 槽位不足时将任务加入等待队列，后续轮询调度
            self.reserve(task, args, kwargs, **options)
            return None

# 定义槽位释放任务
@app.task
def release_slots(slots):
    scheduler = current_app.scheduler
    if hasattr(scheduler, 'used_slots'):
        scheduler.used_slots -= slots

3. 启动worker并指定自定义调度器 #

启动worker时指定我们的GPU感知调度器，同时设置并发数为4：

celery -A your_app_module worker --loglevel=info --concurrency=4 --scheduler=your_app_module.GPUAwareScheduler

关键注意事项 #

• 分布式场景适配：如果部署多个worker，不能用单进程内存变量跟踪used_slots，需改用Redis等共享存储来维护全局资源计数。
• 任务优先级优化：可以结合Celery的任务优先级设置（priority参数），给task1设置更高优先级，确保小任务优先被调度。
• 容错处理：通过task_failure信号捕获任务失败场景，强制释放占用的槽位，避免资源泄漏。

通过以上方案，即可实现你想要的调度逻辑：先运行4个task1（占满4个槽位），任务完成后释放槽位，依次调度剩余的task1和task2（每个task2占用2个槽位），最后调度占用全部4个槽位的task3。

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