#!/bin/bash -e
#
# Copyright (C) Mellanox Technologies Ltd. 2016-2017.  ALL RIGHTS RESERVED.
# See file LICENSE for terms.
#
# Testing script for SHARP, to run OSU benchmark.
#

BASE_DIR=$(dirname $0)
SCRIPT_NAME=$(basename $0)

. $BASE_DIR/sharp_funcs.sh


#
# Defaults
#
SHARP_INI_FILE=${SHARP_INI_FILE:-$USER_INI_FILE}
SHARP_PPN="1"
SHARP_TEST_ITERS="10000"
SHARP_TEST_SKIP_ITERS="1000"
SHARP_TEST_MAX_DATA="4096"
SHARP_MAX_PAYLOAD_SIZE=256
SHARP_SR="SHArP.AggregationManager"
SHARP_TMP_DIR_DEFAULT="/tmp"
SHARP_TMP_DIR=${SHARP_TMP_DIR:-$SHARP_TMP_DIR_DEFAULT}
LOGS_DIR=$SHARP_TMP_DIR/sharp_benchmark_logs.$$
MPI_LOG=$LOGS_DIR/sharp_log_mpirun.log
HOSTFILE=$LOGS_DIR/hostfile

function finish()
{
	local rc=$?

	\rm -rf $LOGS_DIR &>/dev/null

	return $((rc))
}

trap finish EXIT

usage()
{
	echo "This script includes OSU benchmarks for MPI_Allreduce and MPI_Barrier blocking collective operations."
	echo "Both benchmarks run with and without using SHARP technology."
	echo ""
	echo "Usage: $SCRIPT_NAME [-t] [-d] [-h]"
	echo "	-t - tests list (e.g. sharp:barrier)"
	echo "	-d - dry run"
	echo "	-h - display this help and exit"
	echo ""
	echo "Configuration:"
	echo " Runtime:"
	echo "  sharp_ppn - number of processes per compute node (default $SHARP_PPN)"
	echo "  sharp_ib_dev - Infiniband device used for communication. Format <device_name>:<port_number>."
	echo "                 For example: sharp_ib_dev=\"mlx5_0:1\""
	echo "                 This is a mandatory parameter. If it's absent, $SCRIPT_NAME tries to use the first active device on local machine"
	echo "  sharp_hostlist - hostnames of compute nodes used in the benchmark. The list may include normal host names,"
	echo "                   a range of hosts in hostlist format. Under SLURM allocation, SLURM_NODELIST is used as a default"
	echo "  sharp_test_iters - number of test iterations (default $SHARP_TEST_ITERS)"
	echo "  sharp_test_skip_iters - number of test iterations (default $SHARP_TEST_SKIP_ITERS)"
	echo "  sharp_test_max_data - max data size used for testing (default and maximum $SHARP_TEST_MAX_DATA)"
	echo " Environment:"
	echo "  SHARP_INI_FILE - takes configuration from given file instead of $USER_INI_FILE"
	echo "  SHARP_TMP_DIR - store temporary files here instead of $SHARP_TMP_DIR_DEFAULT"
	echo ""
	echo "Examples:"
	echo "  sharp_ib_dev=\"mlx5_0:1\" $SCRIPT_NAME  # run using \"mlx5_0:1\" IB port. Rest parameters are loaded from $USER_INI_FILE or default"
	echo "  SHARP_INI_FILE=~/benchmark.ini  $SCRIPT_NAME # Override default configuration file"
	echo "  SHARP_INI_FILE=~/benchmark.ini  sharp_hostlist=ajna0[2-3]  sharp_ib_dev=\"mlx5_0:1\" $SCRIPT_NAME # Use specific host list"
	echo "  sharp_ppn=1 sharp_hostlist=ajna0[1-8] sharp_ib_dev=\"mlx5_0:1\" $SCRIPT_NAME  -d # Print commands without actual run"
	echo ""
	echo "Dependencies:"
	echo "  $(print_hostlist_dependency_message)"
	exit 2
}

