diff --git a/frontend/app/theme.scss b/frontend/app/theme.scss index 287a004100..4fb4f3b9a3 100644 --- a/frontend/app/theme.scss +++ b/frontend/app/theme.scss @@ -114,6 +114,7 @@ --sysinfo-cpu-color: #58c142; --sysinfo-mem-color: #53b4ea; + --sysinfo-gpu-color: #ffb84d; --bulb-color: rgb(255, 221, 51); diff --git a/frontend/app/view/sysinfo/sysinfo.tsx b/frontend/app/view/sysinfo/sysinfo.tsx index f283096749..25e669aa91 100644 --- a/frontend/app/view/sysinfo/sysinfo.tsx +++ b/frontend/app/view/sysinfo/sysinfo.tsx @@ -26,10 +26,13 @@ export type SysinfoEnv = WaveEnvSubset<{ fullConfigAtom: WaveEnv["atoms"]["fullConfigAtom"]; }; getConnStatusAtom: WaveEnv["getConnStatusAtom"]; - getBlockMetaKeyAtom: MetaKeyAtomFnType<"graph:numpoints" | "sysinfo:type" | "connection" | "count">; + getBlockMetaKeyAtom: MetaKeyAtomFnType<"graph:numpoints" | "graph:metrics" | "sysinfo:type" | "connection" | "count">; }>; const DefaultNumPoints = 120; +const DefaultMetricKeys = ["cpu"]; +const DefaultCpuCoreCount = 32; +const DefaultGpuCount = 16; type DataItem = { ts: number; @@ -58,50 +61,207 @@ function defaultMemMeta(name: string, maxY: string): TimeSeriesMeta { }; } -const PlotTypes: object = { - CPU: function (_dataItem: DataItem): Array { +function defaultGpuMeta(name: string): TimeSeriesMeta { + return { + name: name, + label: "%", + miny: 0, + maxy: 100, + color: "var(--sysinfo-gpu-color)", + decimalPlaces: 0, + }; +} + +function defaultGpuMemMeta(name: string, maxY: string): TimeSeriesMeta { + return { + name: name, + label: "GB", + miny: 0, + maxy: maxY, + color: "var(--sysinfo-gpu-color)", + decimalPlaces: 1, + }; +} + +type PlotTypeFn = (dataItem: DataItem | null) => Array; + +function getIndexedMetrics(dataItem: DataItem | null, prefix: string): Array { + return Object.keys(dataItem ?? {}) + .filter((item) => { + if (!item.startsWith(prefix)) { + return false; + } + const metricIdx = Number(item.slice(prefix.length)); + return Number.isInteger(metricIdx); + }) + .sort((a, b) => { + const valA = Number(a.slice(prefix.length)); + const valB = Number(b.slice(prefix.length)); + return valA - valB; + }); +} + +function getCpuCoreMetrics(dataItem: DataItem | null): Array { + return getIndexedMetrics(dataItem, "cpu:"); +} + +function getGpuMetrics(dataItem: DataItem | null): Array { + return getIndexedMetrics(dataItem, "gpu:"); +} + +const LegacyPlotTypes: Record = { + CPU: function (_dataItem: DataItem | null): Array { return ["cpu"]; }, - Mem: function (_dataItem: DataItem): Array { + Mem: function (_dataItem: DataItem | null): Array { return ["mem:used"]; }, - "CPU + Mem": function (_dataItem: DataItem): Array { + "CPU + Mem": function (_dataItem: DataItem | null): Array { return ["cpu", "mem:used"]; }, - "All CPU": function (dataItem: DataItem): Array { - return Object.keys(dataItem) - .filter((item) => item.startsWith("cpu") && item != "cpu") - .sort((a, b) => { - const valA = parseInt(a.replace("cpu:", "")); - const valB = parseInt(b.replace("cpu:", "")); - return valA - valB; - }); + "All CPU": function (dataItem: DataItem | null): Array { + return getCpuCoreMetrics(dataItem); }, }; -const DefaultPlotMeta = { +const MetricToggles: Array<{ label: string; getMetrics: PlotTypeFn; sublabel?: string }> = [ + { + label: "CPU", + getMetrics: function (_dataItem: DataItem | null): Array { + return ["cpu"]; + }, + }, + { + label: "Mem", + getMetrics: function (_dataItem: DataItem | null): Array { + return ["mem:used"]; + }, + }, + { + label: "GPU", + sublabel: "Installed GPU tools", + getMetrics: function (dataItem: DataItem | null): Array { + if (typeof dataItem?.gpu != "number" || !Number.isFinite(dataItem.gpu)) { + return []; + } + return ["gpu"]; + }, + }, + { + label: "All CPU Cores", + getMetrics: getCpuCoreMetrics, + }, + { + label: "All GPUs", + getMetrics: getGpuMetrics, + }, +]; + +const DefaultPlotMeta: Record = { cpu: defaultCpuMeta("CPU %"), "mem:total": defaultMemMeta("Memory Total", "mem:total"), "mem:used": defaultMemMeta("Memory Used", "mem:total"), "mem:free": defaultMemMeta("Memory Free", "mem:total"), "mem:available": defaultMemMeta("Memory Available", "mem:total"), + gpu: defaultGpuMeta("GPU %"), + "gpumem:total": defaultGpuMemMeta("GPU Memory Total", "gpumem:total"), + "gpumem:used": defaultGpuMemMeta("GPU Memory Used", "gpumem:total"), }; -for (let i = 0; i < 32; i++) { +for (let i = 0; i < DefaultCpuCoreCount; i++) { DefaultPlotMeta[`cpu:${i}`] = defaultCpuMeta(`Core ${i}`); } +for (let i = 0; i < DefaultGpuCount; i++) { + DefaultPlotMeta[`gpu:${i}`] = defaultGpuMeta(`GPU ${i}`); + DefaultPlotMeta[`gpumem:${i}:total`] = defaultGpuMemMeta(`GPU ${i} Memory Total`, `gpumem:${i}:total`); + DefaultPlotMeta[`gpumem:${i}:used`] = defaultGpuMemMeta(`GPU ${i} Memory Used`, `gpumem:${i}:total`); +} + +function dedupeMetricKeys(metrics: Array): Array { + return [...new Set(metrics.filter((metric) => typeof metric == "string" && metric != ""))]; +} + +function getLegacyPlotMetrics(plotType: string, dataItem: DataItem | null): Array { + const plotFn = LegacyPlotTypes[plotType] ?? LegacyPlotTypes.CPU; + return plotFn(dataItem); +} -function convertWaveEventToDataItem(event: Extract): DataItem { +function resolveSelectedMetricKeys(metaMetrics: unknown, plotType: string, dataItem: DataItem | null): Array { + if (Array.isArray(metaMetrics)) { + const metricKeys = dedupeMetricKeys(metaMetrics); + if (metricKeys.length > 0) { + return metricKeys; + } + } + const legacyMetricKeys = getLegacyPlotMetrics(plotType, dataItem); + if (legacyMetricKeys.length > 0) { + return legacyMetricKeys; + } + return DefaultMetricKeys; +} + +function toggleMetricKeys(currentMetrics: Array, toggledMetrics: Array): Array { + if (toggledMetrics.length == 0) { + return currentMetrics; + } + const metricSet = new Set(currentMetrics); + const removeMetrics = toggledMetrics.every((metric) => metricSet.has(metric)); + if (removeMetrics) { + toggledMetrics.forEach((metric) => metricSet.delete(metric)); + } else { + toggledMetrics.forEach((metric) => metricSet.add(metric)); + } + const nextMetrics = Array.from(metricSet); + if (nextMetrics.length == 0) { + return DefaultMetricKeys; + } + return nextMetrics; +} + +function getMetricDisplayName(metric: string): string { + if (metric == "cpu") { + return "CPU"; + } + if (metric == "mem:used") { + return "Mem"; + } + if (metric == "gpu") { + return "GPU"; + } + if (metric.startsWith("cpu:")) { + return `Core ${metric.slice("cpu:".length)}`; + } + if (metric.startsWith("gpu:")) { + return `GPU ${metric.slice("gpu:".length)}`; + } + return metric; +} + +function getMetricsViewName(metrics: Array): string { + if (metrics.length == 0) { + return "CPU"; + } + if (metrics.length <= 3) { + return metrics.map(getMetricDisplayName).join(" + "); + } + return `${metrics.length} Plots`; +} + +function convertWaveEventToDataItem(event: Extract): DataItem | null { const eventData = event.data; if (eventData == null || eventData.ts == null || eventData.values == null) { return null; } - const dataItem = { ts: eventData.ts }; + const dataItem: DataItem = { ts: eventData.ts }; for (const key in eventData.values) { dataItem[key] = eventData.values[key]; } return dataItem; } +function isDataItem(dataItem: DataItem | null): dataItem is DataItem { + return dataItem != null && typeof dataItem.ts == "number"; +} + class SysinfoViewModel implements ViewModel { viewType: string; termMode: jotai.Atom; @@ -197,19 +357,6 @@ class SysinfoViewModel implements ViewModel { } return metaNumPoints; }); - this.metrics = jotai.atom((get) => { - const plotType = get(this.plotTypeSelectedAtom); - const plotData = get(this.dataAtom); - try { - const metrics = PlotTypes[plotType](plotData[plotData.length - 1]); - if (metrics == null || !Array.isArray(metrics)) { - return ["cpu"]; - } - return metrics; - } catch (e) { - return ["cpu"]; - } - }); this.plotTypeSelectedAtom = jotai.atom((get) => { const plotType = get(this.env.getBlockMetaKeyAtom(blockId, "sysinfo:type")); if (plotType == null || typeof plotType != "string") { @@ -217,11 +364,18 @@ class SysinfoViewModel implements ViewModel { } return plotType; }); + this.metrics = jotai.atom((get) => { + const plotData = get(this.dataAtom); + const latestDataItem = plotData[plotData.length - 1]; + const metaMetrics = get(this.env.getBlockMetaKeyAtom(blockId, "graph:metrics")); + const plotType = get(this.plotTypeSelectedAtom); + return resolveSelectedMetricKeys(metaMetrics, plotType, latestDataItem); + }); this.viewIcon = jotai.atom((get) => { return "chart-line"; // should not be hardcoded }); this.viewName = jotai.atom((get) => { - return get(this.plotTypeSelectedAtom); + return getMetricsViewName(get(this.metrics)); }); this.incrementCount = jotai.atom(null, async (get, _set) => { const count = get(this.env.getBlockMetaKeyAtom(blockId, "count")) ?? 0; @@ -264,7 +418,7 @@ class SysinfoViewModel implements ViewModel { return; } this.getDefaultData(); - const initialDataItems: DataItem[] = initialData.map(convertWaveEventToDataItem); + const initialDataItems: DataItem[] = initialData.map(convertWaveEventToDataItem).filter(isDataItem); // splice the initial data into the default data (replacing the newest points) //newData.splice(newData.length - initialDataItems.length, initialDataItems.length, ...initialDataItems); globalStore.set(this.addInitialDataAtom, initialDataItems); @@ -275,6 +429,12 @@ class SysinfoViewModel implements ViewModel { } } + getSelectedMetricKeys(dataItem: DataItem | null): Array { + const metaMetrics = globalStore.get(this.env.getBlockMetaKeyAtom(this.blockId, "graph:metrics")); + const plotType = globalStore.get(this.plotTypeSelectedAtom); + return resolveSelectedMetricKeys(metaMetrics, plotType, dataItem); + } + getSettingsMenuItems(): ContextMenuItem[] { const fullConfig = globalStore.get(this.env.atoms.fullConfigAtom); const termThemes = fullConfig?.termthemes ?? {}; @@ -285,30 +445,30 @@ class SysinfoViewModel implements ViewModel { return (termThemes[a]["display:order"] ?? 0) - (termThemes[b]["display:order"] ?? 0); }); const fullMenu: ContextMenuItem[] = []; - let submenu: ContextMenuItem[]; - if (plotData.length == 0) { - submenu = []; - } else { - submenu = Object.keys(PlotTypes).map((plotType) => { - const dataTypes = PlotTypes[plotType](plotData[plotData.length - 1]); - const currentlySelected = globalStore.get(this.plotTypeSelectedAtom); - const menuItem: ContextMenuItem = { - label: plotType, - type: "radio", - checked: currentlySelected == plotType, - click: async () => { - await this.env.rpc.SetMetaCommand(TabRpcClient, { - oref: makeORef("block", this.blockId), - meta: { "graph:metrics": dataTypes, "sysinfo:type": plotType }, - }); - }, - }; - return menuItem; - }); - } + const latestDataItem = plotData[plotData.length - 1]; + const selectedMetrics = this.getSelectedMetricKeys(latestDataItem); + const submenu = MetricToggles.map((plotType) => { + const dataTypes = plotType.getMetrics(latestDataItem); + const checked = dataTypes.length > 0 && dataTypes.every((metric) => selectedMetrics.includes(metric)); + const menuItem: ContextMenuItem = { + label: plotType.label, + type: "checkbox", + checked: checked, + enabled: dataTypes.length > 0, + sublabel: dataTypes.length > 0 ? plotType.sublabel : "No data yet", + click: async () => { + const nextMetrics = toggleMetricKeys(selectedMetrics, dataTypes); + await this.env.rpc.SetMetaCommand(TabRpcClient, { + oref: makeORef("block", this.blockId), + meta: { "graph:metrics": nextMetrics }, + }); + }, + }; + return menuItem; + }); fullMenu.push({ - label: "Plot Type", + label: "Plots", submenu: submenu, }); fullMenu.push({ type: "separator" }); @@ -370,6 +530,9 @@ function SysinfoView({ model, blockId }: SysinfoViewProps) { return; } const dataItem = convertWaveEventToDataItem(event); + if (dataItem == null) { + return; + } const prevData = globalStore.get(model.dataAtom); const prevLastTs = prevData[prevData.length - 1]?.ts ?? 0; if (dataItem.ts - prevLastTs > 2000) { @@ -396,7 +559,7 @@ function SysinfoView({ model, blockId }: SysinfoViewProps) { type SingleLinePlotProps = { plotData: Array; yval: string; - yvalMeta: TimeSeriesMeta; + yvalMeta?: TimeSeriesMeta; blockId: string; defaultColor: string; title?: boolean; @@ -452,7 +615,7 @@ function SingleLinePlot({ ); if (title) { marks.push( - Plot.text([yvalMeta?.name], { + Plot.text([yvalMeta?.name ?? getMetricDisplayName(yval)], { frameAnchor: "top-left", dx: 4, fill: "var(--grey-text-color)", @@ -481,8 +644,11 @@ function SingleLinePlot({ fill: "var(--main-bg-color)", anchor: "middle", dy: -30, - title: (d) => - `${dayjs.unix(d.ts / 1000).format("HH:mm:ss")} ${Number(d[yval]).toFixed(decimalPlaces)}${labelY}`, + title: (d) => { + const value = Number(d[yval]); + const displayValue = Number.isFinite(value) ? value.toFixed(decimalPlaces) : "n/a"; + return `${dayjs.unix(d.ts / 1000).format("HH:mm:ss")} ${displayValue}${labelY}`; + }, textPadding: 3, }) ) diff --git a/frontend/preview/mock/mockwaveenv.ts b/frontend/preview/mock/mockwaveenv.ts index 123b9d3144..3e23313dca 100644 --- a/frontend/preview/mock/mockwaveenv.ts +++ b/frontend/preview/mock/mockwaveenv.ts @@ -408,6 +408,7 @@ export function makeMockWaveEnv(mockEnv?: MockEnv): MockWaveEnv { view: "sysinfo", connection: MockSysinfoConnection, "sysinfo:type": "CPU + Mem", + "graph:metrics": ["cpu", "mem:used", "gpu"], "graph:numpoints": 90, }, } as Block, diff --git a/frontend/preview/previews/sysinfo.preview-util.ts b/frontend/preview/previews/sysinfo.preview-util.ts index b577d8607b..8572245a22 100644 --- a/frontend/preview/previews/sysinfo.preview-util.ts +++ b/frontend/preview/previews/sysinfo.preview-util.ts @@ -6,6 +6,7 @@ export const MockSysinfoConnection = "local"; const MockMemoryTotal = 32; const MockCoreCount = 6; +const MockGpuMemoryTotal = 12; function clamp(value: number, minValue: number, maxValue: number): number { return Math.min(maxValue, Math.max(minValue, value)); @@ -23,12 +24,20 @@ export function makeMockSysinfoEvent( const baseCpu = clamp(42 + 18 * Math.sin(step / 6) + 8 * Math.cos(step / 3.5), 8, 96); const memUsed = clamp(12 + 4 * Math.sin(step / 10) + 2 * Math.cos(step / 7), 6, MockMemoryTotal - 4); const memAvailable = clamp(MockMemoryTotal - memUsed + 1.5, 0, MockMemoryTotal); + const gpu = clamp(34 + 22 * Math.sin(step / 8) + 10 * Math.cos(step / 5), 0, 100); + const gpuMemUsed = clamp(4 + 2 * Math.sin(step / 9) + 1.5 * Math.cos(step / 4), 1, MockGpuMemoryTotal - 1); const values: Record = { cpu: round1(baseCpu), "mem:total": MockMemoryTotal, "mem:used": round1(memUsed), "mem:free": round1(MockMemoryTotal - memUsed), "mem:available": round1(memAvailable), + gpu: round1(gpu), + "gpu:0": round1(gpu), + "gpumem:used": round1(gpuMemUsed), + "gpumem:total": MockGpuMemoryTotal, + "gpumem:0:used": round1(gpuMemUsed), + "gpumem:0:total": MockGpuMemoryTotal, }; for (let i = 0; i < MockCoreCount; i++) { diff --git a/frontend/preview/previews/sysinfo.preview.test.ts b/frontend/preview/previews/sysinfo.preview.test.ts index 6e696ea2a6..c8c0d8bd08 100644 --- a/frontend/preview/previews/sysinfo.preview.test.ts +++ b/frontend/preview/previews/sysinfo.preview.test.ts @@ -15,6 +15,9 @@ describe("sysinfo preview helpers", () => { expect(event.data.values.cpu).toBeLessThanOrEqual(100); expect(event.data.values["mem:used"]).toBeGreaterThan(0); expect(event.data.values["mem:total"]).toBeGreaterThan(event.data.values["mem:used"]); + expect(event.data.values.gpu).toBeGreaterThanOrEqual(0); + expect(event.data.values.gpu).toBeLessThanOrEqual(100); + expect(event.data.values["gpu:0"]).toBeTypeOf("number"); expect(event.data.values["cpu:0"]).toBeTypeOf("number"); }); diff --git a/pkg/wshrpc/wshremote/sysinfo.go b/pkg/wshrpc/wshremote/sysinfo.go index c573c4d9d1..83af7bf7f3 100644 --- a/pkg/wshrpc/wshremote/sysinfo.go +++ b/pkg/wshrpc/wshremote/sysinfo.go @@ -1,11 +1,21 @@ -// Copyright 2025, Command Line Inc. +// Copyright 2026, Command Line Inc. // SPDX-License-Identifier: Apache-2.0 package wshremote import ( + "context" + "encoding/json" + "fmt" + "io" "log" + "math" + "os/exec" + "runtime" + "sort" "strconv" + "strings" + "sync" "time" "github.com/shirou/gopsutil/v4/cpu" @@ -16,7 +26,31 @@ import ( "github.com/wavetermdev/waveterm/pkg/wshutil" ) -const BYTES_PER_GB = 1073741824 +const ( + BYTES_PER_GB = 1073741824 + mibPerGB = 1024 + gpuQueryMaxOutputBytes = 64 * 1024 + gpuQueryTimeout = 750 * time.Millisecond + gpuDetectionInterval = 30 * time.Second +) + +type gpuSample struct { + idx int + util float64 + memUsedGB float64 + memTotalGB float64 +} + +type gpuCollector struct { + sourceIdx int + getSamples func() []gpuSample +} + +var ( + cachedGpuCollectorsMu sync.Mutex + cachedGpuCollectors []gpuCollector + nextGpuDetectionTime time.Time +) func getCpuData(values map[string]float64) { percentArr, err := cpu.Percent(0, false) @@ -46,11 +80,718 @@ func getMemData(values map[string]float64) { values["mem:free"] = float64(memData.Free) / BYTES_PER_GB } +func parseGpuFloat(raw string) (float64, bool) { + raw = strings.TrimSpace(strings.TrimSuffix(strings.TrimSpace(raw), "%")) + val, err := strconv.ParseFloat(raw, 64) + if err != nil || math.IsNaN(val) || math.IsInf(val, 0) || val < 0 { + return 0, false + } + return val, true +} + +func parseGpuInt(raw string) (int, bool) { + val, err := strconv.Atoi(strings.TrimSpace(raw)) + if err != nil || val < 0 { + return 0, false + } + return val, true +} + +func parseNvidiaSmiOutput(output []byte) []gpuSample { + lines := strings.Split(strings.TrimSpace(string(output)), "\n") + var samples []gpuSample + for _, line := range lines { + line = strings.TrimSpace(line) + if line == "" { + continue + } + parts := strings.Split(line, ",") + if len(parts) != 4 { + continue + } + idx, ok := parseGpuInt(parts[0]) + if !ok { + continue + } + util, ok := parseGpuFloat(parts[1]) + if !ok || util > 100 { + continue + } + memUsedMIB, ok := parseGpuFloat(parts[2]) + if !ok { + continue + } + memTotalMIB, ok := parseGpuFloat(parts[3]) + if !ok || memTotalMIB <= 0 || memUsedMIB > memTotalMIB { + continue + } + samples = append(samples, gpuSample{ + idx: idx, + util: util, + memUsedGB: memUsedMIB / mibPerGB, + memTotalGB: memTotalMIB / mibPerGB, + }) + } + sort.Slice(samples, func(i int, j int) bool { + return samples[i].idx < samples[j].idx + }) + return samples +} + +func parseAmdSmiMemoryUsage(raw string, unit string) (float64, float64, bool) { + parts := strings.Split(raw, "/") + if len(parts) != 2 { + return 0, 0, false + } + used, ok := parseGpuFloat(parts[0]) + if !ok { + return 0, 0, false + } + total, ok := parseGpuFloat(parts[1]) + if !ok || total <= 0 || used > total { + return 0, 0, false + } + switch strings.ToLower(strings.TrimSpace(unit)) { + case "gb", "gib": + return used, total, true + case "mb", "mib": + return used / mibPerGB, total / mibPerGB, true + case "b", "bytes": + return used / BYTES_PER_GB, total / BYTES_PER_GB, true + default: + return 0, 0, false + } +} + +func parseAmdSmiMemoryFields(fields []string) (float64, float64, bool) { + for fieldIdx, field := range fields { + if !strings.Contains(field, "/") { + continue + } + if fieldIdx+1 >= len(fields) { + continue + } + if strings.HasSuffix(field, "/") { + if fieldIdx+2 >= len(fields) { + continue + } + used, total, ok := parseAmdSmiMemoryUsage(field+fields[fieldIdx+1], fields[fieldIdx+2]) + if ok { + return used, total, true + } + continue + } + used, total, ok := parseAmdSmiMemoryUsage(field, fields[fieldIdx+1]) + if ok { + return used, total, true + } + } + return 0, 0, false +} + +func parseAmdSmiUtil(fields []string) (float64, bool) { + for idx := 1; idx < len(fields)-1; idx++ { + if fields[idx+1] != "%" { + continue + } + util, ok := parseGpuFloat(fields[idx]) + if ok && util <= 100 { + return util, true + } + } + return 0, false +} + +func parseAmdSmiMonitorOutput(output []byte) []gpuSample { + lines := strings.Split(strings.TrimSpace(string(output)), "\n") + var samples []gpuSample + for _, line := range lines { + fields := strings.Fields(strings.TrimSpace(line)) + if len(fields) < 4 { + continue + } + idx, ok := parseGpuInt(fields[0]) + if !ok { + continue + } + util, ok := parseAmdSmiUtil(fields) + if !ok { + continue + } + memUsedGB, memTotalGB, ok := parseAmdSmiMemoryFields(fields) + if !ok { + continue + } + samples = append(samples, gpuSample{ + idx: idx, + util: util, + memUsedGB: memUsedGB, + memTotalGB: memTotalGB, + }) + } + sort.Slice(samples, func(i int, j int) bool { + return samples[i].idx < samples[j].idx + }) + return samples +} + +func parseRocmSmiMemoryGB(card map[string]string, keys ...string) (float64, bool) { + for _, key := range keys { + raw, ok := card[key] + if !ok { + continue + } + bytes, ok := parseGpuFloat(raw) + if ok { + return bytes / BYTES_PER_GB, true + } + } + return 0, false +} + +func rocmSmiCardIndex(key string, fallback int) int { + if strings.HasPrefix(key, "card") { + idx, ok := parseGpuInt(strings.TrimPrefix(key, "card")) + if ok { + return idx + } + } + return fallback +} + +func parseRocmSmiJSONOutput(output []byte) []gpuSample { + cardData := make(map[string]map[string]string) + if err := json.Unmarshal(output, &cardData); err != nil { + return nil + } + cardKeys := make([]string, 0, len(cardData)) + for key := range cardData { + cardKeys = append(cardKeys, key) + } + sort.Strings(cardKeys) + var samples []gpuSample + for fallbackIdx, key := range cardKeys { + card := cardData[key] + util, ok := parseGpuFloat(card["GPU use (%)"]) + if !ok || util > 100 { + continue + } + memUsedGB, ok := parseRocmSmiMemoryGB(card, + "VRAM Total Used Memory (B)", + "VIS_VRAM Total Used Memory (B)", + "GTT Total Used Memory (B)", + ) + if !ok { + continue + } + memTotalGB, ok := parseRocmSmiMemoryGB(card, + "VRAM Total Memory (B)", + "VIS_VRAM Total Memory (B)", + "GTT Total Memory (B)", + ) + if !ok || memTotalGB <= 0 || memUsedGB > memTotalGB { + continue + } + samples = append(samples, gpuSample{ + idx: rocmSmiCardIndex(key, fallbackIdx), + util: util, + memUsedGB: memUsedGB, + memTotalGB: memTotalGB, + }) + } + sort.Slice(samples, func(i int, j int) bool { + return samples[i].idx < samples[j].idx + }) + return samples +} + +func parseGpuJSONFloat(raw any) (float64, bool) { + switch v := raw.(type) { + case float64: + if math.IsNaN(v) || math.IsInf(v, 0) || v < 0 { + return 0, false + } + return v, true + case string: + return parseGpuFloat(strings.Trim(v, `"`)) + case json.Number: + val, err := v.Float64() + if err != nil || math.IsNaN(val) || math.IsInf(val, 0) || val < 0 { + return 0, false + } + return val, true + default: + return 0, false + } +} + +func parseMacosIORegNumber(line string, key string) (float64, bool) { + keyIdx := strings.Index(line, `"`+key+`"`) + if keyIdx == -1 { + return 0, false + } + tail := line[keyIdx+len(key)+2:] + eqIdx := strings.Index(tail, "=") + if eqIdx == -1 { + return 0, false + } + tail = strings.TrimSpace(tail[eqIdx+1:]) + if tail == "" { + return 0, false + } + if tail[0] == '"' { + tail = tail[1:] + endIdx := strings.Index(tail, `"`) + if endIdx != -1 { + tail = tail[:endIdx] + } + } else { + endIdx := strings.IndexAny(tail, ",} \t\r\n") + if endIdx != -1 { + tail = tail[:endIdx] + } + } + return parseGpuFloat(tail) +} + +func firstMacosIORegNumber(line string, keys ...string) (float64, bool) { + for _, key := range keys { + val, ok := parseMacosIORegNumber(line, key) + if ok { + return val, true + } + } + return 0, false +} + +func parseMacosIORegUtil(line string) (float64, bool) { + util, ok := firstMacosIORegNumber(line, + "Device Utilization %", + "GPU Device Utilization %", + "GPU HW active residency", + ) + if ok { + return util, util <= 100 + } + var utilSum float64 + for _, key := range []string{"Renderer Utilization %", "Tiler Utilization %", "GPU Core Utilization %"} { + util, ok := parseMacosIORegNumber(line, key) + if !ok || util > 100 { + continue + } + utilSum += util + } + if utilSum == 0 { + return 0, false + } + return min(utilSum, 100), true +} + +func parseMacosIORegMemoryGB(line string) (float64, float64, bool) { + usedBytes, ok := firstMacosIORegNumber(line, + "vramUsedBytes", + "VRAM Used Bytes", + "VRAM Total Used Memory (B)", + ) + if !ok { + return 0, 0, false + } + totalBytes, totalOk := firstMacosIORegNumber(line, + "vramTotalBytes", + "VRAM Total Bytes", + "VRAM Total Memory (B)", + ) + if !totalOk { + freeBytes, freeOk := firstMacosIORegNumber(line, + "vramFreeBytes", + "VRAM Free Bytes", + ) + if freeOk { + totalBytes = usedBytes + freeBytes + totalOk = true + } + } + if !totalOk || totalBytes <= 0 || usedBytes > totalBytes { + return 0, 0, false + } + return usedBytes / BYTES_PER_GB, totalBytes / BYTES_PER_GB, true +} + +func parseMacosIORegGpuOutput(output []byte) []gpuSample { + lines := strings.Split(strings.TrimSpace(string(output)), "\n") + var samples []gpuSample + for _, line := range lines { + line = strings.TrimSpace(line) + if !strings.Contains(line, "PerformanceStatistics") { + continue + } + util, ok := parseMacosIORegUtil(line) + if !ok { + continue + } + sample := gpuSample{ + idx: len(samples), + util: util, + } + memUsedGB, memTotalGB, ok := parseMacosIORegMemoryGB(line) + if ok { + sample.memUsedGB = memUsedGB + sample.memTotalGB = memTotalGB + } + samples = append(samples, sample) + } + return samples +} + +type intelGpuTopEngineStat struct { + Busy any `json:"busy"` + Unit string `json:"unit"` +} + +type intelGpuTopMetric struct { + Value any `json:"value"` + Unit string `json:"unit"` +} + +type intelGpuTopSample struct { + Engines map[string]intelGpuTopEngineStat `json:"engines"` + RC6 intelGpuTopMetric `json:"rc6"` +} + +func parseIntelGpuTopSampleUtil(sample intelGpuTopSample) (float64, bool) { + if sample.RC6.Unit == "" || sample.RC6.Unit == "%" { + rc6, ok := parseGpuJSONFloat(sample.RC6.Value) + if ok && rc6 <= 100 { + return 100 - rc6, true + } + } + var utilSum float64 + var found bool + for _, engine := range sample.Engines { + if engine.Unit != "" && engine.Unit != "%" { + continue + } + util, ok := parseGpuJSONFloat(engine.Busy) + if !ok || util > 100 { + continue + } + utilSum += util + found = true + } + if !found { + return 0, false + } + return min(utilSum, 100), true +} + +func parseIntelGpuTopJSONOutput(output []byte) []gpuSample { + jsonText := strings.TrimSpace(string(output)) + if jsonText == "" { + return nil + } + var samples []intelGpuTopSample + if strings.HasPrefix(jsonText, "{") { + samples = parseIntelGpuTopJSONObjects(jsonText) + } else if strings.HasPrefix(jsonText, "[") { + samples = parseIntelGpuTopJSONObjects(strings.TrimPrefix(jsonText, "[")) + } else { + return nil + } + for idx := len(samples) - 1; idx >= 0; idx-- { + util, ok := parseIntelGpuTopSampleUtil(samples[idx]) + if ok { + return []gpuSample{{idx: 0, util: util}} + } + } + return nil +} + +func parseIntelGpuTopJSONObjects(jsonText string) []intelGpuTopSample { + var samples []intelGpuTopSample + for { + jsonText = strings.TrimLeft(jsonText, " \t\r\n,") + jsonText = strings.TrimRight(jsonText, " \t\r\n,]") + if jsonText == "" { + return samples + } + var sample intelGpuTopSample + decoder := json.NewDecoder(strings.NewReader(jsonText)) + if err := decoder.Decode(&sample); err != nil { + return samples + } + samples = append(samples, sample) + jsonText = jsonText[decoder.InputOffset():] + } +} + +func encodeGpuSampleSource(samples []gpuSample, sourceIdx int, sourceCount int) []gpuSample { + rtn := make([]gpuSample, 0, len(samples)) + for _, sample := range samples { + if sample.idx < 0 || sourceIdx < 0 || sourceIdx >= sourceCount { + continue + } + sample.idx = sample.idx*sourceCount + sourceIdx + rtn = append(rtn, sample) + } + return rtn +} + +func collectGpuSamplesFromCollectors(collectors []gpuCollector, sourceCount int) ([]gpuSample, []gpuCollector) { + var samples []gpuSample + var activeCollectors []gpuCollector + for _, collector := range collectors { + collectorSamples := collector.getSamples() + if len(collectorSamples) == 0 { + continue + } + samples = append(samples, encodeGpuSampleSource(collectorSamples, collector.sourceIdx, sourceCount)...) + activeCollectors = append(activeCollectors, collector) + } + return samples, activeCollectors +} + +func addGpuSamples(values map[string]float64, samples []gpuSample) { + if len(samples) == 0 { + return + } + var utilSum float64 + var memUsedSum float64 + var memTotalSum float64 + for _, sample := range samples { + gpuIdx := strconv.Itoa(sample.idx) + values["gpu:"+gpuIdx] = sample.util + if sample.memTotalGB > 0 { + values["gpumem:"+gpuIdx+":used"] = sample.memUsedGB + values["gpumem:"+gpuIdx+":total"] = sample.memTotalGB + } + utilSum += sample.util + if sample.memTotalGB > 0 { + memUsedSum += sample.memUsedGB + memTotalSum += sample.memTotalGB + } + } + values["gpu"] = utilSum / float64(len(samples)) + if memTotalSum > 0 { + values["gpumem:used"] = memUsedSum + values["gpumem:total"] = memTotalSum + } +} + +func runGpuQuery(ctx context.Context, name string, args ...string) ([]byte, error) { + return runGpuQueryInternal(ctx, false, name, args...) +} + +func runGpuQueryAllowTimeout(ctx context.Context, name string, args ...string) ([]byte, error) { + return runGpuQueryInternal(ctx, true, name, args...) +} + +func runGpuQueryInternal(ctx context.Context, allowTimeoutOutput bool, name string, args ...string) ([]byte, error) { + cmd := exec.CommandContext(ctx, name, args...) + stdout, err := cmd.StdoutPipe() + if err != nil { + return nil, err + } + cmd.Stderr = io.Discard + if err := cmd.Start(); err != nil { + return nil, err + } + output, readErr := io.ReadAll(io.LimitReader(stdout, gpuQueryMaxOutputBytes+1)) + waitErr := cmd.Wait() + if readErr != nil { + return nil, readErr + } + if len(output) > gpuQueryMaxOutputBytes { + return nil, fmt.Errorf("%s output exceeded %d bytes", name, gpuQueryMaxOutputBytes) + } + if waitErr != nil { + if allowTimeoutOutput && ctx.Err() != nil && len(output) > 0 { + return output, nil + } + return nil, waitErr + } + return output, nil +} + +func runNvidiaSmiQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery( + ctx, + "nvidia-smi", + "--query-gpu=index,utilization.gpu,memory.used,memory.total", + "--format=csv,noheader,nounits", + ) +} + +func runAmdSmiQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "amd-smi", "monitor", "--gfx", "--vram-usage") +} + +func runRocmSmiQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "rocm-smi", "--showuse", "--showmeminfo", "vram", "--json") +} + +func runMacosIORegGpuQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "ioreg", "-r", "-d", "1", "-w", "0", "-c", "IOAccelerator") +} + +func runMacosAGXGpuQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "ioreg", "-r", "-d", "1", "-w", "0", "-c", "AGXAccelerator") +} + +func runMacosIntelGpuQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "ioreg", "-r", "-d", "1", "-w", "0", "-c", "IntelAccelerator") +} + +func runIntelGpuTopQuery(ctx context.Context) ([]byte, error) { + return runGpuQuery(ctx, "intel_gpu_top", "-J", "-s", "250", "-n", "2", "-o", "-") +} + +func runIntelGpuTopFallbackQuery(ctx context.Context) ([]byte, error) { + return runGpuQueryAllowTimeout(ctx, "intel_gpu_top", "-J", "-s", "250", "-o", "-") +} + +func getNvidiaGpuSamples() []gpuSample { + ctx, cancel := context.WithTimeout(context.Background(), gpuQueryTimeout) + defer cancel() + output, err := runNvidiaSmiQuery(ctx) + if err != nil { + return nil + } + return parseNvidiaSmiOutput(output) +} + +func getAmdSmiGpuSamples() []gpuSample { + ctx, cancel := context.WithTimeout(context.Background(), gpuQueryTimeout) + defer cancel() + output, err := runAmdSmiQuery(ctx) + if err != nil { + return nil + } + return parseAmdSmiMonitorOutput(output) +} + +func getRocmSmiGpuSamples() []gpuSample { + ctx, cancel := context.WithTimeout(context.Background(), gpuQueryTimeout) + defer cancel() + output, err := runRocmSmiQuery(ctx) + if err != nil { + return nil + } + return parseRocmSmiJSONOutput(output) +} + +func getAmdGpuSamples() []gpuSample { + samples := getAmdSmiGpuSamples() + if len(samples) > 0 { + return samples + } + return getRocmSmiGpuSamples() +} + +func getMacosGpuSamples() []gpuSample { + if runtime.GOOS != "darwin" { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), gpuQueryTimeout) + output, err := runMacosIORegGpuQuery(ctx) + cancel() + if err == nil { + samples := parseMacosIORegGpuOutput(output) + if len(samples) > 0 { + return samples + } + } + for _, queryFn := range []func(context.Context) ([]byte, error){runMacosAGXGpuQuery, runMacosIntelGpuQuery} { + ctx, cancel = context.WithTimeout(context.Background(), gpuQueryTimeout) + output, err = queryFn(ctx) + cancel() + if err != nil { + continue + } + samples := parseMacosIORegGpuOutput(output) + if len(samples) > 0 { + return samples + } + } + return nil +} + +func getIntelGpuSamples() []gpuSample { + if runtime.GOOS != "linux" { + return nil + } + ctx, cancel := context.WithTimeout(context.Background(), gpuQueryTimeout) + output, err := runIntelGpuTopQuery(ctx) + cancel() + if err != nil { + ctx, cancel = context.WithTimeout(context.Background(), gpuQueryTimeout) + output, err = runIntelGpuTopFallbackQuery(ctx) + cancel() + if err != nil { + return nil + } + } + return parseIntelGpuTopJSONOutput(output) +} + +func defaultGpuCollectors() []gpuCollector { + return []gpuCollector{ + {sourceIdx: 0, getSamples: getNvidiaGpuSamples}, + {sourceIdx: 1, getSamples: getAmdGpuSamples}, + {sourceIdx: 2, getSamples: getMacosGpuSamples}, + {sourceIdx: 3, getSamples: getIntelGpuSamples}, + } +} + +func getCachedGpuCollectors(now time.Time) ([]gpuCollector, bool) { + cachedGpuCollectorsMu.Lock() + defer cachedGpuCollectorsMu.Unlock() + if len(cachedGpuCollectors) > 0 { + return append([]gpuCollector(nil), cachedGpuCollectors...), false + } + if now.Before(nextGpuDetectionTime) { + return nil, false + } + return nil, true +} + +func setCachedGpuCollectors(collectors []gpuCollector, now time.Time) { + cachedGpuCollectorsMu.Lock() + defer cachedGpuCollectorsMu.Unlock() + cachedGpuCollectors = append(cachedGpuCollectors[:0], collectors...) + if len(cachedGpuCollectors) == 0 { + nextGpuDetectionTime = now.Add(gpuDetectionInterval) + } else { + nextGpuDetectionTime = time.Time{} + } +} + +func getGpuData(values map[string]float64) { + allCollectors := defaultGpuCollectors() + now := time.Now() + collectors, shouldDetect := getCachedGpuCollectors(now) + if len(collectors) > 0 { + samples, activeCollectors := collectGpuSamplesFromCollectors(collectors, len(allCollectors)) + setCachedGpuCollectors(activeCollectors, now) + if len(samples) > 0 { + addGpuSamples(values, samples) + return + } + shouldDetect = true + } + if !shouldDetect { + return + } + samples, activeCollectors := collectGpuSamplesFromCollectors(allCollectors, len(allCollectors)) + setCachedGpuCollectors(activeCollectors, now) + addGpuSamples(values, samples) +} + func generateSingleServerData(client *wshutil.WshRpc, connName string) { now := time.Now() values := make(map[string]float64) getCpuData(values) getMemData(values) + getGpuData(values) tsData := wshrpc.TimeSeriesData{Ts: now.UnixMilli(), Values: values} event := wps.WaveEvent{ Event: wps.Event_SysInfo, diff --git a/pkg/wshrpc/wshremote/sysinfo_test.go b/pkg/wshrpc/wshremote/sysinfo_test.go new file mode 100644 index 0000000000..a0c5271eb1 --- /dev/null +++ b/pkg/wshrpc/wshremote/sysinfo_test.go @@ -0,0 +1,208 @@ +// Copyright 2026, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +package wshremote + +import "testing" + +func TestParseNvidiaSmiOutput(t *testing.T) { + output := []byte("1, 50, 4096, 8192\n0, 25.5, 1024, 4096\n") + samples := parseNvidiaSmiOutput(output) + + if len(samples) != 2 { + t.Fatalf("expected 2 samples, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 25.5 || samples[0].memUsedGB != 1 || samples[0].memTotalGB != 4 { + t.Fatalf("unexpected first sample: %#v", samples[0]) + } + if samples[1].idx != 1 || samples[1].util != 50 || samples[1].memUsedGB != 4 || samples[1].memTotalGB != 8 { + t.Fatalf("unexpected second sample: %#v", samples[1]) + } + + values := map[string]float64{} + addGpuSamples(values, samples) + if values["gpu"] != 37.75 { + t.Fatalf("expected aggregate gpu utilization 37.75, got %v", values["gpu"]) + } + if values["gpumem:used"] != 5 || values["gpumem:total"] != 12 { + t.Fatalf("unexpected aggregate gpu memory: %#v", values) + } +} + +func TestParseNvidiaSmiOutputSkipsMalformedRows(t *testing.T) { + output := []byte("bad\n0, 101, 1024, 2048\n1, 33, 4096, 2048\n2, 44, 1024, 2048\n") + samples := parseNvidiaSmiOutput(output) + + if len(samples) != 1 { + t.Fatalf("expected 1 valid sample, got %d", len(samples)) + } + if samples[0].idx != 2 || samples[0].util != 44 || samples[0].memUsedGB != 1 || samples[0].memTotalGB != 2 { + t.Fatalf("unexpected sample: %#v", samples[0]) + } +} + +func TestParseAmdSmiMonitorOutput(t *testing.T) { + output := []byte(`GPU XCP POWER GPU_T MEM_T GFX_CLK GFX% MEM% ENC% DEC% VRAM_USAGE +0 0 183 W 49 C 48 C 1427 MHz 12 % 0 % N/A 0 % 0.3/192.0 GB +1 0 42 W 29 C 30 C 47 MHz 2 % 0 % N/A 0 % 512.0/8192.0 MB +2 0 42 W 29 C 30 C 47 MHz 3 % 0 % N/A 0 % 0.5/ 48.0 GB +`) + samples := parseAmdSmiMonitorOutput(output) + + if len(samples) != 3 { + t.Fatalf("expected 3 samples, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 12 || samples[0].memUsedGB != 0.3 || samples[0].memTotalGB != 192 { + t.Fatalf("unexpected first sample: %#v", samples[0]) + } + if samples[1].idx != 1 || samples[1].util != 2 || samples[1].memUsedGB != 0.5 || samples[1].memTotalGB != 8 { + t.Fatalf("unexpected second sample: %#v", samples[1]) + } + if samples[2].idx != 2 || samples[2].util != 3 || samples[2].memUsedGB != 0.5 || samples[2].memTotalGB != 48 { + t.Fatalf("unexpected third sample: %#v", samples[2]) + } +} + +func TestParseRocmSmiJSONOutput(t *testing.T) { + output := []byte(`{ + "card1": { + "GPU use (%)": "7", + "VRAM Total Memory (B)": "2147483648", + "VRAM Total Used Memory (B)": "1073741824" + }, + "card0": { + "GPU use (%)": "50", + "VRAM Total Memory (B)": "4294967296", + "VRAM Total Used Memory (B)": "2147483648" + } +}`) + samples := parseRocmSmiJSONOutput(output) + + if len(samples) != 2 { + t.Fatalf("expected 2 samples, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 50 || samples[0].memUsedGB != 2 || samples[0].memTotalGB != 4 { + t.Fatalf("unexpected first sample: %#v", samples[0]) + } + if samples[1].idx != 1 || samples[1].util != 7 || samples[1].memUsedGB != 1 || samples[1].memTotalGB != 2 { + t.Fatalf("unexpected second sample: %#v", samples[1]) + } +} + +func TestParseMacosIORegGpuOutput(t *testing.T) { + output := []byte(`+-o AGXAccelerator + "PerformanceStatistics" = {"Device Utilization %"=42,"Renderer Utilization %"=25,"Tiler Utilization %"=17,"In use system memory"=268435456} ++-o IntelAccelerator + "PerformanceStatistics" = {"Renderer Utilization %"=20,"Tiler Utilization %"=10,"vramUsedBytes"=1073741824,"vramFreeBytes"=3221225472} +`) + samples := parseMacosIORegGpuOutput(output) + + if len(samples) != 2 { + t.Fatalf("expected 2 samples, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 42 || samples[0].memUsedGB != 0 || samples[0].memTotalGB != 0 { + t.Fatalf("unexpected Apple Silicon sample: %#v", samples[0]) + } + if samples[1].idx != 1 || samples[1].util != 30 || samples[1].memUsedGB != 1 || samples[1].memTotalGB != 4 { + t.Fatalf("unexpected Intel macOS sample: %#v", samples[1]) + } +} + +func TestParseIntelGpuTopJSONOutput(t *testing.T) { + output := []byte(`[ + { + "rc6": {"value": 99, "unit": "%"}, + "engines": { + "Render/3D/0": {"busy": 3.5, "unit": "%"}, + "Video/0": {"busy": "-", "unit": "%"} + } + }, + { + "rc6": {"value": 74.5, "unit": "%"}, + "engines": { + "Render/3D/0": {"busy": 25.5, "unit": "%"}, + "Blitter/0": {"busy": "5", "unit": "%"}, + "Frequency": {"busy": 1200, "unit": "MHz"} + } + } +]`) + samples := parseIntelGpuTopJSONOutput(output) + + if len(samples) != 1 { + t.Fatalf("expected 1 sample, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 25.5 { + t.Fatalf("unexpected Intel GPU sample: %#v", samples[0]) + } +} + +func TestParseIntelGpuTopJSONOutputRepairsUnterminatedArray(t *testing.T) { + output := []byte(`[ + { + "rc6": {"value": 88, "unit": "%"}, + "engines": { + "Render/3D/0": {"busy": 5, "unit": "%"} + } + }, +`) + samples := parseIntelGpuTopJSONOutput(output) + + if len(samples) != 1 { + t.Fatalf("expected 1 sample, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 12 { + t.Fatalf("unexpected Intel GPU sample: %#v", samples[0]) + } +} + +func TestParseIntelGpuTopJSONOutputFallsBackToEngineBusy(t *testing.T) { + output := []byte(`{ + "engines": { + "Render/3D/0": {"busy": 25.5, "unit": "%"}, + "Blitter/0": {"busy": "5", "unit": "%"}, + "Frequency": {"busy": 1200, "unit": "MHz"} + } +}`) + samples := parseIntelGpuTopJSONOutput(output) + + if len(samples) != 1 { + t.Fatalf("expected 1 sample, got %d", len(samples)) + } + if samples[0].idx != 0 || samples[0].util != 30.5 { + t.Fatalf("unexpected Intel GPU fallback sample: %#v", samples[0]) + } +} + +func TestEncodeGpuSampleSourcePreservesCollectorIdentity(t *testing.T) { + sourceCount := len(defaultGpuCollectors()) + samples := append( + encodeGpuSampleSource([]gpuSample{ + {idx: 0, util: 10, memUsedGB: 1, memTotalGB: 2}, + {idx: 1, util: 20, memUsedGB: 2, memTotalGB: 4}, + }, 0, sourceCount), + encodeGpuSampleSource([]gpuSample{ + {idx: 0, util: 30, memUsedGB: 3, memTotalGB: 6}, + }, 1, sourceCount)..., + ) + + if samples[0].idx != 0 || samples[1].idx != sourceCount || samples[2].idx != 1 { + t.Fatalf("expected source-stable indices, got %#v", samples) + } +} + +func TestCollectGpuSamplesFromCollectorsKeepsActiveCollectors(t *testing.T) { + collectors := []gpuCollector{ + {sourceIdx: 0, getSamples: func() []gpuSample { return []gpuSample{{idx: 0, util: 10}} }}, + {sourceIdx: 1, getSamples: func() []gpuSample { return nil }}, + {sourceIdx: 2, getSamples: func() []gpuSample { return []gpuSample{{idx: 0, util: 30}} }}, + } + + samples, activeCollectors := collectGpuSamplesFromCollectors(collectors, len(collectors)) + + if len(samples) != 2 || samples[0].idx != 0 || samples[1].idx != 2 { + t.Fatalf("expected source-stable collected samples, got %#v", samples) + } + if len(activeCollectors) != 2 || activeCollectors[0].sourceIdx != 0 || activeCollectors[1].sourceIdx != 2 { + t.Fatalf("expected only active collectors, got %#v", activeCollectors) + } +}