env_check.yaml

num_nodes: 2

setup: |
  echo "here"
  echo "export TEST_VAR=test" >> ~/.bashrc
  [[ -v CUDA_VISIBLE_DEVICES ]] && exit 1 || exit 0
  export CUDA_VISIBLE_DEVICES=1
  [[ -v CUDA_VISIBLE_DEVICES ]] && exit 0 || exit 1

run: |
  if [[ -z "${TEST_VAR}" ]]; then
    echo "TEST_VAR is not set"
    exit 1
  else
    echo "TEST_VAR is set to ${TEST_VAR}"
  fi

  if [[ -z "$SKYPILOT_NODE_RANK" ]]; then
    echo "SKYPILOT_NODE_RANK is not set"
    exit 1
  else
    echo "SKYPILOT_NODE_RANK is set to ${SKYPILOT_NODE_RANK}"
  fi

  if [[ -z "$SKYPILOT_NODE_IPS" ]]; then
    echo "SKYPILOT_NODE_IPS is not set"
    exit 1
  else
    echo "SKYPILOT_NODE_IPS is set to ${SKYPILOT_NODE_IPS}"
    echo "${SKYPILOT_NODE_IPS}"
    echo "${SKYPILOT_NODE_IPS}" | wc -l | grep 2 || exit 1
  fi
  worker_addr=`echo "$SKYPILOT_NODE_IPS" | sed -n 2p`
  echo Worker IP: $worker_addr

  if [[ -z "$SKYPILOT_TASK_ID" ]]; then
    echo "SKYPILOT_TASK_ID is not set"
    exit 1
  else
    echo "SKYPILOT_TASK_ID is set to ${SKYPILOT_TASK_ID}"
  fi