diff --git a/.env.example b/.env.example index 8f27f4bb..5e945f07 100644 --- a/.env.example +++ b/.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 + diff --git a/platform/src/bridge.fork.ts b/platform/src/bridge.fork.ts new file mode 100644 index 00000000..6427d7be --- /dev/null +++ b/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 | null = null; +let masterProc: ChildProcess | null = null; + +// Singleton; a failed start clears the slot so the next request retries. +function forkMaster(): Promise { + return (master ??= startMaster().catch((err) => { + master = null; + throw err; + })); +} + +function startMaster(): Promise { + 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((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 { + 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((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 | 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)); + } + }); + }); + }); +} diff --git a/platform/src/bridge.ts b/platform/src/bridge.ts index 240e9ade..bd79ad6e 100644 --- a/platform/src/bridge.ts +++ b/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((resolve, reject) => { const proc = spawn( "poetry", diff --git a/platform/src/util/log-line.ts b/platform/src/util/log-line.ts new file mode 100644 index 00000000..b3c69ff7 --- /dev/null +++ b/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):/; diff --git a/services/entry.fork.py b/services/entry.fork.py new file mode 100644 index 00000000..c83cbf71 --- /dev/null +++ b/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: then the service's stdout then EXIT: +""" + +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) diff --git a/services/fork_preload.yaml b/services/fork_preload.yaml new file mode 100644 index 00000000..7aac853a --- /dev/null +++ b/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