#!/bin/bash

NNODES="${NNODES:-4}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
MASTER_ADDR="${MASTER_ADDR:-localhost:29550}"
RDZV_ID=20020
USE_WANDB="${USE_WANDB:-0}"   # set USE_WANDB=1 to enable W&B logging (fill in wandb_key/wandb_entity in the config)

NODE_IPS=("${NODE_IP_0:-127.0.0.1}" "${NODE_IP_1:-your_node_ip_1}" "${NODE_IP_2:-your_node_ip_2}" "${NODE_IP_3:-your_node_ip_3}")

LOCAL_IPS=$(hostname -I 2>/dev/null)
LOCAL_ADDR=""
for node_ip in "${NODE_IPS[@]}"; do
    if echo "$LOCAL_IPS" | grep -qw "$node_ip"; then
        LOCAL_ADDR="$node_ip"
        break
    fi
done
if [ -z "$LOCAL_ADDR" ]; then
    echo "ERROR: failed to match local IP against NODE_IPS; check NODE_IPS configuration"
    echo "local IPs: $LOCAL_IPS"
    exit 1
fi
echo "LOCAL_ADDR=$LOCAL_ADDR"

CONFIG=configs/echo_infinity.yaml
CONFIG_NAME=$(basename "$CONFIG" .yaml)
LOGDIR="logs/${CONFIG_NAME}"
WANDB_SAVE_DIR=wandb

echo "CONFIG=$CONFIG"
echo "LOGDIR=$LOGDIR"

mkdir -p "$LOGDIR"
LOGFILE="${LOGDIR}/train_$(date +%Y%m%d_%H%M%S).log"
echo "LOGFILE=$LOGFILE"

WANDB_ARGS="--disable-wandb"
[ "${USE_WANDB}" != "0" ] && WANDB_ARGS=""

torchrun --nnodes=$NNODES --nproc_per_node=$NPROC_PER_NODE \
    --rdzv_id=$RDZV_ID --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR \
    --local-addr=$LOCAL_ADDR \
  train.py \
  --config_path $CONFIG \
  --logdir $LOGDIR \
  --wandb-save-dir $WANDB_SAVE_DIR \
  ${WANDB_ARGS} \
  2>&1 | tee "$LOGFILE"