#!/bin/bash

version=0.0.1
update_time="June 9, 2023."
output_filename=out_findMoth_$(date '+%Y-%m-%d_%H:%M:%S').txt
cpu_threshold=0.3 #
exec 3>&1 # 打开文件描述符3用于重定向
cpu_moth_filename=.out_cpu_moth_$(date '+%Y-%m-%d_%H:%M:%S').txt
gpu_moth_filename=.out_gpu_moth_$(date '+%Y-%m-%d_%H:%M:%S').txt
original_args="$@"
sleep_time=0.1
PWD=$PWD
#############################
## autor:
## 1.  qiql: hi@qiql.net
## 2. sunty: say@isunty.com
#############################

# 本脚本用于检测slurm集群中的蛀虫
# 判断逻辑为：1. 查询每个节点的资源申请情况，然后与实际的CPU利用率进行匹配，如果匹配不上，说明有问题
#             2. 查询每个节点上的用户，如果该用户没有作业，但仍在节点上，说明有问题

function help(){
  echo ""
  echo " Description:"
  echo -e "   findMoth-${version} is a Pdsh-based batch script suitable for high-performance Slurm clusters.
   It can check whether there are moths in the computing nodes. The judgment criteria are:

   1. For CPU nodes, determine whether the actual CPU utilization of the node is higher than the theoretical
      utilization. The theoretical utilization is calculated by dividing the number of applied cores by the 
      total number of cores of the node
   2. For GPU nodes, determine whether the actual GPU utilization and video memory utilization of the node are 
      higher than the theoretical utilization. For an eight-card GPU node, when the number of applied cards 
      is 8, the theoretical utilization should be 800%."
  echo ""
  echo " Usage: findMoth [options] parameter ... "
  echo ""
  echo " Options:"
  echo "   -h, --help                  Print this message and exit."
  echo "   -v, --version               Print the version number of findMoth and exit."
  echo "   -c, --host-cpu=</to/your/path/nodelist.txt>
                               Set the CPU nodes range through the file, file format: bn001 bn002 bn040..."
  echo "   -g, --host-gpu=</to/your/path/nodelist.txt>
                               Set the CPU nodes range through the file, file format: gpu001 gpu002 gpu003 ...."
  echo "   -o, --output=<filename>     Set the output filename."
  echo "   -t, --threshold=<value>     Set a difference between the allowable theoretical maximum utilization and the
                               actual utilization of the node, the default is 0.3."
  echo ""
  echo " Example:"
  echo "   1. findMoth -c host-cpu.txt "
  echo "   2. findMoth -c ./host-cpu.txt -g ./host-gpu.txt -t 0.35"
  echo "   3. findMoth --host-cpu=./host-cpu.txt --host-gpu=./host-gpu.txt --threshold=0.35 -output=findmoth.txt"
  echo ""
  echo " Bug Report: hi@qiql.net"
  echo ""
}

options=$(getopt -o c:g:t:vh --long host-cpu:,host-gpu:,threshold::version,help -- "$@")
if [ $? -ne 0 ]; then
   echo "提供了无效参数"
   exit 1
fi

eval set -- "$options"

while true; do
  case "$1" in
    -c|--host-cpu)
      nodelist_cpu="$2"
#      echo "while nodelist_cpu=${nodelist_cpu}"
      shift 2
      ;;
    -g|--host-gpu)
      nodelist_gpu="$2"
#      echo "nodelist_gpu=${nodelist_gpu}"
      shift 2
      ;;
    -t|--threshold)
      cpu_threshold="$2"
#      echo "nodelist_gpu=${nodelist_gpu}"
      shift 2
      ;;
    -v|--version)
      echo "finMoth version: ${version}. Published on ${update_time}"
      exit 0
      ;;
    -h|--help)
      help
      exit 0
      ;;
    -o|--output)
      output_filename="$2"
      shift 2
      ;;
    --)
      shift
      break
      ;;
    *)
      echo "无效参数：$1"
      exit 1
      ;;
   esac
done

###########@################
####  开始检查前置条件  ####
############################
echo ""
echo "============================================================================"
echo "*** Startup Prepare"
sleep ${sleep_time}
if ! which pdsh > /dev/null 2>&1; then
  echo "checking for pdsh... no."
  echo -e "\033[31mError: pdsh is not installed. Exit. \033[0m"
  exit 1
fi
echo "checking for pdsh... yes."
sleep ${sleep_time}
if ! which getopt > /dev/null; then
  echo -e "\033[31mError: getopt is not installed. Exit. \033[0m"
  exit 1
fi
echo "checking for getopt... yes."
sleep ${sleep_time}
echo "*** Finish  Prepare"
sleep ${sleep_time}
echo "============================================================================"
echo ""

