#!/usr/bin/env bash
# Launch one keys-2k H3 container on HEAD and WORKER (dual-Spark parallel arms).
# Tony dual-serve order: start DS4 first if co-tenanting; then this script for H3.
set -euo pipefail

IMAGE="${IMAGE:-ghcr.io/drowzeys/keys-2k-minimax-h3-parallel-two-dgx-sparks:0.31.1-pp20260811}"
HEAD="${HEAD:-10.100.10.1}"
WORKER="${WORKER:-10.100.10.5}"
MODELS_DIR="${MODELS_DIR:-$HOME/keys-2k-minimax-h3-parallel/models}"
OUT_DIR="${OUT_DIR:-$HOME/Videos/h3-out}"
NAME="${CONTAINER_NAME:-keys-2k-h3}"
PORT="${H3_PORT:-8188}"

launch() {
  local host="$1"
  echo "--> H3 on $host"
  ssh -o BatchMode=yes "keyspark@${host}" "bash -s" <<REMOTE
set -euo pipefail
mkdir -p '$MODELS_DIR' '$OUT_DIR'
# stop prior
docker rm -f '$NAME' 2>/dev/null || true
docker pull '$IMAGE' 2>/dev/null || true
docker run -d --name '$NAME' --gpus all --restart unless-stopped \
  --network host \
  -e H3_PORT=$PORT \
  -e H3_RESERVE_VRAM=\${H3_RESERVE_VRAM:-48} \
  -e H3_VRAM_HEADROOM=\${H3_VRAM_HEADROOM:-10} \
  -v '$MODELS_DIR:/opt/powerpack-h3/ComfyUI/models' \
  -v '$OUT_DIR:/opt/powerpack-h3/ComfyUI/output' \
  '$IMAGE'
echo "  container \$NAME on $host"
REMOTE
}

launch "$HEAD"
launch "$WORKER"

echo "waiting for system_stats..."
for ip in "$HEAD" "$WORKER"; do
  ok=0
  for _ in $(seq 1 60); do
    if curl -sf -m 3 "http://${ip}:${PORT}/system_stats" >/dev/null; then ok=1; break; fi
    sleep 3
  done
  if [[ "$ok" == "1" ]]; then echo "  OK http://${ip}:${PORT}"; else echo "  DOWN http://${ip}:${PORT}"; fi
done

echo
echo "Parallel 2K bench:"
echo "  HEAD=$HEAD WORKER=$WORKER bash $(dirname "$0")/run_anime_2k_bench.sh"