如何创建可启动支持GPU的Docker容器的Airflow任务

原生DockerOperator目前未直接开放device_requests参数，而终端执行docker run --gpus all命令时，底层正是通过该参数向Docker daemon申请GPU设备资源的，因此无法直接通过原生DockerOperator传参实现。

方案一：自定义扩展DockerOperator（推荐） #

该方案可以保留原生DockerOperator的所有功能（日志收集、XCom推送、资源管理等），仅新增GPU参数支持。

• 第一步：继承原生DockerOperator，添加GPU参数适配逻辑

from airflow.providers.docker.operators.docker import DockerOperator
from docker.types import DeviceRequest

class GPUEnabledDockerOperator(DockerOperator):
    def __init__(self, gpus: str = None, **kwargs):
        super().__init__(**kwargs)
        if gpus:
            # 构造对应--gpus参数的DeviceRequest对象
            self.device_requests.append(
                DeviceRequest(
                    count=-1 if gpus == "all" else int(gpus),
                    capabilities=[["gpu"]]
                )
            )

• 第二步：直接调用自定义Operator即可，示例代码如下

gpu_test_task = GPUEnabledDockerOperator(
    task_id="gpu_container_test",
    image="nvidia/cuda:11.8.0-base-ubuntu22.04",
    command="nvidia-smi",
    gpus="all",
    # 其他原生DockerOperator支持的参数都可正常使用
    network_mode="bridge",
    do_xcom_push=False,
    auto_remove="success"
)

方案二：使用BashOperator直接调用docker命令（适合临时验证） #

如果不需要用到DockerOperator的高级功能，可直接通过BashOperator执行终端的docker run命令：

from airflow.operators.bash import BashOperator

gpu_test_task = BashOperator(
    task_id="gpu_container_test_via_bash",
    bash_command="docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi"
)

• Airflow Worker节点必须已经安装nvidia-docker2运行时，Docker daemon配置正确可正常识别GPU设备
• 使用自定义Operator方案前，需要确保Airflow环境已安装docker Python依赖，可执行pip install docker完成安装
• 部署到Airflow前，先在Worker节点手动执行docker run --rm --gpus all nvidia-smi，确认节点本身可正常启动带GPU的容器

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