From d18a5a6c22716e8cd3d0b9fca9fbbc1e037b2897 Mon Sep 17 00:00:00 2001 From: David Snelling Date: Tue, 24 Jun 2025 10:38:05 -0700 Subject: [PATCH] refactor(worker-threads): pass functions as strings to resolve Promise cloning issues - 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. --- src/augmentationPipeline.ts | 2 +- src/hnsw/hnswIndex.ts | 22 +++---- src/pipeline.ts | 2 +- src/utils/embedding.ts | 27 ++++++--- src/utils/workerUtils.ts | 112 +++++++++++++++++++----------------- 5 files changed, 92 insertions(+), 73 deletions(-) diff --git a/src/augmentationPipeline.ts b/src/augmentationPipeline.ts index bf46f1d5..f3c5f640 100644 --- a/src/augmentationPipeline.ts +++ b/src/augmentationPipeline.ts @@ -601,7 +601,7 @@ export class AugmentationPipeline { return augFn.apply(augmentation, workerArgs); }; - methodPromise = executeInThread>(workerFn, ...args); + methodPromise = executeInThread>(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 diff --git a/src/hnsw/hnswIndex.ts b/src/hnsw/hnswIndex.ts index f489fcf9..ed8adf3f 100644 --- a/src/hnsw/hnswIndex.ts +++ b/src/hnsw/hnswIndex.ts @@ -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>( - distanceCalculator, - queryVector, - vectors, - distanceFnString + distanceCalculator.toString(), + { queryVector, vectors, distanceFnString } ) } catch (error) { console.error('Error in parallel distance calculation, falling back to sequential:', error) diff --git a/src/pipeline.ts b/src/pipeline.ts index e27441c8..36580693 100644 --- a/src/pipeline.ts +++ b/src/pipeline.ts @@ -418,7 +418,7 @@ export class Pipeline implements IPipeline { return augFn.apply(augmentation, workerArgs) } - methodPromise = executeInThread>(workerFn, ...args) + methodPromise = executeInThread>(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 diff --git a/src/utils/embedding.ts b/src/utils/embedding.ts index 6e15e204..495a49fb 100644 --- a/src/utils/embedding.ts +++ b/src/utils/embedding.ts @@ -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(embedInWorker, data) + // Pass the worker implementation as a string to avoid Promise cloning issues + return await executeInThread(workerImplementation.toString(), embedInWorker(data)) } catch (error) { // If threading fails and fallback is enabled, use the standard embedding function if (options.fallbackToMain) { diff --git a/src/utils/workerUtils.ts b/src/utils/workerUtils.ts index 8e500ee3..263d4d30 100644 --- a/src/utils/workerUtils.ts +++ b/src/utils/workerUtils.ts @@ -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(fn: Function, ...args: any[]): Promise { +export function executeInWebWorker(fnString: string, args: any): Promise { 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(fn: Function, ...args: any[]): Promise }); } }; - `; + ` - 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(fn: Function, ...args: any[]): Promise { +export function executeInWorkerThread( + fnString: string, + args: any +): Promise { 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(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(fn: Function, ...args: any[]): Promise { +export function executeInThread(fnString: string, args: any): Promise { if (isBrowser()) { - return executeInWebWorker(fn, ...args); + return executeInWebWorker(fnString, args) } else if (isNode()) { - return executeInWorkerThread(fn, ...args); + return executeInWorkerThread(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) } }