#!/usr/bin/env bash
set -Eeuo pipefail
export TORCHDYNAMO_DISABLE=1

WORKDIR="$(cd "$(dirname "$0")" && pwd)"
WORKER="$WORKDIR/run_sae_steering_qwen2.5_3b.sh"

if [[ ! -f "$WORKER" ]]; then
  echo "[fatal] Worker not found: $WORKER"
  exit 1
fi
chmod +x "$WORKER"

mkdir -p "$WORKDIR/logs/sae_steering_workers"

declare -a pids

# ===== Choose the GPUs to use here =====
# Edit only this line to select which GPUs to run on.
GPUS=(0 1)
NUM_GPUS="${#GPUS[@]}"

# Spawn one worker per selected GPU
# <slot> is the index in GPUS (0..NUM_GPUS-1)
# <gpu_id> is the actual CUDA device number (e.g., 1, 2, 3, 4, 6, 7)
# <num_gpus> is the total number of workers you are launching
for slot in "${!GPUS[@]}"; do
  gpu_id="${GPUS[$slot]}"
  (
    export CUDA_VISIBLE_DEVICES="$gpu_id"
    # pass: <slot> <gpu_id> <num_gpus>
    stdbuf -oL -eL bash -c "$WORKER $slot $gpu_id $NUM_GPUS" 2>&1 \
      | awk -v gpu="$gpu_id" '{ cmd="date +\"%F %T\""; cmd | getline ts; close(cmd); print ts, "[gpu "gpu"]", $0 ; fflush() }' \
      | tee -a "$WORKDIR/logs/workers/worker_${gpu_id}.out"
  ) &
  pids[$slot]=$!
  echo "[spawn] worker on gpu $gpu_id -> pid ${pids[$slot]}"
done

# Wait for all workers
ec=0
for slot in "${!GPUS[@]}"; do
  gpu_id="${GPUS[$slot]}"
  if ! wait "${pids[$slot]}"; then
    echo "[error] worker gpu $gpu_id exited non-zero"
    ec=1
  else
    echo "[done] worker gpu $gpu_id finished"
  fi
done

exit $ec
