#!/bin/bash
# ==============================================================================
# FILE: roles/llm-inference-multimodel/scripts/nvidia-smi-vram-exporter.sh
# DESCRIPTION: NVIDIA VRAM textfile exporter for Prometheus
#              Queries nvidia-smi for GPU VRAM usage and writes Prometheus-
#              formatted metrics to node_exporter's textfile collector
#              (/var/lib/node_exporter/textfile_collector/).
#              
#              Designed for 1-minute cron execution (idempotent; atomic writes).
#              Outputs: llamacpp_vram_used_mib (gauge, MiB)
# 
# CRON ENTRY:  * * * * * /opt/llama-server-monitoring/nvidia-smi-vram-exporter.sh
# OUTPUT FILE: /var/lib/node_exporter/textfile_collector/nvidia.prom
#
# AUTHOR:      Wong (Infrastructure Automation Specialist)
# DATE:        2026-08-18
# ==============================================================================

set -euo pipefail

# Configuration
TEXTFILE_DIR="/var/lib/node_exporter/textfile_collector"
OUTPUT_FILE="${TEXTFILE_DIR}/nvidia.prom"
TMPFILE="${OUTPUT_FILE}.tmp.$$"
GPU_INDEX="${1:-0}"  # Allow override via first positional arg; default GPU 0

# Ensure textfile collector directory exists
if [ ! -d "$TEXTFILE_DIR" ]; then
    echo "ERROR: $TEXTFILE_DIR does not exist. Create it with:" >&2
    echo "  mkdir -p $TEXTFILE_DIR" >&2
    echo "  chown prometheus:prometheus $TEXTFILE_DIR" >&2
    exit 1
fi

# Query nvidia-smi for instantaneous GPU VRAM usage
# Format: plain number (MiB), or empty if nvidia-smi fails
VRAM_MIB=$(nvidia-smi --query-gpu=memory.used \
    --format=csv,noheader,nounits \
    --id="$GPU_INDEX" 2>/dev/null || echo "")

# Validate output is a number; default to 0 if nvidia-smi fails
if [ -z "$VRAM_MIB" ] || ! [[ "$VRAM_MIB" =~ ^[0-9]+$ ]]; then
    VRAM_MIB=0
fi

# Write metric to temp file (atomic swap to avoid partial reads)
cat > "$TMPFILE" << EOF
# HELP llamacpp_vram_used_mib GPU VRAM used in MiB (nvidia-smi)
# TYPE llamacpp_vram_used_mib gauge
llamacpp_vram_used_mib $VRAM_MIB
EOF

# Atomic swap: move temp file to final location
# This ensures node_exporter never reads a partial file
mv "$TMPFILE" "$OUTPUT_FILE"

exit 0