File size: 1,488 Bytes
4357d9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | #!/usr/bin/env bash
set -euo pipefail
GPU_MEM_USED_LIMIT_MB="${GPU_MEM_USED_LIMIT_MB:-2048}"
GPU_UTIL_LIMIT="${GPU_UTIL_LIMIT:-10}"
require_nvidia_smi() {
if ! command -v nvidia-smi >/dev/null 2>&1; then
echo "nvidia-smi is required for GPU guard checks." >&2
return 1
fi
}
_gpu_query_once() {
nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits
}
list_idle_gpus() {
require_nvidia_smi || return 1
local sample idx mem util
sample="$(_gpu_query_once)"
while IFS=',' read -r idx mem util; do
idx="${idx// /}"
mem="${mem// /}"
util="${util// /}"
if [[ -n "$idx" && "$mem" -le "$GPU_MEM_USED_LIMIT_MB" && "$util" -le "$GPU_UTIL_LIMIT" ]]; then
printf '%s\n' "$idx"
fi
done <<< "$sample"
}
count_idle_gpus() {
list_idle_gpus | sed '/^$/d' | wc -l | tr -d ' '
}
join_first_n_idle_gpus() {
local need="$1"
list_idle_gpus | sed '/^$/d' | head -n "$need" | paste -sd, -
}
ensure_idle_gpu_count() {
local need="$1"
local have
have="$(count_idle_gpus)"
if [[ "$have" -lt "$need" ]]; then
echo "Need $need idle GPU(s), but only found $have under guard limits: mem<=${GPU_MEM_USED_LIMIT_MB}MB util<=${GPU_UTIL_LIMIT}%." >&2
return 1
fi
}
print_gpu_guard_summary() {
require_nvidia_smi || return 1
echo "GPU guard limits: mem<=${GPU_MEM_USED_LIMIT_MB}MB util<=${GPU_UTIL_LIMIT}%"
nvidia-smi --query-gpu=index,name,memory.used,memory.free,utilization.gpu --format=csv,noheader
}
|