#!/usr/bin/env bash
# run2 학습 런처: TJ-2~6(5노드×8=40 GPU) FSDP. host=TJ-2, 워커=TJ-3~6. ★TJ-1은 학습 제외(eval watcher 전용).
# TJ-1(오케스트레이터)에서 실행 — 로컬 학습 프로세스 없음. host도 원격(ssh TJ-2)이라 자기참조 함정 없음.
# 사용: scripts/launch_om_run2.sh <rdzv_id> <config_abspath> <outdir> <logdir> [extra train.py args...]
set -u
R=/NHNHOME/WORKSPACE/0526040027_A/OpenPath
RDZV="$1"; CFG="$2"; OUTDIR="$3"; LOGDIR="$4"; shift 4
EXTRA="$*"
HOST="TJ-2"
EP="10.34.15.21:29500"           # TJ-2 IB IP (rendezvous rank0)
WORKERS="TJ-3 TJ-4 TJ-5 TJ-6"
PYC="\$HOME/.cache/op_pyc_$RDZV"  # 노드-로컬. 원격 셸에서 확장되도록 escape
ENV="WANDB_MODE=disabled NCCL_DEBUG=WARN NCCL_IB_HCA=mlx5_1,mlx5_2,mlx5_4,mlx5_9 \
NCCL_SOCKET_IFNAME=bond0 GLOO_SOCKET_IFNAME=bond0 \
NCCL_IB_TIMEOUT=22 NCCL_IB_RETRY_CNT=13 NCCL_IB_QPS_PER_CONNECTION=4 \
PYTHONDONTWRITEBYTECODE=1 PYTHONPYCACHEPREFIX=$PYC PYTHONPATH=$R/third_party/OpenMidnight"
TR="venv/bin/torchrun --nnodes=5 --nproc-per-node=8 --rdzv-backend=c10d \
--rdzv-endpoint=$EP --rdzv-id=$RDZV \
third_party/OpenMidnight/dinov2/train/train.py --config-file $CFG --output-dir $OUTDIR $EXTRA"

mkdir -p "$LOGDIR" "$OUTDIR"

# 1) host(TJ-2) 원격 기동 — ★ssh를 &로 백그라운드(안 그러면 원격 setsid의 상속 fd를 물고 ssh가 반환 안 함=런처 행)
echo "launching host on $HOST ..."
ssh -n -o BatchMode=yes "$HOST" "mkdir -p $PYC; cd $R && setsid env $ENV $TR > $LOGDIR/${HOST}.log 2>&1 < /dev/null & echo ${HOST}-host-fired" &

# 2) host torchrun가 store 바인딩할 시간(고정 대기). c10d는 워커가 재시도로 붙으므로 ss폴링 불필요.
echo "waiting 30s for host store bind on $EP ..."
sleep 30
echo "firing 4 workers"

# 3) 워커(TJ-3~6) 병렬 기동
for n in $WORKERS; do
  ssh -n -o BatchMode=yes "$n" "mkdir -p $PYC; cd $R && setsid env $ENV $TR > $LOGDIR/${n}.log 2>&1 < /dev/null & echo ${n}-fired" &
done
# ★ wait 대신 고정 sleep — bg ssh가 원격 setsid의 fd를 물어 반환 안 할 수 있어 wait는 영원히 멈춤(autoresume blocking 호출도 막힘).
sleep 20
echo "ALL launched: host $HOST + workers $WORKERS | rdzv=$RDZV | logs=$LOGDIR | out=$OUTDIR"