print_configuration()
{
	echo "IB device (sharp_ib_dev): $sharp_ib_dev"
	echo "PPN (sharp_ppn): $sharp_ppn"
	echo "Max data size (sharp_test_max_data): $sharp_test_max_data"
	echo "# test iteration (sharp_test_iters): $sharp_test_iters"
	echo "# skip iterations (sharp_test_skip_iters): $sharp_test_skip_iters"
	echo ""
	echo "Configuration file (SHARP_INI_FILE) : $SHARP_INI_FILE"
	echo "Temporary logs (SHARP_TMP_DIR) : $LOGS_DIR"
}

if [ -f $SHARP_INSTALL/../modulefiles/hpcx ]; then
	module load $SHARP_INSTALL/../modulefiles/hpcx
elif module whatis hpcx-gcc 2>&1 | grep -q HPC-X; then
	module load hpcx-gcc
fi

mpirun="$OMPI_HOME/bin/mpirun"

#
# Command wrapper
# Input parameters: <host name> <command name> <command's prm1 > ... <command's prmN>
# If the host is local host, run the command
# If the host is remote host, run in remote using mpirun
#
command_wrapper()
{
	local host=$1
	local cmd=$2

	if [[ "$host" != "$HOSTNAME" ]]; then
		cmd="$mpirun -H $1 -np 1 $cmd"
	fi

	shift 2
	cmd=$cmd" $@"

	eval $cmd
}

#
# Get first active device
# Input parameters: <hostname>
#
get_active_ib_device()
{
	local port
	local ibdev
	local port_info

	for ibdev in $(command_wrapper $1 ibstat -l)
	do
		port=1
		port_info=$( command_wrapper $1 ibstat $ibdev $port )
		( echo "$port_info" | grep -q Active ) && ( echo "$port_info" | grep -q InfiniBand ) && eval echo "$ibdev:$port" && break ||:
	done
}

#
# Check IB device
#
check_ib_device()
{
	local host=$1
	local ib_dev=$2
	local ca_name
	local port
	local port_info
	local sr_info

	if ! check_ib_dev_name "$ib_dev"; then
		return 1
	fi

	ca_name=${ib_dev%:*}
	port=${ib_dev#*:}

	port_info=$( command_wrapper "$host" ibstat "$ca_name" "$port" )

	if [[ $? != "0" ]]; then
		echo "Failed to check $ib_dev $host"
		return 1
	fi

	if ! [[ "$port_info" =~ Active && "$port_info" =~ InfiniBand ]]; then
		echo "$ib_dev is not active Infiniband device on $host"
		return 1
	fi

	sr_info=$( command_wrapper "$host" saquery -C "$ca_name" -P "$port" SR )

	if [[ $? != "0" ]]; then
		echo "Failed to fetch SR on $host"
		return 1
	fi

	if ! [[ "$sr_info" =~ "$SHARP_SR" ]]; then
		echo "There is no sharp_am running in fabric connected to $ib_dev ($host)"
		return 1
	fi
}

#
# Check ib dev provided by a user. If it's empty run auto detection
#
check_and_update_ib_dev()
{
	local var=$1
	local hostlist=$2
	local first_compute_node

	first_compute_node=$( $hostlist_tool $hostlist "-l 1" )
	if [ -z "${!var}" ]; then
		echo "$var is empty. Try to use the first active device"
		if $($hostlist_tool $hostlist -i "$HOSTNAME" -q -0 ); then
			eval $var=$(get_active_ib_device "$HOSTNAME")
		else
			eval $var=$(get_active_ib_device "$first_compute_node" )
		fi

		if [ -z "${!var}" ]; then
			echo "$var can't be empty. Exit."
			exit 1
		fi
	elif ! check_ib_device "$first_compute_node" "${!var}"; then
		exit 1
	fi
}

allreduce=1
barrier=1
sharp_allreduce=1
sharp_barrier=1

while getopts "dt:h" flag_arg; do
	case $flag_arg in
		d)  dryrun=1   ;;
		t)	allreduce=
			barrier=
			sharp_allreduce=
			sharp_barrier=

			for x in $OPTARG; do
				case $x in
					allreduce)       allreduce=1       ;;
					barrier)         barrier=1         ;;
					sharp:allreduce) sharp_allreduce=1 ;;
					sharp:barrier)   sharp_barrier=1   ;;
					*)               echo "Unknown parameter: $x"
									 usage ;;
				esac
			done ;;
		h) usage ;;
		*) usage ;;

	esac
