@lincoln504/pi-research 0.1.12 → 1.0.0

package/src/infrastructure/embedding/embedding-client.ts

@@ -0,0 +1,205 @@
+/**
+ * Embedding Client
+ *
+ * HTTP client that implements IEmbedder by forwarding embed/embedMany calls
+ * to a remote EmbeddingServer. Mirrors BrowserClient.
+ */
+
+import * as http from 'node:http';
+import { logger } from '../../logger.ts';
+import type { IEmbedder } from '../../core/interfaces/knowledge-interfaces.ts';
+
+export class EmbeddingClient implements IEmbedder {
+  // Public so store.ts can set the dimension from existing table schema
+  // before the embedder has been fully warmed up (see Embedder.setDimension).
+  dimension: number | null = null;
+
+  constructor(
+    private readonly port: number,
+    private _device: string = 'unknown',
+  ) {
+    logger.log(`[EmbeddingClient] Connecting to embedding server at http://127.0.0.1:${port}`);
+  }
+
+  // ---- IEmbedder ----
+
+  getDevice(): string {
+    return this._device;
+  }
+
+  getOriginalDevice(): string {
+    return this._device;
+  }
+
+  getDimension(): number | null {
+    return this.dimension;
+  }
+
+  setDimension(dim: number): void {
+    this.dimension = dim;
+  }
+
+  isInitialized(): boolean {
+    // The server has already initialised; from the client's perspective we are ready.
+    return true;
+  }
+
+  async embed(text: string): Promise<Float32Array> {
+    const resp = await this.request<{ data: number[] }>('/embed', { text });
+    return new Float32Array(resp.data);
+  }
+
+  async embedMany(texts: string[]): Promise<(Float32Array | number[])[]> {
+    const resp = await this.request<{ results: number[][] }>('/embedMany', { texts });
+    return resp.results.map(r => new Float32Array(r));
+  }
+
+  async dispose(): Promise<void> {
+    // Clients don't shut down the server.
+  }
+
+  // ---- Health ----
+
+  async fetchHealth(): Promise<void> {
+    const result = await this.getRequest<{ status: string; device: string; dimension: number | null }>('/health');
+    this._device = result.device;
+    if (result.dimension !== null) {
+      this.dimension = result.dimension;
+    }
+  }
+
+  // ---- Internal HTTP helpers ----
+
+  private getRequest<T>(path: string): Promise<T> {
+    return new Promise((resolve, reject) => {
+      const timeoutMs = 120_000;
+      let resolved = false;
+      let req: http.ClientRequest | undefined;
+      const timer = setTimeout(() => {
+        if (!resolved) {
+          resolved = true;
+          // Destroy the request so a stalled body read is torn down, not left
+          // holding the socket until the server eventually closes it.
+          req?.destroy();
+          reject(new Error(`[EmbeddingClient] GET ${path} timed out after ${timeoutMs}ms`));
+        }
+      }, timeoutMs);
+      // FIX (#34): Don't keep the event loop alive for timeout timers
+      if (timer.unref) timer.unref();
+
+      req = http.request(
+        {
+          hostname: '127.0.0.1',
+          port: this.port,
+          path,
+          method: 'GET',
+        },
+        (res) => {
+          // Keep the timeout armed through the body read — clearing it on headers
+          // would let a stalled body hang for the full process lifetime.
+          let body = '';
+          res.on('data', (chunk: Buffer) => { body += chunk; });
+          res.on('end', () => {
+            if (resolved) return;
+            resolved = true;
+            clearTimeout(timer);
+            try {
+              const parsed = JSON.parse(body) as T;
+              if (res.statusCode !== 200) {
+                reject(new Error(`[EmbeddingClient] GET ${path} returned HTTP ${res.statusCode}`));
+              } else {
+                resolve(parsed);
+              }
+            } catch (_e) {
+              reject(new Error(`[EmbeddingClient] Failed to parse GET ${path} response: ${body}`));
+            }
+          });
+          res.on('error', (err) => {
+            if (resolved) return;
+            resolved = true;
+            clearTimeout(timer);
+            reject(new Error(`[EmbeddingClient] GET ${path} response error: ${err.message}`));
+          });
+        },
+      );
+
+      req.on('error', (err) => {
+        if (resolved) return;
+        resolved = true;
+        clearTimeout(timer);
+        reject(new Error(`[EmbeddingClient] GET ${path} request error: ${err.message}`));
+      });
+
+      req.end();
+    });
+  }
+
+  private request<T>(path: string, data: unknown): Promise<T> {
+    return new Promise((resolve, reject) => {
+      const timeoutMs = 120_000;
+      let resolved = false;
+      let req: http.ClientRequest | undefined;
+      const timer = setTimeout(() => {
+        if (!resolved) {
+          resolved = true;
+          req?.destroy();
+          reject(new Error(`[EmbeddingClient] POST ${path} timed out after ${timeoutMs}ms`));
+        }
+      }, timeoutMs);
+      // FIX (#34): Don't keep the event loop alive for timeout timers
+      if (timer.unref) timer.unref();
+
+      const body = JSON.stringify(data);
+
+      req = http.request(
+        {
+          hostname: '127.0.0.1',
+          port: this.port,
+          path,
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+            'Content-Length': Buffer.byteLength(body),
+          },
+        },
+        (res) => {
+          // Keep the timeout armed through the body read (see getRequest).
+          let responseBody = '';
+          res.on('data', (chunk: Buffer) => { responseBody += chunk; });
+          res.on('end', () => {
+            if (resolved) return;
+            resolved = true;
+            clearTimeout(timer);
+            try {
+              const parsed = JSON.parse(responseBody) as T;
+              if (res.statusCode !== 200) {
+                const msg = (parsed as Record<string, unknown>)?.['error'];
+                reject(new Error(typeof msg === 'string' ? msg : `HTTP ${res.statusCode}`));
+              } else {
+                resolve(parsed);
+              }
+            } catch (_e) {
+              reject(new Error(`[EmbeddingClient] Failed to parse POST ${path} response: ${responseBody}`));
+            }
+          });
+          res.on('error', (err) => {
+            if (resolved) return;
+            resolved = true;
+            clearTimeout(timer);
+            reject(new Error(`[EmbeddingClient] POST ${path} response error: ${err.message}`));
+          });
+        },
+      );
+
+      req.on('error', (err) => {
+        if (resolved) return;
+        resolved = true;
+        clearTimeout(timer);
+        reject(new Error(`[EmbeddingClient] POST ${path} request error: ${err.message}`));
+      });
+
+      req.write(body);
+      req.end();
+    });
+  }
+}

