#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib.sh
source "${SCRIPT_DIR}/lib.sh"

start_worker() {
  local name="$1"
  local node_ip
  node_ip="$(container_ip "${name}")"
  echo "[${name}] ray stop (ignore errors)"
  dexec "${name}" bash -lc "ray stop --force || true"

  local head_ip
  head_ip="$(container_ip "${HEAD_CONTAINER}")"
  echo "[${name}] container ip: ${node_ip}"
  echo "[${name}] ray start -> join ${head_ip}:6379 (num_gpus=4, resources worker_node=100)"
  dexec "${name}" bash -lc "ray start --node-ip-address=${node_ip} --address=${head_ip}:6379 --num-gpus=4 --resources='{\"worker_node\": 100}'"
}

start_worker "${WORKER0_CONTAINER}"
start_worker "${WORKER1_CONTAINER}"

echo "[head] waiting for workers to register"
for _ in $(seq 1 30); do
  if dexec "${HEAD_CONTAINER}" bash -lc "ray status" | grep -q "Active:"; then
    break
  fi
  sleep 2
done

dexec "${HEAD_CONTAINER}" bash -lc "ray status || true"
