diff --git a/scripts/tests/test-qmd-benchmark.mjs b/scripts/tests/test-qmd-benchmark.mjs new file mode 100644 index 0000000..95da139 --- /dev/null +++ b/scripts/tests/test-qmd-benchmark.mjs @@ -0,0 +1,25 @@ +import { test } from 'node:test'; +import assert from 'node:assert/strict'; +import { passes, fixtureDocuments, describeDevices } from '../../shared/qmd/qmd-benchmark.mjs'; + +test('both thresholds must pass', () => { + assert.equal(passes({ querySeconds: 5, chunksPerSecond: 20 }, {}), true); + assert.equal(passes({ querySeconds: 5.01, chunksPerSecond: 50 }, {}), false); + assert.equal(passes({ querySeconds: 1, chunksPerSecond: 19.9 }, {}), false); +}); + +test('thresholds are overridable', () => { + assert.equal(passes({ querySeconds: 8, chunksPerSecond: 5 }, { QMD_BENCHMARK_MAX_QUERY_SECONDS: '10', QMD_BENCHMARK_MIN_CHUNKS_PER_SECOND: '4' }), true); +}); + +test('device description lists every GPU and the thread count', () => { + assert.equal(describeDevices({ gpu: 'cuda', devices: ['RTX 4090', 'RTX 3090'], threads: 64 }), 'CUDA: RTX 4090, RTX 3090; CPU threads: 64'); + assert.equal(describeDevices({ gpu: false, devices: [], threads: 16 }), 'no GPU; CPU threads: 16'); +}); + +test('fixture is deterministic and large enough for throughput', () => { + const a = fixtureDocuments(100); + assert.equal(a.length, 100); + assert.deepEqual(a, fixtureDocuments(100)); + assert.ok(a.every((d) => d.text.length > 2500)); +}); diff --git a/shared/qmd/qmd-benchmark.mjs b/shared/qmd/qmd-benchmark.mjs new file mode 100644 index 0000000..7b5c86b --- /dev/null +++ b/shared/qmd/qmd-benchmark.mjs @@ -0,0 +1,86 @@ +// shared/qmd/qmd-benchmark.mjs — measures warm QMD search latency and embedding throughput on one device. +import { mkdtempSync, mkdirSync, rmSync, writeFileSync } from 'node:fs'; +import { availableParallelism, tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { pathToFileURL } from 'node:url'; +import { parseArgs } from 'node:util'; +import { MODELS, importQmd, qmdPackageDir } from './qmd-lib.mjs'; + +export function passes({ querySeconds, chunksPerSecond }, env = process.env) { + const maxQuery = Number(env.QMD_BENCHMARK_MAX_QUERY_SECONDS) || 5; + const minChunks = Number(env.QMD_BENCHMARK_MIN_CHUNKS_PER_SECOND) || 20; + return querySeconds <= maxQuery && chunksPerSecond >= minChunks; +} + +const TOPICS = ['invoice rounding', 'upload retries', 'trial expiry', 'cache invalidation', 'session tokens']; + +export function fixtureDocuments(count) { + return Array.from({ length: count }, (_, index) => { + const topic = TOPICS[index % TOPICS.length]; + const sections = Array.from({ length: 6 }, (__, section) => + `## ${topic} part ${section}\n\n` + (`This section ${index}-${section} explains how ${topic} behaves when the service handles request ${index * 7 + section}. ` + + `It lists the edge cases, the configuration keys, and the error messages that operators see during incidents.\n`).repeat(2)); + return { name: `doc-${String(index).padStart(3, '0')}.md`, text: `# ${topic} ${index}\n\n${sections.join('\n')}` }; + }); +} + +export function describeDevices({ gpu, devices, threads }) { + const gpus = gpu ? `${String(gpu).toUpperCase()}: ${devices.join(', ') || 'unnamed device'}` : 'no GPU'; + return `${gpus}; CPU threads: ${threads}`; +} + +// llama.cpp's automatic backend uses every GPU of the chosen backend; nothing here narrows the device set. +async function detect() { + const { getLlama } = await import(pathToFileURL(join(qmdPackageDir(), 'node_modules', 'node-llama-cpp', 'dist', 'index.js')).href); + const llama = await getLlama({ gpu: 'auto' }); + const gpu = llama.gpu; + const devices = gpu ? await llama.getGpuDeviceNames() : []; + await llama.dispose(); + return { gpu, devices, threads: availableParallelism() }; +} + +async function measure(device) { + for (const [key, value] of Object.entries({ QMD_EMBED_MODEL: MODELS.embed, QMD_RERANK_MODEL: MODELS.rerank, QMD_GENERATE_MODEL: MODELS.generate })) { + process.env[key] = value; + } + const work = mkdtempSync(join(tmpdir(), 'ai-config-qmd-bench-')); + try { + const docs = join(work, 'docs'); + mkdirSync(docs); + for (const doc of fixtureDocuments(100)) writeFileSync(join(docs, doc.name), doc.text); + const { createStore } = await importQmd(); + const store = await createStore({ dbPath: join(work, 'index.sqlite'), config: { collections: { bench: { path: docs, pattern: '**/*.md' } } } }); + try { + await store.update(); + await store.searchVector('warm up the embedding model', { limit: 1 }); + const embedded = await store.embed({ collection: 'bench' }); + const chunksPerSecond = embedded.chunksEmbedded / Math.max(embedded.durationMs / 1000, 0.001); + const question = 'Which part explains how upload retries behave during incidents?'; + await store.search({ query: question, collections: ['bench'], limit: 5 }); + const times = []; + for (let run = 0; run < 3; run += 1) { + const started = performance.now(); + await store.search({ query: `${question} (${run})`, collections: ['bench'], limit: 5, candidateLimit: 40 }); + times.push((performance.now() - started) / 1000); + } + const querySeconds = times.sort((a, b) => a - b)[1]; + const result = { device, querySeconds: Number(querySeconds.toFixed(3)), chunksPerSecond: Number(chunksPerSecond.toFixed(1)) }; + return { ...result, passed: passes(result) }; + } finally { + await store.close(); + } + } finally { + rmSync(work, { recursive: true, force: true }); + } +} + +async function main() { + const { values } = parseArgs({ options: { detect: { type: 'boolean' }, device: { type: 'string' } } }); + if (values.detect) return console.log(JSON.stringify(await detect())); + if (!['gpu', 'cpu'].includes(values.device)) throw new Error('--device gpu|cpu or --detect is required'); + console.log(JSON.stringify(await measure(values.device))); +} + +if (import.meta.url === pathToFileURL(process.argv[1] || '').href) { + main().catch((error) => { console.error(`qmd-benchmark: ${error.message}`); process.exit(1); }); +}