feat: forked python process

.changeset/spicy-vans-wait.md

@@ -0,0 +1,5 @@
+---
+"apollo": patch
+---
+
+fork python processes to reduce startup time

.env.example

@@ -70,3 +70,9 @@ LANGFUSE_BASE_URL=https://cloud.langfuse.com
 # ZILLIZ_URI =  https://in01-XXXXXXXXXXXXX.aws-us-west-2.vectordb.zillizcloud.com:XXXXX
 # ZILLIZ_TOKEN =db_admin:password (or ApiKey)
 
+# Turn this on to use the fork server mechanism
+# APOLLO_FORK_SERVER=true
+
+# macOS only: required when APOLLO_FORK_SERVER=true
+# OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES
+

platform/src/bridge.fork.ts

@@ -0,0 +1,183 @@
+import readline from "node:readline";
+import net from "node:net";
+import path from "node:path";
+import { spawn, type ChildProcess } from "node:child_process";
+import { rm } from "node:fs/promises";
+import { getInternalToken } from "./auth/internal-token";
+import {
+  emptyResult,
+  malformedResult,
+  subprocessCancelled,
+  subprocessFailed,
+  subprocessKilled,
+  subprocessSpawnFailed,
+} from "./util/errors";
+import { LOG_LINE } from "./util/log-line";
+import pkg from "../../package.json";
+
+export const USE_FORK_SERVER =
+  (process.env.APOLLO_FORK_SERVER ?? "").toLowerCase() === "true";
+
+const onLines = (stream: NodeJS.ReadableStream, fn: (line: string) => void) =>
+  readline.createInterface({ input: stream, crlfDelay: Infinity }).on("line", fn);
+
+// --- the warm master: spawned once, forks a child per request ----------------
+
+type Master = { proc: ChildProcess; socketPath: string };
+let master: Promise<Master> | null = null;
+let masterProc: ChildProcess | null = null;
+
+// Singleton; a failed start clears the slot so the next request retries.
+function forkMaster(): Promise<Master> {
+  return (master ??= startMaster().catch((err) => {
+    master = null;
+    throw err;
+  }));
+}
+
+function startMaster(): Promise<Master> {
+  const socketPath = path.resolve(`tmp/fork_server-${process.pid}.sock`);
+  const proc = spawn(
+    "poetry",
+    ["run", "python", "services/entry.fork.py", "--socket", socketPath],
+    {
+      env: {
+        ...process.env,
+        APOLLO_INTERNAL_TOKEN: getInternalToken(),
+        APOLLO_VERSION: pkg.version,
+      },
+    }
+  );
+  masterProc = proc;
+  onLines(proc.stderr!, (line) => console.error(`fork-server: ${line}`));
+
+  return new Promise<Master>((resolve, reject) => {
+    const fail = (err: Error) => (clearTimeout(timer), reject(subprocessSpawnFailed("fork_server", err)));
+    const timer = setTimeout(() => fail(new Error("not READY in time")), 60_000);
+
+    onLines(proc.stdout!, (line) => {
+      if (line.startsWith("READY")) {
+        clearTimeout(timer);
+        resolve({ proc, socketPath });
+      }
+    });
+    proc.on("error", fail);
+    proc.on("exit", (code, sig) => {
+      console.error(`fork-server exited (code=${code} signal=${sig})`);
+      master = masterProc = null; // next request starts a fresh one
+      rm(socketPath).catch(() => {});
+      fail(new Error(`fork-server exited before ready (code=${code} signal=${sig})`));
+    });
+  });
+}
+
+// Don't leave an orphaned master when the server process goes.
+process.once("exit", () => masterProc?.kill("SIGTERM"));
+
+// --- per-request: ask the master to fork a child, stream it back -------------
+
+export async function runForked(
+  scriptName: string,
+  port: number,
+  inputPath: string,
+  outputPath: string,
+  onLog?: (str: string) => void,
+  onEvent?: (type: string, payload: any) => void,
+  signal?: AbortSignal
+): Promise<JSON | null> {
+  const cleanup = () =>
+    Promise.all([rm(inputPath).catch(() => {}), rm(outputPath).catch(() => {})]);
+
+  let active: Master;
+  try {
+    active = await forkMaster();
+  } catch (err) {
+    await cleanup();
+    throw err;
+  }
+
+  return new Promise<JSON | null>((resolve, reject) => {
+    const conn = net.createConnection({ path: active.socketPath });
+    let childPid: number | null = null;
+    let exitCode: number | null = null;
+    let cancelled = false;
+    let hardKill: ReturnType<typeof setTimeout> | undefined;
+
+    const killChild = () => {
+      if (childPid == null) return; // killed once the PID: line arrives
+      try {
+        process.kill(childPid, "SIGTERM");
+      } catch {}
+      hardKill = setTimeout(() => {
+        try {
+          process.kill(childPid!, "SIGKILL");
+        } catch {}
+      }, 5_000);
+      hardKill.unref?.();
+    };
+
+    const onAbort = () => {
+      cancelled = true;
+      console.warn(`cancelling ${scriptName}: client went away`);
+      killChild();
+      conn.destroy();
+    };
+    if (signal) {
+      if (signal.aborted) onAbort();
+      else signal.addEventListener("abort", onAbort, { once: true });
+    }
+
+    conn.on("connect", () =>
+      conn.write(JSON.stringify({ service: scriptName, input: inputPath, output: outputPath, port }) + "\n")
+    );
+
+    // The child's stream: PID:/EXIT: control lines around the service's own stdout.
+    onLines(conn, (line) => {
+      if (line.startsWith("PID:")) {
+        childPid = Number(line.slice(4)) || null;
+        if (cancelled) killChild();
+      } else if (line.startsWith("EXIT:")) {
+        exitCode = Number(line.slice(5));
+      } else if (LOG_LINE.test(line)) {
+        console.log(line);
+        onLog?.(line);
+      } else if (line.startsWith("EVENT:")) {
+        const [, type, ...rest] = line.split(":");
+        let payload: any = rest.join(":");
+        try {
+          payload = JSON.parse(payload);
+        } catch {
+          // not json, forward the raw string
+        }
+        onEvent?.(type, payload);
+      }
+    });
+
+    const settle = async (finish: () => void) => {
+      if (hardKill) clearTimeout(hardKill);
+      signal?.removeEventListener("abort", onAbort);
+      await cleanup();
+      finish();
+    };
+
+    conn.on("error", (err) =>
+      settle(() => reject(subprocessSpawnFailed(scriptName, err)))
+    );
+
+    conn.on("close", async () => {
+      const text = await Bun.file(outputPath).text().catch(() => "");
+      await settle(() => {
+        if (cancelled) return reject(subprocessCancelled(scriptName, "SIGTERM"));
+        // No EXIT line => the child died mid-flight (crash/OOM/signal).
+        if (exitCode === null) return reject(subprocessKilled(scriptName, "UNKNOWN"));
+        if (exitCode !== 0) return reject(subprocessFailed(scriptName, exitCode));
+        if (!text) return reject(emptyResult(scriptName));
+        try {
+          resolve(JSON.parse(text));
+        } catch {
+          reject(malformedResult(scriptName));
+        }
+      });
+    });
+  });
+}

