223 lines · 7.0 KB · TypeScript
RawDownload
import { existsSync } from "node:fs";import { freemem, totalmem } from "node:os";import type { GpuInfo, RuntimeGpuMonitoringTool } from "../../models/types";import { runCommand, runCommandAsync } from "../../../core/command";import { getGpuInfoFromAmdSmi, getGpuInfoFromRocmSmi } from "./amd-gpu";import { getGpuInfoFromIntelSysfs } from "./intel-gpu";import { resolveRocmSmiTool } from "./rocm-info";import {  resolveAmdSmiBinary,  resolveForcedGpuMonitoringTool,  resolveNvidiaSmiBinary,  resolveRocmSmiBinary,} from "./smi-tools"; const NVIDIA_SMI_GPU_FIELDS = [  "uuid",  "pci.bus_id",  "name",  "memory.total",  "memory.used",  "memory.free",  "utilization.gpu",  "temperature.gpu",  "power.draw",  "power.limit",] as const; // driver_version is appended after the GPU fields so one nvidia-smi invocation// can feed GPU info, CUDA driver info, and the monitoring probe.const NVIDIA_SMI_SNAPSHOT_QUERY = [...NVIDIA_SMI_GPU_FIELDS, "driver_version"].join(",");const NVIDIA_SMI_ARGS = [  `--query-gpu=${NVIDIA_SMI_SNAPSHOT_QUERY}`,  "--format=csv,noheader,nounits",];const NVIDIA_SMI_TIMEOUT_MS = 5_000; const parseNvidiaSmiGpuLine = (line: string, index: number): GpuInfo => {  const parts = line.split(",").map((value) => value.trim());  const [    rawUuid,    rawPciBusId,    rawName,    memoryTotal,    memoryUsed,    memoryFree,    utilization,    temperature,    powerDraw,    powerLimit,  ] = parts;  const name = rawName ?? "Unknown";  const identity = (value: string | undefined): string | undefined => {    if (!value || /^(?:N\/A|\[Not Supported\])$/i.test(value)) return undefined;    return value;  };  const uuid = identity(rawUuid);  const pciBusId = identity(rawPciBusId);  const toFiniteNumber = (value: string | undefined): number => {    const parsed = Number(value ?? 0);    return Number.isFinite(parsed) ? parsed : 0;  };  const toMb = (megabytes: string | undefined): number =>    Math.max(0, Math.round(toFiniteNumber(megabytes)));  const reportedTotalMb = toMb(memoryTotal);  const isUnifiedMemoryNvidia = reportedTotalMb === 0 && /\b(?:GB10|Grace)\b/i.test(name);  const fallbackTotalMb = isUnifiedMemoryNvidia ? Math.round(totalmem() / 1024 / 1024) : 0;  const fallbackFreeMb = isUnifiedMemoryNvidia ? Math.round(freemem() / 1024 / 1024) : 0;  const fallbackUsedMb = Math.max(0, fallbackTotalMb - fallbackFreeMb);  const memoryTotalMb = reportedTotalMb || fallbackTotalMb;  const memoryUsedMb = toMb(memoryUsed) || fallbackUsedMb;  const memoryFreeMb = toMb(memoryFree) || fallbackFreeMb;  return {    ...(uuid ? { uuid } : {}),    ...(pciBusId ? { pci_bus_id: pciBusId } : {}),    index,    name,    memory_total_mb: memoryTotalMb,    memory_used_mb: memoryUsedMb,    memory_free_mb: memoryFreeMb,    utilization_pct: toFiniteNumber(utilization),    temp_c: toFiniteNumber(temperature),    power_draw: toFiniteNumber(powerDraw),    power_limit: toFiniteNumber(powerLimit),  };}; const splitSmiLines = (stdout: string): string[] =>  stdout    .trim()    .split("\n")    .map((line) => line.trim())    .filter(Boolean); const parseNvidiaSmiGpuOutput = (stdout: string): GpuInfo[] =>  splitSmiLines(stdout).map(parseNvidiaSmiGpuLine); const parseNvidiaSmiDriverVersion = (stdout: string): string | null => {  const firstLine = splitSmiLines(stdout)[0];  if (!firstLine) return null;  const driver = firstLine.split(",")[NVIDIA_SMI_GPU_FIELDS.length]?.trim();  return driver || null;}; export type NvidiaSmiSnapshot = {  /** nvidia-smi exists and the query command exited 0. */  available: boolean;  gpus: GpuInfo[];  driverVersion: string | null;}; /** * Single async nvidia-smi invocation whose parsed output feeds GPU info, the * CUDA driver version, and the GPU-monitoring probe. Returns null when * nvidia-smi is not resolvable at all. */export const queryNvidiaSmiSnapshot = async (): Promise<NvidiaSmiSnapshot | null> => {  const nvidiaSmi = resolveNvidiaSmiBinary();  if (!nvidiaSmi) return null;  try {    const result = await runCommandAsync(nvidiaSmi, NVIDIA_SMI_ARGS, {      timeoutMs: NVIDIA_SMI_TIMEOUT_MS,    });    if (result.status !== 0 || !result.stdout) {      return { available: result.status === 0, gpus: [], driverVersion: null };    }    return {      available: true,      gpus: parseNvidiaSmiGpuOutput(result.stdout),      driverVersion: parseNvidiaSmiDriverVersion(result.stdout),    };  } catch {    return { available: false, gpus: [], driverVersion: null };  }}; export const getGpuInfoFromNvidiaSmi = (): GpuInfo[] => {  try {    const nvidiaSmi = resolveNvidiaSmiBinary();    if (!nvidiaSmi) return [];     const result = runCommand(nvidiaSmi, NVIDIA_SMI_ARGS, NVIDIA_SMI_TIMEOUT_MS);    if (result.status !== 0 || !result.stdout) return [];     return parseNvidiaSmiGpuOutput(result.stdout);  } catch {    return [];  }}; /** Tool the cascade in getGpuInfo would use, without running any query commands. */export const detectGpuMonitoringTool = (): RuntimeGpuMonitoringTool | null => {  const forced = resolveForcedGpuMonitoringTool();  if (forced) return forced;  if (resolveNvidiaSmiBinary()) return "nvidia-smi";  const rocmTool = resolveRocmSmiTool();  if (rocmTool) return rocmTool;  if (getGpuInfoFromIntelSysfs().length > 0) return "intel-sysfs";  return null;}; // Logged once per process so CPU-only and missing-driver hosts are distinguishable// from "zero GPUs" without spamming every poll.let warnedNoGpuTooling = false; const warnNoGpuToolingOnce = (): void => {  if (warnedNoGpuTooling) return;  warnedNoGpuTooling = true;  const attempted = [    `nvidia-smi=${resolveNvidiaSmiBinary() ? "found" : "not found"}`,    `amd-smi=${resolveAmdSmiBinary() ? "found" : "not found"}`,    `rocm-smi=${resolveRocmSmiBinary() ? "found" : "not found"}`,    `intel-sysfs=${existsSync("/sys/bus/pci/devices") ? "no compute GPUs" : "unavailable"}`,  ].join(" ");  console.warn(`No GPUs reported by any monitoring tool; attempted: ${attempted}`);}; const collectGpuInfo = (): GpuInfo[] => {  const forced = resolveForcedGpuMonitoringTool();  if (forced === "nvidia-smi") {    return getGpuInfoFromNvidiaSmi();  }  if (forced === "amd-smi") {    return getGpuInfoFromAmdSmi();  }  if (forced === "rocm-smi") {    return getGpuInfoFromRocmSmi();  }  if (forced === "intel-sysfs") {    return getGpuInfoFromIntelSysfs();  }   const nvidia = getGpuInfoFromNvidiaSmi();  if (nvidia.length > 0) {    return nvidia;  }   const rocmTool = resolveRocmSmiTool();  if (rocmTool === "amd-smi") {    const amd = getGpuInfoFromAmdSmi();    if (amd.length > 0) return amd;    return getGpuInfoFromRocmSmi();  }  if (rocmTool === "rocm-smi") {    const rocm = getGpuInfoFromRocmSmi();    if (rocm.length > 0) return rocm;    return getGpuInfoFromAmdSmi();  }   const intel = getGpuInfoFromIntelSysfs();  if (intel.length > 0) {    return intel;  }   return [];}; export const getGpuInfo = (): GpuInfo[] => {  const gpus = collectGpuInfo();  if (gpus.length === 0) {    warnNoGpuToolingOnce();  }  return gpus;};