#!/bin/bash

# Target Processor [cpu|gpu|inf]
export PROCESSOR=inf

# REGISTRY URI, end with '/'
ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
REGION=$(aws configure get region)
export REGISTRY=${ACCOUNT}.dkr.ecr.${REGION}.amazonaws.com/

# Docker IMAGE name
export IMAGE=dlc

# Docker CONTAINER name
export CONTAINER="${IMAGE}-${PROCESSOR}"

# Docker IMAGE tag
export TAG=":v1-${PROCESSOR}"

# Public Neuron repos
export NEURON_YUM_REPO_URL=https://yum.repos.neuron.amazonaws.com
export NEURON_PIP_REPO_URL=https://pip.repos.neuron.amazonaws.com

# Private Beta Neuron repos, uncomment to override default public repos, must specify user & pass
#export NEURON_YUM_REPO_URL=https://user:pass@yum.repos.beta.neuron.annapurna.aws.a2z.com
#export NEURON_PIP_REPO_URL=https://user:pass@pip.repos.beta.neuron.annapurna.aws.a2z.com

# BUILD_OPTS build time options
export BUILD_OPTS="--progress simple --build-arg NEURON_YUM_REPO_URL=$NEURON_YUM_REPO_URL --build-arg NEURON_PIP_REPO_URL=$NEURON_PIP_REPO_URL"

# RUN_OPTS runtime options
if [ "${PROCESSOR}" == "inf" ]; then
	if [ -f /dev/neuron0 ]; then
        	NEURON_DEVICES=$(for d in $(ls /dev/neuron*); do echo -n "--device=$d "; done)
        	export RUN_OPTS="${NEURON_DEVICES}"
	else
        	export RUN_OPTS="-e AWS_NEURON_VISIBLE_DEVICES=ALL --privileged"
	fi
elif [ "$PROCESSOR" == "gpu" ]; then
        export RUN_OPTS="-e CUDA_VISIBLE_DEVICES=ALL"
fi
export RUN_OPTS="$RUN_OPTS --name $CONTAINER"