function read_host_by_hostlist(){
    IFS=$'\n'   # 按换行分割
    file_array=($(cat $1))   # 读取文件内容
    for item in "${file_array[@]}"; do
        IFS=$' '   # 对于每一行再按空格分割
        sub_array=($item)
        for sub_item in "${sub_array[@]}"; do
            echo "$sub_item"
        done
        unset sub_array
    done
    unset IFS
}

if [ ! -z "$nodelist_cpu" ]
then
    cpu_hostname=($(read_host_by_hostlist "${nodelist_cpu}"))
fi

# echo "cpu_hostname=${cpu_hostname[@]}"

if [ ! -z "$nodelist_gpu" ]
then
    gpu_hostname=($(read_host_by_hostlist "${nodelist_gpu}"))
fi

# echo "gpu_hostname=${gpu_hostname[@]}"

function gpu_util(){

    local hostname=$1

    # 查询一个节点的GPU总利用率：
    # cur_gpu_util=$(nvidia-smi --format=csv --query-gpu=utilization.gpu | tail -n+2 | sed 's/ \+%//' | awk '{ sum += $1 } END { print sum }')
    local cur_gpu_util=$(pdsh -w ${hostname} nvidia-smi --format=csv --query-gpu=utilization.gpu | tail -n+2 | sed 's/.*://; s/ \+%//' | awk '{ sum += $1 } END { printf "%.2f\n", sum }')    

    # 查询一个节点的显存总利用率
    # cur_mem_util=$( pdsh -w ${hostname} nvidia-smi --format=csv --query-gpu=utilization.memory | tail -n+2 | sed 's/ \+%//' | awk '{ sum += $1 } END { printf "%.2f\n", sum }')   
    local cur_mem_util=$(pdsh -w ${hostname} nvidia-smi --format=csv --query-gpu=utilization.memory | tail -n+2 | sed 's/.*://; s/ \+%//' | awk '{ sum += $1 } END { printf "%.2f\n", sum }')

    # 查询节点总GPU数：
    local all_gpus=$(scontrol show node ${hostname} | grep -oP 'CfgTRES=.*?(?<=:tesla=)(\d+)' | awk -F= '{print $6}')
    
    # 查询节点已被申请的GPU数
    local alloc_gpus=$(scontrol show node ${hostname} | grep -oP 'AllocTRES=.*?(?<=:tesla=)(\d+)' | awk -F= '{print $4}')
    if [ ! -n "$alloc_gpus" ];then
        alloc_gpus=0
    fi
    # 计算理论利用率
    local max_gpus=$((alloc_gpus * 100))
    
    local is_moth=1
   
    if [ `echo "${cur_gpu_util} <= ${max_gpus} "|bc` -eq 1 ] ; then
        if [ `echo "${cur_mem_util} <= ${max_gpus} "|bc` -eq 1 ] ; then        
            is_moth=0
        fi
    fi
 
    if awk -v cur_util="$cur_gpu_util" -v max_util="$max_gpus" 'BEGIN{ exit !(cur_util < max_util) }'; then
        if awk -v cur_util="$cur_mem_util" -v max_util="$max_mem" 'BEGIN{ exit !(cur_util < max_util) }'; then
            is_moth=0
        fi
    fi

    if [ $is_moth -eq 1 ];then
        printf "%8s  %4d  %4d  %6.2f%% %6.2f%%  %6.2f%%  %2d\n" $hostname $all_cpus $alloc_cpus $max_cpus $cur_cpu_util $cur_mem_util $is_moth >> ${gpu_moth_filename}
    fi

    printf "%8s  %4d  %4d  %6.2f%% %6.2f%%  %6.2f%%  %2d\n" $hostname $all_gpus $alloc_gpus $max_gpus $cur_gpu_util $cur_mem_util $is_moth >> ${output_filename}
}

