diff --git a/src/lib/metrics-poller.ts b/src/lib/metrics-poller.ts new file mode 100644 index 0000000..0c39973 --- /dev/null +++ b/src/lib/metrics-poller.ts @@ -0,0 +1,116 @@ +import { fetchMetrics, type LlamaSwapConfig } from "./llamaswap"; +import { gpuInfos, parseGpuMetrics, type GpuMetricKind, type GpuSample } from "./metrics-parser"; + +const RING_SIZE = 60; + +function avg(values: number[]): number { + return values.reduce((a, b) => a + b, 0) / (values.length || 1); +} +function max(values: number[]): number { + return values.reduce((a, b) => Math.max(a, b), -Infinity); +} +function sum(values: number[]): number { + return values.reduce((a, b) => a + b, 0); +} + +const AGGREGATE: Record number> = { + util_percent: avg, + memory_util_percent: avg, + temperature: max, + power: sum, + fan: avg, +}; + +function key(gpuId: string, kind: GpuMetricKind): string { + return `${gpuId}|${kind}`; +} + +type FetchFn = (cfg: LlamaSwapConfig) => Promise; + +export class MetricsPoller { + private buffers = new Map(); + private gpuName = new Map(); + private gpuIds: string[] = []; + private listeners = new Set<() => void>(); + private timer?: ReturnType; + private lastError?: string; + + constructor( + private cfg: LlamaSwapConfig, + private intervalMs = 5000, + private fetchFn: FetchFn = fetchMetrics, + ) {} + + start(): void { + void this.tick(); + this.timer = setInterval(() => void this.tick(), this.intervalMs); + } + + stop(): void { + if (this.timer) clearInterval(this.timer); + this.timer = undefined; + } + + on(listener: () => void): () => void { + this.listeners.add(listener); + return () => { + this.listeners.delete(listener); + }; + } + + async tick(): Promise { + try { + const text = await this.fetchFn(this.cfg); + this.apply(parseGpuMetrics(text)); + this.lastError = undefined; + } catch (err) { + this.lastError = err instanceof Error ? err.message : String(err); + } finally { + for (const listener of this.listeners) listener(); + } + } + + private apply(samples: GpuSample[]): void { + const byKind = new Map>(); + for (const s of samples) { + if (!byKind.has(s.kind)) byKind.set(s.kind, new Map()); + byKind.get(s.kind)!.set(s.id, s.value); + if (!this.gpuName.has(s.id)) this.gpuName.set(s.id, s.name); + } + this.gpuIds = [...this.gpuName.keys()].sort((a, b) => parseInt(a, 10) - parseInt(b, 10)); + + for (const [kind, values] of byKind) { + for (const [id, value] of values) this.push(key(id, kind), value); + const all = [...values.values()]; + if (all.length > 0) this.push(key("all", kind), AGGREGATE[kind](all)); + } + } + + private push(k: string, value: number): void { + let ring = this.buffers.get(k); + if (!ring) { + ring = []; + this.buffers.set(k, ring); + } + ring.push(value); + if (ring.length > RING_SIZE) ring.shift(); + } + + getValue(gpuId: string, kind: GpuMetricKind): number | undefined { + const ring = this.buffers.get(key(gpuId, kind)); + return ring ? ring[ring.length - 1] : undefined; + } + + getHistory(gpuId: string, kind: GpuMetricKind): number[] { + const ring = this.buffers.get(key(gpuId, kind)); + return ring ? [...ring] : []; + } + + gpus(): { id: string; name: string }[] { + return this.gpuIds.map((id) => ({ id, name: this.gpuName.get(id) ?? "" })); + } + + isOffline(): boolean { + return this.lastError !== undefined; + } +} diff --git a/tests/metrics-poller.test.ts b/tests/metrics-poller.test.ts new file mode 100644 index 0000000..56c4d07 --- /dev/null +++ b/tests/metrics-poller.test.ts @@ -0,0 +1,63 @@ +import assert from "node:assert/strict"; +import { readFileSync } from "node:fs"; +import { test } from "node:test"; +import { type LlamaSwapConfig } from "../src/lib/util"; +import { MetricsPoller } from "../src/lib/metrics-poller"; +import { parseGpuMetrics } from "../src/lib/metrics-parser"; + +const FIXTURE = readFileSync(new URL("./fixtures/metrics.txt", import.meta.url), "utf8"); +const FIXTURE2 = FIXTURE + .replace(/llamaswap_gpu_util_percent\{id="0"[^}]*\} 100/, 'llamaswap_gpu_util_percent{id="0",name="RTX PRO 6000"} 30') + .replace(/llamaswap_gpu_util_percent\{id="1"[^}]*\} 100/, 'llamaswap_gpu_util_percent{id="1",name="RTX PRO 6000"} 50') + .replace(/llamaswap_gpu_util_percent\{id="2"[^}]*\} 100/, 'llamaswap_gpu_util_percent{id="2",name="RTX 5090"} 80'); + +function stubFetch(bodies: string[]) { + let i = 0; + return async () => bodies[Math.min(i++, bodies.length - 1)]; +} + +const CFG: LlamaSwapConfig = { baseUrl: "http://test" }; + +test("poller stores per-GPU history and computes aggregates", async () => { + const poller = new MetricsPoller(CFG, 5000, stubFetch([FIXTURE, FIXTURE2])); + await poller.tick(); + await poller.tick(); + + assert.equal(poller.getValue("0", "util_percent"), 30); + assert.deepEqual(poller.getHistory("0", "util_percent"), [100, 30]); + + const allUtil = poller.getValue("all", "util_percent"); + assert.ok(allUtil !== undefined); + assert.ok(Math.abs(allUtil - (30 + 50 + 80) / 3) < 1e-9); + + const allTemp = poller.getValue("all", "temperature"); + assert.equal(allTemp, 64); + + const allPower = poller.getValue("all", "power"); + assert.ok(allPower !== undefined); + assert.ok(Math.abs(allPower - (316.56 + 336.17 + 377.28)) < 1e-9); + + assert.deepEqual(poller.gpus().map((g) => g.id), ["0", "1", "2"]); + assert.equal(poller.gpus()[2].name, "NVIDIA GeForce RTX 5090"); +}); + +test("poller ring buffer caps at 60 samples", async () => { + const poller = new MetricsPoller(CFG, 5000, stubFetch([FIXTURE])); + for (let i = 0; i < 70; i++) await poller.tick(); + assert.equal(poller.getHistory("0", "fan").length, 60); +}); + +test("poller reports offline after a fetch failure and recovers", async () => { + let fail = true; + const fetchFn = async () => { + if (fail) throw new Error("boom"); + return FIXTURE; + }; + const poller = new MetricsPoller(CFG, 5000, fetchFn); + await poller.tick(); + assert.equal(poller.isOffline(), true); + fail = false; + await poller.tick(); + assert.equal(poller.isOffline(), false); + assert.ok(poller.getValue("0", "util_percent") !== undefined); +});