558 lines
16 KiB
TypeScript
558 lines
16 KiB
TypeScript
import { randomUUID } from 'node:crypto'
|
|
import { localInferenceService } from '../local-inference-service'
|
|
import type { InferenceResources } from '../../shared/local-inference-contracts'
|
|
import {
|
|
EMBEDDING_INFERENCE_MAX_IN_FLIGHT,
|
|
embeddingInferenceProtocolBase,
|
|
parseEmbeddingInferenceResponse,
|
|
validateEmbeddingInferenceTexts,
|
|
validateEmbeddingInferenceVectors,
|
|
type EmbeddingInferenceRequest,
|
|
type EmbeddingInferenceRole
|
|
} from './embedding-inference-contracts'
|
|
|
|
const DEFAULT_TIMEOUT_MS = 120_000
|
|
const MIN_TIMEOUT_MS = 100
|
|
const MAX_TIMEOUT_MS = 300_000
|
|
const DEFAULT_MAX_RESTARTS = 2
|
|
const MAX_RESTARTS = 10
|
|
const SHUTDOWN_GRACE_MS = 1_000
|
|
|
|
export interface EmbeddingInferenceTransport {
|
|
getResources?(): InferenceResources
|
|
postMessage(message: EmbeddingInferenceRequest): void
|
|
onMessage(listener: (message: unknown) => void): () => void
|
|
onClose(listener: () => void): () => void
|
|
close(): void | Promise<void>
|
|
}
|
|
|
|
export type EmbeddingInferenceTransportFactory = () =>
|
|
| EmbeddingInferenceTransport
|
|
| Promise<EmbeddingInferenceTransport>
|
|
|
|
export interface EmbeddingInferenceBrokerOptions {
|
|
readonly createTransport: EmbeddingInferenceTransportFactory
|
|
readonly timeoutMs?: number
|
|
/**
|
|
* Maximum replacement processes over this broker's lifetime. The initial
|
|
* process is not a restart and failed requests are never replayed.
|
|
*/
|
|
readonly maxRestarts?: number
|
|
}
|
|
|
|
type PendingRequest = {
|
|
readonly expectedCount: number
|
|
readonly resolve: (vectors: number[][]) => void
|
|
readonly reject: (error: Error) => void
|
|
readonly rejectCaller: (error: Error) => void
|
|
callerError?: Error
|
|
readonly timeout: ReturnType<typeof setTimeout>
|
|
removeAbortListener?: () => void
|
|
}
|
|
|
|
type ActiveTransport = {
|
|
readonly transport: EmbeddingInferenceTransport
|
|
readonly removeMessageListener: () => void
|
|
readonly removeCloseListener: () => void
|
|
}
|
|
|
|
export class EmbeddingInferenceBrokerError extends Error {
|
|
constructor(
|
|
readonly code:
|
|
| 'CLOSED'
|
|
| 'PROCESS_CLOSED'
|
|
| 'PROTOCOL_VIOLATION'
|
|
| 'RESTART_LIMIT'
|
|
| 'START_FAILURE'
|
|
| 'ENGINE_FAILURE'
|
|
| 'INVALID_RESULT'
|
|
| 'OVERLOADED',
|
|
detail?: string
|
|
) {
|
|
super(`Embedding inference failed (${code})${detail ? `: ${detail}` : ''}`)
|
|
this.name = 'EmbeddingInferenceBrokerError'
|
|
}
|
|
}
|
|
|
|
function boundedInteger(
|
|
value: number,
|
|
field: string,
|
|
minimum: number,
|
|
maximum: number
|
|
): number {
|
|
if (!Number.isSafeInteger(value) || value < minimum || value > maximum) {
|
|
throw new RangeError(
|
|
`${field} must be an integer between ${minimum} and ${maximum}`
|
|
)
|
|
}
|
|
return value
|
|
}
|
|
|
|
function abortError(): Error {
|
|
const error = new Error('Embedding inference was cancelled')
|
|
error.name = 'AbortError'
|
|
return error
|
|
}
|
|
|
|
function timeoutError(): Error {
|
|
const error = new Error('Embedding inference timed out')
|
|
error.name = 'TimeoutError'
|
|
return error
|
|
}
|
|
|
|
export class EmbeddingInferenceBroker {
|
|
private readonly createTransport: EmbeddingInferenceTransportFactory
|
|
private readonly timeoutMs: number
|
|
private readonly maxRestarts: number
|
|
private readonly pending = new Map<string, PendingRequest>()
|
|
private active?: ActiveTransport
|
|
private starting?: Promise<ActiveTransport>
|
|
private restarts = 0
|
|
private everStarted = false
|
|
private closed = false
|
|
private shutdownComplete?: () => void
|
|
private paused = false
|
|
private stopping = false
|
|
private managementError?: string
|
|
|
|
getManagementState(): 'stopped' | 'stopping' | 'starting' | 'running' | 'idle' | 'error' {
|
|
if (this.stopping) return 'stopping'
|
|
if (this.paused || this.closed) return 'stopped'
|
|
if (this.starting) return 'starting'
|
|
if (this.active) return 'running'
|
|
return this.managementError ? 'error' : 'idle'
|
|
}
|
|
|
|
getManagementError(): string | undefined { return this.managementError }
|
|
|
|
getResources(): InferenceResources {
|
|
return this.active?.transport.getResources?.() ?? {
|
|
scope: 'unavailable', reason: this.active ? '当前执行端未提供进程指标' : '服务进程未运行'
|
|
}
|
|
}
|
|
|
|
async start(): Promise<void> {
|
|
if (this.closed) throw new EmbeddingInferenceBrokerError('CLOSED')
|
|
this.paused = false
|
|
this.restarts = 0
|
|
try {
|
|
await this.ensureTransport()
|
|
this.managementError = undefined
|
|
} catch (error) {
|
|
this.managementError = String(error)
|
|
throw error
|
|
}
|
|
}
|
|
|
|
async stop(): Promise<void> {
|
|
this.paused = true
|
|
this.stopping = true
|
|
const error = new Error('向量执行服务已由用户停止；任务不会自动重放')
|
|
this.rejectCallers(error)
|
|
try {
|
|
if (this.starting) await this.starting.catch(() => undefined)
|
|
const active = this.active
|
|
if (!active) return
|
|
await active.transport.close()
|
|
this.disposeTransport(active.transport, false)
|
|
this.rejectAll(error)
|
|
} catch (error) {
|
|
this.paused = false
|
|
this.managementError = error instanceof Error ? error.message : String(error)
|
|
throw error
|
|
} finally {
|
|
this.stopping = false
|
|
}
|
|
}
|
|
|
|
constructor(options: EmbeddingInferenceBrokerOptions) {
|
|
if (typeof options.createTransport !== 'function') {
|
|
throw new TypeError('createTransport must be a function')
|
|
}
|
|
this.createTransport = options.createTransport
|
|
this.timeoutMs = boundedInteger(
|
|
options.timeoutMs ?? DEFAULT_TIMEOUT_MS,
|
|
'timeoutMs',
|
|
MIN_TIMEOUT_MS,
|
|
MAX_TIMEOUT_MS
|
|
)
|
|
this.maxRestarts = boundedInteger(
|
|
options.maxRestarts ?? DEFAULT_MAX_RESTARTS,
|
|
'maxRestarts',
|
|
0,
|
|
MAX_RESTARTS
|
|
)
|
|
}
|
|
|
|
async embed(
|
|
texts: readonly string[],
|
|
role: EmbeddingInferenceRole,
|
|
signal?: AbortSignal
|
|
): Promise<number[][]> {
|
|
const controller = new AbortController()
|
|
return new Promise<number[][]>((resolve, reject) => {
|
|
// Track worker execution separately from prompt caller cancellation/timeout.
|
|
void localInferenceService.run('embedding', role === 'query' ? '向量请求 · 查询' : '向量请求 · 文档',
|
|
(markCancelling) => this.embedRequest(texts, role, signal, controller.signal, (error) => {
|
|
markCancelling()
|
|
reject(error)
|
|
}),
|
|
{ cancel: () => controller.abort() }).then(resolve, reject)
|
|
})
|
|
}
|
|
|
|
private async embedRequest(
|
|
texts: readonly string[], role: EmbeddingInferenceRole, signal: AbortSignal | undefined,
|
|
managementSignal: AbortSignal,
|
|
rejectCaller: (error: Error) => void
|
|
): Promise<number[][]> {
|
|
const normalizedTexts = validateEmbeddingInferenceTexts(texts)
|
|
if (role !== 'query' && role !== 'document') {
|
|
throw new TypeError('role must be query or document')
|
|
}
|
|
if (this.closed || this.paused) {
|
|
throw new EmbeddingInferenceBrokerError('CLOSED')
|
|
}
|
|
if (signal?.aborted) {
|
|
throw abortError()
|
|
}
|
|
if (this.pending.size >= EMBEDDING_INFERENCE_MAX_IN_FLIGHT) {
|
|
throw new EmbeddingInferenceBrokerError('OVERLOADED')
|
|
}
|
|
|
|
const active = await this.ensureTransport()
|
|
if (this.closed || this.paused) {
|
|
throw new EmbeddingInferenceBrokerError('CLOSED')
|
|
}
|
|
if (signal?.aborted) {
|
|
throw abortError()
|
|
}
|
|
if (this.pending.size >= EMBEDDING_INFERENCE_MAX_IN_FLIGHT) {
|
|
throw new EmbeddingInferenceBrokerError('OVERLOADED')
|
|
}
|
|
|
|
const requestId = randomUUID()
|
|
return new Promise<number[][]>((resolve, reject) => {
|
|
const timeout = setTimeout(() => {
|
|
this.cancelPending(requestId, timeoutError())
|
|
}, this.timeoutMs)
|
|
const pending: PendingRequest = {
|
|
expectedCount: normalizedTexts.length,
|
|
resolve,
|
|
reject,
|
|
rejectCaller,
|
|
timeout
|
|
}
|
|
if (signal) {
|
|
const onAbort = (): void => {
|
|
this.cancelPending(requestId, abortError())
|
|
}
|
|
signal.addEventListener('abort', onAbort, { once: true })
|
|
pending.removeAbortListener = () =>
|
|
signal.removeEventListener('abort', onAbort)
|
|
}
|
|
this.pending.set(requestId, pending)
|
|
|
|
const requestCancellation = (): void => {
|
|
// Management cancellation waits for the worker's result or CANCELLED response.
|
|
try {
|
|
active.transport.postMessage({ ...embeddingInferenceProtocolBase(), type: 'cancel', requestId })
|
|
} catch {
|
|
this.failTransport(active.transport, new EmbeddingInferenceBrokerError('PROCESS_CLOSED'))
|
|
}
|
|
}
|
|
const removeSourceAbort = pending.removeAbortListener
|
|
pending.removeAbortListener = () => {
|
|
removeSourceAbort?.()
|
|
managementSignal.removeEventListener('abort', requestCancellation)
|
|
}
|
|
managementSignal.addEventListener('abort', requestCancellation, { once: true })
|
|
|
|
try {
|
|
active.transport.postMessage({
|
|
...embeddingInferenceProtocolBase(),
|
|
type: 'embed',
|
|
requestId,
|
|
role,
|
|
texts: normalizedTexts
|
|
})
|
|
if (managementSignal.aborted) requestCancellation()
|
|
} catch {
|
|
this.failTransport(
|
|
active.transport,
|
|
new EmbeddingInferenceBrokerError('PROCESS_CLOSED')
|
|
)
|
|
}
|
|
})
|
|
}
|
|
|
|
async shutdown(): Promise<void> {
|
|
if (this.closed) {
|
|
return
|
|
}
|
|
this.closed = true
|
|
this.rejectCallers(new EmbeddingInferenceBrokerError('CLOSED'))
|
|
|
|
let active = this.active
|
|
if (!active && this.starting) {
|
|
try {
|
|
active = await this.starting
|
|
} catch {
|
|
return
|
|
}
|
|
}
|
|
if (!active || this.active !== active) {
|
|
return
|
|
}
|
|
|
|
await new Promise<void>((resolve) => {
|
|
let settled = false
|
|
const finish = (): void => {
|
|
if (settled) {
|
|
return
|
|
}
|
|
settled = true
|
|
clearTimeout(timeout)
|
|
this.shutdownComplete = undefined
|
|
this.disposeTransport(active.transport, true)
|
|
resolve()
|
|
}
|
|
const timeout = setTimeout(finish, SHUTDOWN_GRACE_MS)
|
|
this.shutdownComplete = finish
|
|
try {
|
|
active.transport.postMessage({
|
|
...embeddingInferenceProtocolBase(),
|
|
type: 'shutdown'
|
|
})
|
|
} catch {
|
|
finish()
|
|
}
|
|
})
|
|
}
|
|
|
|
private async ensureTransport(): Promise<ActiveTransport> {
|
|
if (this.active) {
|
|
return this.active
|
|
}
|
|
if (this.starting) {
|
|
return this.starting
|
|
}
|
|
if (this.everStarted && this.restarts >= this.maxRestarts) {
|
|
throw new EmbeddingInferenceBrokerError('RESTART_LIMIT')
|
|
}
|
|
if (this.everStarted) {
|
|
this.restarts += 1
|
|
}
|
|
this.everStarted = true
|
|
|
|
const starting = Promise.resolve()
|
|
.then(() => this.createTransport())
|
|
.then((transport) => {
|
|
if (this.closed) {
|
|
void Promise.resolve(transport.close()).catch(() => undefined)
|
|
throw new EmbeddingInferenceBrokerError('CLOSED')
|
|
}
|
|
const active: ActiveTransport = {
|
|
transport,
|
|
removeMessageListener: transport.onMessage((message) =>
|
|
this.handleMessage(transport, message)
|
|
),
|
|
removeCloseListener: transport.onClose(() =>
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROCESS_CLOSED'),
|
|
true
|
|
)
|
|
)
|
|
}
|
|
this.active = active
|
|
return active
|
|
})
|
|
.catch((error: unknown) => {
|
|
this.managementError = error instanceof Error ? error.message : String(error)
|
|
if (error instanceof EmbeddingInferenceBrokerError) {
|
|
throw error
|
|
}
|
|
throw new EmbeddingInferenceBrokerError('START_FAILURE')
|
|
})
|
|
.finally(() => {
|
|
if (this.starting === starting) {
|
|
this.starting = undefined
|
|
}
|
|
})
|
|
this.starting = starting
|
|
return starting
|
|
}
|
|
|
|
private handleMessage(
|
|
transport: EmbeddingInferenceTransport,
|
|
rawMessage: unknown
|
|
): void {
|
|
if (this.active?.transport !== transport) {
|
|
return
|
|
}
|
|
const message = parseEmbeddingInferenceResponse(rawMessage)
|
|
if (!message) {
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
|
|
)
|
|
return
|
|
}
|
|
if (message.type === 'shutdown-complete') {
|
|
if (!this.closed || !this.shutdownComplete) {
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
|
|
)
|
|
return
|
|
}
|
|
this.rejectAll(new EmbeddingInferenceBrokerError('CLOSED'))
|
|
this.shutdownComplete()
|
|
return
|
|
}
|
|
if (message.type === 'fatal') {
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
|
|
)
|
|
return
|
|
}
|
|
|
|
const pending = this.pending.get(message.requestId)
|
|
if (!pending) {
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
|
|
)
|
|
return
|
|
}
|
|
|
|
if (message.type === 'error') {
|
|
this.finishPending(
|
|
message.requestId,
|
|
message.code === 'CANCELLED'
|
|
? abortError()
|
|
: new EmbeddingInferenceBrokerError(message.code, message.message)
|
|
)
|
|
return
|
|
}
|
|
try {
|
|
const vectors = validateEmbeddingInferenceVectors(
|
|
message.vectors,
|
|
pending.expectedCount
|
|
)
|
|
this.finishPending(message.requestId, undefined, vectors)
|
|
} catch {
|
|
this.failTransport(
|
|
transport,
|
|
new EmbeddingInferenceBrokerError('PROTOCOL_VIOLATION')
|
|
)
|
|
}
|
|
}
|
|
|
|
private cancelPending(requestId: string, error: Error): void {
|
|
const pending = this.pending.get(requestId)
|
|
if (!pending || pending.callerError) {
|
|
return
|
|
}
|
|
pending.callerError = error
|
|
clearTimeout(pending.timeout)
|
|
pending.rejectCaller(error)
|
|
try {
|
|
this.active?.transport.postMessage({
|
|
...embeddingInferenceProtocolBase(),
|
|
type: 'cancel',
|
|
requestId
|
|
})
|
|
} catch {
|
|
if (this.active) {
|
|
this.failTransport(
|
|
this.active.transport,
|
|
new EmbeddingInferenceBrokerError('PROCESS_CLOSED')
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
private finishPending(
|
|
requestId: string,
|
|
error?: Error,
|
|
vectors?: number[][]
|
|
): void {
|
|
const pending = this.pending.get(requestId)
|
|
if (!pending) {
|
|
return
|
|
}
|
|
this.pending.delete(requestId)
|
|
clearTimeout(pending.timeout)
|
|
pending.removeAbortListener?.()
|
|
error = pending.callerError ?? error
|
|
if (error) {
|
|
pending.reject(error)
|
|
} else {
|
|
pending.resolve(vectors as number[][])
|
|
}
|
|
}
|
|
|
|
private failTransport(
|
|
transport: EmbeddingInferenceTransport,
|
|
error: Error,
|
|
closed = false
|
|
): void {
|
|
if (this.active?.transport !== transport) {
|
|
return
|
|
}
|
|
this.rejectCallers(error)
|
|
if (closed) {
|
|
this.disposeTransport(transport, false)
|
|
this.rejectAll(error)
|
|
this.shutdownComplete?.()
|
|
return
|
|
}
|
|
this.managementError = error.message
|
|
this.disposeTransport(transport, true)
|
|
}
|
|
|
|
private disposeTransport(
|
|
transport: EmbeddingInferenceTransport,
|
|
close: boolean
|
|
): void {
|
|
const active = this.active
|
|
if (!active || active.transport !== transport) {
|
|
return
|
|
}
|
|
if (close) {
|
|
try {
|
|
void Promise.resolve(transport.close()).then(() => {
|
|
if (this.active !== active) return
|
|
this.disposeTransport(transport, false)
|
|
this.rejectAll(new EmbeddingInferenceBrokerError('PROCESS_CLOSED'))
|
|
}).catch((error: unknown) => {
|
|
this.managementError = error instanceof Error ? error.message : String(error)
|
|
})
|
|
} catch (error) {
|
|
this.managementError = error instanceof Error ? error.message : String(error)
|
|
}
|
|
return
|
|
}
|
|
this.active = undefined
|
|
active.removeMessageListener()
|
|
active.removeCloseListener()
|
|
}
|
|
|
|
private rejectCallers(error: Error): void {
|
|
for (const pending of this.pending.values()) {
|
|
pending.callerError ??= error
|
|
clearTimeout(pending.timeout)
|
|
pending.rejectCaller(pending.callerError)
|
|
}
|
|
}
|
|
|
|
private rejectAll(error: Error): void {
|
|
for (const requestId of [...this.pending.keys()]) {
|
|
this.finishPending(requestId, error)
|
|
}
|
|
}
|
|
}
|