#!/bin/bash


output_filename="/usr/bin/utsrpool"    # 默认输出文件名
PWD=$PWD                                                         # 工作目录路径

# 定义函数，获取正在使用的 GPU 卡数量
get_used_gpu_count() {
    # 运行 nvidia-smi 命令，获取 GPU 显存利用率
    gpu_memory_utilization=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)

    # 统计显存利用率不为零的 GPU 卡数量
    used_count=0
    for utilization in $gpu_memory_utilization; do
        if [ $utilization -gt 0 ]; then
            used_count=$((used_count + 1))
        fi
    done

    # 返回正在使用的 GPU 卡数量
    echo $used_count
}

function util(){
   
    local hostname=$HOSTNAME
    #echo "hostname=$hostname"

    # 查询一个节点的核心总数
    local all_cpus=$(cat /proc/cpuinfo| grep "processor"| wc -l)
    #echo "all_cpus=${all_cpus}"

    # 查询一个节点的CPU利用率
    local cur_cpu_util=$(top -b -n 1 | grep -oE '%Cpu\(s\):.*' | awk '{print $2}')
    #echo "cur_cpu_util=${cur_cpu_util}"
    
    # 查询一个节点的内存利用率
    local cur_mem_util=$(free -m | awk 'NR==2{print $3/$2 * 100}')
    #echo "cur_mem_util=${cur_mem_util}"    

    # 查询节点总GPU数
    local all_gpus=$(nvidia-smi --query-gpu=count --format=csv,noheader | awk 'NR==1{print; exit}')
    #echo "all_gpus=${all_gpus}"

    # 调用函数获取正在使用的 GPU 卡数量
    local used_gpu_count=$(get_used_gpu_count)
    #echo "used_gpu_count=${used_gpu_count}"

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

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

    # 计算理论利用率
    local max_gpus=$((all_gpus * 100))
    #echo "max_gpus=${max_gpus}"

    # local cur_date=$(date '+%Y-%m-%d %H:%M:%S')
    local cur_date=$(date '+%s')

    printf "%8s %8d %7.2f%% %6.2f%% %8d %8d %6.2f%% %7.2f%% %(%Y-%m-%d-%H:%M:%S)T\n" $hostname $all_cpus $cur_cpu_util $cur_mem_util $all_gpus $used_gpu_count $cur_gpu_util $cur_gpu_mem_util $cur_date >> ${output_filename}
}

function main(){
    echo "Full path to this output file: ${PWD}/${output_filename}" >> ${output_filename}
    printf "hostname all_cpus cpu_util cpu_mem all_gpus used_gpu gpu_util gpu_mem date\n" >> ${output_filename}
    while true
    do
        util
        sleep 30
    done

}

main &
