diff --git a/dist/embeddings.d.ts b/dist/embeddings.d.ts index c244884f..eaf5fa54 100644 --- a/dist/embeddings.d.ts +++ b/dist/embeddings.d.ts @@ -1,4 +1,28 @@ export declare const BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "; +/** + * Resolve an intra-op thread cap for the embedding session, or null to leave + * onnxruntime at its default (one worker per core). + * + * Embedding is done one exchange at a time (batch=1) throughout indexing, sync, + * and the re-embed migration. At that size onnxruntime-node's default intra-op + * thread pool fans each tiny int8 matmul across every core, which on Apple + * Silicon pegs ~5 cores for almost no throughput gain and makes a bulk re-embed + * hog the whole machine. Capping intra-op threads keeps embedding off the + * user's cores; measured on an M-series Mac, a cap of 2 is as fast or faster + * than the uncapped default while using ~1.7 cores instead of ~5 (and pulls + * further ahead when the machine is otherwise busy, since it stops the pool + * from oversubscribing). + * + * This stays on the same CPU / q8 execution provider, so embeddings are + * bit-identical to the uncapped output — no re-index is triggered. (CoreML and + * WebGPU were evaluated and rejected: on this quantized model they either run + * far slower and never leave the CPU, or require an fp32 model whose vectors + * differ from the stored q8 ones and would force a full re-index.) + * + * Override with EPISODIC_MEMORY_EMBED_THREADS: a positive integer sets the cap; + * 0 (or any non-positive/invalid value) restores onnxruntime's default. + */ +export declare function resolveIntraOpThreads(): number | null; export declare function initEmbeddings(): Promise; export declare function generateEmbedding(text: string): Promise; /** diff --git a/dist/embeddings.js b/dist/embeddings.js index b3881d64..d5f4c84a 100644 --- a/dist/embeddings.js +++ b/dist/embeddings.js @@ -18,11 +18,56 @@ env.useBrowserCache = false; const MODEL_ID = 'Xenova/bge-small-en-v1.5'; const MODEL_DTYPE = 'q8'; export const BGE_QUERY_PREFIX = 'Represent this sentence for searching relevant passages: '; +/** + * Resolve an intra-op thread cap for the embedding session, or null to leave + * onnxruntime at its default (one worker per core). + * + * Embedding is done one exchange at a time (batch=1) throughout indexing, sync, + * and the re-embed migration. At that size onnxruntime-node's default intra-op + * thread pool fans each tiny int8 matmul across every core, which on Apple + * Silicon pegs ~5 cores for almost no throughput gain and makes a bulk re-embed + * hog the whole machine. Capping intra-op threads keeps embedding off the + * user's cores; measured on an M-series Mac, a cap of 2 is as fast or faster + * than the uncapped default while using ~1.7 cores instead of ~5 (and pulls + * further ahead when the machine is otherwise busy, since it stops the pool + * from oversubscribing). + * + * This stays on the same CPU / q8 execution provider, so embeddings are + * bit-identical to the uncapped output — no re-index is triggered. (CoreML and + * WebGPU were evaluated and rejected: on this quantized model they either run + * far slower and never leave the CPU, or require an fp32 model whose vectors + * differ from the stored q8 ones and would force a full re-index.) + * + * Override with EPISODIC_MEMORY_EMBED_THREADS: a positive integer sets the cap; + * 0 (or any non-positive/invalid value) restores onnxruntime's default. + */ +export function resolveIntraOpThreads() { + const override = process.env.EPISODIC_MEMORY_EMBED_THREADS; + if (override !== undefined) { + const n = Number.parseInt(override, 10); + return Number.isFinite(n) && n > 0 ? n : null; + } + if (process.platform === 'darwin' && process.arch === 'arm64') { + return 2; + } + return null; +} let embeddingPipeline = null; export async function initEmbeddings() { if (!embeddingPipeline) { console.error('Loading embedding model (first run may take time)...'); - embeddingPipeline = await pipeline('feature-extraction', MODEL_ID, { dtype: MODEL_DTYPE, progress_callback: () => { } }); + const options = { + dtype: MODEL_DTYPE, + progress_callback: () => { }, + }; + const intraOpThreads = resolveIntraOpThreads(); + if (intraOpThreads !== null) { + options.session_options = { + intraOpNumThreads: intraOpThreads, + interOpNumThreads: 1, + }; + } + embeddingPipeline = await pipeline('feature-extraction', MODEL_ID, options); console.error('Embedding model loaded'); } } diff --git a/dist/mcp-server.js b/dist/mcp-server.js index b46fabb7..a74ab8d6 100644 --- a/dist/mcp-server.js +++ b/dist/mcp-server.js @@ -3837,7 +3837,7 @@ var require_fast_uri = __commonJS({ if (!options.unicodeSupport && (!schemeHandler || !schemeHandler.unicodeSupport)) { if (parsed.host && (options.domainHost || schemeHandler && schemeHandler.domainHost) && isIP === false && nonSimpleDomain(parsed.host)) { try { - parsed.host = URL.domainToASCII(parsed.host.toLowerCase()); + parsed.host = new URL("http://" + parsed.host).hostname; } catch (e) { parsed.error = parsed.error || "Host's domain name can not be converted to ASCII: " + e; } @@ -24966,16 +24966,34 @@ env.useBrowserCache = false; var MODEL_ID = "Xenova/bge-small-en-v1.5"; var MODEL_DTYPE = "q8"; var BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "; +function resolveIntraOpThreads() { + const override = process.env.EPISODIC_MEMORY_EMBED_THREADS; + if (override !== void 0) { + const n = Number.parseInt(override, 10); + return Number.isFinite(n) && n > 0 ? n : null; + } + if (process.platform === "darwin" && process.arch === "arm64") { + return 2; + } + return null; +} var embeddingPipeline = null; async function initEmbeddings() { if (!embeddingPipeline) { console.error("Loading embedding model (first run may take time)..."); - embeddingPipeline = await pipeline( - "feature-extraction", - MODEL_ID, - { dtype: MODEL_DTYPE, progress_callback: () => { - } } - ); + const options = { + dtype: MODEL_DTYPE, + progress_callback: () => { + } + }; + const intraOpThreads = resolveIntraOpThreads(); + if (intraOpThreads !== null) { + options.session_options = { + intraOpNumThreads: intraOpThreads, + interOpNumThreads: 1 + }; + } + embeddingPipeline = await pipeline("feature-extraction", MODEL_ID, options); console.error("Embedding model loaded"); } } @@ -26865,7 +26883,7 @@ ${result} } // src/version.ts -var VERSION = "1.4.1"; +var VERSION = "1.4.2"; // src/mcp-server.ts import fs4 from "fs"; diff --git a/src/embeddings.ts b/src/embeddings.ts index 2ff0d74a..782baf44 100644 --- a/src/embeddings.ts +++ b/src/embeddings.ts @@ -21,16 +21,58 @@ const MODEL_ID = 'Xenova/bge-small-en-v1.5'; const MODEL_DTYPE = 'q8'; export const BGE_QUERY_PREFIX = 'Represent this sentence for searching relevant passages: '; +/** + * Resolve an intra-op thread cap for the embedding session, or null to leave + * onnxruntime at its default (one worker per core). + * + * Embedding is done one exchange at a time (batch=1) throughout indexing, sync, + * and the re-embed migration. At that size onnxruntime-node's default intra-op + * thread pool fans each tiny int8 matmul across every core, which on Apple + * Silicon pegs ~5 cores for almost no throughput gain and makes a bulk re-embed + * hog the whole machine. Capping intra-op threads keeps embedding off the + * user's cores; measured on an M-series Mac, a cap of 2 is as fast or faster + * than the uncapped default while using ~1.7 cores instead of ~5 (and pulls + * further ahead when the machine is otherwise busy, since it stops the pool + * from oversubscribing). + * + * This stays on the same CPU / q8 execution provider, so embeddings are + * bit-identical to the uncapped output — no re-index is triggered. (CoreML and + * WebGPU were evaluated and rejected: on this quantized model they either run + * far slower and never leave the CPU, or require an fp32 model whose vectors + * differ from the stored q8 ones and would force a full re-index.) + * + * Override with EPISODIC_MEMORY_EMBED_THREADS: a positive integer sets the cap; + * 0 (or any non-positive/invalid value) restores onnxruntime's default. + */ +export function resolveIntraOpThreads(): number | null { + const override = process.env.EPISODIC_MEMORY_EMBED_THREADS; + if (override !== undefined) { + const n = Number.parseInt(override, 10); + return Number.isFinite(n) && n > 0 ? n : null; + } + if (process.platform === 'darwin' && process.arch === 'arm64') { + return 2; + } + return null; +} + let embeddingPipeline: FeatureExtractionPipeline | null = null; export async function initEmbeddings(): Promise { if (!embeddingPipeline) { console.error('Loading embedding model (first run may take time)...'); - embeddingPipeline = await pipeline( - 'feature-extraction', - MODEL_ID, - { dtype: MODEL_DTYPE, progress_callback: () => {} } - ); + const options: Parameters[2] = { + dtype: MODEL_DTYPE, + progress_callback: () => {}, + }; + const intraOpThreads = resolveIntraOpThreads(); + if (intraOpThreads !== null) { + options.session_options = { + intraOpNumThreads: intraOpThreads, + interOpNumThreads: 1, + }; + } + embeddingPipeline = await pipeline('feature-extraction', MODEL_ID, options); console.error('Embedding model loaded'); } } diff --git a/test/embed-threads.test.ts b/test/embed-threads.test.ts new file mode 100644 index 00000000..6996c000 --- /dev/null +++ b/test/embed-threads.test.ts @@ -0,0 +1,38 @@ +import { describe, it, expect, afterEach } from 'vitest'; +import { resolveIntraOpThreads } from '../src/embeddings.js'; + +const ENV_KEY = 'EPISODIC_MEMORY_EMBED_THREADS'; + +describe('resolveIntraOpThreads', () => { + const original = process.env[ENV_KEY]; + + afterEach(() => { + if (original === undefined) delete process.env[ENV_KEY]; + else process.env[ENV_KEY] = original; + }); + + it('honors a positive integer override', () => { + process.env[ENV_KEY] = '4'; + expect(resolveIntraOpThreads()).toBe(4); + }); + + it('treats 0 as "use onnxruntime default" (null)', () => { + process.env[ENV_KEY] = '0'; + expect(resolveIntraOpThreads()).toBeNull(); + }); + + it('treats a non-numeric override as default (null)', () => { + process.env[ENV_KEY] = 'nope'; + expect(resolveIntraOpThreads()).toBeNull(); + }); + + it('defaults to a small cap on Apple Silicon and null elsewhere when unset', () => { + delete process.env[ENV_KEY]; + const result = resolveIntraOpThreads(); + if (process.platform === 'darwin' && process.arch === 'arm64') { + expect(result).toBe(2); + } else { + expect(result).toBeNull(); + } + }); +});