#!/usr/bin/env bash
# Stop workers A and B duplicating the last chunks.
#
# A walks 0..8 in order; B owns 5..8. Once A finishes chunk 4 it would reach chunk 7
# while B is still on it, so both would compute 7 and then both would compute 8 -
# about an hour of wasted GPU. This waits for chunk 4 to land, then restarts A pinned
# to chunk 8 only, leaving 7 to B. A keeps the global decision (no SKIP_DECIDE), which
# is guarded to refuse until all 9 chunks exist.
set -uo pipefail
cd "$(dirname "$0")/.."
ROOT=$(pwd); LOGS="$ROOT/logs"; CK="$ROOT/work/results_A/test_chunks"
WD="$LOGS/coordinate.log"
say() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" | tee -a "$WD"; }

say "=== waiting for chunk_000800000.parquet (worker A's chunk 4) ==="
while [ ! -f "$CK/chunk_000800000.parquet" ]; do
  n=$(ls "$CK"/chunk_*.parquet 2>/dev/null | wc -l)
  [ "$n" -ge 9 ] && { say "all chunks done already - nothing to coordinate"; exit 0; }
  sleep 60
done
say "chunk 4 landed. Pinning worker A to chunk 8 so it stops shadowing B."

# find worker A: model-3 artifacts, and WITHOUT the chunk-slice env that marks worker B
pid=""
for p in $(pgrep -f "[r]un_test.py"); do
  ps -p "$p" -o cmd --no-headers 2>/dev/null | grep -q "model-3/work/results_A" || continue
  tr '\0' '\n' < "/proc/$p/environ" 2>/dev/null | grep -q "M3_CHUNK_LO" && continue
  pid=$p; break
done
[ -n "$pid" ] && { say "killing worker A pid $pid"; kill -9 "$pid"; }

# stop the supervisor relaunching an unpinned worker A
for s in $(pgrep -f "[s]upervise.sh test"); do say "stopping supervisor $s"; kill -9 "$s"; done
sleep 5

say "relaunching worker A pinned to chunk 8 (index 8:9), decision ENABLED"
setsid nohup env CUDA_VISIBLE_DEVICES=1 HF_HOME=/scratch/user4/mtp/cache/huggingface \
  PYTHONPATH="$ROOT/src" TOKENIZERS_PARALLELISM=false M3_CHUNK_LO=8 M3_CHUNK_HI=9 \
  /home/user4/Jyo/DAI/er_pipeline/.venv/bin/python "$ROOT/src/run_test.py" \
  --gpu 0 --artifacts "$ROOT/work/results_A" >> "$LOGS/run_test.log" 2>&1 &
say "worker A relaunched pid $!"
say "=== done; A=chunk8+decision, B=chunk7 ==="