#!/bin/bash
# Wait for GPUs to become idle (no other compute processes), then run DSpark profiles.

set -e

ROOT=/data/user1/yy
PROFILE_PID_FILE=/tmp/profile_dspark.pid
MONITOR_LOG=$ROOT/bench_results/dspark_profile_monitor.log
RESULT_BASE=$ROOT/bench_results/dspark_profile_$(date +%Y%m%d_%H%M%S)

mkdir -p "$RESULT_BASE" "$ROOT/tmp"

log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$MONITOR_LOG"
}

is_gpu_idle() {
    # Consider idle if no process other than ours uses > 1GB on any GPU.
    # Returns 0 (true) if idle, 1 (false) if busy.
    nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits | \
    awk -F',' '
    {
        pid=$1; mem=$2;
        # Exclude our own known PIDs (monitor script, shell, etc.)
        if (mem > 1024) { busy++; }
    }
    END { exit (busy > 0 ? 1 : 0) }
    '
}

log "Starting GPU idle monitor. Will check every 5 minutes."
log "Results will go to $RESULT_BASE"

wait_for_idle() {
    while true; do
        if is_gpu_idle; then
            log "GPUs appear idle."
            return 0
        else
            log "GPUs still busy. Sleeping 5 minutes."
            sleep 300
        fi
    done
}

# Run profiles sequentially, checking GPU idle state before each run.
CONFIGS=("dspark-st3:64" "dspark-st5:64" "dspark-st3:1" "nospec:64")

for entry in "${CONFIGS[@]}"; do
    IFS=':' read -r config concurrency <<< "$entry"
    log "Waiting before profile: config=$config concurrency=$concurrency"
    wait_for_idle
    log "Running profile: config=$config concurrency=$concurrency"
    "$ROOT/envs/vllm-dspark/bin/python" "$ROOT/scripts/profile_dspark.py" \
        --config "$config" --concurrency "$concurrency" --tool nsys \
        >> "$RESULT_BASE/profile_${config}_c${concurrency}.log" 2>&1
    log "Finished profile: config=$config concurrency=$concurrency"
done

log "All profiles complete. Results in $RESULT_BASE"