#!/usr/bin/env bash
# Launch a second test worker on GPU 0 the moment that card is genuinely free.
#
# Worker A (GPU 1) walks chunks 0..8 in order. Worker B takes 5..8 from the other
# end, so they only ever collide near the finish, and a collision is harmless -
# chunk parquets are written atomically and completed chunks are skipped.
#
# GOOD-CITIZEN RULE: this box is shared. Earlier, adding a worker while bge-out and
# two other users' run_099.py jobs were resident pushed GPU 0 to 93/98 GB and
# triggered an OOM in someone else's process. So we launch ONLY when GPU 0 is
# almost entirely free, and require it to stay that way across several checks
# rather than jumping on a momentary dip between batches.
set -uo pipefail

cd "$(dirname "$0")/.."
ROOT=$(pwd)
LOGS=${M3_LOGS:-$ROOT/logs}
WD="$LOGS/gpu0_watcher.log"

FREE_BELOW=${FREE_BELOW:-20000}     # MiB in use on GPU 0 to count as "free"
NEED_CLEAR=${NEED_CLEAR:-4}         # consecutive clear checks required
POLL=${POLL:-120}
PY=${PY:-/home/user4/Jyo/DAI/er_pipeline/.venv/bin/python}

say() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$WD"; }
say "=== gpu0 watcher start (need GPU0 < ${FREE_BELOW}MiB for ${NEED_CLEAR} checks) ==="

clear_count=0
while true; do
  # already done?
  done_n=$(ls "$ROOT"/work/results_A/test_chunks/chunk_*.parquet 2>/dev/null | wc -l)
  if [ "$done_n" -ge 9 ]; then
    say "all 9 chunks present - watcher exiting"; exit 0
  fi
  # worker B already up?
  if pgrep -f "M3_CHUNK_LO" >/dev/null 2>&1 || \
     ps -u "$USER" -o cmd --no-headers 2>/dev/null | grep -q "[r]un_test.py.*results_A" && \
     [ "$(pgrep -cf '[r]un_test.py')" -gt 2 ]; then
    :
  fi
  if [ -f "$LOGS/worker_b_launched" ]; then
    if pgrep -f "[r]un_test.py" >/dev/null && \
       grep -q . "$LOGS/run_test_w2.log" 2>/dev/null; then
      # still alive?
      if ! ps -p "$(cat "$LOGS/worker_b_launched")" >/dev/null 2>&1; then
        say "worker B exited; clearing marker so it can relaunch"
        rm -f "$LOGS/worker_b_launched"
      else
        sleep "$POLL"; continue
      fi
    fi
  fi

  used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits -i 0 2>/dev/null | tr -d ' ')
  [ -z "$used" ] && { sleep "$POLL"; continue; }

  if [ "$used" -lt "$FREE_BELOW" ]; then
    clear_count=$((clear_count + 1))
    say "GPU0 clear (${used} MiB) ${clear_count}/${NEED_CLEAR}"
  else
    [ "$clear_count" -gt 0 ] && say "GPU0 busy again (${used} MiB) - resetting"
    clear_count=0
  fi

  if [ "$clear_count" -ge "$NEED_CLEAR" ]; then
    say "launching worker B on GPU 0 (chunks 5-8, global decision skipped)"
    setsid nohup env CUDA_VISIBLE_DEVICES=0 \
      HF_HOME=/scratch/user4/mtp/cache/huggingface PYTHONPATH="$ROOT/src" \
      TOKENIZERS_PARALLELISM=false M3_CHUNK_LO=5 M3_CHUNK_HI=9 M3_SKIP_DECIDE=1 \
      "$PY" "$ROOT/src/run_test.py" --gpu 0 --artifacts "$ROOT/work/results_A" \
      > "$LOGS/run_test_w2.log" 2>&1 &
    echo $! > "$LOGS/worker_b_launched"
    say "worker B pid $(cat "$LOGS/worker_b_launched")"
    clear_count=0
    sleep 600
  fi
  sleep "$POLL"
done