如何用Bash编写GPU工作池实现多任务并行利用CUDA设备

需求回顾 #

• 4台CUDA设备（ID: 0、1、2、3）
• 10个任务脚本：run01.sh 至 run10.sh
• 单GPU同一时间仅执行1个任务
• 最大化GPU利用率，有剩余任务时无空闲设备

实现脚本 #

#!/bin/bash

# 定义CUDA设备列表
CUDA_DEVICES=(0 1 2 3)
# 定义任务脚本列表
TASKS=(run01.sh run02.sh run03.sh run04.sh run05.sh run06.sh run07.sh run08.sh run09.sh run10.sh)

# 为每个GPU创建临时FIFO管道用于任务调度
declare -A GPU_PIPES
for dev in "${CUDA_DEVICES[@]}"; do
    PIPE="/tmp/gpu_${dev}_pipe"
    mkfifo "$PIPE"
    GPU_PIPES[$dev]="$PIPE"

    # 启动GPU专属的任务处理后台进程
    (
        while true; do
            read task_script
            [[ -z "$task_script" ]] && break

            echo "GPU $dev 启动任务: $task_script"
            CUDA_VISIBLE_DEVICES=$dev bash "$task_script"
            echo "GPU $dev 完成任务: $task_script"
        done
        rm -f "$PIPE"
    ) &
done

# 轮询分发所有任务到GPU管道
task_idx=0
num_devices=${#CUDA_DEVICES[@]}
for task in "${TASKS[@]}"; do
    dev=${CUDA_DEVICES[$((task_idx % num_devices))]}
    echo "$task" > "${GPU_PIPES[$dev]}"
    ((task_idx++))
done

# 向各GPU管道发送结束信号
for dev in "${CUDA_DEVICES[@]}"; do
    echo "" > "${GPU_PIPES[$dev]}"
done

# 等待所有后台任务进程完成
wait

echo "所有任务执行完毕"

脚本说明 #

1. 设备与任务定义：明确列出可用CUDA设备ID和任务脚本，便于后续修改维护。
2. FIFO管道调度：为每个GPU创建独立管道，对应一个后台进程持续读取任务执行，确保单GPU同一时间仅处理一个任务。
3. 轮询任务分配：通过索引取模实现任务均匀分发，保证所有GPU尽可能同时工作，最大化并行效率。
4. 收尾处理：任务分发完成后向管道发送空信号触发进程退出，最后等待所有后台进程结束，确保无遗漏任务。

使用方法 #

1. 将脚本保存为 gpu_task_scheduler.sh
2. 赋予执行权限：chmod +x gpu_task_scheduler.sh
3. 在任务脚本所在目录执行：./gpu_task_scheduler.sh

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