tryharness 0.7.22 → 0.7.25

package/lib/inference.mjs

@@ -0,0 +1,663 @@
+import fs from "node:fs";
+import net from "node:net";
+import os from "node:os";
+import path from "node:path";
+import { spawn, spawnSync } from "node:child_process";
+import {
+  MODELS,
+  modelById,
+  modelDir,
+  modelFile,
+  modelsPresent,
+  ensureLlama,
+  llamaServer,
+} from "./localmodels.mjs";
+import { ensureDir } from "./files.mjs";
+
+// The relay's local inference host. GGUF models (the search model) run in one llama.cpp server in
+// router mode, on the GPU when the machine has one that is faster, the CPU otherwise; the small ONNX classifiers run in
+// native ONNX Runtime inside the daemon, one thread each (measured: about 13x less CPU energy per item than the
+// WebAssembly build). Every request goes through one scheduler with three lanes: interactive (the person is waiting),
+// handoff (a brief or post is in flight) and background (indexing, the nightly learner). Background work runs at the
+// OS's background priority and waits while the machine is on low battery, in Low Power Mode or hot.
+
+const LANES = { interactive: 0, handoff: 1, background: 2 };
+// Models a sleeping llama.cpp instance keeps about 100 MB of; idle ONNX sessions are dropped outright.
+const SMALL_MACHINE = os.totalmem() <= 8.5 * 2 ** 30;
+// The smoke local-runtime shortens both to watch models go (HARNESS_ONNX_IDLE_MS, HARNESS_LLAMA_SLEEP_SECONDS).
+const ONNX_IDLE_MS =
+  Number(process.env.HARNESS_ONNX_IDLE_MS) ||
+  (SMALL_MACHINE ? 60_000 : 5 * 60_000);
+const SLEEP_IDLE_SECONDS =
+  Number(process.env.HARNESS_LLAMA_SLEEP_SECONDS) || (SMALL_MACHINE ? 30 : 120);
+const POWER_CHECK_MS = 60_000;
+
+const freePort = () =>
+  new Promise((resolve, reject) => {
+    const server = net.createServer();
+    server.unref();
+    server.on("error", reject);
+    server.listen(0, "127.0.0.1", () => {
+      const { port } = server.address();
+      server.close(() => resolve(port));
+    });
+  });
+
+/** Battery, Low Power Mode and heat, read from the OS without any permission. paused() is true when background work should wait. */
+export function createPower({
+  platform = process.platform,
+  run = spawnSync,
+} = {}) {
+  let state = {
+    onBattery: false,
+    percent: null,
+    lowPower: false,
+    hot: false,
+    at: 0,
+  };
+  const read = () => {
+    const next = {
+      onBattery: false,
+      percent: null,
+      lowPower: false,
+      hot: false,
+      at: Date.now(),
+    };
+    if (platform === "darwin") {
+      const batt =
+        run("pmset", ["-g", "batt"], { encoding: "utf8", timeout: 3_000 })
+          .stdout || "";
+      next.onBattery = /Battery Power/.test(batt);
+      next.percent = Number(batt.match(/(\d+)%/)?.[1] ?? NaN);
+      next.lowPower = /lowpowermode\s+1/.test(
+        run("pmset", ["-g"], { encoding: "utf8", timeout: 3_000 }).stdout || ""
+      );
+      const therm =
+        run("pmset", ["-g", "therm"], { encoding: "utf8", timeout: 3_000 })
+          .stdout || "";
+      next.hot =
+        /CPU_Speed_Limit\s*=\s*([0-9]+)/.test(therm) &&
+        Number(therm.match(/CPU_Speed_Limit\s*=\s*([0-9]+)/)[1]) < 80;
+    } else if (platform === "linux") {
+      try {
+        const battery = fs
+          .readdirSync("/sys/class/power_supply")
+          .find((name) => name.startsWith("BAT"));
+        if (battery) {
+          next.onBattery =
+            fs
+              .readFileSync(`/sys/class/power_supply/${battery}/status`, "utf8")
+              .trim() === "Discharging";
+          next.percent = Number(
+            fs.readFileSync(
+              `/sys/class/power_supply/${battery}/capacity`,
+              "utf8"
+            )
+          );
+        }
+      } catch {}
+    }
+    if (!Number.isFinite(next.percent)) next.percent = null;
+    state = next;
+    return state;
+  };
+  return {
+    state: () => (Date.now() - state.at > POWER_CHECK_MS ? read() : state),
+    paused() {
+      const now = this.state();
+      return (
+        now.lowPower ||
+        now.hot ||
+        (now.onBattery && now.percent !== null && now.percent < 20)
+      );
+    },
+    /** First index passes wait for the charger as well. */
+    onCharger() {
+      return !this.state().onBattery;
+    },
+  };
+}
+
+export function createInference({
+  root,
+  log = () => {},
+  power = createPower(),
+  list = MODELS,
+}) {
+  ensureDir(root);
+  const deviceFile = path.join(root, "device.json");
+  let server = null;
+  let starting = null;
+  const queues = { llama: [], onnx: [] };
+  const busy = { llama: false, onnx: false };
+  const sessions = new Map();
+
+  const device = () => {
+    try {
+      return JSON.parse(fs.readFileSync(deviceFile, "utf8"));
+    } catch {
+      return null;
+    }
+  };
+
+  const preset = (gpu) => {
+    const lines = [
+      "version = 1",
+      "[*]",
+      `sleep-idle-seconds = ${SLEEP_IDLE_SECONDS}`,
+      `n-gpu-layers = ${gpu ? 99 : 0}`,
+    ];
+    for (const entry of list.filter(
+      (model) =>
+        model.runtime === "llama" &&
+        modelsPresent(
+          root,
+          [model.id, ...(model.llama?.weights ? [model.llama.weights] : [])],
+          list
+        )
+    )) {
+      // A projector runs on another entry's weights (media on search's): its own preset, loaded only for media.
+      const weights = entry.llama?.weights
+        ? modelById(entry.llama.weights, list)
+        : null;
+      lines.push(
+        `[${entry.id}]`,
+        `model = ${modelFile(root, weights ?? entry)}`,
+        ...(weights ? [`mmproj = ${modelFile(root, entry)}`] : [])
+      );
+      for (const [key, value] of Object.entries(entry.llama ?? {}))
+        if (key !== "weights") lines.push(`${key} = ${value}`);
+    }
+    const file = path.join(root, "router.ini");
+    fs.writeFileSync(file, lines.join("\n") + "\n");
+    return file;
+  };
+
+  const launch = async (gpu) => {
+    const binary = llamaServer(root) ?? (await ensureLlama({ root, gpu, log }));
+    const port = await freePort();
+    const child = spawn(
+      binary,
+      [
+        "--models-preset",
+        preset(gpu),
+        "--models-max",
+        SMALL_MACHINE ? "1" : "2",
+        "--host",
+        "127.0.0.1",
+        "--port",
+        String(port),
+        "--no-webui",
+      ],
+      { cwd: path.dirname(binary), stdio: ["ignore", "ignore", "pipe"] }
+    );
+    let tail = "";
+    child.stderr.on("data", (chunk) => {
+      tail = (tail + chunk).slice(-4_000);
+    });
+    const base = `http://127.0.0.1:${port}`;
+    for (let i = 0; i < 200; i += 1) {
+      if (child.exitCode !== null)
+        throw new Error(
+          `the local model server stopped: ${tail.trim().split("\n").at(-1) || child.exitCode}`
+        );
+      try {
+        if ((await fetch(`${base}/health`)).ok)
+          return {
+            child,
+            base,
+            gpu,
+            media: modelsPresent(root, ["search", "media"], list),
+            marker: null,
+          };
+      } catch {}
+      await new Promise((resolve) => setTimeout(resolve, 100));
+    }
+    child.kill();
+    throw new Error("the local model server did not start");
+  };
+
+  /** One-time check on this machine: the GPU build is kept unless it fails or is slower than the CPU on the search model. */
+  const chooseDevice = async () => {
+    const known = device();
+    if (known) return known.gpu;
+    const timed = async (gpu) => {
+      const trial = await launch(gpu);
+      try {
+        const texts = Array.from(
+          { length: 8 },
+          (_, i) =>
+            `title: none | text: device check ${i} ${"lorem ipsum dolor sit amet ".repeat(40)}`
+        );
+        await post(trial.base, "/v1/embeddings", {
+          model: "search",
+          input: texts.slice(0, 1),
+        });
+        const started = performance.now();
+        await post(trial.base, "/v1/embeddings", {
+          model: "search",
+          input: texts,
+        });
+        return performance.now() - started;
+      } finally {
+        trial.child.kill();
+      }
+    };
+    let gpuMs = null;
+    try {
+      gpuMs = await timed(true);
+    } catch (error) {
+      log(`local models: GPU check failed (${error.message}); using the CPU`);
+    }
+    const cpuMs = await timed(false).catch(() => null);
+    const gpu = gpuMs !== null && (cpuMs === null || gpuMs <= cpuMs);
+    fs.writeFileSync(
+      deviceFile,
+      JSON.stringify(
+        {
+          gpu,
+          gpuMs,
+          cpuMs,
+          checkedAt: new Date().toISOString(),
+          platform: `${process.platform}-${process.arch}`,
+        },
+        null,
+        2
+      )
+    );
+    log(
+      `local models: running on the ${gpu ? "GPU" : "CPU"} (search check ${gpuMs?.toFixed(0) ?? "failed"} ms GPU, ${cpuMs?.toFixed(0) ?? "failed"} ms CPU)`
+    );
+    return gpu;
+  };
+
+  const ensureServer = () => {
+    if (server && server.child.exitCode === null)
+      return Promise.resolve(server);
+    starting ||= (async () => {
+      if (!modelsPresent(root, ["search"], list))
+        throw new Error("the search model is not downloaded yet");
+      server = await launch(await chooseDevice());
+      server.child.on("exit", () => {
+        server = null;
+      });
+      return server;
+    })().finally(() => {
+      starting = null;
+    });
+    return starting;
+  };
+
+  async function post(base, pathname, body) {
+    const response = await fetch(`${base}${pathname}`, {
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      body: JSON.stringify(body),
+    });
+    const text = await response.text();
+    if (!response.ok)
+      throw new Error(
+        `local model server ${pathname}: HTTP ${response.status} ${text.slice(0, 300)}`
+      );
+    return JSON.parse(text);
+  }
+
+  // macOS lets a process move itself or its children in and out of the background band; the llama.cpp children are
+  // moved only while a background job runs so interactive requests are never slowed.
+  const background = (on) => {
+    if (process.platform !== "darwin" || !server) return;
+    const children = (
+      spawnSync("pgrep", ["-P", String(server.child.pid)], { encoding: "utf8" })
+        .stdout || ""
+    )
+      .split("\n")
+      .filter(Boolean);
+    for (const pid of [String(server.child.pid), ...children])
+      spawnSync("taskpolicy", [on ? "-b" : "-B", "-p", pid]);
+  };
+
+  const pump = (backend) => {
+    if (busy[backend]) return;
+    const queue = queues[backend];
+    queue.sort((a, b) => a.rank - b.rank || a.seq - b.seq);
+    const index = queue.findIndex(
+      (job) => job.rank < LANES.background || !power.paused()
+    );
+    if (index < 0) {
+      if (queue.length)
+        setTimeout(() => pump(backend), POWER_CHECK_MS).unref?.();
+      return;
+    }
+    const [job] = queue.splice(index, 1);
+    busy[backend] = true;
+    const bg = job.rank === LANES.background && backend === "llama";
+    Promise.resolve()
+      .then(() => {
+        if (bg) background(true);
+        return job.run();
+      })
+      .then(job.resolve, job.reject)
+      .finally(() => {
+        if (bg) background(false);
+        busy[backend] = false;
+        pump(backend);
+      });
+  };
+  let seq = 0;
+  const schedule = (backend, lane, run) =>
+    new Promise((resolve, reject) => {
+      queues[backend].push({
+        rank: LANES[lane] ?? LANES.handoff,
+        seq: seq++,
+        run,
+        resolve,
+        reject,
+      });
+      pump(backend);
+    });
+
+  // Models on one vocabulary share its tokenizer while any of them is open.
+  const tokenizers = new Map();
+  const tokenizerOf = (Tokenizer, source) => {
+    const dir = modelDir(root, source);
+    let tokenizer = tokenizers.get(dir)?.deref();
+    if (!tokenizer) {
+      const tokenizerPath = source.files.find((file) =>
+        file.path.endsWith("tokenizer.json")
+      ).path;
+      const configPath = source.files.find((file) =>
+        file.path.endsWith("tokenizer_config.json")
+      )?.path;
+      tokenizer = new Tokenizer(
+        JSON.parse(fs.readFileSync(path.join(dir, tokenizerPath), "utf8")),
+        configPath
+          ? JSON.parse(fs.readFileSync(path.join(dir, configPath), "utf8"))
+          : {}
+      );
+      tokenizers.set(dir, new WeakRef(tokenizer));
+    }
+    return tokenizer;
+  };
+
+  const session = async (id) => {
+    const entry = modelById(id, list);
+    if (!entry || entry.runtime !== "onnx")
+      throw new Error(`no ONNX model ${id}`);
+    let held = sessions.get(id);
+    if (!held) {
+      held = (async () => {
+        const source = entry.tokenizer
+          ? modelById(entry.tokenizer.from, list)
+          : entry;
+        if (!modelsPresent(root, [...new Set([id, source.id])], list))
+          throw new Error(`${entry.name} is not downloaded yet`);
+        const [{ Tokenizer }, ort] = await Promise.all([
+          import("@huggingface/tokenizers"),
+          import("onnxruntime-node"),
+        ]);
+        const onnx = await ort.InferenceSession.create(modelFile(root, entry), {
+          executionProviders: ["cpu"],
+          intraOpNumThreads: 1,
+          interOpNumThreads: 1,
+          logSeverityLevel: 3,
+        });
+        return {
+          ort,
+          tokenizer: tokenizerOf(Tokenizer, source),
+          prepare: entry.tokenizer?.prepare ?? ((text) => text),
+          onnx,
+          timer: null,
+        };
+      })();
+      sessions.set(id, held);
+      held.catch(() => sessions.delete(id));
+    }
+    const ready = await held;
+    clearTimeout(ready.timer);
+    ready.timer = setTimeout(() => {
+      sessions.delete(id);
+      ready.onnx.release?.();
+    }, ONNX_IDLE_MS);
+    ready.timer.unref?.();
+    return ready;
+  };
+
+  // Models made of several ONNX graphs (OCR's detector and recogniser, speech's encoder and decoder) open each graph
+  // on first use and drop all of them together once idle.
+  const graphs = new Map();
+  const graphsOf = async (id) => {
+    const entry = modelById(id, list);
+    if (!entry || entry.runtime !== "onnx")
+      throw new Error(`no ONNX model ${id}`);
+    if (!modelsPresent(root, [id], list))
+      throw new Error(`${entry.name} is not downloaded yet`);
+    let held = graphs.get(id);
+    if (!held)
+      graphs.set(
+        id,
+        (held = {
+          ort: await import("onnxruntime-node"),
+          open: new Map(),
+          timer: null,
+          threads: entry.threads ?? 1,
+        })
+      );
+    clearTimeout(held.timer);
+    held.timer = setTimeout(() => {
+      graphs.delete(id);
+      for (const graph of held.open.values())
+        graph.then((onnx) => onnx.release?.()).catch(() => {});
+    }, ONNX_IDLE_MS);
+    held.timer.unref?.();
+    const graph = (file) => {
+      if (!held.open.has(file)) {
+        const opened = held.ort.InferenceSession.create(
+          path.join(modelDir(root, entry), ...file.split("/")),
+          {
+            executionProviders: ["cpu"],
+            intraOpNumThreads: held.threads,
+            interOpNumThreads: 1,
+            logSeverityLevel: 3,
+          }
+        );
+        held.open.set(file, opened);
+        opened.catch(() => held.open.delete(file));
+      }
+      return held.open.get(file);
+    };
+    return { ort: held.ort, graph, dir: modelDir(root, entry) };
+  };
+
+  /** Runs an ONNX model on each encoded input; returns the first output's leading row (its logits). */
+  const logits = async (id, pairs, maxLength = 512) => {
+    const { ort, tokenizer, prepare, onnx } = await session(id);
+    const rows = [];
+    for (const [text, pair] of pairs) {
+      const encoded = tokenizer.encode(
+        prepare(text),
+        pair === undefined
+          ? { add_special_tokens: true }
+          : { text_pair: prepare(pair), add_special_tokens: true }
+      );
+      const ids = encoded.ids.slice(0, maxLength);
+      const feeds = {};
+      for (const name of onnx.inputNames) {
+        if (name === "input_ids")
+          feeds[name] = new ort.Tensor(
+            "int64",
+            BigInt64Array.from(ids, BigInt),
+            [1, ids.length]
+          );
+        else if (name === "attention_mask")
+          feeds[name] = new ort.Tensor(
+            "int64",
+            new BigInt64Array(ids.length).fill(1n),
+            [1, ids.length]
+          );
+        else if (name === "token_type_ids")
+          feeds[name] = new ort.Tensor(
+            "int64",
+            BigInt64Array.from(
+              (encoded.type_ids ?? ids.map(() => 0)).slice(0, ids.length),
+              BigInt
+            ),
+            [1, ids.length]
+          );
+      }
+      const out = await onnx.run(feeds);
+      rows.push(Array.from(out[onnx.outputNames[0]].data, Number));
+    }
+    return rows;
+  };
+
+  return {
+    device,
+    power,
+    /** Embeds texts with the search model. kind: "query" | "code-query" | "document" (title is a path or "none"). */
+    embed: (items, { lane = "handoff" } = {}) =>
+      schedule("llama", lane, async () => {
+        const { base } = await ensureServer();
+        const input = items.map((item) =>
+          typeof item === "string"
+            ? item
+            : item.kind === "document"
+              ? `title: ${item.title || "none"} | text: ${item.text}`
+              : `task: ${item.kind === "code-query" ? "code retrieval" : "search result"} | query: ${item.text}`
+        );
+        // A text too long for the model's window is retried alone, halved until it fits.
+        const one = async (text) => {
+          for (
+            let cut = text;
+            ;
+            cut = cut.slice(0, Math.floor(cut.length / 2))
+          ) {
+            try {
+              return (
+                await post(base, "/v1/embeddings", {
+                  model: "search",
+                  input: [cut],
+                })
+              ).data[0];
+            } catch (error) {
+              if (cut.length < 200) throw error;
+            }
+          }
+        };
+        const out = [];
+        for (let i = 0; i < input.length; i += 16) {
+          const slice = input.slice(i, i + 16);
+          const data = await post(base, "/v1/embeddings", {
+            model: "search",
+            input: slice,
+          })
+            .then((reply) => reply.data.sort((a, b) => a.index - b.index))
+            .catch(async () => {
+              const rows = [];
+              for (const text of slice) rows.push(await one(text));
+              return rows;
+            });
+          for (const row of data) {
+            const vector = Float32Array.from(row.embedding);
+            let norm = 0;
+            for (const value of vector) norm += value * value;
+            norm = Math.sqrt(norm) || 1;
+            for (let j = 0; j < vector.length; j += 1) vector[j] /= norm;
+            out.push(vector);
+          }
+        }
+        return out;
+      }),
+    /**
+     * Embeds recordings and images with the search model's projector, in the same space as typed questions. items:
+     * { title, data (a Buffer of WAV, MP3, FLAC, PNG or JPEG) }. A server started before the projector came down is
+     * restarted once with it.
+     */
+    embedMedia: (items, { lane = "background" } = {}) =>
+      schedule("llama", lane, async () => {
+        let ready = await ensureServer();
+        if (!ready.media) {
+          ready.child.kill();
+          server = null;
+          ready = await ensureServer();
+        }
+        if (!ready.media)
+          throw new Error(
+            "the image and audio part of the search model is not downloaded yet"
+          );
+        // The marker that stands for a recording or an image in the text is the server's own, read once it has loaded.
+        if (!ready.marker) {
+          await post(ready.base, "/v1/embeddings", {
+            model: "media",
+            input: ["warm up"],
+          });
+          ready.marker = (
+            await (await fetch(`${ready.base}/props?model=media`)).json()
+          ).media_marker;
+        }
+        const out = [];
+        for (const item of items) {
+          const reply = await post(ready.base, "/v1/embeddings", {
+            model: "media",
+            input: [
+              {
+                prompt_string: `title: ${item.title || "none"} | text: ${ready.marker}`,
+                multimodal_data: [Buffer.from(item.data).toString("base64")],
+              },
+            ],
+          });
+          const vector = Float32Array.from(reply.data[0].embedding);
+          let norm = 0;
+          for (const value of vector) norm += value * value;
+          norm = Math.sqrt(norm) || 1;
+          out.push(vector.map((value) => value / norm));
+        }
+        return out;
+      }),
+    /** Raw logits rows from an ONNX sequence classifier, for models calibrated on logit margins. */
+    logits: (id, texts, { lane = "handoff", maxLength } = {}) =>
+      schedule("onnx", lane, () =>
+        logits(
+          id,
+          texts.map((text) => [text]),
+          maxLength
+        )
+      ),
+    /** Runs work against a multi-graph ONNX model in the scheduler: run({ ort, graph(file), dir }) opens graphs by file path. */
+    onnx: (id, run, { lane = "handoff" } = {}) =>
+      schedule("onnx", lane, async () => run(await graphsOf(id))),
+    running: () => Boolean(server && server.child.exitCode === null),
+    /** What is in memory now: the router's models and their state, and the ONNX models open (for `tryharness local`). */
+    loaded: async () => {
+      let llama = [];
+      if (server && server.child.exitCode === null) {
+        try {
+          llama = (
+            (await (await fetch(`${server.base}/models`)).json()).data ?? []
+          ).map((model) => ({
+            id: model.id,
+            state: model.status?.value ?? null,
+          }));
+        } catch {}
+      }
+      return { llama, onnx: [...sessions.keys(), ...graphs.keys()] };
+    },
+    stop() {
+      if (server?.child.exitCode === null) server.child.kill();
+      server = null;
+      for (const held of sessions.values())
+        held
+          .then((ready) => {
+            clearTimeout(ready.timer);
+            ready.onnx.release?.();
+          })
+          .catch(() => {});
+      sessions.clear();
+      for (const held of graphs.values()) {
+        clearTimeout(held.timer);
+        for (const graph of held.open.values())
+          graph.then((onnx) => onnx.release?.()).catch(() => {});
+      }
+      graphs.clear();
+    },
+  };
+}