#!/bin/bash

# Description: Build AMI for Slurm Workers
# Arguments:
#   $1 (required) :  cpu|gpu|a100

AMI_TYPE="$1"
export AWS_MAX_ATTEMPTS=90
export AWS_POLL_DELAY_SECONDS=60

function usage() {
    echo "usage: ${BASH_SOURCE[0]} <cpu|gpu|a100>"
}

function setup_directories() {
    export SCRIPT_PATH="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
    export DRIVING_ROOT="$(readlink -f $SCRIPT_PATH/../../../)"
    export PACKER_ROOT="$DRIVING_ROOT/infra/packer/slurm"
    export ANSIBLE_ROOT="$DRIVING_ROOT/infra/ansible"
}

function main() {
    cd "$ANSIBLE_ROOT"
    case "$AMI_TYPE" in
        cpu|CPU|gpu|GPU|a100)
            cpu_or_gpu=$(echo "$AMI_TYPE" | awk '{print tolower($0)}')

            packer build \
            -var-file="$PACKER_ROOT/$cpu_or_gpu-worker-variables.json" \
            "$PACKER_ROOT/worker-ami.json"
            ;;
        *)
            usage
            exit 1
            ;;
    esac
}

setup_directories
main
