188 lines
7.4 KiB
TypeScript
188 lines
7.4 KiB
TypeScript
import { describe, it, expect } from 'vitest';
|
||
import {
|
||
parseNvidiaSmi,
|
||
detectNvidia,
|
||
appleVramFraction,
|
||
detectAppleSilicon,
|
||
isAppleSilicon,
|
||
detectHardware,
|
||
type CommandRunner,
|
||
type SystemProbe,
|
||
} from '../src/local/hardware-detect.js';
|
||
|
||
const baseSystem = (over: Partial<SystemProbe> = {}): SystemProbe => ({
|
||
platform: 'linux',
|
||
arch: 'x64',
|
||
totalRamGb: 32,
|
||
freeRamGb: 20,
|
||
cpuCores: 16,
|
||
cpuName: 'AMD Ryzen 9 5950X',
|
||
...over,
|
||
});
|
||
|
||
/** A runner that answers only for the bare 'nvidia-smi'/'nvidia-smi.exe' name. */
|
||
const runnerYielding = (out: string | null): CommandRunner =>
|
||
async (command) =>
|
||
command === 'nvidia-smi' || command === 'nvidia-smi.exe' ? out : null;
|
||
|
||
const runnerNever: CommandRunner = async () => null;
|
||
|
||
describe('parseNvidiaSmi', () => {
|
||
it('parses a single discrete GPU (12282 MiB → 12.0 GB)', () => {
|
||
const p = parseNvidiaSmi('NVIDIA GeForce RTX 4070, 12282\n');
|
||
expect(p.driverError).toBeNull();
|
||
expect(p.unified).toEqual([]);
|
||
expect(p.discrete).toEqual([{ index: 0, name: 'NVIDIA GeForce RTX 4070', vramGb: 12.0 }]);
|
||
});
|
||
|
||
it('parses multi-GPU (2× RTX 3090, 24576 MiB each → 24.0 GB each)', () => {
|
||
const p = parseNvidiaSmi('NVIDIA GeForce RTX 3090, 24576\nNVIDIA GeForce RTX 3090, 24576\n');
|
||
expect(p.discrete).toHaveLength(2);
|
||
expect(p.discrete[0]).toEqual({ index: 0, name: 'NVIDIA GeForce RTX 3090', vramGb: 24.0 });
|
||
expect(p.discrete[1].index).toBe(1);
|
||
expect(p.discrete[1].vramGb).toBe(24.0);
|
||
});
|
||
|
||
it('flags a driver/library mismatch as driverError, not a GPU', () => {
|
||
const p = parseNvidiaSmi('Failed to initialize NVML: Driver/library version mismatch\n');
|
||
expect(p.discrete).toEqual([]);
|
||
expect(p.driverError).toBe('Failed to initialize NVML: Driver/library version mismatch');
|
||
});
|
||
|
||
it('treats a non-numeric memory.total ([N/A]) as a unified-memory part', () => {
|
||
const p = parseNvidiaSmi('NVIDIA GB10, [N/A]\n');
|
||
expect(p.discrete).toEqual([]);
|
||
expect(p.unified).toEqual([{ index: 0, name: 'NVIDIA GB10' }]);
|
||
});
|
||
|
||
it('is safe on malformed/garbage output (no crash, no false GPU)', () => {
|
||
const p = parseNvidiaSmi('garbage with no comma\n\n \n,\n');
|
||
expect(p.discrete).toEqual([]);
|
||
expect(p.unified).toEqual([]);
|
||
expect(p.driverError).toBeNull();
|
||
});
|
||
|
||
it('is safe on empty string', () => {
|
||
expect(parseNvidiaSmi('')).toEqual({ discrete: [], unified: [], driverError: null });
|
||
});
|
||
});
|
||
|
||
describe('detectNvidia', () => {
|
||
it('resolves a single GPU from the bare nvidia-smi name', async () => {
|
||
const res = await detectNvidia(runnerYielding('NVIDIA GeForce RTX 4070, 12282'), 'linux', 32);
|
||
expect(res.gpus).toEqual([{ index: 0, name: 'NVIDIA GeForce RTX 4070', vramGb: 12.0 }]);
|
||
expect(res.unifiedMemory).toBe(false);
|
||
expect(res.driverError).toBeNull();
|
||
});
|
||
|
||
it('resolves a unified-memory part by reporting system RAM as its VRAM', async () => {
|
||
const res = await detectNvidia(runnerYielding('NVIDIA GB10, [N/A]'), 'linux', 128);
|
||
expect(res.unifiedMemory).toBe(true);
|
||
expect(res.gpus).toEqual([{ index: 0, name: 'NVIDIA GB10', vramGb: 128.0 }]);
|
||
});
|
||
|
||
it('returns the driver error (and no GPUs) when nvidia-smi cannot reach the driver', async () => {
|
||
const res = await detectNvidia(
|
||
runnerYielding('Failed to initialize NVML: Driver/library version mismatch'), 'linux', 32);
|
||
expect(res.gpus).toEqual([]);
|
||
expect(res.driverError).toContain('NVML');
|
||
});
|
||
|
||
it('falls through every candidate to a clean no-GPU result when nvidia-smi is absent', async () => {
|
||
const res = await detectNvidia(runnerNever, 'linux', 32);
|
||
expect(res).toEqual({ gpus: [], driverError: null, unifiedMemory: false });
|
||
});
|
||
|
||
it('finds nvidia-smi on a WSL absolute-path candidate when the bare name misses', async () => {
|
||
const wslRunner: CommandRunner = async (command) =>
|
||
command === '/usr/lib/wsl/lib/nvidia-smi' ? 'NVIDIA RTX A6000, 49140' : null;
|
||
const res = await detectNvidia(wslRunner, 'linux', 64);
|
||
expect(res.gpus).toEqual([{ index: 0, name: 'NVIDIA RTX A6000', vramGb: 48.0 }]); // 49140/1024=47.98→48.0
|
||
});
|
||
});
|
||
|
||
describe('appleVramFraction', () => {
|
||
it('uses the macOS working-set tiers', () => {
|
||
expect(appleVramFraction(8)).toBe(0.67);
|
||
expect(appleVramFraction(16)).toBe(0.67);
|
||
expect(appleVramFraction(32)).toBe(0.75);
|
||
expect(appleVramFraction(64)).toBe(0.75);
|
||
expect(appleVramFraction(128)).toBe(0.80);
|
||
});
|
||
});
|
||
|
||
describe('detectAppleSilicon', () => {
|
||
it('budgets 0.67 of RAM on a 16 GB M2 (→ 10.7 GB)', () => {
|
||
const gpu = detectAppleSilicon(baseSystem({ platform: 'darwin', arch: 'arm64', totalRamGb: 16, cpuName: 'Apple M2' }));
|
||
expect(gpu).toEqual({ index: 0, name: 'Apple M2', vramGb: 10.7 }); // 16*0.67=10.72→10.7
|
||
});
|
||
|
||
it('budgets 0.75 of RAM on a 64 GB M4 Max (→ 48.0 GB)', () => {
|
||
const gpu = detectAppleSilicon(baseSystem({ platform: 'darwin', arch: 'arm64', totalRamGb: 64, cpuName: 'Apple M4 Max' }));
|
||
expect(gpu?.vramGb).toBe(48.0);
|
||
});
|
||
|
||
it('returns null on an Intel Mac (x64 Darwin)', () => {
|
||
expect(isAppleSilicon(baseSystem({ platform: 'darwin', arch: 'x64' }))).toBe(false);
|
||
expect(detectAppleSilicon(baseSystem({ platform: 'darwin', arch: 'x64' }))).toBeNull();
|
||
});
|
||
});
|
||
|
||
describe('detectHardware (orchestrator)', () => {
|
||
it('Apple path: arm64 Darwin → metal, unified, no subprocess spawned', async () => {
|
||
const spawned: string[] = [];
|
||
const spy: CommandRunner = async (c) => { spawned.push(c); return null; };
|
||
const hw = await detectHardware({
|
||
run: spy,
|
||
system: baseSystem({ platform: 'darwin', arch: 'arm64', totalRamGb: 16, freeRamGb: 9, cpuName: 'Apple M2' }),
|
||
});
|
||
expect(spawned).toEqual([]); // Macs never carry nvidia-smi — skip the spawn
|
||
expect(hw.hasGpu).toBe(true);
|
||
expect(hw.backend).toBe('metal');
|
||
expect(hw.unifiedMemory).toBe(true);
|
||
expect(hw.gpuVramGb).toBe(10.7);
|
||
expect(hw.gpus).toEqual([{ name: 'Apple M2', vramGb: 10.7, backend: 'metal' }]);
|
||
});
|
||
|
||
it('NVIDIA path: single GPU → cuda, summed VRAM, gpuName set', async () => {
|
||
const hw = await detectHardware({
|
||
run: runnerYielding('NVIDIA GeForce RTX 4070, 12282'),
|
||
system: baseSystem({ totalRamGb: 32, freeRamGb: 20 }),
|
||
});
|
||
expect(hw.hasGpu).toBe(true);
|
||
expect(hw.backend).toBe('cuda');
|
||
expect(hw.gpuName).toBe('NVIDIA GeForce RTX 4070');
|
||
expect(hw.gpuVramGb).toBe(12.0);
|
||
expect(hw.gpuCount).toBe(1);
|
||
expect(hw.totalRamGb).toBe(32);
|
||
expect(hw.gpuError).toBeNull();
|
||
});
|
||
|
||
it('NVIDIA multi-GPU: VRAM summed across the pool', async () => {
|
||
const hw = await detectHardware({
|
||
run: runnerYielding('NVIDIA GeForce RTX 3090, 24576\nNVIDIA GeForce RTX 3090, 24576'),
|
||
system: baseSystem(),
|
||
});
|
||
expect(hw.gpuCount).toBe(2);
|
||
expect(hw.gpuVramGb).toBe(48.0);
|
||
});
|
||
|
||
it('CPU fallback: nvidia-smi absent → hasGpu false, CPU backend', async () => {
|
||
const hw = await detectHardware({ run: runnerNever, system: baseSystem({ arch: 'x64' }) });
|
||
expect(hw.hasGpu).toBe(false);
|
||
expect(hw.gpuName).toBeNull();
|
||
expect(hw.gpus).toEqual([]);
|
||
expect(hw.backend).toBe('CPU (x64)');
|
||
expect(hw.gpuError).toBeNull();
|
||
});
|
||
|
||
it('CPU fallback carries the driver error string (driver mismatch, not "No GPU")', async () => {
|
||
const hw = await detectHardware({
|
||
run: runnerYielding('Failed to initialize NVML: Driver/library version mismatch'),
|
||
system: baseSystem(),
|
||
});
|
||
expect(hw.hasGpu).toBe(false);
|
||
expect(hw.gpuError).toContain('NVML');
|
||
});
|
||
});
|