GPU Task Assignment Script

改进脚本

我正在整理评估任务的bash脚本，确保GPU管理得当，并且包含相关日志。

以下是一个基于您提供的训练脚本的评估脚本。该脚本将并行运行 eval.py 和 eval_sac.py，并且保持与训练脚本相同的 GPU 分配和日志管理机制。

#!/bin/bash

# ================================
# Configuration
# ================================

# 数据集列表
data_list=("3" "5" "7" "butterfly" "dolphin" "golden_fish" "griffen" "mermaid" "seahorse" "shark" "unicorn")

# Conda 环境名称
CONDA_ENV="kirl"

# GPU 数量
NUM_GPUS=8

# 每个 GPU 的最大任务数
MAX_TASKS_PER_GPU=2

# 日志目录
LOG_DIR="./logs"

# ================================
# Function Definitions
# ================================

# 初始化日志目录
init_logs() {
    if [ ! -d "$LOG_DIR" ]; then
        mkdir -p "$LOG_DIR"
        if [ $? -ne 0 ]; then
            echo "Error: 无法创建日志目录在 $LOG_DIR"
            exit 1
        fi
    fi
}

# 激活 conda 环境
activate_conda_env() {
    # 尝试找到 conda 的安装路径
    if [ -f "$HOME/miniconda3/etc/profile.d/conda.sh" ]; then
        source "$HOME/miniconda3/etc/profile.d/conda.sh"
    elif [ -f "$HOME/anaconda3/etc/profile.d/conda.sh" ]; then
        source "$HOME/anaconda3/etc/profile.d/conda.sh"
    else
        echo "Error: 未找到 Conda。请确保已安装 Conda。"
        exit 1
    fi

    # 激活指定的 conda 环境
    conda activate "$CONDA_ENV"
    if [ $? -ne 0 ]; then
        echo "Error: 无法激活 conda 环境 '$CONDA_ENV'"
        exit 1
    fi
}

# 分配 GPU
assign_gpus() {
    # 创建一个数组来跟踪每个 GPU 的当前任务数
    for ((i=0; i<NUM_GPUS; i++)); do
        gpu_tasks[$i]=0
    done
}

# 获取下一个可用的 GPU
get_available_gpu() {
    for ((i=0; i<NUM_GPUS; i++)); do
        if [ "${gpu_tasks[$i]}" -lt "$MAX_TASKS_PER_GPU" ]; then
            echo "$i"
            return
        fi
    done
    # 如果所有 GPU 都满了，返回 -1
    echo "-1"
}

# ================================
# Main Script
# ================================

# 初始化日志目录
init_logs

# 激活 conda 环境
# 如果需要激活 conda 环境，请取消下面一行的注释
# activate_conda_env

# 分配 GPU
assign_gpus

# 遍历数据集列表并分配任务
for data in "${data_list[@]}"; do
    assigned=false
    while [ "$assigned" = false ]; do
        gpu_id=$(get_available_gpu)
        if [ "$gpu_id" -ge "0" ]; then
            # 分配 GPU
            gpu_tasks[$gpu_id]=$((gpu_tasks[$gpu_id] + 1))
            assigned=true

            # 设置 CUDA_VISIBLE_DEVICES 并运行评估脚本
            (
                export CUDA_VISIBLE_DEVICES=$gpu_id

                # 运行 eval.py
                python eval.py --data "$data" > "$LOG_DIR/eval_${data}.log" 2>&1
                if [ $? -ne 0 ]; then
                    echo "Error: eval.py 执行失败，数据集 '$data'。请检查日志: $LOG_DIR/eval_${data}.log"
                fi

                # 运行 eval_sac.py
                python eval_sac.py --data "$data" > "$LOG_DIR/eval_sac_${data}.log" 2>&1
                if [ $? -ne 0 ]; then
                    echo "Error: eval_sac.py 执行失败，数据集 '$data'。请检查日志: $LOG_DIR/eval_sac_${data}.log"
                fi

                # 任务完成，减少 GPU 任务计数
                gpu_tasks[$gpu_id]=$((gpu_tasks[$gpu_id] - 1))
            ) &

        else
            # 所有 GPU 都满了，等待一秒后重试
            sleep 1
        fi
    done
done

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

echo "所有评估任务已完成。"

# ================================
# End of Script
# ================================

说明

1. 配置部分：

   • data_list：定义了要评估的数据集列表。
   • CONDA_ENV：指定了要使用的 Conda 环境名称。
   • NUM_GPUS 和 MAX_TASKS_PER_GPU：定义了可用 GPU 的数量以及每个 GPU 上允许的最大并行任务数。
   • LOG_DIR：指定了日志文件的存储目录。

2. 函数定义：

   • init_logs：检查并创建日志目录。
   • activate_conda_env：查找并激活指定的 Conda 环境。如果您的环境需要激活，请取消主脚本中对应行的注释。
   • assign_gpus 和 get_available_gpu：管理 GPU 资源，确保不会超过每个 GPU 的任务限制。

3. 主脚本流程：

   • 初始化日志目录。
   • （可选）激活 Conda 环境。
   • 分配 GPU 资源。
   • 遍历每个数据集，分配可用的 GPU 并启动评估任务。
   • 每个评估任务包括运行 eval.py 和 eval_sac.py，并将输出重定向到对应的日志文件。
   • 等待所有后台任务完成后，输出完成信息。

使用说明

1. 确保环境配置正确：

   • 确保 Conda 已安装，并且脚本中 CONDA_ENV 指定的环境存在。
   • 如果需要激活 Conda 环境，请取消主脚本中 activate_conda_env 的注释。

2. 脚本权限：

   • 确保脚本具有执行权限。您可以使用以下命令赋予执行权限：

     chmod +x your_evaluation_script.sh
     

3. 运行脚本：

   • 通过以下命令运行脚本：

     ./your_evaluation_script.sh
     

4. 日志查看：

   • 所有评估任务的输出和错误日志将保存在 ./logs 目录中，文件名格式为 eval_<data>.log 和 eval_sac_<data>.log。

5. GPU 监控：

   • 您可以使用 nvidia-smi 命令实时监控 GPU 的使用情况，确保任务分配合理。

注意事项

• 资源管理：确保系统中有足够的 GPU 资源来运行并行任务，以避免资源竞争和性能下降。
• 错误处理：脚本中已经包含了基本的错误处理机制，但建议在运行前手动测试每个评估脚本以确保其正常运行。
• 日志管理：定期检查和清理日志目录，以防止日志文件占用过多磁盘空间。

希望这个评估脚本能够满足您的需求，并有效地管理多 GPU 环境下的评估任务。