@oh-my-pi/pi-mnemopi 18.13.1 → 18.13.2

package/CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## [Unreleased]
 
+## [18.13.2] - 2026-10-27
+
+### Fixed
+
+- Fixed local embeddings using one spinning CPU thread per core, so a single recall query cost several CPU-seconds on a many-core machine ([#17109](https://github.com/can1357/oh-my-pi/pull/17109) by [@sjawhar](https://github.com/sjawhar))
+
 ## [18.13.0] - 2026-10-26
 
 ### Fixed

package/dist/types/core/embeddings.d.ts

@@ -37,6 +37,11 @@ export declare function quarantineCorruptModelFile(message: string, cacheDir?: s
  * retries once; fastembed re-downloads any file missing from the model's
  * cache directory. Also the initializer the embed worker uses in its
  * subprocess; the in-process seam stays {@link setLocalModelInitializer}.
+ *
+ * The ONNX sessions fastembed builds here from the model cache carry the
+ * embedding thread budget ({@link withEmbedSessionBudget}); it is off again by
+ * the time this returns, and never touches a session built from anywhere else,
+ * so a consumer sharing the same onnxruntime copy keeps its own defaults.
  */
 export declare function defaultLocalModelInitializer(options: LocalModelInitOptions): Promise<LocalEmbeddingModel>;
 export declare function embeddingsDisabled(): boolean;
@@ -51,6 +56,8 @@ export declare function embeddingsDisabled(): boolean;
 export declare function currentEmbeddingModel(): string;
 export declare function isApiModel(modelName: string): boolean;
 export declare function embeddingDimFor(modelName: string): number;
+/** The bundled fastembed model id for a model name, or `null` when fastembed has no local model for it. */
+export declare function fastembedModelName(modelName: string): StandardEmbeddingModel | null;
 export declare function setEmbeddingProviderForTests(provider: EmbeddingProvider | null | undefined): void;
 export declare const setEmbeddingProvider: typeof setEmbeddingProviderForTests;
 export declare function setLocalModelInitializerForTests(initializer: LocalModelInitializer | null | undefined): void;

package/dist/types/core/fastembed-runtime.d.ts

@@ -1,6 +1,9 @@
 import { type RuntimeInstallSpec } from "@oh-my-pi/pi-utils";
 import type * as Fastembed from "fastembed";
+import type * as Ort from "onnxruntime-node";
 type FastembedModule = typeof Fastembed;
+type OrtModule = typeof Ort;
+type OrtSessionOptions = Ort.InferenceSession.SessionOptions;
 /** Runtime install inputs for the optional fastembed embedding stack. */
 export interface FastembedRuntimeInstallPlan {
     /** Cache directory key; changes when runtime resolution policy changes. */
@@ -37,5 +40,35 @@ export interface WindowsFastembedRuntime {
  * unrelated `onnxruntime.dll` from the inherited system path.
  */
 export declare function prepareWindowsFastembedRuntime({ fastembedEntry, fastembedPackageDir, arch, env, }: WindowsFastembedRuntimeOptions): Promise<WindowsFastembedRuntime>;
+/**
+ * Session options for one embedding session: a bounded intra-op pool, a single
+ * inter-op thread, and no thread spinning. Every value the caller set wins,
+ * including inside `extra.session`, so a consumer that tunes onnxruntime itself
+ * keeps its own numbers.
+ */
+export declare function embedSessionOptions(base: OrtSessionOptions, parallelism: number): OrtSessionOptions;
+/**
+ * Apply {@link embedSessionOptions} to sessions created through this ORT module
+ * from a model file under `cacheRoot`, until the returned function releases the
+ * scope. fastembed's `FlagEmbedding.init` passes only `executionProviders` and
+ * `graphOptimizationLevel` and exposes no way to reach the session options, so
+ * the module its dependency graph loads is the only seam. That module is
+ * onnxruntime-common's exported class, which an unrelated consumer shares when
+ * ORT is hoisted beside fastembed rather than nested inside it, so the budget
+ * is confined twice: only to model files under `cacheRoot`, and only while a
+ * scope is open. One wrapper serves every open scope (the newest open scope
+ * whose cache holds the model applies); releasing the last scope, in any
+ * order, puts the original `create` back unless something else has since
+ * wrapped it.
+ */
+export declare function applyEmbedSessionBudget(ort: OrtModule, parallelism: number, cacheRoot: string): () => void;
+/**
+ * Run `work` with the embedding thread budget on the ORT module paired with the
+ * loaded fastembed, for sessions built from model files under `cacheRoot`, and
+ * take it off again afterwards. Budgeting is best-effort: an ORT that cannot be
+ * found or loaded leaves the defaults in place rather than failing the
+ * embedding that asked for it. `loadOrt` is the test seam.
+ */
+export declare function withEmbedSessionBudget<T>(cacheRoot: string, work: () => Promise<T>, loadOrt?: () => OrtModule): Promise<T>;
 export declare function loadFastembed(): Promise<FastembedModule>;
 export {};

package/package.json

@@ -1,7 +1,7 @@
 {
 	"type": "module",
 	"name": "@oh-my-pi/pi-mnemopi",
-	"version": "18.13.1",
+	"version": "18.13.2",
 	"description": "Local SQLite memory engine for omp agents",
 	"homepage": "https://omp.sh",
 	"author": {
@@ -42,10 +42,10 @@
 		"fmt": "oxfmt --no-error-on-unmatched-pattern 'src/**/*.{ts,tsx}' '{test,bench,examples,scripts}/**/*.ts' '*.ts'"
 	},
 	"dependencies": {
-		"@oh-my-pi/pi-ai": "18.13.1",
-		"@oh-my-pi/pi-catalog": "18.13.1",
-		"@oh-my-pi/pi-natives": "18.13.1",
-		"@oh-my-pi/pi-utils": "18.13.1"
+		"@oh-my-pi/pi-ai": "18.13.2",
+		"@oh-my-pi/pi-catalog": "18.13.2",
+		"@oh-my-pi/pi-natives": "18.13.2",
+		"@oh-my-pi/pi-utils": "18.13.2"
 	},
 	"peerDependencies": {
 		"fastembed": "3.0.0",

package/src/core/embeddings.ts

@@ -15,7 +15,7 @@ import {
 } from "@oh-my-pi/pi-utils";
 import { LRUCache } from "@oh-my-pi/pi-utils/lru";
 import type { EmbeddingModel } from "fastembed";
-import { loadFastembed } from "./fastembed-runtime";
+import { loadFastembed, withEmbedSessionBudget } from "./fastembed-runtime";
 import {
 	type EmbeddingOutput,
 	getMnemopiRuntimeOptions,
@@ -99,20 +99,27 @@ export async function quarantineCorruptModelFile(message: string, cacheDir?: str
  * retries once; fastembed re-downloads any file missing from the model's
  * cache directory. Also the initializer the embed worker uses in its
  * subprocess; the in-process seam stays {@link setLocalModelInitializer}.
+ *
+ * The ONNX sessions fastembed builds here from the model cache carry the
+ * embedding thread budget ({@link withEmbedSessionBudget}); it is off again by
+ * the time this returns, and never touches a session built from anywhere else,
+ * so a consumer sharing the same onnxruntime copy keeps its own defaults.
  */
 export async function defaultLocalModelInitializer(options: LocalModelInitOptions): Promise<LocalEmbeddingModel> {
 	const cacheDir = options.cacheDir ?? getFastembedCacheDir();
 	const initOptions = options.cacheDir === undefined ? { ...options, cacheDir } : options;
 	const { FlagEmbedding } = await loadFastembed();
-	try {
-		return await FlagEmbedding.init(initOptions);
-	} catch (error) {
-		const message = error instanceof Error ? error.message : "";
-		if (/Protobuf parsing failed/i.test(message) && (await quarantineCorruptModelFile(message, cacheDir))) {
-			return FlagEmbedding.init(initOptions);
+	return withEmbedSessionBudget(cacheDir, async () => {
+		try {
+			return await FlagEmbedding.init(initOptions);
+		} catch (error) {
+			const message = error instanceof Error ? error.message : "";
+			if (/Protobuf parsing failed/i.test(message) && (await quarantineCorruptModelFile(message, cacheDir))) {
+				return FlagEmbedding.init(initOptions);
+			}
+			throw error;
 		}
-		throw error;
-	}
+	});
 }
 
 function activeEmbeddingOptions() {
@@ -313,7 +320,8 @@ const KNOWN_MODEL_NAMES: Record<string, string> = {
 	"intfloat/multilingual-e5-large": "fast-multilingual-e5-large",
 	"sentence-transformers/all-MiniLM-L6-v2": "fast-all-MiniLM-L6-v2",
 };
-function fastembedModelName(modelName: string): StandardEmbeddingModel | null {
+/** The bundled fastembed model id for a model name, or `null` when fastembed has no local model for it. */
+export function fastembedModelName(modelName: string): StandardEmbeddingModel | null {
 	// Fastembed `EmbeddingModel` enum string values, inlined so resolving a model name
 	// (and `available()`) never imports `fastembed` — its module eagerly loads the
 	// `onnxruntime-node` native addon, which segfaults in some runtimes.

package/src/core/fastembed-runtime.ts

@@ -1,4 +1,5 @@
 import { createRequire } from "node:module";
+import * as os from "node:os";
 import * as path from "node:path";
 import {
 	ensureRuntimeInstalled,
@@ -9,9 +10,12 @@ import {
 	resolveRuntimeModule,
 } from "@oh-my-pi/pi-utils";
 import type * as Fastembed from "fastembed";
+import type * as Ort from "onnxruntime-node";
 import packageManifest from "../../package.json" with { type: "json" };
 
 type FastembedModule = typeof Fastembed;
+type OrtModule = typeof Ort;
+type OrtSessionOptions = Ort.InferenceSession.SessionOptions;
 
 /** Runtime install inputs for the optional fastembed embedding stack. */
 export interface FastembedRuntimeInstallPlan {
@@ -58,6 +62,8 @@ export function fastembedRuntimeInstallPlan(): FastembedRuntimeInstallPlan {
 	};
 }
 let fastembedLoad: Promise<FastembedModule> | null = null;
+/** The fastembed installation {@link loadFastembed} resolved, and the `require` that sees its dependency graph. */
+let loadedFastembed: { entry: string; fastembedPackageDir: string; require: NodeRequire } | null = null;
 
 /** Inputs for selecting the Windows DLL directory paired with a fastembed installation. */
 export interface WindowsFastembedRuntimeOptions {
@@ -81,6 +87,18 @@ export interface WindowsFastembedRuntime {
 	dllDir: string;
 }
 
+/** The `onnxruntime-node` installation fastembed's own dependency graph loads: nested first, then the hoisted copy beside it. */
+function resolveFastembedOrt(fastembedPackageDir: string): { ortEntry: string | null; ortPackageDir: string } {
+	const nestedNodeModules = path.join(fastembedPackageDir, "node_modules");
+	const rootNodeModules = path.dirname(fastembedPackageDir);
+	const nestedOrtEntry = resolveRuntimeModule(nestedNodeModules, "onnxruntime-node");
+	const ortEntry = nestedOrtEntry ?? resolveRuntimeModule(rootNodeModules, "onnxruntime-node");
+	return {
+		ortEntry,
+		ortPackageDir: path.join(nestedOrtEntry ? nestedNodeModules : rootNodeModules, "onnxruntime-node"),
+	};
+}
+
 /**
  * Prepend the ORT DLL directory paired with fastembed before Bun loads its
  * native binding. Compiled Windows binaries extract `.node` files to a
@@ -93,11 +111,7 @@ export async function prepareWindowsFastembedRuntime({
 	arch = process.arch,
 	env = process.env,
 }: WindowsFastembedRuntimeOptions): Promise<WindowsFastembedRuntime> {
-	const nestedNodeModules = path.join(fastembedPackageDir, "node_modules");
-	const rootNodeModules = path.dirname(fastembedPackageDir);
-	const nestedOrtEntry = resolveRuntimeModule(nestedNodeModules, "onnxruntime-node");
-	const ortEntry = nestedOrtEntry ?? resolveRuntimeModule(rootNodeModules, "onnxruntime-node");
-	const ortPackageDir = path.join(nestedOrtEntry ? nestedNodeModules : rootNodeModules, "onnxruntime-node");
+	const { ortEntry, ortPackageDir } = resolveFastembedOrt(fastembedPackageDir);
 	if (!ortEntry) {
 		throw new Error(`Cannot find module onnxruntime-node beside ${fastembedEntry}`);
 	}
@@ -120,6 +134,172 @@ export async function prepareWindowsFastembedRuntime({
 	return { ortEntry, ortPackageDir, dllDir };
 }
 
+/**
+ * Upper bound on the intra-op thread pool of an embedding session.
+ *
+ * fastembed creates its ONNX session without any thread option, so
+ * onnxruntime applies its own defaults: one intra-op thread per physical core,
+ * with those threads spinning between operators. The real work of embedding a
+ * recall query is about half a CPU-second (fastembed pads every text to the
+ * model's maximum length), so on a many-core host most of what the default
+ * pool burns beyond that is spin and scheduling. Measured on a 32-core host,
+ * one short query cost 3.85 CPU-seconds and kept 31 cores busy for 124 ms;
+ * with the budget below it costs 0.56 CPU-seconds and finishes in 178 ms, and its
+ * vector points the same way (cosine 0.9999997 against the unbudgeted one;
+ * `bench/embed-session-threads.bench.ts` measures and checks both).
+ *
+ * The cap is a trade, not a free win: on an otherwise idle many-core host a
+ * single embedding finishes later (124 ms to 178 ms above, about 40% slower)
+ * for about 7x less CPU. It is fixed because the default pool's cost
+ * multiplies with every process embedding at the same time, while the extra
+ * latency is tens of milliseconds per retain or recall.
+ */
+const EMBED_INTRA_OP_THREAD_CAP = 4;
+
+/** Narrow an `extra` sub-object for merging; anything else is replaced wholesale. */
+function extraSection(value: unknown): Record<string, unknown> | undefined {
+	return typeof value === "object" && value !== null ? (value as Record<string, unknown>) : undefined;
+}
+
+/**
+ * Session options for one embedding session: a bounded intra-op pool, a single
+ * inter-op thread, and no thread spinning. Every value the caller set wins,
+ * including inside `extra.session`, so a consumer that tunes onnxruntime itself
+ * keeps its own numbers.
+ */
+export function embedSessionOptions(base: OrtSessionOptions, parallelism: number): OrtSessionOptions {
+	const extra = extraSection(base.extra);
+	const session = extraSection(extra?.session);
+	return {
+		...base,
+		intraOpNumThreads: base.intraOpNumThreads ?? Math.max(1, Math.min(EMBED_INTRA_OP_THREAD_CAP, parallelism)),
+		interOpNumThreads: base.interOpNumThreads ?? 1,
+		extra: {
+			...extra,
+			session: {
+				...session,
+				intra_op: { allow_spinning: "0", ...extraSection(session?.intra_op) },
+				inter_op: { allow_spinning: "0", ...extraSection(session?.inter_op) },
+			},
+		},
+	};
+}
+
+/** One open budget: the host's core count it sizes for, and the model directory its sessions load from. */
+interface EmbedBudgetScope {
+	parallelism: number;
+	cacheRoot: string;
+}
+
+/** The single budgeting wrapper installed on an ORT session factory, and the scopes it currently serves. */
+interface BudgetedFactory {
+	original: typeof Ort.InferenceSession.create;
+	wrapper: typeof Ort.InferenceSession.create;
+	scopes: EmbedBudgetScope[];
+}
+
+const budgetedFactories = new WeakMap<object, BudgetedFactory>();
+
+/** Whether `model` is a model file under `cacheRoot`, the only sessions an embedding budget may touch. */
+function isEmbeddingModelPath(model: unknown, cacheRoot: string): boolean {
+	return typeof model === "string" && path.resolve(model).startsWith(path.resolve(cacheRoot) + path.sep);
+}
+
+/**
+ * Apply {@link embedSessionOptions} to sessions created through this ORT module
+ * from a model file under `cacheRoot`, until the returned function releases the
+ * scope. fastembed's `FlagEmbedding.init` passes only `executionProviders` and
+ * `graphOptimizationLevel` and exposes no way to reach the session options, so
+ * the module its dependency graph loads is the only seam. That module is
+ * onnxruntime-common's exported class, which an unrelated consumer shares when
+ * ORT is hoisted beside fastembed rather than nested inside it, so the budget
+ * is confined twice: only to model files under `cacheRoot`, and only while a
+ * scope is open. One wrapper serves every open scope (the newest open scope
+ * whose cache holds the model applies); releasing the last scope, in any
+ * order, puts the original `create` back unless something else has since
+ * wrapped it.
+ */
+export function applyEmbedSessionBudget(ort: OrtModule, parallelism: number, cacheRoot: string): () => void {
+	const factory = ort.InferenceSession;
+	let budgeted = budgetedFactories.get(factory);
+	if (budgeted === undefined) {
+		const original = factory.create;
+		const create = original.bind(factory);
+		const scopes: EmbedBudgetScope[] = [];
+		const wrapper = ((
+			model: string | ArrayBufferLike | Uint8Array,
+			...rest: unknown[]
+		): Promise<Ort.InferenceSession> => {
+			// A budgeted call is always the `(path, options?)` form: only a path
+			// can name a model file under a scope's cache. Every other call,
+			// including the buffer forms, reaches ORT exactly as it was made.
+			// The newest scope that owns this model wins, so two scopes on
+			// different caches each budget their own sessions.
+			const scope = scopes.findLast(open => isEmbeddingModelPath(model, open.cacheRoot));
+			if (scope === undefined) {
+				return create(model as string, ...(rest as [OrtSessionOptions]));
+			}
+			return create(
+				model as string,
+				embedSessionOptions((rest[0] as OrtSessionOptions | undefined) ?? {}, scope.parallelism),
+			);
+		}) as typeof factory.create;
+		budgeted = { original, wrapper, scopes };
+		budgetedFactories.set(factory, budgeted);
+		factory.create = wrapper;
+	}
+	const held = budgeted;
+	const scope: EmbedBudgetScope = { parallelism, cacheRoot };
+	held.scopes.push(scope);
+	return () => {
+		const index = held.scopes.indexOf(scope);
+		if (index === -1) return;
+		held.scopes.splice(index, 1);
+		if (held.scopes.length > 0) return;
+		// The last scope is gone: forget the factory so the next scope installs
+		// a fresh wrapper, and restore the original only when ours is still the
+		// installed `create` — something wrapped on top of it owns the chain.
+		budgetedFactories.delete(factory);
+		if (factory.create === held.wrapper) factory.create = held.original;
+	};
+}
+
+/** The ORT module paired with the fastembed {@link loadFastembed} resolved. */
+function pairedFastembedOrt(): OrtModule {
+	const loaded = loadedFastembed;
+	if (loaded === null) throw new Error("fastembed is not loaded");
+	const { ortEntry } = resolveFastembedOrt(loaded.fastembedPackageDir);
+	if (ortEntry === null) throw new Error(`Cannot find module onnxruntime-node beside ${loaded.entry}`);
+	return loaded.require(ortEntry);
+}
+
+/**
+ * Run `work` with the embedding thread budget on the ORT module paired with the
+ * loaded fastembed, for sessions built from model files under `cacheRoot`, and
+ * take it off again afterwards. Budgeting is best-effort: an ORT that cannot be
+ * found or loaded leaves the defaults in place rather than failing the
+ * embedding that asked for it. `loadOrt` is the test seam.
+ */
+export async function withEmbedSessionBudget<T>(
+	cacheRoot: string,
+	work: () => Promise<T>,
+	loadOrt: () => OrtModule = pairedFastembedOrt,
+): Promise<T> {
+	let release: (() => void) | undefined;
+	try {
+		release = applyEmbedSessionBudget(loadOrt(), os.availableParallelism(), cacheRoot);
+	} catch (error) {
+		logger.debug("mnemopi: onnxruntime not budgetable; embedding sessions keep ORT thread defaults", {
+			error: String(error),
+		});
+	}
+	try {
+		return await work();
+	} finally {
+		release?.();
+	}
+}
+
 export function loadFastembed(): Promise<FastembedModule> {
 	fastembedLoad ??= loadFastembedOnce().catch(error => {
 		fastembedLoad = null;
@@ -153,6 +333,8 @@ async function loadResolvedFastembed(entry: string, fastembedPackageDir: string)
 		requireFastembed(ortEntry);
 	}
 	const loaded: FastembedModule = requireFastembed(entry);
+	// Where the ORT paired with this fastembed lives, for `withEmbedSessionBudget`.
+	loadedFastembed = { entry, fastembedPackageDir, require: requireFastembed };
 	return loaded;
 }