platform/src/bridge.ts

@@ -12,10 +12,8 @@ import {
   subprocessSpawnFailed,
 } from "./util/errors";
 import pkg from "../../package.json";
-
-// A line a service logged on purpose, as opposed to whatever else lands on a
-// stream. Only these are forwarded to the caller.
-const LOG_LINE = /^(INFO|DEBUG|ERROR|WARNING):/;
+import { USE_FORK_SERVER, runForked } from "./bridge.fork";
+import { LOG_LINE } from "./util/log-line";
 
 /**
   Run a python script
@@ -61,6 +59,10 @@ export const run = async (
     throw subprocessSpawnFailed(scriptName, error);
   }
 
+  if (USE_FORK_SERVER) {
+    return runForked(scriptName, port, inputPath, outputPath, onLog, onEvent, signal);
+  }
+
   return new Promise<JSON | null>((resolve, reject) => {
     const proc = spawn(
       "poetry",

platform/src/util/log-line.ts

@@ -0,0 +1,4 @@
+// A line a service logged on purpose, as opposed to whatever else lands on a
+// stream. Only these are forwarded to the caller. Shared by bridge.ts and
+// bridge.fork.ts.
+export const LOG_LINE = /^(INFO|DEBUG|ERROR|WARNING):/;

services/entry.fork.py

@@ -0,0 +1,137 @@
+"""Fork entry for the Python side (used when APOLLO_FORK_SERVER=true; see
+platform/src/bridge.fork.ts). Run as a script, never imported.
+
+Pays the heavy imports once, then forks a fresh child per request. Each child
+inherits the imports via copy-on-write and imports entry.py AFTER forking, so
+entry.py's own init runs in the child where it's fork-safe - entry.py is left
+untouched. The parent only imports and never inits, so it stays forkable.
+
+macOS: set OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES in your env (see .env.example),
+or a forked child that touches an Objective-C framework aborts. No-op on Linux.
+
+Protocol (NDJSON over a unix socket, one connection per request):
+  bridge -> master : {"service","input","output","port"}
+  child  -> bridge : PID:<pid>  then the service's stdout  then EXIT:<code>
+"""
+
+import argparse
+import importlib
+import json
+import os
+import signal
+import socket
+import sys
+import traceback
+from pathlib import Path
+
+import sentry_sdk  # for the per-child service tag + flush (see run_child)
+import yaml
+
+# Warm entry.py's heavy deps so forked children inherit them (copy-on-write).
+# See fork_preload.yaml for the list and why each module is there.
+PRELOAD_PATH = Path(__file__).parent / "fork_preload.yaml"
+
+
+def preload() -> int:
+    """Warm the shared deps, plus any opt-in services. Services are opt-in
+    (APOLLO_FORK_PRELOAD_SERVICES) because some do I/O at import and would hang
+    the parent; the child imports the rest on demand."""
+    with PRELOAD_PATH.open() as f:
+        modules = yaml.safe_load(f)["preload"]
+    for module in modules:
+        importlib.import_module(module)
+
+    names = [s.strip() for s in os.environ.get("APOLLO_FORK_PRELOAD_SERVICES", "").split(",") if s.strip()]
+    loaded = 0
+    for name in names:
+        try:
+            importlib.import_module(f"{name}.{name}")
+            loaded += 1
+        except Exception as e:
+            print(f"entry.fork: could not preload {name}: {e}", file=sys.stderr, flush=True)
+    return loaded
+
+
+def run_child(conn: socket.socket, req: dict) -> None:
+    """In the forked child: stream through the socket, run one job, exit."""
+    # Route stdout to the socket so the service's logging reaches the bridge.
+    # Leave stderr on fd 2 (the master's stderr pipe) so Python tracebacks still
+    # reach the server logs, the way the spawn path forwards them.
+    os.dup2(conn.fileno(), 1)
+    sys.stdout = os.fdopen(1, "w", buffering=1)
+
+    print(f"PID:{os.getpid()}")  # first, so the bridge can cancel this child
+
+    code = 0
+    try:
+        import entry  # post-fork: entry.py's init runs here (fork-safe), deps warm
+
+        sentry_sdk.set_tag("service", req["service"])  # entry.main() tags it; call() doesn't
+        entry.call(
+            service=req["service"],
+            input_path=req.get("input"),
+            output_path=req.get("output"),
+            apollo_port=req.get("port"),
+        )
+    except Exception:
+        traceback.print_exc()
+        code = 1
+    finally:
+        try:
+            sentry_sdk.flush(timeout=2)
+        except Exception:
+            pass
+        langfuse = getattr(sys.modules.get("entry"), "langfuse", None)
+        if langfuse is not None:
+            try:
+                langfuse.flush()
+            except Exception:
+                pass
+        print(f"EXIT:{code}")
+    os._exit(code)
+
+
+def read_request(conn: socket.socket) -> dict | None:
+    """Read one newline-terminated JSON request; None if there isn't a usable one."""
+    line = b""
+    while not line.endswith(b"\n"):
+        chunk = conn.recv(65536)
+        if not chunk:
+            break
+        line += chunk
+    try:
+        return json.loads(line) if line else None
+    except json.JSONDecodeError:
+        return None
+
+
+def serve(sock_path: str) -> None:
+    signal.signal(signal.SIGCHLD, signal.SIG_IGN)  # auto-reap children
+    count = preload()
+
+    if os.path.exists(sock_path):
+        os.unlink(sock_path)
+    server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    server.bind(sock_path)
+    server.listen(128)
+    print(f"READY services={count}", flush=True)  # handshake the bridge waits for
+
+    while True:
+        try:
+            conn, _ = server.accept()
+        except InterruptedError:
+            continue
+        req = read_request(conn)
+        if req is None:
+            conn.close()
+            continue
+        if os.fork() == 0:
+            server.close()  # the child must not hold the listening socket
+            run_child(conn, req)  # never returns
+        conn.close()  # parent drops its copy; the child owns the reply
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--socket", required=True)
+    serve(parser.parse_args().socket)

services/fork_preload.yaml

@@ -0,0 +1,22 @@
+# Modules the fork master (entry.fork.py) imports once at startup, before
+# accepting any requests. A forked child inherits these via copy-on-write, so
+# the cost of importing them is paid once per server lifetime instead of once
+# per request. A module missing from this list isn't an error - the child
+# just imports it itself on demand, which is slower but never wrong.
+
+preload:
+  # Mirrors entry.py's own top-level imports (Sentry init, Langfuse, otel).
+  - sentry_sdk
+  - anthropic
+  - langfuse
+  - langfuse_util
+  - util
+  - opentelemetry.instrumentation.anthropic
+  - opentelemetry.instrumentation.threading
+
+  # global_chat is ~95% of fork traffic; warm its RAG path (search_docsite)
+  # so the common case doesn't pay for these on the first request.
+  - pinecone
+  - openai
+  - langchain_openai
+  - langchain_pinecone