chore(vendor): import pi 1.3.0
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
This commit is contained in:
253
third_party/pi/packages/durable/test/nested-tools-bench.ts
vendored
Normal file
253
third_party/pi/packages/durable/test/nested-tools-bench.ts
vendored
Normal file
@@ -0,0 +1,253 @@
|
||||
/**
|
||||
* Nested tool call storage benchmark. Per run, one model-issued tool makes `calls` parallel nested calls of `blob`,
|
||||
* each streaming `bytes` of output in 8 chunks `gapMs` apart; at depth 2 it makes them through `calls / 10` parallel
|
||||
* `fan` calls of 10 each. With `progress: false` the nested calls commit no running output. Reports commits, logical
|
||||
* bytes written (JSON of the storage writes) by write type, the stored size after each run, and the heap: its peak during a run, sampled every 5 ms, and what stays after a
|
||||
* run and a GC, both over the heap after a GC before the first run. A 2 ms gap is a fast tool such as `read`: it ends
|
||||
* before a second progress commit. A 50 ms gap streams for 400 ms, as a short `bash` command does, and commits progress
|
||||
* every 100 ms.
|
||||
*
|
||||
* node --conditions=source --experimental-strip-types --expose-gc test/nested-tools-bench.ts [streaming|nesting]
|
||||
*/
|
||||
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 { createModels } from "@earendil-works/pi-ai/models";
|
||||
import { fauxAssistantMessage, fauxProvider, fauxText, fauxToolCall } from "@earendil-works/pi-ai/providers/faux";
|
||||
import { Type } from "typebox";
|
||||
import {
|
||||
type Conversation,
|
||||
createRegistry,
|
||||
defineExtension,
|
||||
defineTool,
|
||||
Harness,
|
||||
type Storage,
|
||||
type StorageWrite,
|
||||
ToolResultEntry,
|
||||
} 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 Scenario = {
|
||||
backend: Backend;
|
||||
calls: number;
|
||||
depth: 1 | 2;
|
||||
bytes: number;
|
||||
gapMs: number;
|
||||
/** Whether nested calls stream their running output to `pi.live`. */
|
||||
progress: boolean;
|
||||
runs: number;
|
||||
};
|
||||
type Written = Record<StorageWrite["type"], number>;
|
||||
type Run = {
|
||||
commits: number;
|
||||
written: Written;
|
||||
storedBytes: number;
|
||||
ms: number;
|
||||
/** Over the heap after a GC before the first run. */
|
||||
peakHeap: number;
|
||||
heldHeap: 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;
|
||||
}
|
||||
|
||||
const CHUNKS = 8;
|
||||
/** Calls each inner `fan` makes at depth 2. */
|
||||
const FAN_OUT = 10;
|
||||
|
||||
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,
|
||||
};
|
||||
}
|
||||
|
||||
async function runScenario(scenario: Scenario): Promise<{ perRun: Run[]; closedBytes: number }> {
|
||||
const directory = await mkdtemp(join(tmpdir(), "pi-durable-nested-bench-"));
|
||||
try {
|
||||
const faux = fauxProvider();
|
||||
const models = createModels();
|
||||
models.setProvider(faux.provider);
|
||||
const registry = createRegistry();
|
||||
const chunk = `${"x".repeat(scenario.bytes / CHUNKS - 1)}\n`;
|
||||
// Streams its output and returns none, so its result is the retained output, as bash's is.
|
||||
const blob = defineTool({
|
||||
name: "blob",
|
||||
description: "Stream a blob",
|
||||
parameters: Type.Object({ n: Type.Number() }),
|
||||
execute: async (_args, api) => {
|
||||
for (let n = 0; n < CHUNKS; n++) {
|
||||
api.output(chunk);
|
||||
await new Promise((resolve) => setTimeout(resolve, scenario.gapMs));
|
||||
}
|
||||
return {};
|
||||
},
|
||||
});
|
||||
const options = scenario.progress ? {} : { progress: false as const };
|
||||
/** A tool that makes `count` parallel nested calls of `target`, which get `{ n, count: FAN_OUT }`. */
|
||||
const fanOut = (name: string, target: string) =>
|
||||
defineTool({
|
||||
name,
|
||||
description: `Call ${target} in parallel`,
|
||||
parameters: Type.Object({ count: Type.Number() }),
|
||||
execute: async (args, api, callContext) => {
|
||||
const calls = Array.from({ length: args.count }, (_, n) =>
|
||||
api.executeTool(target, { n, count: FAN_OUT }, callContext, options),
|
||||
);
|
||||
const failed = (await Promise.all(calls)).filter((result) => result.isError).length;
|
||||
return { output: [{ type: "text", text: `${failed} failed` }], isError: failed > 0 };
|
||||
},
|
||||
});
|
||||
registry.install(defineExtension({ name: "bench", tools: [blob, fanOut("fan", "blob"), fanOut("fan2", "fan")] }));
|
||||
|
||||
const storage = await open(scenario.backend, directory);
|
||||
let commits = 0;
|
||||
const written = noWrites();
|
||||
const commit = storage.commit.bind(storage);
|
||||
storage.commit = (writes, commitContext) => {
|
||||
commits++;
|
||||
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" } } });
|
||||
const perRun: Run[] = [];
|
||||
const top = scenario.depth === 1 ? scenario.calls : scenario.calls / FAN_OUT;
|
||||
const baseHeap = heapAfterGc();
|
||||
for (let run = 0; run < scenario.runs; run++) {
|
||||
faux.setResponses([
|
||||
fauxAssistantMessage(
|
||||
[fauxToolCall(scenario.depth === 1 ? "fan" : "fan2", { count: top }, { id: `r${run}` })],
|
||||
{
|
||||
stopReason: "toolUse",
|
||||
},
|
||||
),
|
||||
fauxAssistantMessage([fauxText("done")]),
|
||||
]);
|
||||
const before = { commits, written: { ...written } };
|
||||
let peakHeap = 0;
|
||||
const sampler = setInterval(() => {
|
||||
peakHeap = Math.max(peakHeap, process.memoryUsage().heapUsed);
|
||||
}, 5);
|
||||
const started = performance.now();
|
||||
const settled = await (await root.submit({ type: "input", content: `run ${run}` }, context)).wait(context);
|
||||
if (settled.status !== "done") throw new Error(`run ${run}: ${settled.status}`);
|
||||
await assertToolsSucceeded(root, run);
|
||||
await harness.waitForIdle(context);
|
||||
const ms = Math.round(performance.now() - started);
|
||||
clearInterval(sampler);
|
||||
const delta = noWrites();
|
||||
for (const type of Object.keys(delta) as StorageWrite["type"][]) {
|
||||
delta[type] = written[type] - before.written[type];
|
||||
}
|
||||
perRun.push({
|
||||
commits: commits - before.commits,
|
||||
written: delta,
|
||||
storedBytes: await directorySize(directory),
|
||||
ms,
|
||||
peakHeap: peakHeap - baseHeap,
|
||||
heldHeap: heapAfterGc() - baseHeap,
|
||||
});
|
||||
}
|
||||
await harness.close(context);
|
||||
return { perRun, closedBytes: await directorySize(directory) };
|
||||
} finally {
|
||||
await rm(directory, { recursive: true, force: true });
|
||||
}
|
||||
}
|
||||
|
||||
/** A run that answers `done` may still have failed tools; the model-issued call's result says whether any did. */
|
||||
async function assertToolsSucceeded(root: Conversation, run: number): Promise<void> {
|
||||
const page = await root.entries({}, 10, undefined, context);
|
||||
const result = page.items.find((entry) => ToolResultEntry.is(entry))?.model?.[0];
|
||||
if (result?.role !== "toolResult" || result.isError) throw new Error(`run ${run}: a nested call failed`);
|
||||
}
|
||||
|
||||
function mib(bytes: number): string {
|
||||
return `${(bytes / 1024 / 1024).toFixed(2)} MiB`;
|
||||
}
|
||||
|
||||
async function report(title: string, scenarios: Scenario[]): Promise<void> {
|
||||
console.log(`\n${title}`);
|
||||
console.log(
|
||||
"backend | calls | depth | output | gap | stream | run 1: commits, ms, written (all; tasks / doc changes / doc creates) | stored after run 1 / 2 / 3 | after close | heap: run 1 peak; held after run 1 / 2 / 3",
|
||||
);
|
||||
for (const scenario of scenarios) {
|
||||
const { perRun, closedBytes } = await runScenario(scenario);
|
||||
const first = perRun[0]!;
|
||||
const all = Object.values(first.written).reduce((sum, bytes) => sum + bytes, 0);
|
||||
console.log(
|
||||
[
|
||||
scenario.backend.padEnd(6),
|
||||
String(scenario.calls).padEnd(5),
|
||||
String(scenario.depth).padEnd(5),
|
||||
`${scenario.bytes / 1024} KiB`.padEnd(6),
|
||||
`${scenario.gapMs} ms`.padEnd(5),
|
||||
(scenario.progress ? "yes" : "no").padEnd(6),
|
||||
`${first.commits}, ${first.ms} ms, ${mib(all)} (${mib(first.written.task)} / ${mib(first.written["document.change"])} / ${mib(first.written["document.create"])})`,
|
||||
perRun.map((run) => mib(run.storedBytes)).join(" / "),
|
||||
mib(closedBytes),
|
||||
`${mib(first.peakHeap)}; ${perRun.map((run) => mib(run.heldHeap)).join(" / ")}`,
|
||||
].join(" | "),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async function main(): Promise<void> {
|
||||
const section = process.argv[2];
|
||||
if (section === undefined || section === "streaming") {
|
||||
const scenarios: Scenario[] = [];
|
||||
for (const backend of ["sqlite", "jsonl"] as const) {
|
||||
for (const bytes of [1024, 16 * 1024]) {
|
||||
for (const gapMs of [2, 50]) {
|
||||
for (const progress of [true, false]) {
|
||||
scenarios.push({ backend, calls: 200, depth: 1, bytes, gapMs, progress, runs: 3 });
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
await report("Streaming: 200 calls, with and without progress", scenarios);
|
||||
}
|
||||
if (section === undefined || section === "nesting") {
|
||||
const scenarios: Scenario[] = [];
|
||||
for (const backend of ["sqlite", "jsonl"] as const) {
|
||||
for (const calls of [200, 2000]) {
|
||||
for (const depth of [1, 2] as const) {
|
||||
scenarios.push({ backend, calls, depth, bytes: 1024, gapMs: 2, progress: false, runs: 3 });
|
||||
}
|
||||
}
|
||||
}
|
||||
await report("Nesting: fast calls without progress, at depth 1 and 2", scenarios);
|
||||
}
|
||||
}
|
||||
|
||||
await main();
|
||||
Reference in New Issue
Block a user