soulcraft/brainy
2
0
Fork 0
Code Issues 1 Pull requests Projects Releases 335 Packages Wiki Activity Actions

refactor(worker-threads): pass functions as strings to resolve Promise cloning issues

Browse source 
- Updated worker execution logic to pass functions as stringified versions instead of reference types to avoid cloning issues in worker threads.
- Refactored utility functions in `workerUtils` to accept serialized function strings.
- Adjusted embedding and augmentation pipeline workflows to align with the updated worker handling mechanism.
- Enhanced parallel distance calculator to support structured arguments with stringified distance functions for better compatibility.
This commit is contained in:
David Snelling 2025-06-24 10:38:05 -07:00
parent 887a08b689
commit d18a5a6c22
5 changed files with 92 additions and 73 deletions
Download patch file Download diff file Expand all files Collapse all files

2
src/augmentationPipeline.ts
View file

@ -601,7 +601,7 @@ export class AugmentationPipeline {
return augFn.apply(augmentation, workerArgs);
};
methodPromise = executeInThread<AugmentationResponse<R>>(workerFn, ...args);
methodPromise = executeInThread<AugmentationResponse<R>>(workerFn.toString(), args);
} catch (threadError) {
console.warn(`Failed to execute in thread, falling back to main thread: ${threadError}`);
// Fall back to executing in the main thread

22
src/hnsw/hnswIndex.ts
View file

@ -69,17 +69,21 @@ export class HNSWIndex {
// Function to be executed in a worker thread
const distanceCalculator = (
query: Vector,
items: Array<{ id: string; vector: Vector }>,
distanceFn: string
args: {
queryVector: Vector,
vectors: Array<{ id: string; vector: Vector }>,
distanceFnString: string
}
) => {
const { queryVector, vectors, distanceFnString } = args;
// Recreate the distance function from its string representation
const distanceFunction = new Function('return ' + distanceFn)() as DistanceFunction
const distanceFunction = new Function('return ' + distanceFnString)() as DistanceFunction
// Calculate distances for all items
return items.map(item => ({
return vectors.map(item => ({
id: item.id,
distance: distanceFunction(query, item.vector)
distance: distanceFunction(queryVector, item.vector)
}))
}
@ -89,10 +93,8 @@ export class HNSWIndex {
// Execute the distance calculation in a separate thread
return await executeInThread<Array<{ id: string; distance: number }>>(
distanceCalculator,
queryVector,
vectors,
distanceFnString
distanceCalculator.toString(),
{ queryVector, vectors, distanceFnString }
)
} catch (error) {
console.error('Error in parallel distance calculation, falling back to sequential:', error)

2
src/pipeline.ts
View file

@ -418,7 +418,7 @@ export class Pipeline implements IPipeline {
return augFn.apply(augmentation, workerArgs)
}
methodPromise = executeInThread<AugmentationResponse<T>>(workerFn, ...args)
methodPromise = executeInThread<AugmentationResponse<T>>(workerFn.toString(), args)
} catch (threadError) {
console.warn(`Failed to execute in thread, falling back to main thread: ${threadError}`)
// Fall back to executing in the main thread

27
src/utils/embedding.ts
View file

@ -214,7 +214,15 @@ export function createThreadedEmbeddingFunction(
try {
// Function to be executed in a worker thread
const embedInWorker = async (inputData: any) => {
// This must be a regular function (not async) to avoid Promise cloning issues
const embedInWorker = (inputData: any) => {
// Return a plain object with the input data
// All async operations will be performed inside the worker
return { data: inputData }
}
// Worker implementation function that will be stringified and run in the worker
const workerImplementation = async ({ data }: { data: any }) => {
// We need to dynamically import TensorFlow.js and USE in the worker
const tf = await import('@tensorflow/tfjs')
const use = await import('@tensorflow-models/universal-sentence-encoder')
@ -224,19 +232,19 @@ export function createThreadedEmbeddingFunction(
// Handle different input types
let textToEmbed: string[]
if (typeof inputData === 'string') {
if (inputData.trim() === '') {
if (typeof data === 'string') {
if (data.trim() === '') {
return new Array(512).fill(0)
}
textToEmbed = [inputData]
textToEmbed = [data]
} else if (
Array.isArray(inputData) &&
inputData.every((item) => typeof item === 'string')
Array.isArray(data) &&
data.every((item) => typeof item === 'string')
) {
if (inputData.length === 0 || inputData.every((item) => item.trim() === '')) {
if (data.length === 0 || data.every((item) => item.trim() === '')) {
return new Array(512).fill(0)
}
textToEmbed = inputData.filter((item) => item.trim() !== '')
textToEmbed = data.filter((item) => item.trim() !== '')
if (textToEmbed.length === 0) {
return new Array(512).fill(0)
}
@ -259,7 +267,8 @@ export function createThreadedEmbeddingFunction(
}
// Execute the embedding function in a separate thread
return await executeInThread<Vector>(embedInWorker, data)
// Pass the worker implementation as a string to avoid Promise cloning issues
return await executeInThread<Vector>(workerImplementation.toString(), embedInWorker(data))
} catch (error) {
// If threading fails and fallback is enabled, use the standard embedding function
if (options.fallbackToMain) {

112
src/utils/workerUtils.ts
View file

@ -2,24 +2,23 @@
* Utility functions for working with Web Workers and Worker Threads
*/
import { isNode, isBrowser } from './environment.js';
import { isNode, isBrowser } from './environment.js'
/**
* Execute a function in a Web Worker (browser environment)
*
* @param fn The function to execute
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInWebWorker<T>(fn: Function, ...args: any[]): Promise<T> {
export function executeInWebWorker<T>(fnString: string, args: any): Promise<T> {
return new Promise((resolve, reject) => {
// Create a blob URL for the worker script
const fnString = fn.toString();
const workerScript = `
self.onmessage = function(e) {
self.onmessage = async function(e) {
try {
const fn = ${fnString};
const result = fn(...e.data);
const result = await fn(e.data);
self.postMessage({ success: true, data: result });
} catch (error) {
self.postMessage({
@ -28,59 +27,63 @@ export function executeInWebWorker<T>(fn: Function, ...args: any[]): Promise<T>
});
}
};
`;
`
const blob = new Blob([workerScript], { type: 'application/javascript' });
const blobURL = URL.createObjectURL(blob);
const blob = new Blob([workerScript], { type: 'application/javascript' })
const blobURL = URL.createObjectURL(blob)
// Create a worker
const worker = new Worker(blobURL);
const worker = new Worker(blobURL)
// Set up message handling
worker.onmessage = function(e: MessageEvent<{ success: boolean; data: T; error?: string }>) {
URL.revokeObjectURL(blobURL); // Clean up
worker.terminate(); // Terminate the worker
worker.onmessage = function (
e: MessageEvent<{ success: boolean; data: T; error?: string }>
) {
URL.revokeObjectURL(blobURL) // Clean up
worker.terminate() // Terminate the worker
if (e.data.success) {
resolve(e.data.data);
resolve(e.data.data)
} else {
reject(new Error(e.data.error));
reject(new Error(e.data.error))
}
};
}
worker.onerror = function(error: ErrorEvent) {
URL.revokeObjectURL(blobURL); // Clean up
worker.terminate(); // Terminate the worker
reject(error);
};
worker.onerror = function (error: ErrorEvent) {
URL.revokeObjectURL(blobURL) // Clean up
worker.terminate() // Terminate the worker
reject(error)
}
// Start the worker
worker.postMessage(args);
});
worker.postMessage(args)
})
}
/**
* Execute a function in a Worker Thread (Node.js environment)
*
* @param fn The function to execute
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInWorkerThread<T>(fn: Function, ...args: any[]): Promise<T> {
export function executeInWorkerThread<T>(
fnString: string,
args: any
): Promise<T> {
return new Promise((resolve, reject) => {
try {
// Dynamic import to avoid errors in browser environments
const { Worker } = require('worker_threads');
const { Worker } = require('worker_threads')
// Create a worker script
const fnString = fn.toString();
const workerScript = `
const { parentPort } = require('worker_threads');
parentPort.once('message', (data) => {
parentPort.once('message', async (data) => {
try {
const fn = ${fnString};
const result = fn(...data);
const result = await fn(data);
parentPort.postMessage({ success: true, data: result });
} catch (error) {
parentPort.postMessage({
@ -89,49 +92,54 @@ export function executeInWorkerThread<T>(fn: Function, ...args: any[]): Promise<
});
}
});
`;
`
// Create a worker
const worker = new Worker(workerScript, { eval: true });
const worker = new Worker(workerScript, { eval: true })
// Set up message handling
worker.on('message', (data: { success: boolean; data: T; error?: string }) => {
worker.terminate(); // Terminate the worker
worker.on(
'message',
(data: { success: boolean; data: T; error?: string }) => {
worker.terminate() // Terminate the worker
if (data.success) {
resolve(data.data);
} else {
reject(new Error(data.error));
if (data.success) {
resolve(data.data)
} else {
reject(new Error(data.error))
}
}
});
)
worker.on('error', (error: Error) => {
worker.terminate(); // Terminate the worker
reject(error);
});
worker.terminate() // Terminate the worker
reject(error)
})
// Start the worker
worker.postMessage(args);
worker.postMessage(args)
} catch (error) {
reject(error);
reject(error)
}
});
})
}
/**
* Execute a function in a separate thread based on the environment
*
* @param fn The function to execute
*
* @param fnString The function to execute as a string
* @param args The arguments to pass to the function
* @returns A promise that resolves with the result of the function
*/
export function executeInThread<T>(fn: Function, ...args: any[]): Promise<T> {
export function executeInThread<T>(fnString: string, args: any): Promise<T> {
if (isBrowser()) {
return executeInWebWorker<T>(fn, ...args);
return executeInWebWorker<T>(fnString, args)
} else if (isNode()) {
return executeInWorkerThread<T>(fn, ...args);
return executeInWorkerThread<T>(fnString, args)
} else {
// Fall back to executing in the main thread
return Promise.resolve(fn(...args) as T);
// Parse the function from string and execute it
const fn = new Function('return ' + fnString)()
return Promise.resolve(fn(args) as T)
}
}