alexion
6081e3e325
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: 71023aec1546d6d677a65669c5665c0927ce0dd0
305 lines
12 KiB
TypeScript
305 lines
12 KiB
TypeScript
/**
|
|
* Nested image storage benchmark. One model-issued `script` call, as a code mode call would, loops `count` times: a
|
|
* nested `read_image` call returns a unique base64 image of `bytes` characters, then a nested `classify` call reads it.
|
|
* The `inline` variant passes the image data in `classify`'s arguments; the `reference` variant passes the `taskId` of
|
|
* the `read_image` result, whose value is the image, and `classify` resolves it from its caller's `NestedCallDoc`.
|
|
* Reports commits, wall time, logical bytes written (JSON of the storage writes) by write type, the stored size before
|
|
* `script` returns, after the run settles, after close, and after reopening, the heap (peak during the run, sampled
|
|
* every 5 ms and at each commit; after a GC once the run settles, after `waitForIdle()`, and after reopening; all over
|
|
* the heap after a GC before opening), the number of nested task records, and the size of the largest one.
|
|
*
|
|
* node --conditions=source --experimental-strip-types --expose-gc test/nested-images-bench.ts [count]
|
|
*/
|
|
import { randomBytes } from "node:crypto";
|
|
import { mkdtemp, readdir, rm, stat } from "node:fs/promises";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
|
|
import type { ImageContent } from "@earendil-works/pi-ai";
|
|
import { createModels } from "@earendil-works/pi-ai/models";
|
|
import { fauxAssistantMessage, fauxProvider, fauxText, fauxToolCall } from "@earendil-works/pi-ai/providers/faux";
|
|
import { Type } from "typebox";
|
|
import {
|
|
createRegistry,
|
|
defineExtension,
|
|
defineTool,
|
|
Harness,
|
|
type JsonObject,
|
|
NestedCallDoc,
|
|
type Storage,
|
|
type StorageWrite,
|
|
type TaskId,
|
|
ToolResultEntry,
|
|
type ToolTaskInput,
|
|
} from "../src/index.ts";
|
|
import { openNodeJsonlStorage } from "../src/storage/jsonl/node.ts";
|
|
import { openNodeSqliteStorage } from "../src/storage/sqlite/node.ts";
|
|
|
|
type Backend = "sqlite" | "jsonl";
|
|
type Variant = "inline" | "reference";
|
|
type Scenario = { backend: Backend; variant: Variant; count: number; bytes: number };
|
|
/** Hooks the tools call: a heap sample, and the stored size measured before `script` returns. */
|
|
type Probe = { sample: () => void; beforeReturn: () => Promise<void> };
|
|
type Written = Record<StorageWrite["type"], number>;
|
|
type Result = {
|
|
commits: number;
|
|
ms: number;
|
|
written: Written;
|
|
/** Before `script` returns, while its nested results are live. */
|
|
storedLive: number;
|
|
storedRun: number;
|
|
storedClosed: number;
|
|
storedReopened: number;
|
|
/** Over the heap after a GC before opening. */
|
|
peakHeap: number;
|
|
settledHeap: number;
|
|
idleHeap: number;
|
|
reopenedHeap: number;
|
|
nestedTasks: number;
|
|
largestTask: number;
|
|
};
|
|
|
|
const gc = (globalThis as { gc?: () => void }).gc;
|
|
if (gc === undefined) throw new Error("Run with --expose-gc");
|
|
|
|
function heapAfterGc(): number {
|
|
gc!();
|
|
return process.memoryUsage().heapUsed;
|
|
}
|
|
|
|
async function directorySize(path: string): Promise<number> {
|
|
let total = 0;
|
|
for (const entry of await readdir(path, { withFileTypes: true })) {
|
|
const full = join(path, entry.name);
|
|
total += entry.isDirectory() ? await directorySize(full) : (await stat(full)).size;
|
|
}
|
|
return total;
|
|
}
|
|
|
|
async function open(backend: Backend, directory: string): Promise<Storage> {
|
|
return backend === "sqlite"
|
|
? openNodeSqliteStorage(join(directory, "session.sqlite"))
|
|
: openNodeJsonlStorage(directory, context);
|
|
}
|
|
|
|
function noWrites(): Written {
|
|
return {
|
|
conversation: 0,
|
|
entry: 0,
|
|
task: 0,
|
|
submission: 0,
|
|
"document.create": 0,
|
|
"document.copy": 0,
|
|
"document.change": 0,
|
|
"document.retire": 0,
|
|
};
|
|
}
|
|
|
|
const ImageArgument = Type.Object({
|
|
type: Type.Literal("image"),
|
|
mimeType: Type.String(),
|
|
data: Type.Optional(Type.String()),
|
|
source: Type.Optional(Type.Object({ taskId: Type.Number() })),
|
|
});
|
|
|
|
function registryFor(scenario: Scenario, probe: Probe) {
|
|
const registry = createRegistry();
|
|
const readImage = defineTool({
|
|
name: "read_image",
|
|
description: "Read an image",
|
|
parameters: Type.Object({ n: Type.Number() }),
|
|
// One image, so programs get it as an `ImageContent`. Unique per call: base64 of random bytes.
|
|
execute: async () => ({
|
|
output: [
|
|
{ type: "image", data: randomBytes((scenario.bytes / 4) * 3).toString("base64"), mimeType: "image/png" },
|
|
],
|
|
}),
|
|
});
|
|
const classify = defineTool({
|
|
name: "classify",
|
|
description: "Classify images",
|
|
parameters: Type.Object({ images: Type.Array(ImageArgument) }),
|
|
execute: async (args, api, callContext) => {
|
|
for (const image of args.images) {
|
|
let data = image.data;
|
|
if (image.source !== undefined) {
|
|
// The caller is the nested call's `parent`; its `NestedCallDoc` holds the referenced result.
|
|
const input = (await api.getTask(api.taskId, callContext))?.input as ToolTaskInput | undefined;
|
|
if (input?.kind !== "nested") throw new Error("classify resolves references only as a nested call");
|
|
// Results are keyed by call key, which the referenced call's input carries.
|
|
const source = (await api.getTask(image.source.taskId as TaskId, callContext))?.input as ToolTaskInput;
|
|
if (source.kind !== "nested") throw new Error("A reference names a nested call");
|
|
const stored = await api.snapshot(NestedCallDoc, input.parent, source.key, callContext);
|
|
const block = stored?.result?.structuredOutput as ImageContent | undefined;
|
|
data = block?.type === "image" ? block.data : undefined;
|
|
}
|
|
if (data?.length !== scenario.bytes) throw new Error(`Image has ${data?.length} characters`);
|
|
}
|
|
return { output: [{ type: "text", text: "ok" }] };
|
|
},
|
|
});
|
|
const script = defineTool({
|
|
name: "script",
|
|
description: "Read and classify images",
|
|
parameters: Type.Object({ count: Type.Number() }),
|
|
callers: ["model"],
|
|
execute: async (args, api, callContext) => {
|
|
for (let n = 0; n < args.count; n++) {
|
|
const key = `read${n}`;
|
|
const read = await api.executeTool("read_image", { n }, callContext, { key });
|
|
const image = read.structuredOutput as unknown as ImageContent;
|
|
const argument: JsonObject =
|
|
scenario.variant === "inline"
|
|
? { type: "image", data: image.data, mimeType: image.mimeType }
|
|
: { type: "image", source: { taskId: read.taskId }, mimeType: image.mimeType };
|
|
const classified = await api.executeTool("classify", { images: [argument] }, callContext, {
|
|
key: `classify${n}`,
|
|
});
|
|
if (classified.isError) throw new Error(`classify ${n} failed: ${JSON.stringify(classified.diagnostics)}`);
|
|
probe.sample();
|
|
}
|
|
await probe.beforeReturn();
|
|
return { output: [{ type: "text", text: `classified ${args.count}` }] };
|
|
},
|
|
});
|
|
registry.install(defineExtension({ name: "bench", tools: [readImage, classify, script] }));
|
|
return registry;
|
|
}
|
|
|
|
async function runScenario(scenario: Scenario): Promise<Result> {
|
|
const directory = await mkdtemp(join(tmpdir(), "pi-durable-images-bench-"));
|
|
try {
|
|
const faux = fauxProvider();
|
|
const models = createModels();
|
|
models.setProvider(faux.provider);
|
|
let peakHeap = 0;
|
|
let storedLive = 0;
|
|
const sample = () => {
|
|
peakHeap = Math.max(peakHeap, process.memoryUsage().heapUsed);
|
|
};
|
|
const registry = registryFor(scenario, {
|
|
sample,
|
|
beforeReturn: async () => {
|
|
storedLive = await directorySize(directory);
|
|
},
|
|
});
|
|
const baseHeap = heapAfterGc();
|
|
|
|
const storage = await open(scenario.backend, directory);
|
|
let commits = 0;
|
|
const written = noWrites();
|
|
const commit = storage.commit.bind(storage);
|
|
storage.commit = (writes, commitContext) => {
|
|
commits++;
|
|
// SQLite commits synchronously, which can starve the timer; sample here too.
|
|
sample();
|
|
for (const write of writes) written[write.type] += JSON.stringify(write).length;
|
|
return commit(writes, commitContext);
|
|
};
|
|
const harness = await Harness.open(storage, { models, registry }, context);
|
|
const root = await harness.root(context, { agent: { model: { provider: "faux", modelId: "faux-1" } } });
|
|
faux.setResponses([
|
|
fauxAssistantMessage([fauxToolCall("script", { count: scenario.count }, { id: "s" })], {
|
|
stopReason: "toolUse",
|
|
}),
|
|
fauxAssistantMessage([fauxText("done")]),
|
|
]);
|
|
const before = { commits, written: { ...written } };
|
|
peakHeap = 0;
|
|
const sampler = setInterval(sample, 5);
|
|
const started = performance.now();
|
|
const settled = await (await root.submit({ type: "input", content: "run" }, context)).wait(context);
|
|
if (settled.status !== "done") throw new Error(`run: ${settled.status}`);
|
|
// `done` only says the model answered; the script's result says whether it ran through.
|
|
const latest = (await root.entries({}, 10, undefined, context)).items;
|
|
const result = latest.find((entry) => ToolResultEntry.is(entry))?.model?.[0];
|
|
if (result?.role !== "toolResult" || result.isError) throw new Error("run: the script failed");
|
|
const ms = Math.round(performance.now() - started);
|
|
clearInterval(sampler);
|
|
const settledHeap = heapAfterGc() - baseHeap;
|
|
await harness.waitForIdle(context);
|
|
const idleHeap = heapAfterGc() - baseHeap;
|
|
const delta = noWrites();
|
|
for (const type of Object.keys(delta) as StorageWrite["type"][]) {
|
|
delta[type] = written[type] - before.written[type];
|
|
}
|
|
const storedRun = await directorySize(directory);
|
|
await harness.close(context);
|
|
const storedClosed = await directorySize(directory);
|
|
|
|
const reopenedStorage = await open(scenario.backend, directory);
|
|
const reopened = await Harness.open(reopenedStorage, { models, registry }, context);
|
|
const reopenedHeap = heapAfterGc() - baseHeap;
|
|
const storedReopened = await directorySize(directory);
|
|
let nestedTasks = 0;
|
|
let largestTask = 0;
|
|
let cursor: Parameters<Storage["scanTasks"]>[2];
|
|
do {
|
|
const page = await reopenedStorage.scanTasks({ kind: "pi.tool" }, 50, cursor, context);
|
|
for (const record of page.items) {
|
|
if ((record.input as ToolTaskInput).kind === "nested") nestedTasks++;
|
|
largestTask = Math.max(largestTask, JSON.stringify(record).length);
|
|
}
|
|
cursor = page.next;
|
|
} while (cursor !== undefined);
|
|
await reopened.close(context);
|
|
return {
|
|
commits: commits - before.commits,
|
|
ms,
|
|
written: delta,
|
|
storedLive,
|
|
storedRun,
|
|
storedClosed,
|
|
storedReopened,
|
|
peakHeap: peakHeap - baseHeap,
|
|
settledHeap,
|
|
idleHeap,
|
|
reopenedHeap,
|
|
nestedTasks,
|
|
largestTask,
|
|
};
|
|
} finally {
|
|
await rm(directory, { recursive: true, force: true });
|
|
}
|
|
}
|
|
|
|
function mib(bytes: number): string {
|
|
return `${(bytes / 1024 / 1024).toFixed(1)} MiB`;
|
|
}
|
|
|
|
function kib(bytes: number): string {
|
|
return `${(bytes / 1024).toFixed(1)} KiB`;
|
|
}
|
|
|
|
async function main(): Promise<void> {
|
|
const counts = process.argv[2] === undefined ? [50, 200] : [Number(process.argv[2])];
|
|
const bytes = 1024 * 1024;
|
|
console.log(`Images: ${bytes / 1024} KiB base64 each, read then classified sequentially in one script call`);
|
|
console.log(
|
|
"backend | variant | count | commits, ms | written: all (task / doc create / doc change / entry) | stored: live / run / closed / reopened | heap: peak; settled / idle / reopened | nested tasks, largest task",
|
|
);
|
|
for (const count of counts) {
|
|
for (const backend of ["sqlite", "jsonl"] as const) {
|
|
for (const variant of ["inline", "reference"] as const) {
|
|
const result = await runScenario({ backend, variant, count, bytes });
|
|
const all = Object.values(result.written).reduce((sum, value) => sum + value, 0);
|
|
const written = result.written;
|
|
console.log(
|
|
[
|
|
backend.padEnd(7),
|
|
variant.padEnd(9),
|
|
String(count).padEnd(5),
|
|
`${result.commits}, ${result.ms} ms`,
|
|
`${mib(all)} (${mib(written.task)} / ${mib(written["document.create"])} / ${mib(written["document.change"])} / ${mib(written.entry)})`,
|
|
`${mib(result.storedLive)} / ${mib(result.storedRun)} / ${mib(result.storedClosed)} / ${mib(result.storedReopened)}`,
|
|
`${mib(result.peakHeap)}; ${mib(result.settledHeap)} / ${mib(result.idleHeap)} / ${mib(result.reopenedHeap)}`,
|
|
`${result.nestedTasks}, ${kib(result.largestTask)}`,
|
|
].join(" | "),
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
await main();
|