#!/usr/bin/env bash
# Usage: ./docker.sh <gpu count> <image name>
set -e; set -o pipefail

GPU_COUNT=${1:-1}
IMG=${2:-"nvcr.io/hpc/qmcpack:v3.5.0"}

echo "Downloading S32 Dataset..."
wget -O - https://gitlab.com/NVHCP/ngc-examples/raw/master/qmcpack/v3.5.0/get_S32.sh | bash
INPUT="/host_pwd/NiO-fcc-S32-dmc.xml"

echo "Changing to ./S32_example directory"
cd S32_example

if [ -f /dev/infiniband ]; then
    echo "Enabling Infiniband support"
    IB_DEV="--device=/dev/infiniband --cap-add=IPC_LOCK --net=host"
fi

DOCKER="nvidia-docker run ${IB_DEV} -it --rm -v ${PWD}:/host_pwd ${IMG}"
MPIRUN="mpirun --allow-run-as-root -n ${GPU_COUNT}"

echo "Running NiO S32 example in ${IMG} on ${GPU_COUNT} GPUS..."
${DOCKER} ${MPIRUN} qmcpack ${INPUT}