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
75 lines
3.3 KiB
TypeScript
75 lines
3.3 KiB
TypeScript
/**
|
|
* Scheduler fan-out benchmark. One model-issued tool makes `calls` parallel nested calls of a tool that returns at
|
|
* once, on in-memory storage, so the time is the Harness's and the scheduler's own. Quadratic work shows as four times
|
|
* the time per doubling of `calls`.
|
|
*
|
|
* node --conditions=source --experimental-strip-types test/scheduler-fanout-bench.ts [calls...]
|
|
*/
|
|
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 { createRegistry, defineExtension, defineTool, Harness, ToolResultEntry } from "../src/index.ts";
|
|
import { MemoryStorage } from "../src/storage/memory.ts";
|
|
|
|
async function run(calls: number): Promise<{ ms: number; commits: number }> {
|
|
const faux = fauxProvider();
|
|
const models = createModels();
|
|
models.setProvider(faux.provider);
|
|
const registry = createRegistry();
|
|
const leaf = defineTool({
|
|
name: "leaf",
|
|
description: "Return at once",
|
|
parameters: Type.Object({ n: Type.Number() }),
|
|
execute: async () => ({ output: [{ type: "text", text: "ok" }] }),
|
|
});
|
|
const fan = defineTool({
|
|
name: "fan",
|
|
description: "Call leaf in parallel",
|
|
parameters: Type.Object({ count: Type.Number() }),
|
|
execute: async (args, api, callContext) => {
|
|
const results = await Promise.all(
|
|
Array.from({ length: args.count }, (_, n) =>
|
|
api.executeTool("leaf", { n }, callContext, { progress: false }),
|
|
),
|
|
);
|
|
const failed = results.filter((result) => result.isError).length;
|
|
return { output: [{ type: "text", text: `${failed} failed` }], isError: failed > 0 };
|
|
},
|
|
});
|
|
registry.install(defineExtension({ name: "bench", tools: [leaf, fan] }));
|
|
const storage = new MemoryStorage();
|
|
let commits = 0;
|
|
const commit = storage.commit.bind(storage);
|
|
storage.commit = (writes, commitContext) => {
|
|
commits++;
|
|
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("fan", { count: calls }, { id: "fan" })], { stopReason: "toolUse" }),
|
|
fauxAssistantMessage([fauxText("done")]),
|
|
]);
|
|
const started = performance.now();
|
|
const settled = await (await root.submit({ type: "input", content: "go" }, context)).wait(context);
|
|
if (settled.status !== "done") throw new Error(`run: ${settled.status}`);
|
|
await harness.waitForIdle(context);
|
|
const ms = Math.round(performance.now() - started);
|
|
// A run whose nested calls failed measures nothing.
|
|
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: nested calls failed`);
|
|
await harness.close(context);
|
|
return { ms, commits };
|
|
}
|
|
|
|
const sizes = process.argv.slice(2).map(Number);
|
|
console.log("calls | ms | commits | ms per call");
|
|
for (const calls of sizes.length > 0 ? sizes : [250, 500, 1000, 2000, 4000]) {
|
|
const { ms, commits } = await run(calls);
|
|
console.log(
|
|
`${String(calls).padEnd(5)} | ${String(ms).padEnd(6)} | ${String(commits).padEnd(7)} | ${(ms / calls).toFixed(2)}`,
|
|
);
|
|
}
|