#!/bin/bash
# ============================================================
# 多机器 benchmark 分发脚本
# 根据当前进度把 official suite 拆到 3 台机器跑
# 用法: bash scripts/run_multi_machine.sh
# ============================================================

set -euo pipefail

PASS="${SSHPASS:-sskj2025}"
SSH="sshpass -p $PASS ssh -o StrictHostKeyChecking=no -o ConnectTimeout=10"

# 机器配置
MACHINES=(
    "174.1.51.1:/data1/sora"   # gpu048 / P800-01
    "174.1.51.2:/data1/sora"   # gpu049 / P800-02
    "174.1.51.4:/data1/sora"   # gpu051 / P800-04
)

COMMON_ARGS="--model DeepSeek-V4-Flash-Int8 --api-url http://localhost:30000/v1 --dataset-dir /data1/sora/evalscope --output-dir /data1/sora/evalscope/output --batch-size 4 --thinking --limit none --seed 42"

# 分发方案：尽量让 3 台机器同时跑完
# 估算：aime25/aime26 12 run，live_code_bench 5 run，其余 1 run
# Machine 1: 推理/数学 multi-run（约 30h）
M1_BENCHES="aime25,aime26,live_code_bench"
M1_FOLDER="DP4-flash-int8-thinking-m1"

# Machine 2: 知识与长上下文（约 18h）
M2_BENCHES="hle,mmlu_pro,gpqa_diamond,longbench_v2"
M2_FOLDER="DP4-flash-int8-thinking-m2"

# Machine 3: SWE + Agent（约 30-50h，取决于 SWE 镜像）
M3_BENCHES="swe_bench_verified,tau2_bench"
M3_FOLDER="DP4-flash-int8-thinking-m3"

run_on_machine() {
    local ip_base="$1"
    local ip="${ip_base%%:*}"
    local base_dir="${ip_base##*:}"
    local datasets="$2"
    local folder="$3"
    local cmd="cd $base_dir/evalscope && nohup python bash/run.py --datasets $datasets --folder-name $folder $COMMON_ARGS > $base_dir/evalscope/logs/run_${folder}.log 2>&1 &"

    echo "==> 在 $ip 启动: $folder ($datasets)"
    $SSH root@$ip "$cmd"
}

run_on_machine "${MACHINES[0]}" "$M1_BENCHES" "$M1_FOLDER"
run_on_machine "${MACHINES[1]}" "$M2_BENCHES" "$M2_FOLDER"
run_on_machine "${MACHINES[2]}" "$M3_BENCHES" "$M3_FOLDER"

echo ""
echo "==================================================="
echo "已分发到 3 台机器，各自日志在 logs/run_*.log"
echo "==================================================="