Show k8s nodes with GPU model, allocated and free resources.

view-gpu-with-model.sh

#!/usr/bin/env bash
# -----------------------------------------------------------------------------
# view-gpu-with-model.sh
# -----------------------------------------------------------------------------
# Wrapper for `kubectl-view-allocations -r gpu`:
#   1. Retrieves each Node's NVIDIA GPU model
#   2. Appends the model name at end of each line (preserves column alignment)
#   3. Highlights models by performance class (H100, A100, 5090, 4090, 3090)
#   4. Colorizes utilization percentages "(XX%)" in a gradient (green→yellow→red)
#
# Usage:
#   chmod +x view-gpu-with-model.sh
#   ./view-gpu-with-model.sh
# -----------------------------------------------------------------------------
set -e

# -----------------------------------------------------------------------------
# ANSI color codes (ANSI 颜色码)
# -----------------------------------------------------------------------------
RESET=$'\e[0m'       # Reset all colors (重置所有颜色)

# GPU model → color map (型号→颜色映射)
declare -A COLOR_MAP=(
  [H100]=$'\e[31m'    # Red for H100 series
  [A100]=$'\e[33m'    # Yellow for A100 series
  [6000]=$'\e[34m'
  [5090]=$'\e[34m'    # Blue for 5090 (high-mid tier)
  [4090]=$'\e[36m'    # Cyan for 4090 (mid tier)
  [3090]=$'\e[36m'    # Cyan for 3090 (mid tier)
)
DEFAULT_COLOR=$'\e[37m'  # White fallback (默认白色)

# Utilization % → color thresholds (利用率→颜色区间)
util_color() {
  local pct=$1
  if   (( pct >= 80 )); then echo $'\e[31m'   # ≥80% → red
  elif (( pct >= 50 )); then echo $'\e[33m'   # ≥50% → yellow
  elif (( pct >= 20 )); then echo $'\e[32m'   # ≥20% → green
  else                      echo $'\e[36m'; # <20% → cyan
  fi
}

# -----------------------------------------------------------------------------
# Step 0: Ensure `kubectl-view-allocations` is installed
# 步骤0：检查 kubectl-view-allocations 是否安装
# -----------------------------------------------------------------------------
if ! command -v kubectl-view-allocations &>/dev/null; then
  echo "Error: 'kubectl-view-allocations' not found."
  echo "错误：找不到 kubectl-view-allocations 命令。"
  echo
  echo "Install with:"
  echo "  curl https://raw.githubusercontent.com/davidB/kubectl-view-allocations/master/scripts/getLatest.sh | bash"
  exit 1
fi

# -----------------------------------------------------------------------------
# Step 1: Build map of node → GPU model
# 步骤1：构建 节点 → GPU 型号 映射
# -----------------------------------------------------------------------------
declare -A MODEL_MAP
while IFS=$'\t' read -r node model; do
  MODEL_MAP["$node"]="$model"
done < <(
  kubectl get nodes \
    -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.metadata.labels.nvidia\.com/gpu\.product}{"\n"}{end}'
)

# -----------------------------------------------------------------------------
# Step 2: Run allocations view and transform lines
# 步骤2：执行资源视图并转换每行
# -----------------------------------------------------------------------------
kubectl-view-allocations -r gpu | while IFS= read -r line; do
  # 1) Header or blank lines: print unchanged
  if [[ $line =~ ^[[:space:]]*$ ]] || [[ $line =~ ^[[:space:]]*Resource ]]; then
    echo -e "$line"
    continue
  fi

  # 2) Colorize utilization “(XX%)” using util_color()
  if [[ $line =~ \(([0-9]{1,3})%\) ]]; then
    pct=${BASH_REMATCH[1]}
    uc=$(util_color "$pct")
    line="${line//(${pct}%)/(${uc}${pct}%${RESET})}"
  fi

  # 3) Append colored GPU model at end of each node line
  if [[ $line =~ ([[:space:]]*[├└]─[[:space:]]*)([A-Za-z0-9._-]+) ]] || \
     [[ $line =~ ^[[:space:]]*([A-Za-z0-9._-]+) ]]; then

    node="${BASH_REMATCH[2]:-${BASH_REMATCH[1]}}"
    model="${MODEL_MAP[$node]}"

    if [[ -n $model ]]; then
      # select model color by matching keywords in uppercase model string
      up_model="${model^^}"
      color="$DEFAULT_COLOR"
      for key in "${!COLOR_MAP[@]}"; do
        if [[ $up_model == *"$key"* ]]; then
          color="${COLOR_MAP[$key]}"
          break
        fi
      done
      echo -e "${line} (${color}${model}${RESET})"
      continue
    fi
  fi

  # 4) All other lines: print unchanged
  echo -e "$line"
done