done


if [ -w $SHARP_INSTALL ]; then
	_local_dir=$SHARP_INSTALL
else
	_local_dir=$HOME/sharp_conf
fi

load_params $_local_dir

OUT_AM=/tmp/d_sharp_am.log
OUT_SD=/tmp/d_sharpd.log
AM_CFG=$DCONF_DIR/sharp_am.cfg
SD_CFG=$DCONF_DIR/sharpd.cfg

#
# Check dependency
#
check_hostlist_tool

#
# Check/Update mandatory parameters
#
check_and_update_hostlist sharp_hostlist
check_and_update_ib_dev sharp_ib_dev $sharp_hostlist

#
# Run OMPI tests
#
sharp_test_iters=${sharp_test_iters:-$SHARP_TEST_ITERS}
sharp_test_skip_iters=${sharp_test_skip_iters:-$SHARP_TEST_SKIP_ITERS}
sharp_test_max_data=${sharp_test_max_data:-$SHARP_TEST_MAX_DATA}
sharp_ppn=${sharp_ppn:-$SHARP_PPN}

bind_to_core=" taskset -c 1 numactl --membind=0 "
nnodes=$($hostlist_tool -e $sharp_hostlist | wc -l)

np=$((sharp_ppn * nnodes))

if (( "$sharp_test_max_data" > "$SHARP_TEST_MAX_DATA" )); then
       echo "For data packets bigger than $SHARP_TEST_MAX_DATA SHARP is not used"
fi

#
# Prepare logs
#
mkdir -p $LOGS_DIR &>/dev/null

print_configuration

#
# If local host is in the hosts list, place it as the first host
# it is performed by concatenating two hosts lists:
# 1. Intersecting local host with hostlist
# 2. Removing local host from hostlist
#
$hostlist_tool -e -i $(hostname) $sharp_hostlist > $HOSTFILE
$hostlist_tool -e -d $sharp_hostlist $(hostname) >> $HOSTFILE


#
# Prepare common options for both all runs: with/without SHARP
#
opt="--bind-to core --map-by node -hostfile $HOSTFILE -np $np "
opt+=" --display-map "
opt+=" -mca pml yalla "
opt+=" -mca btl_openib_warn_default_gid_prefix 0 "
opt+=" -mca rmaps_dist_device $sharp_ib_dev -mca rmaps_base_mapping_policy dist:span "
opt+=" -mca btl_openib_if_include $sharp_ib_dev "
opt+=" -x MXM_RDMA_PORTS=$sharp_ib_dev -x HCOLL_MAIN_IB=$sharp_ib_dev "
opt+=" -x MXM_LOG_LEVEL=ERROR "
#   opt+=" -x HCOLL_ENABLE_MCAST_ALL=1 "  # default: Enabled
#   opt+=" -x HCOLL_MCAST_NP=1 "
#	opt+=" -x HCOLL_BCOL=basesmuma,mlnx_p2p " # HCOLL_BCOL and HCOLL_SBGP enable using multi-channel feature
#	opt+=" -x HCOLL_SBGP=basesmsocket,p2p "
#   opt+=" -x LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:$SHARP_INSTALL/lib "
opt+=" -x LD_PRELOAD=$SHARP_INSTALL/lib/libsharp.so:$SHARP_INSTALL/lib/libsharp_coll.so "


#
# Find executables
#
osu_allreduce_exe=$(find $OMPI_HOME/tests/osu* -name osu_allreduce | tail -1)
osu_barrier_exe=$(find $OMPI_HOME/tests/osu* -name osu_barrier | tail -1)

if [ "$sharp_ppn" -eq "1" ]; then
	osu_allreduce_exe="$bind_to_core $osu_allreduce_exe"
	osu_barrier_exe="$bind_to_core $osu_barrier_exe"
