#!/bin/bash

INS_TYPE_URL="http://metadata.tencentyun.com/latest/meta-data/instance/instance-type"
SUPPORT_INS_TYPE="HCCPNV4H HCCPNV4S HCCG5V HCCTG5v HCCPNV5 HCCPNV4ane HCCPNV6"

log_file="/tmp/gpudirect_rdma_setup.log"


log() {
    local msg=$1
    [ -z "$msg" ] && return 1
    local now=`date +"%F %H:%M:%S"`
    local ms=`date +%N`
    ms=`echo | awk "{printf(\"%03d\", $ms/1000000)}"`
    echo "[$now,$ms]: $msg" | tee -a $log_file
}

exitlog() {
    log "$1"
    exit 0
}


type=`curl $INS_TYPE_URL`
[ -z "$type" ] && exitlog "instance type($type) is invalid"


found=""
for ins_type in $SUPPORT_INS_TYPE
do
    found=`echo $type | grep -i $ins_type`
    [ -n "$found" ] && break
done

[ -z "$found" ] && exitlog "ins-type:$type no need setup"

#add modprobe in nv_gpu_conf.sh which run with rc.local
service_dir="/usr/local/qcloud/gpu"
sudo echo "sudo modprobe nvidia-peermem" >> $service_dir/nv_gpu_conf.sh

log "setup auto modprobe"

setup() {
#gpu driver install need 10min(600s)
for ((i=0;i<60;i++))
do
	which nvidia-smi
	if [ $? -eq 0 ]; then
		nv_mod=`lsmod | grep -i nvidia`
		if [ -n "$nv_mod" ]; then
			log "gpu driver is ok"
			pm_enabled=`nvidia-smi -q -i 0 | grep -i "Persistence Mode" | grep -i "Enabled"`
			if [ -n "$pm_enabled" ]; then
			    log "persistence mode enabled"
			    break;
			fi
		fi
	fi
	log "waiting for gpu driver install"
	sleep 10
done

#470 and newer driver need modprobe peermem
driver_version=`nvidia-smi -q -i 0 | grep "Driver Version" | awk '{print$4}' | awk -F '.' '{print$1}'`

[ -z "$driver_version" ] && exitlog "no nvidia-smi command"

[ $driver_version -lt 470 ] && exitlog "driver:$driver_version no need setup"

sudo modprobe nvidia-peermem
[ $? -eq 0 ] && log "modprobe successfully"

}

setup &