#!/usr/bin/env bash

set -euo pipefail

CHECKOUT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
CONFIG="${CHECKOUT}/runs/frontier-workload-regime-taxonomy-v0/fleet_simulator.toml"
JOBS="${CHECKOUT}/runs/frontier-workload-regime-taxonomy-v0/jobs_simulator_per_gpu_v1.toml"
FLEET="${CHECKOUT}/infra/gpu_fleet/gpu_fleet.py"
INTERVAL_SECONDS="${INTERVAL_SECONDS:-30}"

CHECKS=(
  "dash1:/home/admin/cpfs/wjh/aituner/gpu-fleet-workload-regime-state/dash1/runs/workload-regime-sim-tp1-v4-20260720-20260720T101247381816Z"
  "dash2:/home/admin/cpfs/wjh/aituner/gpu-fleet-workload-regime-state/dash2/runs/workload-regime-sim-tp2-v4-20260720-20260720T101252134828Z"
  "dash3:/home/admin/cpfs/wjh/aituner/gpu-fleet-workload-regime-state/dash3/runs/workload-regime-sim-tp4-noprefix-v4-20260720-20260720T101256981633Z"
  "dash4:/home/admin/cpfs/wjh/aituner/gpu-fleet-workload-regime-state/dash4/runs/workload-regime-sim-tp4-prefix-v4-20260720-20260720T101301837408Z"
)

while true; do
  ready=true
  for entry in "${CHECKS[@]}"; do
    host="${entry%%:*}"
    path="${entry#*:}"
    if ! ssh -o BatchMode=yes -o ConnectTimeout=8 "${host}" \
      "test -f '${path}/finished_at'"; then
      ready=false
    fi
  done
  if [[ "${ready}" == true ]]; then
    break
  fi
  sleep "${INTERVAL_SECONDS}"
done

for entry in "${CHECKS[@]}"; do
  host="${entry%%:*}"
  path="${entry#*:}"
  rc="$(ssh -o BatchMode=yes "${host}" "cat '${path}/exit_code'")"
  if [[ "${rc}" != 0 ]]; then
    echo "V4_FAILED host=${host} rc=${rc}"
    exit 1
  fi
done

python3 "${FLEET}" probe --config "${CONFIG}"
python3 "${FLEET}" dispatch --config "${CONFIG}" --jobs "${JOBS}"