function cpu_util(){
   
    local hostname=$1
    local threshold=$2 # 设置阈值，比如某个节点上，即使是空载的，CPU利用率可能为0.3%
    # 查询一个节点的CPU利用率
    local cur_cpu_util=$(pdsh -w ${hostname} top -b -n 1 | grep -oE '%Cpu\(s\):.*' | awk '{print $2}')
    # 查询一个节点的内存利用率
    local cur_mem_util=$(pdsh -w ${hostname} free -m | sed 's/[^ ]*://' | awk 'NR==2{print $3/$2 * 100}')
    local all_cpus=$(scontrol show node ${hostname}| grep -oP 'CfgTRES=.*?cpu=\K\d+')
    local alloc_cpus=$(scontrol show node ${hostname} | grep -oP 'AllocTRES=cpu=\K\d+')
    if [ ! -n "$alloc_cpus" ];then
        alloc_cpus=0
    fi
    # 计算理论最大CPU利用率
    # printf "alloc_cpus=%3d  all_cpus=%4d\n\n" ${alloc_cpus} ${all_cpus}
    # local max_cpus=$(echo "scale=2; ${alloc_cpus} / ${all_cpus} * 100" | bc -l)
    # local max_cpus=$(printf "%.2f" $(echo "$alloc_cpus / $all_cpus" | bc -l))
    local max_cpus=$(awk -v a="${alloc_cpus}" -v b="${all_cpus}" 'BEGIN{ printf "%.2f", a/b }')
    max_cpus=$(awk -v num="${max_cpus}" 'BEGIN{printf "%.2f\n", num * 100}')
    local is_moth=1

    if [ `echo "${cur_cpu_util} <= ${max_cpus} "|bc` -eq 1 ] ; then
        is_moth=0
    fi
    if [ $(echo "${cur_cpu_util} - ${max_cpus} <= ${threshold}" | bc) -eq 1 ]; then
        is_moth=0
    fi
 
    if [ $is_moth -eq 1 ];then
        printf "%8s  %4d  %4d  %6.2f%% %6.2f%%  %6.2f%%  %2d\n" $hostname $all_cpus $alloc_cpus $max_cpus $cur_cpu_util $cur_mem_util $is_moth >> ${cpu_moth_filename}
    fi

    printf "%8s  %4d  %4d  %6.2f%% %6.2f%%  %6.2f%%  %2d\n" $hostname $all_cpus $alloc_cpus $max_cpus $cur_cpu_util $cur_mem_util $is_moth >> ${output_filename}
}

echo "Full path to this output file: ${PWD}/${output_filename}" >> ${output_filename}
script_name="$0"
command_args="$@"

echo "Full Command:${script_name} ${original_args}" >> ${output_filename}