package/src/infrastructure/embedding/embedding-factory.ts

@@ -0,0 +1,324 @@
+/**
+ * Embedding Factory
+ *
+ * Leader-election for the embedding server. Mirrors getScheduler() in
+ * scheduler-factory.ts: one process wins and runs EmbeddingServer; all
+ * others become EmbeddingClient instances pointing at the winner's port.
+ */
+
+import * as crypto from 'node:crypto';
+import * as net from 'node:net';
+import type { Config } from '../../config.ts';
+import { getConfig } from '../../config.ts';
+import { logger } from '../../logger.ts';
+import { getService } from '../../core/service-registry.ts';
+import { ServiceNames } from '../../core/service-interfaces.ts';
+import type { IStateManager } from '../../core/interfaces/state-manager-interfaces.ts';
+import type { IProcessLifecycle } from '../../core/interfaces/process-interfaces.ts';
+import type { IEmbedder } from '../../core/interfaces/knowledge-interfaces.ts';
+import { Embedder } from '../../knowledge/embedder.ts';
+import { getModelEmbedderConfig } from '../../knowledge/model-config.ts';
+
+import { EmbeddingServer } from './embedding-server.ts';
+import { EmbeddingClient } from './embedding-client.ts';
+
+// ---------------------------------------------------------------------------
+// Port probe — same helper as scheduler-factory.ts
+// ---------------------------------------------------------------------------
+
+async function isPortListening(port: number, timeoutMs = 2000): Promise<boolean> {
+  return new Promise((resolve) => {
+    const socket = new net.Socket();
+    let settled = false;
+    const settle = (ok: boolean) => {
+      if (settled) return;
+      settled = true;
+      socket.destroy();
+      resolve(ok);
+    };
+    socket.setTimeout(timeoutMs);
+    socket.once('connect', () => settle(true));
+    socket.once('error', () => settle(false));
+    socket.once('timeout', () => settle(false));
+    socket.connect(port, '127.0.0.1');
+  });
+}
+
+// ---------------------------------------------------------------------------
+// Leader liveness — PID paired with start-time to defeat PID reuse
+// ---------------------------------------------------------------------------
+
+/**
+ * Decide whether a persisted embedding-leader entry still refers to a live process.
+ *
+ * Pairs the recorded PID with its start-time (when present) so a recycled PID — the OS
+ * handing a dead leader's PID to an unrelated new process, which a bare signal-0 check
+ * would mistake for the still-running leader — is correctly reported dead. This matches
+ * the PID-reuse handling the GPU lock and session cleanup already use. Entries written
+ * before start-time was recorded (no `startTime`) degrade to a bare liveness check, so
+ * older on-disk state keeps working with no migration.
+ *
+ * The reachable failure this closes is narrow: a leader that dies mid-model-init (its
+ * entry still carries the sentinel port -1, so the port probe is skipped) whose PID the
+ * OS then recycles. Without start-time the waiter would spin the full poll timeout on a
+ * leader that will never come; with it the stale entry is cleared immediately.
+ */
+export async function isEmbeddingLeaderAlive(
+  info: { pid: number; startTime?: number },
+  processLifecycle: Pick<IProcessLifecycle, 'isProcessAlive'>,
+): Promise<boolean> {
+  return processLifecycle.isProcessAlive(info.pid, info.startTime);
+}
+
+// ---------------------------------------------------------------------------
+// Module-level singleton cache (mirrors scheduler-factory.ts pattern)
+// ---------------------------------------------------------------------------
+
+let _embeddingInstance: IEmbedder | null = null;
+let _embeddingInitPromise: Promise<IEmbedder> | null = null;
+let _cachedModel: string | null = null;
+let _cachedDevice: string | null = null;
+
+// ---------------------------------------------------------------------------
+// getEmbedder — public entry point
+// ---------------------------------------------------------------------------
+
+// ---------------------------------------------------------------------------
+// In-process Embedder construction — shared by the leader path (wrapped in an
+// EmbeddingServer) and the model-mismatch fallback (used bare, never elected).
+// ---------------------------------------------------------------------------
+
+function buildInProcessEmbedder(cfg: Config, stateManager: IStateManager): Embedder {
+  const modelCfg = getModelEmbedderConfig(cfg.EMBEDDING_MODEL);
+  return new Embedder({
+    model: cfg.EMBEDDING_MODEL,
+    pooling: modelCfg.pooling,
+    queryPrefix: modelCfg.queryPrefix,
+    initializationTimeoutMs: cfg.EMBEDDING_MODEL_INIT_TIMEOUT_MS,
+    device: cfg.EMBEDDING_DEVICE,
+    maxTokens: modelCfg.maxTokens,
+    batchSize: modelCfg.batchSize,
+    charsPerToken: modelCfg.charsPerToken,
+    documentPrefix: modelCfg.documentPrefix,
+    stateManager,
+    useCache: modelCfg.useCache,
+  });
+}
+
+// Bounds the leader-wait retry tail-call below. Without a cap, a cluster where no
+// process can ever become leader (e.g. every model init OOMs) would recurse
+// forever, each hop doing a state-lock round-trip and accumulating async frames.
+const MAX_LEADER_WAIT_RETRIES = 5;
+
+export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbedder> {
+  const cfg = config ?? getConfig();
+
+  // If we have a cached instance, check if it matches the current config
+  if (_embeddingInstance) {
+    if (_cachedModel === cfg.EMBEDDING_MODEL && _cachedDevice === cfg.EMBEDDING_DEVICE) {
+      return _embeddingInstance;
+    }
+    logger.info(`[EmbeddingFactory] Configuration change detected (${_cachedModel} on ${_cachedDevice} -> ${cfg.EMBEDDING_MODEL} on ${cfg.EMBEDDING_DEVICE}). Disposing stale instance.`);
+    await _embeddingInstance.dispose?.();
+    _embeddingInstance = null;
+  }
+
+  if (_embeddingInitPromise) return _embeddingInitPromise;
+
+  let p: Promise<IEmbedder>;
+
+  const init = async (): Promise<IEmbedder> => {
+    const serverId = crypto.randomUUID();
+    const stateManager = await getService<IStateManager>(ServiceNames.STATE_MANAGER);
+    const processLifecycle = await getService<IProcessLifecycle>(ServiceNames.PROCESS_LIFECYCLE);
+    // Recorded alongside our PID so another process can distinguish "leader still
+    // alive" from "our PID was recycled after we died" (getCurrentProcessStartTime
+    // caches, so this is a single cheap read). Null on a platform where it cannot be
+    // read → the entry omits it and liveness falls back to a bare signal-0 check.
+    const myStartTime = await processLifecycle.getCurrentProcessStartTime();
+
+    // ---- Phase 1: Atomically claim candidacy under the state lock ----
+    // Writing port=-1 acts as a mutex: other processes that see it will wait
+    // for the real port rather than starting their own model initialization.
+    // This prevents concurrent GPU sessions and wasted init work.
+    let iAmCandidate = false;
+    let fastPathPort: number | null = null;
+    let fastPathModel: string | null = null;
+
+    await stateManager.updateState(async (state) => {
+      if (state.embeddingServer) {
+        const alive = await isEmbeddingLeaderAlive(state.embeddingServer, processLifecycle);
+        if (alive) {
+          // Either a real server (port > 0) or another process is initializing (port = -1)
+          fastPathPort = state.embeddingServer.port;
+          fastPathModel = state.embeddingServer.model ?? null;
+          return state;
+        }
+        // Stale entry from a dead process — clear it
+        delete state.embeddingServer;
+      }
+      // Claim the slot with sentinel port -1 to signal "initializing". Record the
+      // model we will serve so followers with a DIFFERENT configured model can
+      // refuse adoption (see mismatch check below) instead of silently
+      // cross-embedding into a shared vector space.
+      state.embeddingServer = {
+        port: -1,
+        pid: process.pid,
+        serverId,
+        model: cfg.EMBEDDING_MODEL,
+        ...(myStartTime !== null ? { startTime: myStartTime } : {}),
+      };
+      iAmCandidate = true;
+      return state;
+    });
+
+    // A live leader serving a DIFFERENT model must not be adopted: two sessions
+    // with different EMBEDDING_MODELs would cross-embed (same-dimension models
+    // corrupt the vector space silently). Fall back to a local in-process
+    // embedder for this session — no competing election, the leader's entry is
+    // left untouched. Entries without a recorded model (older process) are
+    // tolerated as matching.
+    const localFallbackForMismatch = (leaderModel: string): IEmbedder => {
+      logger.warn(
+        `[EmbeddingFactory] Live embedding leader serves model "${leaderModel}" but this session is configured for ` +
+        `"${cfg.EMBEDDING_MODEL}". Not adopting it — using a local in-process embedder for this session.`,
+      );
+      const embedder = buildInProcessEmbedder(cfg, stateManager);
+      _embeddingInstance = embedder;
+      _cachedModel = cfg.EMBEDDING_MODEL;
+      _cachedDevice = cfg.EMBEDDING_DEVICE;
+      return embedder;
+    };
+
+    if (!iAmCandidate) {
+      if (fastPathModel !== null && fastPathModel !== cfg.EMBEDDING_MODEL) {
+        return localFallbackForMismatch(fastPathModel);
+      }
+      // Another process is already leader or initializing — wait for real port
+      const POLL_INTERVAL_MS = 500;
+      const POLL_TIMEOUT_MS = 120_000;
+      const deadline = Date.now() + POLL_TIMEOUT_MS;
+
+      logger.info(`[EmbeddingFactory] Another process is initializing (port=${fastPathPort}), waiting...`);
+
+      while (Date.now() < deadline) {
+        await new Promise<void>((r) => setTimeout(r, POLL_INTERVAL_MS));
+        const info = await stateManager.getEmbeddingServer();
+        if (!info) {
+          // Candidate died and cleared state — fall through to restart below
+          break;
+        }
+        if (!(await isEmbeddingLeaderAlive(info, processLifecycle))) {
+          // Candidate process died without cleaning up (or its PID was recycled by
+          // an unrelated process — the start-time pairing catches that too).
+          await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+          break;
+        }
+        // Re-check model identity every poll: the leader may have changed while
+        // we waited (old one died, a new one with a different model claimed).
+        if (info.model !== undefined && info.model !== cfg.EMBEDDING_MODEL) {
+          return localFallbackForMismatch(info.model);
+        }
+        if (info.port > 0) {
+          const portOk = await isPortListening(info.port);
+          if (portOk) {
+            logger.info(`[EmbeddingFactory] Connecting to embedding server on port ${info.port}`);
+            const client = new EmbeddingClient(info.port);
+            await client.fetchHealth();
+            _embeddingInstance = client;
+            _cachedModel = cfg.EMBEDDING_MODEL;
+            _cachedDevice = cfg.EMBEDDING_DEVICE;
+            return client;
+          }
+          // Port registered but not responding — wait one more interval and re-check
+          await new Promise<void>((r) => setTimeout(r, POLL_INTERVAL_MS));
+          const secondCheck = await isPortListening(info.port);
+          if (!secondCheck) {
+            logger.warn(`[EmbeddingFactory] Registered port ${info.port} is unreachable after two checks — clearing stale state`);
+            await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+            break;
+          }
+        }
+      }
+
+      // Timed out or candidate died — attempt to claim candidacy ourselves (tail-call)
+      _embeddingInitPromise = null;
+      if (_attempt + 1 >= MAX_LEADER_WAIT_RETRIES) {
+        throw new Error(
+          `[EmbeddingFactory] Could not obtain an embedding server after ${MAX_LEADER_WAIT_RETRIES} attempts ` +
+          `(no process became leader; model init may be failing repeatedly).`,
+        );
+      }
+      logger.warn(`[EmbeddingFactory] Wait for embedding leader timed out or leader died, retrying (attempt ${_attempt + 2}/${MAX_LEADER_WAIT_RETRIES})...`);
+      return getEmbedder(cfg, _attempt + 1);
+    }
+
+    // ---- Phase 2: We are the exclusive candidate — do model init ----
+    const embedder = buildInProcessEmbedder(cfg, stateManager);
+
+    const server = new EmbeddingServer(embedder, stateManager, serverId, () => {
+      // This leader poisoned itself after a permanently hung inference; drop the
+      // cached singleton so this process's next getEmbedder() re-elects. The store's
+      // reconnectFactory also triggers on the first fast-failed embed, so cluster
+      // recovery does not depend solely on this hook.
+      void clearEmbeddingInstance();
+    });
+    let port: number;
+    try {
+      port = await server.startServer();
+    } catch (err) {
+      // Init failed — clear our placeholder so others can try
+      await stateManager.clearEmbeddingServer().catch((err) => logger.debug('Swallowed clear embedding server error:', err));
+      throw err;
+    }
+
+    // ---- Phase 3: Update state with real port ----
+    await stateManager.updateState(async (state) => {
+      if (state.embeddingServer?.serverId === serverId) {
+        state.embeddingServer.port = port;
+      }
+      return state;
+    });
+
+    logger.info(`[EmbeddingFactory] Won election, serving as embedding leader on port ${port} (PID ${process.pid})`);
+    // Start leadership check only after winning — model init (inside startServer)
+    // can exceed the 30s interval, so starting earlier would produce false misses.
+    server.startLeadershipCheck();
+    _embeddingInstance = server;
+    _cachedModel = cfg.EMBEDDING_MODEL;
+    _cachedDevice = cfg.EMBEDDING_DEVICE;
+    return server;
+  };
+
+  p = init();
+  _embeddingInitPromise = p;
+  // init() already assigns _embeddingInstance before it resolves, so do NOT
+  // re-assign it here. If clearEmbeddingInstance() ran (and disposed the server)
+  // between init resolving and this microtask, re-assigning would resurrect the
+  // disposed instance and hand it to the next getEmbedder() caller. Only clear
+  // the in-flight promise marker.
+  p.then(() => {
+    _embeddingInitPromise = null;
+  }).catch(() => {
+    _embeddingInitPromise = null;
+  });
+  return p;
+}
+
+// ---------------------------------------------------------------------------
+// clearEmbeddingInstance — for testing / restart scenarios
+// ---------------------------------------------------------------------------
+
+export async function clearEmbeddingInstance(): Promise<void> {
+  if (_embeddingInstance) {
+    try {
+      await _embeddingInstance.dispose?.();
+    } catch (err) {
+      logger.warn('[EmbeddingFactory] Error disposing old embedding instance:', err);
+    }
+  }
+  _embeddingInstance = null;
+  _embeddingInitPromise = null;
+  _cachedModel = null;
+  _cachedDevice = null;
+}