#!/bin/bash

# 设置最大并发数
MAX_CONCURRENT=100

# 读取主机文件
HOSTFILE="hostfile-test2"

# 创建日志目录
time=$(date +%Y%m%d%H%M%S)
mkdir -pv mpirun-logs/$time
sleep 3

# 读取文件中的主机 IP 并存储到数组中
hosts=()
while IFS= read -r line; do
    # 跳过空行和注释行
    [[ -z "$line" || "$line" =~ ^#.* ]] && continue
    hosts+=("$line")
done < "$HOSTFILE"

# 每两个主机组成一组
host_pairs=()
for ((i = 0; i < ${#hosts[@]}; i+=2)); do
    if [[ $((i+1)) -lt ${#hosts[@]} ]]; then
        host_pair="${hosts[i]}:8,${hosts[i+1]}:8"
    else
        host_pair="${hosts[i]}:8"
    fi
    host_pairs+=("$host_pair")
done

# 每个主机的槽位数
slots_per_host=8

# 处理并发执行
execute_parallel() {
    # 任务函数
    host_string="$1"
    total_processes=$((slots_per_host * $(echo "$host_string" | tr ',' '\n' | wc -l)))

    cmd=(
        "mpirun"
        "--allow-run-as-root"
        "-np" "$total_processes"
        "-host" "$host_string"
        "-x" "NCCL_DEBUG=INFO"
        "-x" "NCCL_IB_HCA=mlx5_100,mlx5_101,mlx5_102,mlx5_103,mlx5_104,mlx5_105,mlx5_106,mlx5_107"
        "-x" "NCCL_ALGO=Ring"
        "-x" "NCCL_SOCKET_IFNAME=bond0"
        "-x" "NCCL_IB_GID_INDEX=3"
        "-x" "NCCL_TIMEOUT=60"
        "./build/all_reduce_perf"
        "-b" "16G"
        "-e" "16G"
        "-g" "1"
        "-f" "2"
        "-T" "60"
    )
    echo "Executing: ${cmd[*]}"
    output=$("${cmd[@]}" | tee /dev/tty)
    bandwidth=$(echo "$output" | grep "Avg bus bandwidth" | awk '{print $NF}')
    filename="${host_string}_$(date +%Y%m%d%H%M%S)_[${bandwidth}].log"
    echo -e "$output" > "./mpirun-logs/$time/$filename.log"
}

# 初始化并发计数
running_jobs=0

# 使用进程控制最大并发数
for host_pair in "${host_pairs[@]}"; do
    # 如果当前正在运行的并发数达到最大限制
    while ((running_jobs >= MAX_CONCURRENT)); do
        # 等待其中一个进程结束
        wait -n
        ((running_jobs--))
    done

    # 执行命令
    execute_parallel "$host_pair" &

    # 增加并发计数
    ((running_jobs++))
done

# 等待所有并发任务完成
wait