###############################################
### 检查GPU计算节点
###############################################
length_gpu_hostname=${#gpu_hostname[@]}
if [ ${length_gpu_hostname} -ne 0 ]; then
    echo "" | tee >(cat - >&3) >> "${output_filename}" 
    echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
#    echo "############################################################################" | tee >(cat - >&3) >> "${output_filename}"
    echo "=== Check GPU Start: $(date '+%Y-%m-%d %H:%M:%S')" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    pre_start_time=$(date +%s)
    schedule=''
    schedule_index=0
    k=${length_gpu_hostname}
    printf "progress:[%-25s]%d%%\r" ${schedule} ${schedule_index}
    # echo "length_gpu_hostname=${length_gpu_hostname}"
    printf "hostname  gpus alloc max_gpus gpu_util mem_util is_moth\n" >> ${output_filename}
    for (( i=0;i<${length_gpu_hostname};i++ ))
    do
        start_time=$(date +%s)
        host=${gpu_hostname[$i]}
        gpu_util "${host}"
        if [[ ${i} == $(( length_gpu_hostname-1 )) ]];then
            schedule_index=100
        else
#        schedule_index=`expr $(( 100/length_cpu_hostname )) + ${schedule_index}`
            schedule_index=$(printf "%.3f" "$(echo "scale=3; 100 / $length_gpu_hostname + $schedule_index" | bc)")
        fi
        schedule=''
        tmp_schedule_index=$(echo "$schedule_index/4" | bc)
        for (( sch_index=0;sch_index<${tmp_schedule_index};sch_index++ ))
        do
            schedule='$'${schedule}
        done
        cen_time=$(date +%s); time_diff=$((cen_time - start_time)); k=$((k-1)); end_time=$((k*time_diff))
        hours=$((end_time / 3600)); minutes=$((end_time % 3600 / 60)); seconds=$((end_time % 60))
        printf "progress:[%-25s]%2.3f%% eta: %d hours %02d min %02d s\r" ${schedule} ${schedule_index} $hours $minutes $seconds

#        if [[ ${i} == $(( length_gpu_hostname-1 )) ]];then
#            schedule_index=100
#        else
#            schedule_index=`expr $(( 100/length_gpu_hostname )) + ${schedule_index}`
#        fi
#        schedule=''
#        for (( sch_index=0;sch_index<$(( schedule_index/4 ));sch_index++ ))
#        do
#            schedule='$'${schedule}
#        done
#        cen_time=$(date +%s); time_diff=$((cen_time - start_time)); k=$((k-1)); end_time=$((k*time_diff))
#        hours=$((end_time / 3600)); minutes=$((end_time % 3600 / 60)); seconds=$((end_time % 60))
#        printf "progress:[%-25s]%d%% eta: %d hours %02d min %02d s\r" ${schedule} ${schedule_index} $hours $minutes $seconds
    done
    end_time=$(date +%s)
    cost_time=$((end_time - pre_start_time))
    hours=$((cost_time / 3600)); minutes=$((cost_time % 3600 / 60)); seconds=$((cost_time % 60))
    echo ""
    sleep ${sleep_time}
    echo "===   Check GPU End: $(date '+%Y-%m-%d %H:%M:%S')" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    printf "===  Check GPU Cost: %d hours %02d min %02d s" $hours $minutes $seconds | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    echo ""  | tee >(cat - >&3) >> "${output_filename}"
    echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
#    echo "############################################################################" | tee >(cat - >&3) >> "${output_filename}"
    echo "" >> "${output_filename}"
fi

##############################################
### 检查CPU计算节点
##############################################
length_cpu_hostname=${#cpu_hostname[@]}
if [ ${length_cpu_hostname} -ne 0 ]; then
    echo "" | tee >(cat - >&3) >> "${output_filename}"
    echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
#    echo "############################################################################" | tee >(cat - >&3) >> "${output_filename}"
    echo "=== Check CPU Start: $(date '+%Y-%m-%d %H:%M:%S')" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    pre_start_time=$(date +%s)
    schedule=''
    schedule_index=0
    k=${length_cpu_hostname}
    printf "progress:[%-25s]%d%%\r" ${schedule} ${schedule_index}
    length_cpu_hostname=${#cpu_hostname[@]}
    # echo "length_cpu_hostname=${length_cpu_hostname}"
    printf "hostname cpus alloc max_cpus cpu_util mem_util is_moth\n" >> ${output_filename}
    for (( i=0;i<${length_cpu_hostname};i++ ))
    do
        start_time=$(date +%s)
        host=${cpu_hostname[$i]}
        cpu_util "${host}" "${cpu_threshold}"
        if [[ ${i} == $(( length_cpu_hostname-1 )) ]];then
            schedule_index=100
        else
#        schedule_index=`expr $(( 100/length_cpu_hostname )) + ${schedule_index}`
            schedule_index=$(printf "%.3f" "$(echo "scale=3; 100 / $length_cpu_hostname + $schedule_index" | bc)")
        fi
        schedule=''
        tmp_schedule_index=$(echo "$schedule_index/4" | bc)
        for (( sch_index=0;sch_index<${tmp_schedule_index};sch_index++ ))
        do
            schedule='$'${schedule}
        done
        cen_time=$(date +%s); time_diff=$((cen_time - start_time)); k=$((k-1)); end_time=$((k*time_diff))
        hours=$((end_time / 3600)); minutes=$((end_time % 3600 / 60)); seconds=$((end_time % 60))
        printf "progress:[%-25s]%2.3f%% eta: %d hours %02d min %02d s\r" ${schedule} ${schedule_index} $hours $minutes $seconds
    done
    end_time=$(date +%s)
    cost_time=$((end_time - pre_start_time))
    hours=$((cost_time / 3600)); minutes=$((cost_time % 3600 / 60)); seconds=$((cost_time % 60))
    echo ""
    sleep ${sleep_time}
    echo "===   Check CPU End: $(date '+%Y-%m-%d %H:%M:%S')" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    printf "===  Check CPU Cost: %d hours %02d min %02d s" $hours $minutes $seconds | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    echo ""  | tee >(cat - >&3) >> "${output_filename}"
    echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
#    echo "############################################################################" | tee >(cat - >&3) >> "${output_filename}"
    echo "" >> "${output_filename}"
fi
echo ""
echo ""
echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
sleep ${sleep_time}
if [ -s "${cpu_moth_filename}" ]; then
    echo "The following are CPU Moth nodes: " | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    printf "hostname  cpus  alloc max_cpus cpu_util mem_util is_moth\n" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    cat ${cpu_moth_filename} | tee >(cat - >&3) >> "${output_filename}"
elif [ -s "${gpu_moth_filename}" ];then
    echo "The following are GPU Moth nodes: " | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    printf "hostname  gpus alloc max_gpus gpu_util mem_util is_moth\n" | tee >(cat - >&3) >> "${output_filename}"
    sleep ${sleep_time}
    cat ${gpu_moth_filename} | tee >(cat - >&3) >> "${output_filename}"
else
    sleep ${sleep_time}
    echo "*** No Moths!"
fi
sleep ${sleep_time}
echo "*** Work DIR: ${PWD}"
sleep ${sleep_time}
echo "*** Log File: ${output_filename}"
sleep ${sleep_time}
echo "Bug report: hi@qiql.net" >> "${output_filename}"
echo "============================================================================" | tee >(cat - >&3) >> "${output_filename}"
echo ""
rm -f ${cpu_moth_filename}
rm -f ${gpu_moth_filename}
# 关闭文件描述符3
exec 3>&-
