diff --git a/scripts/tests/test-qmd-benchmark.mjs b/scripts/tests/test-qmd-benchmark.mjs index 95da139..b56ac2d 100644 --- a/scripts/tests/test-qmd-benchmark.mjs +++ b/scripts/tests/test-qmd-benchmark.mjs @@ -1,6 +1,6 @@ import { test } from 'node:test'; import assert from 'node:assert/strict'; -import { passes, fixtureDocuments, describeDevices } from '../../shared/qmd/qmd-benchmark.mjs'; +import { passes, fixtureDocuments, describeDevices, embedThroughput } from '../../shared/qmd/qmd-benchmark.mjs'; test('both thresholds must pass', () => { assert.equal(passes({ querySeconds: 5, chunksPerSecond: 20 }, {}), true); @@ -23,3 +23,9 @@ test('fixture is deterministic and large enough for throughput', () => { assert.deepEqual(a, fixtureDocuments(100)); assert.ok(a.every((d) => d.text.length > 2500)); }); + +test('embed throughput rejects errors and empty runs', () => { + assert.deepEqual(embedThroughput({ chunksEmbedded: 200, durationMs: 4000, errors: 0 }), { chunksPerSecond: 50, ok: true }); + assert.equal(embedThroughput({ chunksEmbedded: 200, durationMs: 4000, errors: 1 }).ok, false); + assert.equal(embedThroughput({ chunksEmbedded: 0, durationMs: 4000, errors: 0 }).ok, false); +}); diff --git a/shared/qmd/qmd-benchmark.mjs b/shared/qmd/qmd-benchmark.mjs index 7b5c86b..a2fdfa4 100644 --- a/shared/qmd/qmd-benchmark.mjs +++ b/shared/qmd/qmd-benchmark.mjs @@ -24,6 +24,10 @@ export function fixtureDocuments(count) { }); } +export function embedThroughput({ chunksEmbedded, durationMs, errors }) { + return { chunksPerSecond: chunksEmbedded / Math.max(durationMs / 1000, 0.001), ok: !errors && chunksEmbedded > 0 }; +} + export function describeDevices({ gpu, devices, threads }) { const gpus = gpu ? `${String(gpu).toUpperCase()}: ${devices.join(', ') || 'unnamed device'}` : 'no GPU'; return `${gpus}; CPU threads: ${threads}`; @@ -33,10 +37,13 @@ export function describeDevices({ gpu, devices, threads }) { 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() }; + try { + const gpu = llama.gpu; + const devices = gpu ? await llama.getGpuDeviceNames() : []; + return { gpu, devices, threads: availableParallelism() }; + } finally { + await llama.dispose(); + } } async function measure(device) { @@ -52,9 +59,9 @@ async function measure(device) { 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); + await store.embed({ collection: 'bench' }); // warm-up: loads (and on a first run downloads) the model + const embedded = await store.embed({ collection: 'bench', force: true }); + const { chunksPerSecond, ok } = embedThroughput(embedded); const question = 'Which part explains how upload retries behave during incidents?'; await store.search({ query: question, collections: ['bench'], limit: 5 }); const times = []; @@ -64,8 +71,8 @@ async function measure(device) { 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) }; + const result = { device, querySeconds: Number(querySeconds.toFixed(3)), chunksPerSecond: Number(chunksPerSecond.toFixed(1)), embedErrors: embedded.errors }; + return { ...result, passed: ok && passes(result) }; } finally { await store.close(); }