fi


#
# Prepare OSU benchmark parameters
#
osu_opt="-i $sharp_test_iters -x $sharp_test_skip_iters -f -m :${sharp_test_max_data}"


#
# Prepare SHARP specific parameters
#
hcoll_sharp_opt=" -x HCOLL_ENABLE_SHARP=2 " # Enables sharp without fallback
hcoll_sharp_opt+=" -x SHARP_COLL_LOG_LEVEL=3  " # Enables libsharp logging
#   hcoll_sharp_opt+=" -x SHARP_COLL_GROUP_RESOURCE_POLICY=1 " # Equal distribution of OSTs among the groups. By default : 1
#	hcoll_sharp_opt+=" -x SHARP_COLL_ENABLE_MCAST_TARGET=1 " # UD MCAST. By default : enabled
#	hcoll_sharp_opt+=" -x SHARP_COLL_GROUP_IS_TARGET=1 " # Enable RC distribution. By default : enabled
#	hcoll_sharp_opt+=" -x SHARP_COLL_ENABLE_GROUP_TRIM=1 " # SHARP group trim. By default : enabled
#   hcoll_sharp_opt+=" -x SHARP_COLL_MAX_PAYLOAD_SIZE=$SHARP_MAX_PAYLOAD_SIZE" # Maximum payload size of sharp collective request. By default: 256(max).
#   hcoll_sharp_opt+=" -x HCOLL_SHARP_UPROGRESS_NUM_POLLS=999 "
hcoll_sharp_opt+=" -x HCOLL_BCOL_P2P_ALLREDUCE_SHARP_MAX=$sharp_test_max_data "
#   hcoll_sharp_opt+=" -x SHARP_COLL_PIPELINE_DEPTH=32 " # Size of fragmentation pipeline for larger collective payload
#	hcoll_sharp_opt+=" -x SHARP_COLL_POLL_BATCH=1 " # How many CQ completions to poll on at once. Maximum:16
# Job quota
hcoll_sharp_opt+=" -x SHARP_COLL_JOB_QUOTA_OSTS=32 " # OST quota request. value 0 mean allocate default quota.
hcoll_sharp_opt+=" -x SHARP_COLL_JOB_QUOTA_MAX_GROUPS=4 "
hcoll_sharp_opt+=" -x SHARP_COLL_JOB_QUOTA_PAYLOAD_PER_OST=$SHARP_MAX_PAYLOAD_SIZE " # Maximum payload per OST quota request. value 256 is max

run_cmd()
{
	local _test=$1
	local _cmd=$2
	local _var=$3

	if [ -n "$_test" ]; then
		echo -e "$_cmd"
		if [ -z "$dryrun" ]; then
			$_cmd 2>&1 | tee -a $MPI_LOG
			eval $_var=${PIPESTATUS[0]}
		fi
		echo ""
	fi
}

var1=0
var2=0
var3=0
var4=0

if [ -n "$sharp_allreduce" ] || [ -n "$sharp_barrier" ]; then
	printf "\n\n....................................................\n"
	printf "   w/ SHARP\n"
	printf "....................................................\n"
fi
run_cmd "$sharp_allreduce" "$mpirun $opt $hcoll_sharp_opt $osu_allreduce_exe $osu_opt" "var1"
run_cmd "$sharp_barrier"   "$mpirun $opt $hcoll_sharp_opt $osu_barrier_exe $osu_opt" "var2"

if [ -n "$allreduce" ] || [ -n "$barrier" ]; then
	printf "\n\n....................................................\n"
	printf "   w/o SHARP\n"
	printf "....................................................\n"
fi
run_cmd "$allreduce" "$mpirun $opt $osu_allreduce_exe $osu_opt" "var3"
run_cmd "$barrier"   "$mpirun $opt $osu_barrier_exe $osu_opt" "var4"


if [ $var1 -ne 0 ] || [ $var2 -ne 0 ] || [ $var3 -ne 0 ] || [ $var4 -ne 0 ]; then
	exit 1
else
	exit 0
fi
