peardata/server/services/collectors/nvidia.js

/**
 * NVIDIA GPU metrics via nvidia-smi.
 *
 * Enable: PEARDATA_NVIDIA=1
 *
 * Charts: gpu.util.{i}, gpu.mem.{i}, gpu.temp.{i}, gpu.power.{i}
 */
import { execFile } from 'child_process'
import { promisify } from 'util'
import { CollectorPlugin } from './plugin.js'
import { registerChart } from '../../../shared/metrics.js'
import logger from '../../utils/logger.js'

const log = logger.child('nvidia')
const execFileAsync = promisify(execFile)

const QUERY =
  'index,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw'

export function isNvidiaEnabled() {
  const v = process.env.PEARDATA_NVIDIA
  return v === '1' || v === 'on' || v === 'true'
}

/**
 * @param {number} index
 * @param {'util'|'mem'|'temp'|'power'} kind
 */
function makeGpuChart(index, kind) {
  const id = `gpu.${kind}.${index}`
  const titles = {
    util: `GPU ${index} utilization`,
    mem: `GPU ${index} memory`,
    temp: `GPU ${index} temperature`,
    power: `GPU ${index} power`,
  }
  const units = {
    util: 'percentage',
    mem: 'MiB',
    temp: 'Celsius',
    power: 'Watts',
  }
  const contexts = {
    util: 'gpu.utilization',
    mem: 'gpu.memory',
    temp: 'gpu.temperature',
    power: 'gpu.power',
  }
  const dims =
    kind === 'mem'
      ? [
          { id: 'used', name: 'used', algorithm: 'absolute' },
          { id: 'total', name: 'total', algorithm: 'absolute' },
        ]
      : [{ id: 'value', name: 'value', algorithm: 'absolute' }]
  return {
    id,
    name: id,
    context: contexts[kind],
    title: titles[kind],
    units: units[kind],
    family: `gpu${index}`,
    chartType: kind === 'mem' ? 'area' : 'line',
    priority: 7200 + index,
    plugin: 'nvidia',
    dimensions: dims,
  }
}

/**
 * @param {string} stdout
 * @returns {Array<{ index: number, util: number, memUsed: number, memTotal: number, temp: number, power: number }>}
 */
export function parseNvidiaSmiCsv(stdout) {
  /** @type {Array<{ index: number, util: number, memUsed: number, memTotal: number, temp: number, power: number }>} */
  const out = []
  for (const line of String(stdout || '').split('\n')) {
    const t = line.trim()
    if (!t) continue
    const parts = t.split(',').map((s) => s.trim())
    if (parts.length < 6) continue
    out.push({
      index: Number(parts[0]) || out.length,
      util: Number(parts[1]) || 0,
      memUsed: Number(parts[2]) || 0,
      memTotal: Number(parts[3]) || 0,
      temp: Number(parts[4]) || 0,
      power: Number(parts[5]) || 0,
    })
  }
  return out
}

export async function queryNvidiaGpus(timeoutMs = 5000) {
  const { stdout } = await execFileAsync(
    'nvidia-smi',
    [`--query-gpu=${QUERY}`, '--format=csv,noheader,nounits'],
    { timeout: timeoutMs, encoding: 'utf8' }
  )
  return parseNvidiaSmiCsv(stdout)
}

export class NvidiaCollector extends CollectorPlugin {
  constructor(opts = {}) {
    super({ name: 'nvidia', intervalMs: opts.intervalMs })
    /** @type {boolean|null} */
    this._available = null
  }

  isEnabled() {
    return isNvidiaEnabled()
  }

  start() {
    if (!this.isEnabled()) return
    log.info('NVIDIA collector started')
    super.start()
  }

  async collect() {
    let gpus
    try {
      gpus = await queryNvidiaGpus()
      this._available = true
    } catch (err) {
      if (this._available !== false) {
        log.info('nvidia-smi unavailable — collector idle', { error: err.message })
        this._available = false
      }
      return []
    }
    const ts = Date.now()
    /** @type {Array<{ chart: string, context: string, ts: number, values: Record<string, number|null> }>} */
    const batch = []
    for (const g of gpus) {
      const utilDef = makeGpuChart(g.index, 'util')
      const memDef = makeGpuChart(g.index, 'mem')
      const tempDef = makeGpuChart(g.index, 'temp')
      const powerDef = makeGpuChart(g.index, 'power')
      registerChart(utilDef)
      registerChart(memDef)
      registerChart(tempDef)
      registerChart(powerDef)
      batch.push(
        { chart: utilDef.id, context: utilDef.context, ts, values: { value: g.util } },
        {
          chart: memDef.id,
          context: memDef.context,
          ts,
          values: { used: g.memUsed, total: g.memTotal },
        },
        { chart: tempDef.id, context: tempDef.context, ts, values: { value: g.temp } },
        { chart: powerDef.id, context: powerDef.context, ts, values: { value: g.power } }
      )
    }
    return batch
  }
}

/** @type {NvidiaCollector|null} */
let singleton = null

export function getNvidiaCollector() {
  if (!singleton) singleton = new NvidiaCollector()
  return singleton
}