@lincoln504/pi-research 1.3.12 → 1.4.3

package/src/infrastructure/browser/worker-pool-manager.ts

@@ -28,6 +28,55 @@ const __dirname = dirname(__filename);
 /**
  * Worker pool manager for browser operations.
  */
+/**
+ * The pool's dispatch shape: how a worker is chosen, and how much each may hold.
+ *
+ * Separated out so it is testable. The claim this configuration makes is behavioural
+ * — that no task is parked behind a busy worker while another idles — and the first
+ * test written for it asserted that this file contained the string
+ * "WorkerChoiceStrategies.LEAST_USED", which passes against an implementation that
+ * dispatches nothing.
+ *
+ * LEAST_USED picks the node minimising executing + queued. ROUND_ROBIN is a bare
+ * modular cursor whose only filter is a one-time `ready` flag — it cannot see load,
+ * and poolifier consults per-node capacity only AFTER the strategy has chosen. A task
+ * handed to a busy node is queued ON that node and never re-routed, so one degraded
+ * worker takes a share of every burst with it: on 2026-08-15 a worker whose browser
+ * launch hung was still handed roughly one turn in six, and the run lost exactly the
+ * three queries that landed on it while five workers sat idle.
+ *
+ * The guarantee is conditional, not absolute, and worth stating precisely: a
+ * load-reading strategy never parks a task WHILE A FREE NODE EXISTS. The
+ * PriorityTaskQueue above normally makes sure one does, by capping in-flight work at
+ * exactly WORKER_THREADS x WORKER_CONCURRENCY — which is why a WORKER_CONCURRENCY
+ * change has to rebuild the pool, or the two layers disagree about capacity. But that
+ * cap is released early when a caller aborts a task still running on a worker, and a
+ * respawning worker is not `ready`, so there are windows where every node is genuinely
+ * full. In those, ties resolve to the lowest-index node and a task does park.
+ *
+ * Which is also why poolifier's stolen-task accounting defect is made rare rather than
+ * unreachable: stealing fires only when something is queued on a node, and a stolen
+ * task's completion decrements the ORIGINAL node's executing count instead of the one
+ * that ran it, freeing a slot on a worker that is still working. That produced 1,425
+ * genuinely overlapping task pairs across 37 days at a configured concurrency of 1.
+ *
+ * LEAST_BUSY was the other candidate and is worse here: it minimises accumulated
+ * run+wait time, which turns on per-task performance measurement for every task and
+ * still lets a busy fast worker outrank an idle slow one. LEAST_USED overrides no
+ * statistics requirements, so it costs nothing extra.
+ */
+export function buildWorkerPoolOptions(config: Config): {
+    workerChoiceStrategy: typeof WorkerChoiceStrategies[keyof typeof WorkerChoiceStrategies];
+    enableTasksQueue: boolean;
+    tasksQueueOptions: { concurrency: number };
+} {
+    return {
+        workerChoiceStrategy: WorkerChoiceStrategies.LEAST_USED,
+        enableTasksQueue: true,
+        tasksQueueOptions: { concurrency: config.WORKER_CONCURRENCY },
+    };
+}
+
 export class WorkerPoolManager implements IService {
     readonly name = ServiceNames.WORKER_POOL_MANAGER;
     lifecycle = ServiceLifecycle.UNINITIALIZED;
@@ -35,6 +84,12 @@ export class WorkerPoolManager implements IService {
     private pool: any | null = null;
     private poolInitializationPromise: Promise<any> | null = null;
     private currentWorkerCount: number | null = null;
+    /** Per-worker task concurrency the LIVE pool was built with. Baked into
+     *  poolifier's tasksQueueOptions at construction and not changeable after, so a
+     *  change has to recreate the pool — the upper PriorityTaskQueue recomputes its
+     *  own capacity from the same two values on every task, and the two disagreeing
+     *  is what lets a task be queued on a worker that the queue believes is free. */
+    private currentWorkerConcurrency: number | null = null;
     private consecutiveErrors: number = 0;
     private isShuttingDown: boolean = false;
     // Monotonic shutdown generation. Each shutdown() bumps it, so an in-flight
@@ -95,7 +150,8 @@ export class WorkerPoolManager implements IService {
 
         // If pool exists, worker count matches, and we are not shutting down,
         // return the existing pool immediately.
-        if (this.pool && this.currentWorkerCount === maxWorkers) {
+        const workerConcurrency = (config || getConfig()).WORKER_CONCURRENCY;
+        if (this.pool && this.currentWorkerCount === maxWorkers && this.currentWorkerConcurrency === workerConcurrency) {
             return this.pool;
         }
 
@@ -110,14 +166,15 @@ export class WorkerPoolManager implements IService {
         const myGeneration = this.generation;
         this.poolInitializationPromise = (async () => {
             try {
-                if (this.pool && this.currentWorkerCount !== maxWorkers) {
-                    logger.log(`[WorkerPoolManager] Worker count changed from ${this.currentWorkerCount} to ${maxWorkers}, recreating pool...`);
+                if (this.pool && (this.currentWorkerCount !== maxWorkers || this.currentWorkerConcurrency !== workerConcurrency)) {
+                    logger.log(`[WorkerPoolManager] Pool shape changed (workers ${this.currentWorkerCount} -> ${maxWorkers}, concurrency ${this.currentWorkerConcurrency} -> ${workerConcurrency}), recreating pool...`);
                     this.poolEpoch++; // intentional destroy: silence the dying pool's exit events
                     await this.pool.destroy();
                     this.pool = null;
                 }
 
                 this.currentWorkerCount = maxWorkers;
+                this.currentWorkerConcurrency = workerConcurrency;
 
                 // Install the master-side IPC guard BEFORE any worker is forked, so a
                 // `write EPIPE` on a closing worker channel (poolifier dispatch racing a
@@ -161,7 +218,6 @@ export class WorkerPoolManager implements IService {
                 void cleanupOrphanedCamoufoxProcesses()
                     .catch((e) => logger.debug('[WorkerPoolManager] Startup orphan sweep failed:', e));
 
-                const workerConcurrency = (config || getConfig()).WORKER_CONCURRENCY;
 
                 // thread-worker.mjs is the esbuild-compiled bundle of thread-worker.ts and its
                 // local imports. Using a pre-compiled JS file eliminates any TypeScript loading
@@ -205,10 +261,15 @@ export class WorkerPoolManager implements IService {
                             return;
                         }
                         this.consecutiveErrors++;
-                        metrics.increment('browser_pool_errors_total', 1);
+                        // Session-scoped on purpose. This handler is installed on the pool,
+                        // and the pool is built lazily inside whichever research run first
+                        // needed a browser — so the run-registry context every later
+                        // invocation inherits is that FIRST run's, long after it was
+                        // discarded. Pool health is process state, not run accounting.
+                        metrics.session.increment('browser_pool_errors_total', 1);
                         logger.error('[WorkerPoolManager] Cluster Error:', e);
                         if (this.consecutiveErrors >= 3) {
-                            metrics.increment('browser_pool_unhealthy_events_total', 1);
+                            metrics.session.increment('browser_pool_unhealthy_events_total', 1);
                             logger.error(`[WorkerPoolManager] Worker pool may be unhealthy: ${this.consecutiveErrors} consecutive errors. Consider restarting.`);
                             if (this.onPoolError) {
                                 this.onPoolError(e, this.consecutiveErrors);
@@ -235,11 +296,7 @@ export class WorkerPoolManager implements IService {
                             }
                         }
                     },
-                    workerChoiceStrategy: WorkerChoiceStrategies.ROUND_ROBIN,
-                    enableTasksQueue: true,
-                    tasksQueueOptions: {
-                        concurrency: workerConcurrency, // configurable via PI_RESEARCH_WORKER_CONCURRENCY
-                    }
+                    ...buildWorkerPoolOptions(config || getConfig()),
                 });
 
                 // Secondary race check: if a shutdown() was called concurrent with
@@ -315,6 +372,60 @@ export class WorkerPoolManager implements IService {
         if (this.consecutiveErrors > 0) this.consecutiveErrors--;
     }
 
+    /**
+     * Record how the pool is actually distributing work.
+     *
+     * The head-of-line failure this pool used to have was invisible to every metric
+     * and every log line: tasks parked on one node while other nodes idled, and the
+     * only trace was a PriorityTaskQueue reporting free slots and an empty backlog
+     * while queries timed out. poolifier computes all of it on `pool.info` and
+     * nothing read it.
+     *
+     * `browser_pool_parked_tasks` is the number that should stay at zero: a task
+     * queued on a worker node while some other node is idle. Non-zero means the
+     * choice strategy handed work to a busy node when a free one existed — which is
+     * also the precondition for poolifier's task stealing, and therefore for its
+     * stolen-task accounting defect. `stolenTasks` is recorded for the same reason.
+     *
+     * Two limits, stated because a gauge that reads zero is otherwise taken as proof:
+     * `idleWorkerNodes` counts nodes that are READY and free, so during a respawn
+     * window a genuinely parked task is attributed to legitimate saturation; and with
+     * every node busy the figure is zero by construction, which is correct but is not
+     * evidence the strategy is choosing well. `browser_pool_idle_workers` is exported
+     * alongside so the denominator is visible rather than inferred.
+     *
+     * With no pool the gauges are ZEROED rather than left alone. Returning early held
+     * the last live reading forever — so a burst that ended with tasks parked left a
+     * standing non-zero gauge for a pool that no longer exists, and every later read
+     * of it was a report about a dead process.
+     *
+     * Session-scoped, for the same reason as the pool's error counters above. The only
+     * caller is the leader's 5s leadership tick — a self-rescheduling timer chain
+     * started when this process won the election, i.e. inside whichever research run
+     * first needed a browser. AsyncLocalStorage resolves at emit time from the async
+     * resource that fired, so a plain `metrics.setGauge` here would write into that
+     * FIRST run's registry for the life of the process: invisible to every later run's
+     * summary, and holding a discarded registry reachable through the timer. The pool's
+     * dispatch shape is process state, not run accounting.
+     */
+    recordDispatchHealth(): void {
+        const info = (this.pool as unknown as { info?: Record<string, number> } | null)?.info;
+        if (!info) {
+            metrics.session.setGauge('browser_pool_parked_tasks', 0);
+            metrics.session.setGauge('browser_pool_queued_tasks', 0);
+            metrics.session.setGauge('browser_pool_idle_workers', 0);
+            return;
+        }
+        const queued = info['queuedTasks'] ?? 0;
+        const idle = info['idleWorkerNodes'] ?? 0;
+        metrics.session.setGauge('browser_pool_parked_tasks', idle > 0 ? queued : 0);
+        metrics.session.setGauge('browser_pool_queued_tasks', queued);
+        metrics.session.setGauge('browser_pool_idle_workers', idle);
+        // Cumulative, not a level: never zeroed on teardown, or a restart would read
+        // as the defect having been fixed.
+        metrics.session.setGauge('browser_pool_stolen_tasks', info['stolenTasks'] ?? 0);
+    }
+
     /**
      * Schedule an out-of-band pool reset. Called from poolifier event handlers
      * where destroying the pool synchronously would deadlock. Waits 1 s so the
@@ -334,12 +445,15 @@ export class WorkerPoolManager implements IService {
         // is still being destroyed. Instead, use a flag to signal reset is in progress.
         this._resetInProgress = true;
         this.currentWorkerCount = null;
+        this.currentWorkerConcurrency = null;
         this.consecutiveErrors = 0;
         // Pin the generation this reset belongs to. shutdown() bumps generation, so if a
         // shutdown (and a subsequent re-init) happens while this timer is pending, the
         // generation no longer matches and the timer must NOT touch the new pool's state.
         const myGen = this.generation;
-        metrics.increment('browser_pool_auto_recoveries_total', 1);
+        // Session-scoped: every caller of schedulePoolReset is one of the pool's own
+        // lifecycle handlers, which carry the async context of pool construction.
+        metrics.session.increment('browser_pool_auto_recoveries_total', 1);
         logger.info('[WorkerPoolManager] Pool scheduled for auto-recovery; next ensurePool() will wait for old pool destruction.');
         // Destroy the old pool asynchronously after the event handler returns.
         const t = setTimeout(async () => {
@@ -427,6 +541,7 @@ export class WorkerPoolManager implements IService {
 
         this.poolInitializationPromise = null;
         this.currentWorkerCount = null;
+        this.currentWorkerConcurrency = null;
         this.consecutiveErrors = 0;
         // Clear any in-flight auto-recovery flag too: if a schedulePoolReset() timer
         // was pending when shutdown ran, leaving this true makes the next ensurePool()

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

@@ -72,8 +72,19 @@ export class EmbeddingClient implements IEmbedder {
 
   // ---- Health ----
 
-  async fetchHealth(): Promise<void> {
-    const result = await this.getRequest<{ status: string; device: string; dimension: number | null }>('/health');
+  /**
+   * `timeoutMs` defaults to the general-purpose 120s used for real embed requests, but the
+   * embedding-leader poll loop (embedding-factory.ts) needs a much shorter deadline: it uses
+   * a failed health check as one "unreachable poll" toward declaring a leader's registration
+   * stale, on the documented assumption that each poll costs "the poll interval plus up to
+   * the [probe] timeout" (~2.5s). At the 120s default that assumption is off by 48x — a
+   * leader whose event loop is merely busy (not dead) can take the full 120s to answer, so
+   * the poll loop's outer wait times out and retries long before its own unreachable-poll
+   * counter would ever fire. Callers that need the fast-fail poll semantics must pass an
+   * explicit short timeout.
+   */
+  async fetchHealth(timeoutMs?: number): Promise<void> {
+    const result = await this.getRequest<{ status: string; device: string; dimension: number | null }>('/health', timeoutMs);
     this._device = result.device;
     if (result.dimension !== null) {
       this.dimension = result.dimension;
@@ -82,9 +93,8 @@ export class EmbeddingClient implements IEmbedder {
 
   // ---- Internal HTTP helpers ----
 
-  private getRequest<T>(path: string): Promise<T> {
+  private getRequest<T>(path: string, timeoutMs = 120_000): Promise<T> {
     return new Promise((resolve, reject) => {
-      const timeoutMs = 120_000;
       let resolved = false;
       let req: http.ClientRequest | undefined;
       const timer = setTimeout(() => {

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

@@ -227,6 +227,33 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
       const POLL_TIMEOUT_MS = 120_000;
       const deadline = Date.now() + POLL_TIMEOUT_MS;
 
+      /**
+       * How many CONSECUTIVE polls must find a confirmed-live leader's endpoint
+       * unusable before its registration is judged stale.
+       *
+       * Deleting a live leader's entry is the expensive mistake in this loop: the
+       * leader is not told, treats a missing entry as benign, and never steps down,
+       * so the follower's subsequent election produces a second model init and a
+       * second GPU context while the first server keeps running, invisible. The old
+       * code reached that from two failed TCP probes ~2.5s apart — well inside a
+       * blip from a saturated accept backlog, a paused event loop during a large
+       * batch embed, or CPU starvation on a shared box.
+       *
+       * Not clearing at all is not the answer either: a leader whose HTTP server
+       * died while the process lives on leaves a registration nothing else collects,
+       * and every later process then falls back to its own in-process embedder. So
+       * the threshold is raised rather than removed. Each iteration costs the poll
+       * interval plus up to the 2s socket probe (isPortListening) plus up to a
+       * further 2s health-check probe (fetchHealth, called below with an explicit
+       * short timeout — NOT the client's 120s default, which would let a merely-busy
+       * leader blow the outer POLL_TIMEOUT_MS long before this counter ever fired),
+       * making this roughly 45s of uninterrupted unreachability — far past any blip,
+       * far short of the 120s wait.
+       */
+      const UNREACHABLE_POLLS_BEFORE_STALE = 10;
+      let unreachablePolls = 0;
+      let unreachableEndpoint: string | null = null;
+
       logger.info(`[EmbeddingFactory] Another process is initializing (port=${fastPathPort}), waiting...`);
 
       while (Date.now() < deadline) {
@@ -253,24 +280,50 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
           return localFallbackForMismatch(info.model);
         }
         if (info.port > 0) {
-          const portOk = await isPortListening(info.port);
-          if (portOk) {
+          // A new claimant resets the streak — the count must describe ONE endpoint's
+          // continuous unreachability, not a tally across successive leaders.
+          const endpoint = `${info.serverId}:${info.port}`;
+          if (endpoint !== unreachableEndpoint) {
+            unreachableEndpoint = endpoint;
+            unreachablePolls = 0;
+          }
+
+          // Reaching the counter below means "not usable this poll" by construction:
+          // the only usable outcome returns.
+          if (await isPortListening(info.port)) {
             logger.info(`[EmbeddingFactory] Connecting to embedding server on port ${info.port}`);
             const client = new EmbeddingClient(info.port, cfg.EMBEDDING_DEVICE, info.authSecret);
-            await client.fetchHealth();
-            _embeddingInstance = client;
-            _cachedModel = cfg.EMBEDDING_MODEL;
-            _cachedDevice = cfg.EMBEDDING_DEVICE;
-            return client;
+            try {
+              // A short, poll-scale timeout — NOT EmbeddingClient's 120s default. The
+              // unreachable-poll count above assumes each iteration costs roughly the poll
+              // interval plus this probe's timeout (~2.5s total); at the 120s default, a
+              // leader that is merely busy (not dead) can take the full 120s to answer,
+              // blowing the outer POLL_TIMEOUT_MS well before the stale-registration
+              // threshold below would ever fire. Matches isPortListening's own 2s probe.
+              await client.fetchHealth(2_000);
+              _embeddingInstance = client;
+              _cachedModel = cfg.EMBEDDING_MODEL;
+              _cachedDevice = cfg.EMBEDDING_DEVICE;
+              return client;
+            } catch (err) {
+              // The socket accepted but the server did not answer. Previously this
+              // rejected out of getEmbedder entirely, so one failed health request
+              // against an otherwise-recoverable leader (mid model swap, momentarily
+              // wedged) took embedding down for the whole session instead of costing
+              // one poll. Count it as an unreachable poll and keep waiting.
+              logger.debug(`[EmbeddingFactory] Health check against port ${info.port} failed:`, err);
+            }
           }
-          // 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`);
-            // serverId-scoped: this snapshot is up to ~2.5s old (two probes + sleep);
-            // the entry may already belong to a new claimant. See the dead-leader
-            // clear above for the failure this prevents.
+
+          unreachablePolls++;
+          if (unreachablePolls >= UNREACHABLE_POLLS_BEFORE_STALE) {
+            logger.warn(
+              `[EmbeddingFactory] Registered port ${info.port} unreachable for ${unreachablePolls} consecutive polls ` +
+              `(PID ${info.pid} is alive) — treating the registration as stale and electing a replacement.`,
+            );
+            // serverId-scoped: this snapshot is many seconds old by now and the entry
+            // may already belong to a new claimant. See the dead-leader clear above
+            // for the failure an unconditional delete causes.
             await stateManager.clearEmbeddingServer({ serverId: info.serverId }).catch((err) => logger.debug('Swallowed clear embedding server error:', err));
             break;
           }
@@ -355,6 +408,21 @@ export async function getEmbedder(config?: Config, _attempt = 0): Promise<IEmbed
 // clearEmbeddingInstance — for testing / restart scenarios
 // ---------------------------------------------------------------------------
 
+/**
+ * Does THIS process currently hold a live embedding server (as opposed to a client
+ * pointing at someone else's, or nothing at all)?
+ *
+ * Asked by the teardown-time state clear, which is pid-scoped and therefore cannot
+ * tell "my leftover registration" from "my server that is still listening". That
+ * clear runs before disposeAllServices, so at that moment our own server usually
+ * IS still up — and deleting its entry makes it an invisible leader, since its
+ * leadership check reads a missing entry as benign and never steps down. Its own
+ * dispose clears the entry properly, keyed on serverId.
+ */
+export function ownsLiveEmbeddingServer(): boolean {
+    return _embeddingInstance !== null && _embeddingInstance instanceof EmbeddingServer;
+}
+
 export async function clearEmbeddingInstance(): Promise<void> {
   if (_embeddingInstance) {
     try {

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

@@ -80,6 +80,11 @@ class SerialQueue {
   private readonly hardDeadlineMs: number;
   private readonly onPoison?: (err: Error) => void;
   private poisoned: Error | null = null;
+  // The currently-executing task's real work promise (not the soft-timeout race) —
+  // set only while a native call is actually in flight. Lets a caller that needs
+  // to dispose the underlying pipeline (EmbeddingServer.shutdown) wait for true
+  // idleness instead of racing disposal against a still-running inference.
+  private currentWork: Promise<unknown> | null = null;
 
   constructor(
     maxDepth = 200,
@@ -132,6 +137,7 @@ class SerialQueue {
         let hardTimer: ReturnType<typeof setTimeout> | undefined;
         try {
           work = fn();
+          this.currentWork = work;
           // Arm the hard-deadline watchdog from the moment the real GPU work starts, so
           // the cap is an honest wall-clock deadline measured from work-start rather than
           // (soft-timeout + hardDeadlineMs). A single inference still unsettled after
@@ -162,6 +168,7 @@ class SerialQueue {
           } finally {
             settled = true;
             if (hardTimer !== undefined) clearTimeout(hardTimer);
+            if (this.currentWork === work) this.currentWork = null;
           }
         }
       };
@@ -170,6 +177,23 @@ class SerialQueue {
     });
   }
 
+  /**
+   * Resolve once no native call is in flight. Used by EmbeddingServer.shutdown()
+   * to wait out a still-running inference before disposing the pipeline —
+   * without this, disposal could run concurrently with a legitimately slow (not
+   * hung) embed/embedMany call on the same re-entrant-locked GPU context, the
+   * exact native-segfault class this queue exists to prevent. No timeout: a
+   * promise that truly never settles is already an accepted, pre-existing
+   * leak-to-process-exit scenario elsewhere in this class (see the poison path's
+   * disposeEmbedder=false handling) — better to leak than risk concurrent
+   * native access.
+   */
+  async waitForIdle(): Promise<void> {
+    while (this.currentWork) {
+      await this.currentWork.catch(() => {});
+    }
+  }
+
   private poison(err: Error): void {
     if (this.poisoned) return;
     this.poisoned = err;
@@ -414,6 +438,16 @@ export class EmbeddingServer implements IEmbedder {
     // is enough for another process to re-elect; the leaked native resources die with
     // eventual process exit.
     if (disposeEmbedder) {
+      // Wait for the SerialQueue to go idle BEFORE disposing. This shutdown() path
+      // also runs on ordinary leadership loss — not just the poison path above — and
+      // a legitimately slow (not hung) embed/embedMany can still be mid-flight on
+      // this same GPU context. Embedder.dispose()'s own drainActiveEmbeddings only
+      // waits 5s, an order of magnitude under the queue's own 120s soft timeout, so
+      // without this a normal large batch could have its pipeline disposed out from
+      // under it — the same concurrent-native-call hazard the queue exists to
+      // prevent, just triggered by disposal racing a live call instead of two calls
+      // racing each other.
+      await this.queue.waitForIdle();
       await this.embedder.dispose().catch((err) => {
         if (this.isShuttingDown) {
           logger.debug('[EmbeddingServer] Error disposing embedder during shutdown (expected):', err);