From 79e1a48d973f741cf38a8ae6673e5503995d68d9 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Thu, 3 Sep 2026 02:19:57 +0200 Subject: [PATCH 1/4] Add an in-browser optimization runtime to petrinaut-core --- .changeset/browser-optimization-runtime.md | 5 + .../petrinaut-cli/src/runtime/optimization.ts | 159 +--- libs/@hashintel/petrinaut-core/.oxlintrc.json | 6 +- libs/@hashintel/petrinaut-core/package.json | 10 + .../src/browser-optimization.ts | 20 + .../browser-optimization.test.ts | 692 ++++++++++++++++++ .../browser-optimization.ts | 438 +++++++++++ .../create-optimizer-worker.ts | 30 + .../src/browser-optimization/messages.ts | 109 +++ .../browser-optimization/optimizer.worker.ts | 166 +++++ .../pyodide-config.test.ts | 33 + .../browser-optimization/pyodide-config.ts | 23 + .../src/browser-optimization/pyodide-like.ts | 32 + .../browser-optimization/python-sources.ts | 14 + .../src/browser-optimization/run-log.test.ts | 104 +++ .../src/browser-optimization/run-log.ts | 112 +++ .../study-runner.pyodide.test.ts | 257 +++++++ .../src/browser-optimization/study-runner.ts | 222 ++++++ .../petrinaut-core/src/environment.ts | 11 + .../petrinaut-core/src/optimization.ts | 93 +++ .../src/optimization/describe.test.ts | 140 ++++ .../src/optimization/describe.ts | 192 +++++ .../shared/optimization-manifest.fixtures.ts | 80 ++ .../petrinaut-core/src/vite-types.d.ts | 5 + .../petrinaut-core/src/workers/README.md | 1 + .../petrinaut-core/src/workers/optimizer.ts | 23 + libs/@hashintel/petrinaut-core/vite.config.ts | 11 +- yarn.lock | 21 +- 28 files changed, 2862 insertions(+), 147 deletions(-) create mode 100644 .changeset/browser-optimization-runtime.md create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/create-optimizer-worker.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-like.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/python-sources.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts create mode 100644 libs/@hashintel/petrinaut-core/src/optimization/describe.test.ts create mode 100644 libs/@hashintel/petrinaut-core/src/optimization/describe.ts create mode 100644 libs/@hashintel/petrinaut-core/src/shared/optimization-manifest.fixtures.ts create mode 100644 libs/@hashintel/petrinaut-core/src/workers/optimizer.ts diff --git a/.changeset/browser-optimization-runtime.md b/.changeset/browser-optimization-runtime.md new file mode 100644 index 00000000000..6bd95875612 --- /dev/null +++ b/.changeset/browser-optimization-runtime.md @@ -0,0 +1,5 @@ +--- +"@hashintel/petrinaut-core": patch +--- + +Adds an in-browser optimization capability that runs the Optuna study in a Pyodide worker and evaluates trials through a host channel. diff --git a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts index 3eb73674aa0..0ce2be79777 100644 --- a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts +++ b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts @@ -3,32 +3,28 @@ import { readFile } from "node:fs/promises"; import { compileScenario, createMonteCarloExperiment, - deriveRunSeed, parseDocumentText, petrinautOptimizationEvaluateParamsSchema, petrinautOptimizationManifestSchema, } from "@hashintel/petrinaut-core"; import { lowerScenarioToHir } from "@hashintel/petrinaut-core/hir"; +import { + deriveOptimizationTrialSeeds, + describeOptimization, + resolveTrialScenarioParameterValues, +} from "@hashintel/petrinaut-core/optimization"; import { createInProcessMonteCarloWorker } from "@hashintel/petrinaut-core/workers/monte-carlo"; import type { MonteCarloExperiment, - PetrinautOptimizationDescribeParameter, PetrinautOptimizationDescribeResult, PetrinautOptimizationEvaluateResult, PetrinautOptimizationManifest, - Scenario, WorkerFactory, } from "@hashintel/petrinaut-core"; import type { PetrinautCompiledModel } from "@hashintel/petrinaut-core/compiled-model"; -type OptimizationScalar = number | boolean; -type ScenarioParameter = Scenario["scenarioParameters"][number]; -type OptimizedBinding = Extract< - PetrinautOptimizationManifest["scenario"]["parameterBindings"][string], - { kind: "optimize" } ->; -type OptimizationDomain = OptimizedBinding["domain"]; +export { deriveOptimizationTrialSeeds as deriveTrialSeeds } from "@hashintel/petrinaut-core/optimization"; function formatManifestIssues( prefix: string, @@ -67,95 +63,11 @@ export async function loadOptimizationManifest( return parseOptimizationManifest(document.data); } -function describeParameter( - parameter: ScenarioParameter, - domain: OptimizationDomain, -): PetrinautOptimizationDescribeParameter { - switch (domain.kind) { - case "continuous": - return { - identifier: parameter.identifier, - type: "float", - default: parameter.default, - minimum: domain.minimum, - maximum: domain.maximum, - scale: domain.scale, - }; - case "integer": - return { - identifier: parameter.identifier, - type: "int", - default: parameter.default, - minimum: domain.minimum, - maximum: domain.maximum, - step: domain.step, - scale: domain.scale, - }; - case "boolean": - return { - identifier: parameter.identifier, - type: "boolean", - default: parameter.default !== 0, - }; - } -} - -function validateSuggestedValue( - parameter: ScenarioParameter, - domain: OptimizationDomain, - value: OptimizationScalar, -): void { - if (domain.kind === "boolean") { - if (typeof value !== "boolean") { - throw new Error( - `Optimization parameter "${parameter.identifier}" must be boolean`, - ); - } - return; - } - if (typeof value !== "number") { - throw new Error( - `Optimization parameter "${parameter.identifier}" must be numeric`, - ); - } - if (value < domain.minimum || value > domain.maximum) { - throw new Error( - `Optimization parameter "${parameter.identifier}" must be between ${domain.minimum} and ${domain.maximum}`, - ); - } - if (domain.kind === "integer") { - if (!Number.isInteger(value)) { - throw new Error( - `Optimization parameter "${parameter.identifier}" must be an integer`, - ); - } - if ((value - domain.minimum) % domain.step !== 0) { - throw new Error( - `Optimization parameter "${parameter.identifier}" must align with step ${domain.step} from ${domain.minimum}`, - ); - } - } -} - export type OptimizationProtocol = { describe(): PetrinautOptimizationDescribeResult; evaluate(params: unknown): Promise; }; -/** - * Derives one trial's run seeds. Run 0 keeps the base seed, so a single-seed - * trial matches the old fixed-seed behaviour; later runs use the Monte Carlo - * derivation. Every trial gets the same sequence: common random numbers. - */ -export function deriveTrialSeeds( - baseSeed: number, - seedsPerTrial: number, -): number[] { - return Array.from({ length: seedsPerTrial }, (_, index) => - index === 0 ? baseSeed : deriveRunSeed(baseSeed, index), - ); -} - /** Resolves when the experiment reports its terminal event. */ function waitForCompletion(experiment: MonteCarloExperiment): Promise { return new Promise((resolve, reject) => { @@ -204,7 +116,10 @@ export function createOptimizationProtocol(args: { const createWorker = args.createWorker ?? createInProcessMonteCarloWorker; const createExperiment = args.createExperiment ?? createMonteCarloExperiment; const seedsPerTrial = manifest.execution.seedsPerTrial ?? 1; - const trialSeeds = deriveTrialSeeds(manifest.execution.seed, seedsPerTrial); + const trialSeeds = deriveOptimizationTrialSeeds( + manifest.execution.seed, + seedsPerTrial, + ); const scenario = manifest.model.definition.scenarios?.[0]; const metric = manifest.model.definition.metrics?.[0]; if (!scenario || !metric) { @@ -212,17 +127,6 @@ export function createOptimizationProtocol(args: { "An optimization manifest requires exactly one scenario and one metric", ); } - const optimizedParameters = scenario.scenarioParameters.flatMap( - (parameter) => { - const binding = manifest.scenario.parameterBindings[parameter.identifier]; - return binding?.kind === "optimize" - ? [{ parameter, domain: binding.domain }] - : []; - }, - ); - const optimizedIdentifiers = new Set( - optimizedParameters.map(({ parameter }) => parameter.identifier), - ); // Lower the scenario's expressions once per study; each trial re-runs only // the type-check and the interpreter with that trial's parameter values. @@ -237,17 +141,7 @@ export function createOptimizationProtocol(args: { return { describe() { - return { - direction: manifest.objective.direction, - study: { - ...manifest.study, - seed: manifest.execution.seed, - seedsPerTrial, - }, - parameters: optimizedParameters.map(({ parameter, domain }) => - describeParameter(parameter, domain), - ), - }; + return describeOptimization(manifest); }, async evaluate(params) { const parsed = @@ -258,33 +152,10 @@ export function createOptimizationProtocol(args: { parsed.error.issues, ); } - const values = parsed.data.parameterValues; - for (const { parameter } of optimizedParameters) { - const { identifier } = parameter; - if (!Object.hasOwn(values, identifier)) { - throw new Error(`Missing optimized parameter "${identifier}"`); - } - } - for (const identifier of Object.keys(values)) { - if (!optimizedIdentifiers.has(identifier)) { - throw new Error(`Unexpected optimization parameter "${identifier}"`); - } - } - - const scenarioParameterValues: Record = {}; - for (const parameter of scenario.scenarioParameters) { - const binding = - manifest.scenario.parameterBindings[parameter.identifier]!; - const value = - binding.kind === "fixed" - ? binding.value - : values[parameter.identifier]!; - if (binding.kind === "optimize") { - validateSuggestedValue(parameter, binding.domain, value); - } - scenarioParameterValues[parameter.identifier] = - typeof value === "boolean" ? (value ? 1 : 0) : value; - } + const scenarioParameterValues = resolveTrialScenarioParameterValues( + manifest, + parsed.data.parameterValues, + ); const compiledScenario = compileScenario( scenario, diff --git a/libs/@hashintel/petrinaut-core/.oxlintrc.json b/libs/@hashintel/petrinaut-core/.oxlintrc.json index 5a4b9641d42..901d9f085b3 100644 --- a/libs/@hashintel/petrinaut-core/.oxlintrc.json +++ b/libs/@hashintel/petrinaut-core/.oxlintrc.json @@ -89,7 +89,11 @@ { "patterns": [ { - "group": ["@local/*"], + "group": [ + "@local/*", + "!@local/petrinaut-optimizer-core", + "!@local/petrinaut-optimizer-core/**" + ], "message": "You cannot use unpublished local packages in a published package." }, { diff --git a/libs/@hashintel/petrinaut-core/package.json b/libs/@hashintel/petrinaut-core/package.json index 32b4826b368..8ff070f4f02 100644 --- a/libs/@hashintel/petrinaut-core/package.json +++ b/libs/@hashintel/petrinaut-core/package.json @@ -34,6 +34,10 @@ "types": "./dist/ai.d.ts", "import": "./dist/ai.js" }, + "./browser-optimization": { + "types": "./dist/browser-optimization.d.ts", + "import": "./dist/browser-optimization.js" + }, "./compiled-model": { "types": "./dist/compiled-model.d.ts", "import": "./dist/compiled-model.js" @@ -66,6 +70,10 @@ "types": "./dist/workers/monte-carlo.d.ts", "import": "./dist/workers/monte-carlo.js" }, + "./workers/optimizer": { + "types": "./dist/workers/optimizer.d.ts", + "import": "./dist/workers/optimizer.js" + }, "./workers/simulation": { "types": "./dist/workers/simulation.d.ts", "import": "./dist/workers/simulation.js" @@ -98,12 +106,14 @@ "zod": "4.4.3" }, "devDependencies": { + "@local/petrinaut-optimizer-core": "workspace:*", "@types/js-yaml": "^4", "@types/node": "22.18.13", "@typescript/native-preview": "7.0.0-dev.20260511.1", "@webgpu/types": "0.1.71", "oxlint": "1.63.0", "oxlint-tsgolint": "0.22.1", + "pyodide": "314.0.6", "rolldown": "1.2.6", "rolldown-plugin-dts": "0.28.3", "typescript": "5.9.3", diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization.ts new file mode 100644 index 00000000000..a6a02805749 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization.ts @@ -0,0 +1,20 @@ +export { + createBrowserOptimization, + type CreateBrowserOptimizationOptions, +} from "./browser-optimization/browser-optimization"; +export type { + OptimizerWorkerErrorEvent, + OptimizerWorkerLike, +} from "./browser-optimization/create-optimizer-worker"; +export { + defaultOptimizerPyodideConfig, + type OptimizerPyodideConfig, +} from "./browser-optimization/pyodide-config"; +export type { + OptimizationScalar, + PetrinautConnectedOptimization, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, + PetrinautOptimizationTrialOutcome, + PetrinautOptimizationTrialRequest, +} from "./optimization"; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts new file mode 100644 index 00000000000..6482ffe9ad0 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts @@ -0,0 +1,692 @@ +import { describe, expect, it, vi } from "vitest"; + +import { PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE } from "../optimization"; +import { createOptimizationManifestInput } from "../shared/optimization-manifest.fixtures"; +import { createBrowserOptimization } from "./browser-optimization"; + +import type { WorkerMessageHandler } from "../environment"; +import type { + PetrinautOptimizationChannel, + PetrinautOptimizationEvent, + PetrinautOptimizationTrialRequest, +} from "../optimization"; +import type { + OptimizerWorkerErrorEvent, + OptimizerWorkerLike, +} from "./create-optimizer-worker"; +import type { + OptimizerStudySummary, + OptimizerToMainMessage, + OptimizerToWorkerMessage, + OptimizerTrialPayload, +} from "./messages"; + +type FakeWorkerErrorHandler = (event: OptimizerWorkerErrorEvent) => void; + +type FakeWorker = OptimizerWorkerLike & { + readonly sent: OptimizerToWorkerMessage[]; + readonly terminated: boolean; + /** Deliver a message as the worker would post it. */ + emit(message: OptimizerToMainMessage): void; + /** Fire the `error` event a worker whose script fails to load fires. */ + emitError(event: OptimizerWorkerErrorEvent): void; + /** The messages of one type posted so far. */ + sentOfType( + type: TType, + ): Extract[]; +}; + +const createFakeWorker = (): FakeWorker => { + const listeners = new Set>(); + const errorListeners = new Set(); + const sent: OptimizerToWorkerMessage[] = []; + let terminated = false; + return { + sent, + get terminated() { + return terminated; + }, + postMessage(message) { + sent.push(message); + }, + addEventListener( + type: "message" | "error", + listener: + | WorkerMessageHandler + | FakeWorkerErrorHandler, + ) { + if (type === "message") { + listeners.add(listener as WorkerMessageHandler); + } else { + errorListeners.add(listener as FakeWorkerErrorHandler); + } + }, + removeEventListener(_type, listener) { + listeners.delete(listener); + }, + terminate() { + terminated = true; + }, + emit(message) { + for (const listener of listeners) { + listener({ data: message }); + } + }, + emitError(event) { + for (const listener of errorListeners) { + listener(event); + } + }, + sentOfType(type) { + return sent.filter( + (message): message is Extract => + message.type === type, + ); + }, + }; +}; + +const flush = async (): Promise => { + for (let index = 0; index < 5; index++) { + await Promise.resolve(); + } +}; + +const summary: OptimizerStudySummary = { + requestedTrials: 20, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: { trial: 0, parameters: { rate: 0.5 }, objective: 1 }, +}; + +const completedTrial: OptimizerTrialPayload = { + trial: 0, + parameters: { rate: 0.5, count: 6, enabled: true }, + objective: 1, + state: "complete", + best: { + trial: 0, + parameters: { rate: 0.5, count: 6, enabled: true }, + objective: 1, + }, +}; + +const collectEvents = ( + iterable: AsyncIterable, +): Promise => + (async () => { + const events: PetrinautOptimizationEvent[] = []; + for await (const event of iterable) { + events.push(event); + } + return events; + })(); + +const setUp = (options?: { + evaluateTrial?: PetrinautOptimizationChannel["evaluateTrial"]; + createWorker?: (attempt: number) => FakeWorker; +}) => { + const workers: FakeWorker[] = []; + let attempts = 0; + const evaluateTrial = vi.fn( + options?.evaluateTrial ?? + (async (request) => ({ + kind: "objective", + objective: Number(request.suggestedValues.rate) * 2, + })), + ); + const capability = createBrowserOptimization({ + pyodide: { indexURL: "https://example.test/pyodide/" }, + createWorker: () => { + attempts += 1; + const worker = (options?.createWorker ?? createFakeWorker)(attempts); + workers.push(worker); + return worker; + }, + }).connect({ evaluateTrial }); + return { + capability, + evaluateTrial, + workers, + get worker() { + const worker = workers.at(-1); + if (!worker) { + throw new Error("no worker was created"); + } + return worker; + }, + }; +}; + +/** Creates a run and brings the worker to the point of having received `start`. */ +const startRun = async (context: ReturnType) => { + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + context.worker.emit({ type: "ready" }); + await flush(); + return runId; +}; + +describe("createBrowserOptimization", () => { + it("initialises the worker lazily with the runtime config and the Python sources", async () => { + const context = setUp(); + expect(context.workers).toHaveLength(0); + + await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + + const [init] = context.worker.sent; + if (init?.type !== "init") { + throw new Error("the first worker message must be init"); + } + expect(init.pyodide.indexURL).toBe("https://example.test/pyodide/"); + expect(init.pyodide.packages.optuna).toMatch(/^\d/); + expect(Object.keys(init.pythonSources)).toEqual([ + "petrinaut_optimizer_core/__init__.py", + "petrinaut_optimizer_core/description.py", + "petrinaut_optimizer_core/study.py", + "petrinaut_optimizer_core/ask_tell.py", + "petrinaut_optimizer_core/pyodide_entry.py", + ]); + expect(context.worker.sentOfType("start")).toHaveLength(0); + }); + + it("streams started, trial and complete events with dense sequence numbers", async () => { + const context = setUp(); + const runId = await startRun(context); + const events = collectEvents( + context.capability.attachOptimizationRun(runId), + ); + + const [start] = context.worker.sentOfType("start"); + expect(start).toMatchObject({ + runId, + description: { + direction: "maximize", + study: { trials: 20, sampler: "tpe", seed: 42, seedsPerTrial: 1 }, + }, + }); + + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + + expect(context.evaluateTrial).toHaveBeenCalledTimes(1); + const request = context.evaluateTrial.mock + .calls[0]?.[0] as PetrinautOptimizationTrialRequest; + expect(request).toMatchObject({ + runId, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + scenarioParameterValues: { rate: 0.5, count: 6, enabled: 1, share: 0.25 }, + seeds: [42], + }); + expect(request.manifest.name).toBe("Find the best rate"); + expect(request.signal.aborted).toBe(false); + expect(context.worker.sentOfType("evaluated")).toEqual([ + { + type: "evaluated", + requestId: 1, + outcome: { kind: "objective", objective: 1 }, + }, + ]); + + context.worker.emit({ type: "trial", runId, event: completedTrial }); + context.worker.emit({ type: "complete", runId, summary }); + + expect(await events).toEqual([ + { type: "started", requestedTrials: 20, seq: 1 }, + { + type: "trial", + trial: 0, + parameters: { rate: 0.5, count: 6, enabled: true }, + objective: 1, + state: "complete", + best: completedTrial.best, + seq: 2, + }, + { + type: "complete", + requestedTrials: 20, + completedTrials: 1, + prunedTrials: 0, + failedTrials: 0, + best: summary.best, + seq: 3, + }, + ]); + }); + + it("reports pruned trials with a null objective and forwards pruned outcomes", async () => { + const context = setUp({ + evaluateTrial: async () => ({ kind: "pruned", reason: "no frames" }), + }); + const runId = await startRun(context); + + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + expect(context.worker.sentOfType("evaluated")[0]?.outcome).toEqual({ + kind: "pruned", + reason: "no frames", + }); + + context.worker.emit({ + type: "trial", + runId, + event: { ...completedTrial, objective: 3, state: "pruned", best: null }, + }); + context.worker.emit({ + type: "complete", + runId, + summary: { ...summary, completedTrials: 0, prunedTrials: 1, best: null }, + }); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events[1]).toMatchObject({ + type: "trial", + objective: null, + state: "pruned", + best: null, + }); + expect(events[2]).toMatchObject({ type: "complete", prunedTrials: 1 }); + }); + + it("cancels a running study through the worker and ends with the cancelled error code", async () => { + const context = setUp({ + evaluateTrial: (request) => + new Promise((_resolve, reject) => { + request.signal.addEventListener("abort", () => { + const error = new Error("aborted"); + error.name = "AbortError"; + reject(error); + }); + }), + }); + const runId = await startRun(context); + const events = collectEvents( + context.capability.attachOptimizationRun(runId), + ); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + + await context.capability.cancelOptimizationRun(runId); + await flush(); + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + expect(context.worker.sentOfType("evaluated")).toEqual([ + { + type: "evaluated", + requestId: 1, + outcome: { kind: "pruned", reason: "cancelled" }, + }, + ]); + + context.worker.emit({ type: "cancelled", runId }); + await context.capability.cancelOptimizationRun(runId); + + expect((await events).at(-1)).toEqual({ + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + message: "optimization cancelled", + retryable: false, + seq: 2, + }); + }); + + it("cancels a run that is still waiting for the runtime without starting it", async () => { + const context = setUp(); + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + + await context.capability.cancelOptimizationRun(runId); + context.worker.emit({ type: "ready" }); + await flush(); + + expect(context.worker.sentOfType("start")).toHaveLength(0); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.map((event) => event.type)).toEqual(["started", "error"]); + expect(events[1]).toMatchObject({ + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + }); + }); + + it("rejects attaching to an unknown run with the not-found shape", () => { + const context = setUp(); + + expect(() => context.capability.attachOptimizationRun("missing")).toThrow( + expect.objectContaining({ category: "http", httpStatus: 404 }), + ); + }); + + it("fails the run when the channel throws and stops the study", async () => { + const context = setUp({ + evaluateTrial: async () => { + throw new Error("backend unavailable"); + }, + }); + const runId = await startRun(context); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + expect(context.worker.sentOfType("evaluated")).toHaveLength(0); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "trial_evaluation_failed", + message: "backend unavailable", + retryable: false, + seq: 2, + }); + + // The worker winds the study down afterwards; its late messages change nothing. + context.worker.emit({ type: "cancelled", runId }); + expect( + await collectEvents(context.capability.attachOptimizationRun(runId)), + ).toHaveLength(2); + }); + + it("fails the run when the channel throws synchronously", async () => { + const context = setUp({ + evaluateTrial: () => { + throw new Error("channel misconfigured"); + }, + }); + const runId = await startRun(context); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + expect(context.worker.sentOfType("evaluated")).toHaveLength(0); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toMatchObject({ + type: "error", + code: "trial_evaluation_failed", + message: "channel misconfigured", + }); + }); + + it("fails the run when the optimizer suggests values outside the manifest", async () => { + const context = setUp(); + const runId = await startRun(context); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 7, enabled: true }, + }); + await flush(); + + expect(context.evaluateTrial).not.toHaveBeenCalled(); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toMatchObject({ + type: "error", + code: "trial_evaluation_failed", + message: 'Optimization parameter "count" must align with step 2 from 2', + }); + }); + + it("reports a study error from the worker", async () => { + const context = setUp(); + const runId = await startRun(context); + + context.worker.emit({ type: "error", runId, message: "ValueError: nope" }); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "study_failed", + message: "ValueError: nope", + retryable: false, + seq: 2, + }); + }); + + it("runs studies one at a time on a shared worker", async () => { + const context = setUp(); + const first = await startRun(context); + const { runId: second } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + await flush(); + + expect(context.workers).toHaveLength(1); + expect( + context.worker.sentOfType("start").map(({ runId }) => runId), + ).toEqual([first]); + + context.worker.emit({ type: "complete", runId: first, summary }); + await flush(); + + expect( + context.worker.sentOfType("start").map(({ runId }) => runId), + ).toEqual([first, second]); + expect( + (await collectEvents(context.capability.attachOptimizationRun(first))).at( + -1, + )?.type, + ).toBe("complete"); + }); + + it("fails a run the runtime could not load and retries with a fresh worker", async () => { + const context = setUp(); + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + const firstWorker = context.worker; + + firstWorker.emit({ type: "init-error", message: "offline" }); + await flush(); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + const last = events.at(-1); + expect(last).toMatchObject({ + type: "error", + code: "optimizer_unavailable", + retryable: true, + }); + expect(last?.type === "error" ? last.message : "").toContain("offline"); + expect(firstWorker.terminated).toBe(true); + + await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + expect(context.workers).toHaveLength(2); + expect(context.worker.sentOfType("init")).toHaveLength(1); + }); + + it("fails a run whose worker script does not load and retries with a fresh worker", async () => { + const context = setUp(); + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + const firstWorker = context.worker; + + firstWorker.emitError({}); + await flush(); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "optimizer_unavailable", + message: + "The in-browser optimizer could not start: The optimizer worker failed to load", + retryable: true, + seq: 2, + }); + expect(firstWorker.terminated).toBe(true); + + await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + expect(context.workers).toHaveLength(2); + expect(context.worker.sentOfType("init")).toHaveLength(1); + }); + + it("fails a run whose worker cannot be created and retries on the next run", async () => { + const context = setUp({ + createWorker: (attempt) => { + if (attempt === 1) { + throw new Error("SecurityError: cross-origin worker script"); + } + return createFakeWorker(); + }, + }); + + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "optimizer_unavailable", + message: + "The in-browser optimizer could not start: SecurityError: cross-origin worker script", + retryable: true, + seq: 2, + }); + expect(context.workers).toHaveLength(0); + + const second = await startRun(context); + expect(context.workers).toHaveLength(1); + expect( + context.worker.sentOfType("start").map(({ runId: started }) => started), + ).toEqual([second]); + }); + + it("replays past a cursor and aborts a tailing attachment", async () => { + const context = setUp(); + const runId = await startRun(context); + context.worker.emit({ type: "trial", runId, event: completedTrial }); + + const replayed = await collectEvents( + (async function* takeFirstEvent() { + const iterator = context.capability + .attachOptimizationRun(runId, { cursor: 1 }) + [Symbol.asyncIterator](); + const next = await iterator.next(); + if (!next.done) { + yield next.value; + } + await iterator.return?.(undefined); + })(), + ); + expect(replayed.map((event) => event.seq)).toEqual([2]); + + const onAttached = vi.fn(); + let abort = (): void => {}; + const tailing = collectEvents( + context.capability.attachOptimizationRun(runId, { + cursor: 2, + onAttached, + signal: { + aborted: false, + addEventListener: (_type, listener) => { + abort = listener; + }, + removeEventListener: () => {}, + }, + }), + ); + await flush(); + expect(onAttached).toHaveBeenCalledTimes(1); + abort(); + await expect(tailing).rejects.toMatchObject({ name: "AbortError" }); + }); + + it("dispose cancels every run and terminates the worker", async () => { + const context = setUp(); + const running = await startRun(context); + const { runId: queued } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + + context.capability.dispose(); + + expect(context.worker.terminated).toBe(true); + for (const runId of [running, queued]) { + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toMatchObject({ + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + }); + } + await expect( + context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ), + ).rejects.toThrow("disposed"); + }); + + it("rejects an invalid manifest before allocating a run", async () => { + const context = setUp(); + + await expect( + context.capability.createOptimizationRun({ + ...createOptimizationManifestInput(), + study: { trials: 0, sampler: "tpe" }, + }), + ).rejects.toThrow(/Invalid optimization manifest: study\.trials/); + expect(context.workers).toHaveLength(0); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts new file mode 100644 index 00000000000..c6b2c11cfe0 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts @@ -0,0 +1,438 @@ +/** + * @layerRoot core.optimization.browser + * @role Runs the Optuna study in a Pyodide worker and evaluates trials through the host channel + */ +import { v4 as generateUuid } from "uuid"; + +import { createAbortController } from "../environment"; +import { + deriveOptimizationTrialSeeds, + describeOptimization, + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + petrinautOptimizationManifestSchema, + resolveTrialScenarioParameterValues, +} from "../optimization"; +import { + createOptimizerWorker, + type OptimizerWorkerErrorEvent, + type OptimizerWorkerLike, +} from "./create-optimizer-worker"; +import { + defaultOptimizerPyodideConfig, + type OptimizerPyodideConfig, +} from "./pyodide-config"; +import { optimizerPythonSources } from "./python-sources"; +import { + createOptimizationRunLog, + type OptimizationRunLog, + type OptimizationRunLogEvent, +} from "./run-log"; + +import type { AbortSignalLike } from "../environment"; +import type { + PetrinautConnectedOptimization, + PetrinautOptimization, + PetrinautOptimizationChannel, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationEvent, + PetrinautOptimizationManifest, + PetrinautOptimizationTrialOutcome, + PetrinautOptimizationTrialRequest, +} from "../optimization"; +import type { + OptimizerEvaluateMessage, + OptimizerToMainMessage, +} from "./messages"; + +export type CreateBrowserOptimizationOptions = { + pyodide?: Partial; + createWorker?: () => OptimizerWorkerLike; +}; + +type RunStatus = "queued" | "starting" | "running" | "finished"; + +type RunRecord = { + readonly runId: string; + readonly manifest: PetrinautOptimizationManifest; + readonly description: PetrinautOptimizationDescribeResult; + readonly seeds: readonly number[]; + readonly log: OptimizationRunLog; + readonly signal: AbortSignalLike; + readonly abort: () => void; + status: RunStatus; +}; + +type WorkerSession = { + readonly worker: OptimizerWorkerLike; + readonly ready: Promise; +}; + +const cancelledEvent: OptimizationRunLogEvent = { + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + message: "optimization cancelled", + retryable: false, +}; + +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : String(error); + +const workerLoadError = (event: OptimizerWorkerErrorEvent): Error => + new Error( + event.message === undefined || event.message === "" + ? "The optimizer worker failed to load" + : event.message, + ); + +const unavailableEvent = (error: unknown): OptimizationRunLogEvent => ({ + type: "error", + code: "optimizer_unavailable", + message: `The in-browser optimizer could not start: ${errorMessage(error)}`, + retryable: true, +}); + +const isAbortError = (error: unknown): boolean => + typeof error === "object" && + error !== null && + "name" in error && + error.name === "AbortError"; + +const invalidManifestError = ( + issues: readonly { path: PropertyKey[]; message: string }[], +): Error => + new Error( + `Invalid optimization manifest: ${issues + .map( + ({ path, message }) => + `${path.length > 0 ? path.join(".") : "manifest"}: ${message}`, + ) + .join("; ")}`, + ); + +/** The shape the optimizations provider drops a stale stored run on. */ +const unknownRunError = (runId: string): Error => + Object.assign(new Error(`Unknown optimization run "${runId}"`), { + category: "http", + httpStatus: 404, + }); + +async function* attachToLog( + log: OptimizationRunLog, + options: { + cursor?: number; + signal?: AbortSignalLike; + onAttached?: () => void; + }, +): AsyncGenerator { + options.onAttached?.(); + yield* log.replay({ cursor: options.cursor, signal: options.signal }); +} + +const connectBrowserOptimization = (options: { + channel: PetrinautOptimizationChannel; + pyodide: OptimizerPyodideConfig; + createWorker: () => OptimizerWorkerLike; +}): PetrinautOptimization & { dispose(this: void): void } => { + const { channel } = options; + const runs = new Map(); + const queue: RunRecord[] = []; + let active: RunRecord | null = null; + let session: WorkerSession | null = null; + let disposed = false; + + const activeRunFor = (runId: string): RunRecord | null => { + const run = runs.get(runId); + return run && run.status === "running" ? run : null; + }; + + const resetSession = (stale: WorkerSession): void => { + if (session === stale) { + session = null; + stale.worker.terminate(); + } + }; + + const finish = (run: RunRecord, event: OptimizationRunLogEvent): void => { + if (run.status === "finished") { + return; + } + // eslint-disable-next-line no-param-reassign -- the record's status is the session state this helper advances + run.status = "finished"; + run.log.append(event); + const queuedAt = queue.indexOf(run); + if (queuedAt !== -1) { + queue.splice(queuedAt, 1); + } + if (active === run) { + active = null; + // eslint-disable-next-line no-use-before-define -- mutual recursion + startNext(); + } + }; + + const reply = ( + requestId: number, + outcome: PetrinautOptimizationTrialOutcome, + ): void => { + session?.worker.postMessage({ type: "evaluated", requestId, outcome }); + }; + + const failTrialEvaluation = (run: RunRecord, error: unknown): void => { + run.abort(); + session?.worker.postMessage({ type: "cancel", runId: run.runId }); + finish(run, { + type: "error", + code: "trial_evaluation_failed", + message: errorMessage(error), + retryable: false, + }); + }; + + const handleEvaluate = (message: OptimizerEvaluateMessage): void => { + const run = activeRunFor(message.runId); + if (!run) { + return; + } + let request: PetrinautOptimizationTrialRequest; + try { + request = { + runId: run.runId, + trial: message.trial, + manifest: run.manifest, + suggestedValues: message.suggestedValues, + scenarioParameterValues: resolveTrialScenarioParameterValues( + run.manifest, + message.suggestedValues, + ), + seeds: run.seeds, + signal: run.signal, + }; + } catch (error) { + failTrialEvaluation(run, error); + return; + } + const evaluated = (outcome: PetrinautOptimizationTrialOutcome): void => { + if (run.status === "running") { + reply(message.requestId, outcome); + } + }; + const evaluationFailed = (error: unknown): void => { + if (run.status !== "running") { + return; + } + if (isAbortError(error)) { + reply(message.requestId, { kind: "pruned", reason: "cancelled" }); + } else { + failTrialEvaluation(run, error); + } + }; + let evaluation: Promise; + try { + evaluation = channel.evaluateTrial(request); + } catch (error) { + evaluationFailed(error); + return; + } + evaluation.then(evaluated, evaluationFailed); + }; + + const handleWorkerMessage = (message: OptimizerToMainMessage): void => { + switch (message.type) { + case "ready": + case "init-error": + return; + case "evaluate": + handleEvaluate(message); + return; + case "trial": { + const run = activeRunFor(message.runId); + const { event } = message; + run?.log.append({ + type: "trial", + trial: event.trial, + parameters: event.parameters, + objective: event.state === "complete" ? event.objective : null, + state: event.state, + best: event.best, + }); + return; + } + case "complete": { + const run = activeRunFor(message.runId); + const { summary } = message; + if (run) { + finish(run, { + type: "complete", + requestedTrials: summary.requestedTrials, + completedTrials: summary.completedTrials, + prunedTrials: summary.prunedTrials, + failedTrials: summary.failedTrials, + best: summary.best, + }); + } + return; + } + case "cancelled": { + const run = activeRunFor(message.runId); + if (run) { + finish(run, cancelledEvent); + } + return; + } + case "error": { + const run = activeRunFor(message.runId); + if (run) { + finish(run, { + type: "error", + code: "study_failed", + message: message.message, + retryable: false, + }); + } + } + } + }; + + const ensureSession = (): WorkerSession => { + if (session) { + return session; + } + const worker = options.createWorker(); + const ready = new Promise((resolve, reject) => { + worker.addEventListener("message", ({ data }) => { + if (data.type === "ready") { + resolve(); + } else if (data.type === "init-error") { + reject(new Error(data.message)); + } else { + handleWorkerMessage(data); + } + }); + worker.addEventListener("error", (event) => { + reject(workerLoadError(event)); + }); + }); + worker.postMessage({ + type: "init", + pyodide: options.pyodide, + pythonSources: optimizerPythonSources, + }); + session = { worker, ready }; + return session; + }; + + const startNext = (): void => { + if (disposed || active) { + return; + } + const run = queue.shift(); + if (!run) { + return; + } + active = run; + run.status = "starting"; + let current: WorkerSession; + try { + current = ensureSession(); + } catch (error) { + finish(run, unavailableEvent(error)); + return; + } + current.ready.then( + () => { + if (run.status === "starting" && session === current) { + run.status = "running"; + current.worker.postMessage({ + type: "start", + runId: run.runId, + description: run.description, + }); + } + }, + (error: unknown) => { + resetSession(current); + finish(run, unavailableEvent(error)); + }, + ); + }; + + return { + async createOptimizationRun(input) { + if (disposed) { + throw new Error("The in-browser optimizer was disposed"); + } + const parsed = petrinautOptimizationManifestSchema.safeParse(input); + if (!parsed.success) { + throw invalidManifestError(parsed.error.issues); + } + const manifest = parsed.data; + const controller = createAbortController(); + const run: RunRecord = { + runId: generateUuid(), + manifest, + description: describeOptimization(manifest), + seeds: deriveOptimizationTrialSeeds( + manifest.execution.seed, + manifest.execution.seedsPerTrial ?? 1, + ), + log: createOptimizationRunLog(), + signal: controller.signal, + abort: () => controller.abort(), + status: "queued", + }; + runs.set(run.runId, run); + run.log.append({ + type: "started", + requestedTrials: manifest.study.trials, + }); + queue.push(run); + startNext(); + return { runId: run.runId }; + }, + attachOptimizationRun(runId, attachOptions = {}) { + const run = runs.get(runId); + if (!run) { + throw unknownRunError(runId); + } + return attachToLog(run.log, attachOptions); + }, + async cancelOptimizationRun(runId) { + const run = runs.get(runId); + if (!run || run.status === "finished") { + return; + } + run.abort(); + if (run.status === "running") { + session?.worker.postMessage({ type: "cancel", runId }); + return; + } + finish(run, cancelledEvent); + }, + dispose() { + disposed = true; + for (const run of runs.values()) { + if (run.status !== "finished") { + run.abort(); + run.status = "finished"; + run.log.append(cancelledEvent); + } + } + queue.length = 0; + active = null; + session?.worker.terminate(); + session = null; + }, + }; +}; + +export const createBrowserOptimization = ( + options: CreateBrowserOptimizationOptions = {}, +): PetrinautConnectedOptimization => ({ + kind: "connected", + connect: (channel) => + connectBrowserOptimization({ + channel, + pyodide: { ...defaultOptimizerPyodideConfig(), ...options.pyodide }, + createWorker: options.createWorker ?? createOptimizerWorker, + }), +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/create-optimizer-worker.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/create-optimizer-worker.ts new file mode 100644 index 00000000000..1b19996e0ef --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/create-optimizer-worker.ts @@ -0,0 +1,30 @@ +// eslint-disable-next-line import/default -- Vite resolves the `?worker&url` query to the bundled worker script's URL +import workerUrl from "./optimizer.worker.ts?worker&url"; + +import type { WorkerLike } from "../environment"; +import type { + OptimizerToMainMessage, + OptimizerToWorkerMessage, +} from "./messages"; + +/** A worker's `error` event; a script that fails to load fires one without a message. */ +export type OptimizerWorkerErrorEvent = { readonly message?: string }; + +export type OptimizerWorkerLike = WorkerLike< + OptimizerToWorkerMessage, + OptimizerToMainMessage +> & { + addEventListener( + type: "error", + listener: (event: OptimizerWorkerErrorEvent) => void, + ): void; +}; + +declare const Worker: new ( + scriptUrl: string, + options: { type: "module" }, +) => OptimizerWorkerLike; + +/** Pyodide loads through a dynamic `import()`, which only a module worker can run. */ +export const createOptimizerWorker = (): OptimizerWorkerLike => + new Worker(workerUrl, { type: "module" }); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts new file mode 100644 index 00000000000..9aaa7f11197 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts @@ -0,0 +1,109 @@ +import type { + OptimizationScalar, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationTrialOutcome, +} from "../optimization"; +import type { OptimizerPyodideConfig } from "./pyodide-config"; + +export type OptimizerBestTrial = { + trial: number; + parameters: Record; + objective: number; +}; + +/** One finished Optuna trial, as the Python study reports it. */ +export type OptimizerTrialPayload = { + trial: number; + parameters: Record; + objective: number | null; + state: "complete" | "pruned" | "failed"; + best: OptimizerBestTrial | null; +}; + +export type OptimizerStudySummary = { + requestedTrials: number; + completedTrials: number; + prunedTrials: number; + failedTrials: number; + best: OptimizerBestTrial | null; + /** Set when the study stopped early because the run was cancelled. */ + cancelled?: boolean; +}; + +export type OptimizerInitMessage = { + type: "init"; + pyodide: OptimizerPyodideConfig; + pythonSources: Readonly>; +}; + +export type OptimizerStartMessage = { + type: "start"; + runId: string; + description: PetrinautOptimizationDescribeResult; +}; + +export type OptimizerEvaluatedMessage = { + type: "evaluated"; + requestId: number; + outcome: PetrinautOptimizationTrialOutcome; +}; + +export type OptimizerCancelMessage = { + type: "cancel"; + runId: string; +}; + +export type OptimizerToWorkerMessage = + | OptimizerInitMessage + | OptimizerStartMessage + | OptimizerEvaluatedMessage + | OptimizerCancelMessage; + +export type OptimizerReadyMessage = { + type: "ready"; +}; + +export type OptimizerInitErrorMessage = { + type: "init-error"; + message: string; +}; + +export type OptimizerEvaluateMessage = { + type: "evaluate"; + runId: string; + requestId: number; + trial: number; + suggestedValues: Record; +}; + +export type OptimizerTrialMessage = { + type: "trial"; + runId: string; + event: OptimizerTrialPayload; +}; + +export type OptimizerCompleteMessage = { + type: "complete"; + runId: string; + summary: OptimizerStudySummary; +}; + +export type OptimizerCancelledMessage = { + type: "cancelled"; + runId: string; +}; + +export type OptimizerErrorMessage = { + type: "error"; + runId: string; + message: string; +}; + +export type OptimizerToMainMessage = + | OptimizerReadyMessage + | OptimizerInitErrorMessage + | OptimizerEvaluateMessage + | OptimizerTrialMessage + | OptimizerCompleteMessage + | OptimizerCancelledMessage + | OptimizerErrorMessage; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts new file mode 100644 index 00000000000..bc768c2a95f --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts @@ -0,0 +1,166 @@ +import { createWorkerThreadRuntime } from "../environment"; +import { createOptimizerStudyRunner } from "./study-runner"; + +import type { PetrinautOptimizationTrialOutcome } from "../optimization"; +import type { + OptimizerCancelMessage, + OptimizerEvaluatedMessage, + OptimizerInitMessage, + OptimizerStartMessage, + OptimizerToMainMessage, + OptimizerToWorkerMessage, +} from "./messages"; +import type { LoadPyodide } from "./pyodide-like"; +import type { OptimizerStudyRunner } from "./study-runner"; + +type ActiveRun = { + cancelled: boolean; + pendingRequestIds: Set; +}; + +const workerRuntime = createWorkerThreadRuntime< + OptimizerToWorkerMessage, + OptimizerToMainMessage +>(); + +let runner: OptimizerStudyRunner | null = null; +const runs = new Map(); +const pendingEvaluations = new Map< + number, + (outcome: PetrinautOptimizationTrialOutcome) => void +>(); +let nextRequestId = 1; + +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : String(error); + +const ensureTrailingSlash = (url: string): string => + url.endsWith("/") ? url : `${url}/`; + +const importLoadPyodide = async (indexURL: string): Promise => { + const module: unknown = await import( + /* @vite-ignore */ `${indexURL}pyodide.mjs` + ); + if ( + typeof module !== "object" || + module === null || + !("loadPyodide" in module) || + typeof module.loadPyodide !== "function" + ) { + throw new Error(`Pyodide at ${indexURL} exposes no loadPyodide`); + } + return module.loadPyodide as LoadPyodide; +}; + +const handleInit = (message: OptimizerInitMessage): void => { + const indexURL = ensureTrailingSlash(message.pyodide.indexURL); + runner = createOptimizerStudyRunner({ + loadPyodide: async (options) => + (await importLoadPyodide(indexURL))(options), + config: { ...message.pyodide, indexURL }, + pythonSources: message.pythonSources, + }); + runner.ready.then( + () => workerRuntime.postMessage({ type: "ready" }), + (error: unknown) => + workerRuntime.postMessage({ + type: "init-error", + message: errorMessage(error), + }), + ); +}; + +const handleStart = (message: OptimizerStartMessage): void => { + const { runId } = message; + if (!runner) { + workerRuntime.postMessage({ + type: "error", + runId, + message: "The optimizer worker received a study before its runtime", + }); + return; + } + const run: ActiveRun = { cancelled: false, pendingRequestIds: new Set() }; + runs.set(runId, run); + runner + .run({ + description: message.description, + evaluate: (trial, suggestedValues) => + new Promise((resolve) => { + const requestId = nextRequestId; + nextRequestId += 1; + pendingEvaluations.set(requestId, resolve); + run.pendingRequestIds.add(requestId); + workerRuntime.postMessage({ + type: "evaluate", + runId, + requestId, + trial, + suggestedValues, + }); + }), + onTrial: (event) => + workerRuntime.postMessage({ type: "trial", runId, event }), + isCancelled: () => run.cancelled, + }) + .then( + (summary) => + workerRuntime.postMessage( + summary.cancelled === true + ? { type: "cancelled", runId } + : { type: "complete", runId, summary }, + ), + (error: unknown) => + workerRuntime.postMessage({ + type: "error", + runId, + message: errorMessage(error), + }), + ) + .finally(() => { + runs.delete(runId); + }); +}; + +const handleEvaluated = (message: OptimizerEvaluatedMessage): void => { + const resolve = pendingEvaluations.get(message.requestId); + if (!resolve) { + return; + } + pendingEvaluations.delete(message.requestId); + for (const run of runs.values()) { + run.pendingRequestIds.delete(message.requestId); + } + resolve(message.outcome); +}; + +const handleCancel = (message: OptimizerCancelMessage): void => { + const run = runs.get(message.runId); + if (!run) { + return; + } + run.cancelled = true; + for (const requestId of run.pendingRequestIds) { + const resolve = pendingEvaluations.get(requestId); + pendingEvaluations.delete(requestId); + resolve?.({ kind: "pruned", reason: "cancelled" }); + } + run.pendingRequestIds.clear(); +}; + +workerRuntime.onMessage((message) => { + switch (message.type) { + case "init": + handleInit(message); + break; + case "start": + handleStart(message); + break; + case "evaluated": + handleEvaluated(message); + break; + case "cancel": + handleCancel(message); + break; + } +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.test.ts new file mode 100644 index 00000000000..a5e96b5de54 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.test.ts @@ -0,0 +1,33 @@ +import { describe, expect, it } from "vitest"; + +import runtimeLock from "@local/petrinaut-optimizer-core/runtime-lock.json"; + +import { + defaultOptimizerPyodideConfig, + micropipRequirements, +} from "./pyodide-config"; + +describe("defaultOptimizerPyodideConfig", () => { + it("pins the CDN runtime and the packages to the shared runtime lock", () => { + const config = defaultOptimizerPyodideConfig(); + + expect(config.indexURL).toBe( + `https://cdn.jsdelivr.net/pyodide/v${runtimeLock.pyodide}/full/`, + ); + expect(config.packages).toEqual(runtimeLock.packages); + expect(config.distributionPackages).toEqual( + runtimeLock.pyodideDistributionPackages, + ); + expect(config.packages.optuna).toMatch(/^\d+\.\d+\.\d+$/); + }); +}); + +describe("micropipRequirements", () => { + it("formats exact-version requirements", () => { + expect( + micropipRequirements({ + packages: { optuna: "4.9.0", colorlog: "6.10.1" }, + }), + ).toEqual(["optuna==4.9.0", "colorlog==6.10.1"]); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.ts new file mode 100644 index 00000000000..40f5acc1bb7 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-config.ts @@ -0,0 +1,23 @@ +import runtimeLock from "@local/petrinaut-optimizer-core/runtime-lock.json"; + +export type OptimizerPyodideConfig = { + /** Where `pyodide.mjs`, the runtime assets and the distribution packages load from. */ + readonly indexURL: string; + /** PyPI package name → exact version, installed with micropip. */ + readonly packages: Readonly>; + /** Packages loaded from the Pyodide distribution rather than from PyPI. */ + readonly distributionPackages: readonly string[]; +}; + +export const defaultOptimizerPyodideConfig = (): OptimizerPyodideConfig => ({ + indexURL: `https://cdn.jsdelivr.net/pyodide/v${runtimeLock.pyodide}/full/`, + packages: runtimeLock.packages, + distributionPackages: runtimeLock.pyodideDistributionPackages, +}); + +export const micropipRequirements = ( + config: Pick, +): string[] => + Object.entries(config.packages).map( + ([name, version]) => `${name}==${version}`, + ); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-like.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-like.ts new file mode 100644 index 00000000000..3e0d6456c88 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/pyodide-like.ts @@ -0,0 +1,32 @@ +/** + * The slice of the Pyodide API the optimizer runtime uses, typed structurally + * so published code never depends on the `pyodide` package's declarations. + */ +export type PyProxyLike = { + toJs(options?: { + dict_converter?: (entries: Iterable<[string, unknown]>) => unknown; + }): unknown; + destroy(): void; +}; + +export type PyodideLike = { + loadPackage(names: string | string[]): Promise; + pyimport(moduleName: string): unknown; + runPythonAsync(code: string): Promise; + FS: { + mkdirTree(path: string): void; + writeFile(path: string, data: string): void; + }; +}; + +export type LoadPyodide = (options: { + indexURL: string; +}) => Promise; + +export const isPyProxyLike = (value: unknown): value is PyProxyLike => + typeof value === "object" && + value !== null && + "toJs" in value && + typeof value.toJs === "function" && + "destroy" in value && + typeof value.destroy === "function"; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/python-sources.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/python-sources.ts new file mode 100644 index 00000000000..45e1cc98a0c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/python-sources.ts @@ -0,0 +1,14 @@ +import initSource from "@local/petrinaut-optimizer-core/python/__init__.py?raw"; +import askTellSource from "@local/petrinaut-optimizer-core/python/ask_tell.py?raw"; +import descriptionSource from "@local/petrinaut-optimizer-core/python/description.py?raw"; +import pyodideEntrySource from "@local/petrinaut-optimizer-core/python/pyodide_entry.py?raw"; +import studySource from "@local/petrinaut-optimizer-core/python/study.py?raw"; + +/** The optimizer's Python package, keyed by path relative to the import root. */ +export const optimizerPythonSources: Readonly> = { + "petrinaut_optimizer_core/__init__.py": initSource, + "petrinaut_optimizer_core/description.py": descriptionSource, + "petrinaut_optimizer_core/study.py": studySource, + "petrinaut_optimizer_core/ask_tell.py": askTellSource, + "petrinaut_optimizer_core/pyodide_entry.py": pyodideEntrySource, +}; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts new file mode 100644 index 00000000000..df6fa014170 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, it } from "vitest"; + +import { createAbortController } from "../environment"; +import { createOptimizationRunLog } from "./run-log"; + +import type { PetrinautOptimizationEvent } from "../optimization"; + +const collect = async ( + iterable: AsyncIterable, +): Promise => { + const events: PetrinautOptimizationEvent[] = []; + for await (const event of iterable) { + events.push(event); + } + return events; +}; + +const trial = ( + index: number, +): Parameters["append"]>[0] => ({ + type: "trial", + trial: index, + parameters: { rate: index }, + objective: index, + state: "complete", + best: { trial: index, parameters: { rate: index }, objective: index }, +}); + +const complete = { + type: "complete", + requestedTrials: 2, + completedTrials: 2, + prunedTrials: 0, + failedTrials: 0, + best: null, +} as const; + +describe("createOptimizationRunLog", () => { + it("stamps dense sequence numbers from 1 and closes on a terminal event", () => { + const log = createOptimizationRunLog(); + + expect(log.append({ type: "started", requestedTrials: 2 }).seq).toBe(1); + expect(log.append(trial(0)).seq).toBe(2); + expect(log.closed).toBe(false); + expect(log.append(complete).seq).toBe(3); + expect(log.closed).toBe(true); + expect(log.events.map((event) => event.seq)).toEqual([1, 2, 3]); + expect(() => log.append(trial(1))).toThrow( + "optimization run log is closed", + ); + }); + + it("replays the stored events past the cursor, then tails until the terminal event", async () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 2 }); + log.append(trial(0)); + + const replay = collect(log.replay({ cursor: 1 })); + await Promise.resolve(); + log.append(trial(1)); + log.append(complete); + + expect((await replay).map((event) => [event.type, event.seq])).toEqual([ + ["trial", 2], + ["trial", 3], + ["complete", 4], + ]); + }); + + it("ends at once when attached past the end of a closed log", async () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 2 }); + log.append(complete); + + expect(await collect(log.replay({ cursor: 2 }))).toEqual([]); + expect(await collect(log.replay())).toHaveLength(2); + }); + + it("notifies subscribers of each appended event", () => { + const log = createOptimizationRunLog(); + const seen: number[] = []; + const unsubscribe = log.subscribe((event) => { + seen.push(event.seq ?? -1); + }); + + log.append({ type: "started", requestedTrials: 2 }); + unsubscribe(); + log.append(trial(0)); + + expect(seen).toEqual([1]); + }); + + it("aborts a tailing replay with an AbortError", async () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 2 }); + const controller = createAbortController(); + + const replay = collect(log.replay({ signal: controller.signal })); + await Promise.resolve(); + controller.abort(); + + await expect(replay).rejects.toMatchObject({ name: "AbortError" }); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts new file mode 100644 index 00000000000..6bc67e50a6c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts @@ -0,0 +1,112 @@ +import type { AbortSignalLike } from "../environment"; +import type { PetrinautOptimizationEvent } from "../optimization"; + +type WithoutSeq = TEvent extends unknown ? Omit : never; + +/** An optimization event before the log stamps its sequence number. */ +export type OptimizationRunLogEvent = WithoutSeq; + +export type OptimizationRunLog = { + readonly events: readonly PetrinautOptimizationEvent[]; + /** True once a terminal `complete`/`error` event was appended. */ + readonly closed: boolean; + /** Stamps the next dense `seq` (from 1) and stores the event. Throws once closed. */ + append(event: OptimizationRunLogEvent): PetrinautOptimizationEvent; + subscribe(listener: (event: PetrinautOptimizationEvent) => void): () => void; + /** + * Yields the stored events with `seq` greater than `cursor`, then tails live + * events until the terminal one. Aborting the signal ends the iteration with + * an `AbortError`. + */ + replay(options?: { + cursor?: number; + signal?: AbortSignalLike; + }): AsyncIterable; +}; + +const isTerminalEvent = (event: PetrinautOptimizationEvent): boolean => + event.type === "complete" || event.type === "error"; + +const createAbortError = (): Error => { + const error = new Error("optimization run attachment aborted"); + error.name = "AbortError"; + return error; +}; + +export const createOptimizationRunLog = (): OptimizationRunLog => { + const events: PetrinautOptimizationEvent[] = []; + const listeners = new Set<(event: PetrinautOptimizationEvent) => void>(); + let closed = false; + + const subscribe = ( + listener: (event: PetrinautOptimizationEvent) => void, + ): (() => void) => { + listeners.add(listener); + return () => { + listeners.delete(listener); + }; + }; + + const waitForNextEvent = (signal?: AbortSignalLike): Promise => + new Promise((resolve, reject) => { + let unsubscribe = (): void => {}; + const onAbort = (): void => { + unsubscribe(); + reject(createAbortError()); + }; + unsubscribe = subscribe(() => { + unsubscribe(); + signal?.removeEventListener("abort", onAbort); + resolve(); + }); + signal?.addEventListener("abort", onAbort, { once: true }); + }); + + return { + get events() { + return events; + }, + get closed() { + return closed; + }, + append(event) { + if (closed) { + throw new Error("optimization run log is closed"); + } + const stamped = { ...event, seq: events.length + 1 }; + events.push(stamped); + if (isTerminalEvent(stamped)) { + closed = true; + } + for (const listener of listeners) { + listener(stamped); + } + return stamped; + }, + subscribe, + async *replay(options) { + const signal = options?.signal; + // Sequence numbers are dense from 1, so the first event past the cursor + // sits at index `cursor`. + let index = Math.max(0, Math.min(options?.cursor ?? 0, events.length)); + for (;;) { + if (signal?.aborted) { + throw createAbortError(); + } + const event = events[index]; + if (event) { + index++; + yield event; + if (isTerminalEvent(event)) { + return; + } + continue; + } + if (closed) { + return; + } + await waitForNextEvent(signal); + } + }, + }; +}; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts new file mode 100644 index 00000000000..1b2be8c3e7e --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts @@ -0,0 +1,257 @@ +/** + * Runs the real Optuna study under Pyodide in Node. The Pyodide runtime comes + * from the `pyodide` devDependency; its packages and the PyPI wheels download + * from the CDN on the first run (Pyodide caches the distribution packages + * next to the runtime). When `CI` is set and that download fails, the tests + * skip instead of failing, so an offline CI runner does not turn a network + * outage into a red build. Every other load failure, a Python syntax or + * import error included, fails the suite. + */ +import { loadPyodide } from "pyodide"; +import { beforeAll, describe, expect, test } from "vitest"; + +import { defaultOptimizerPyodideConfig } from "./pyodide-config"; +import { optimizerPythonSources } from "./python-sources"; +import { + createOptimizerStudyRunner, + type OptimizerStudyRunner, +} from "./study-runner"; + +import type { + OptimizationScalar, + PetrinautOptimizationDescribeResult, +} from "../optimization"; +import type { OptimizerTrialPayload } from "./messages"; + +declare const process: { + readonly env: Readonly>; +}; + +const loadTimeout = 180_000; +const runningInCi = (process.env.CI ?? "") !== ""; + +const downloadFailurePattern = + /fetch|network|request failed|ENOTFOUND|ECONN|EAI_AGAIN|timed out|Can't fetch/i; + +const isDownloadFailure = (error: unknown): boolean => + error instanceof Error && + (downloadFailurePattern.test(error.message) || + (error.cause instanceof Error && + downloadFailurePattern.test(error.cause.message))); + +const description: PetrinautOptimizationDescribeResult = { + direction: "minimize", + study: { trials: 30, sampler: "tpe", seed: 7, seedsPerTrial: 1 }, + parameters: [ + { + identifier: "rate", + type: "float", + default: 1, + minimum: 0.1, + maximum: 10, + scale: "log", + }, + { + identifier: "offset", + type: "float", + default: 0, + minimum: -5, + maximum: 5, + scale: "linear", + }, + { + identifier: "count", + type: "int", + default: 4, + minimum: 2, + maximum: 20, + step: 2, + scale: "linear", + }, + { identifier: "enabled", type: "boolean", default: false }, + ], +}; + +const asNumber = (value: OptimizationScalar | undefined): number => { + if (typeof value !== "number") { + throw new Error(`expected a number, received ${String(value)}`); + } + return value; +}; + +const objectiveOf = (values: Record): number => + (Math.log(asNumber(values.rate)) - Math.log(2)) ** 2 + + (asNumber(values.offset) - 1) ** 2 + + (asNumber(values.count) - 8) ** 2 / 16 + + (values.enabled === true ? 0 : 1); + +let runner: OptimizerStudyRunner; +let loadFailure: string | null = null; + +beforeAll(async () => { + runner = createOptimizerStudyRunner({ + // The runtime itself comes from node_modules; only the packages download. + loadPyodide: () => loadPyodide(), + config: defaultOptimizerPyodideConfig(), + pythonSources: optimizerPythonSources, + }); + try { + await runner.ready; + } catch (error) { + if (!runningInCi || !isDownloadFailure(error)) { + throw error; + } + loadFailure = error instanceof Error ? error.message : String(error); + } +}, loadTimeout); + +const skipWhenOffline = (skip: (note?: string) => never): void => { + if (loadFailure !== null) { + skip(`Pyodide packages could not be downloaded in CI: ${loadFailure}`); + } +}; + +describe("createOptimizerStudyRunner", () => { + test( + "drives a seeded TPE study through the host evaluate callback", + async ({ skip }) => { + skipWhenOffline(skip); + const trials: OptimizerTrialPayload[] = []; + const evaluated: Record[] = []; + + const summary = await runner.run({ + description, + evaluate: async (trial, suggestedValues) => { + expect(trial).toBe(evaluated.length); + evaluated.push(suggestedValues); + return trial === 3 + ? { kind: "pruned", reason: "no frames" } + : { kind: "objective", objective: objectiveOf(suggestedValues) }; + }, + onTrial: (event) => { + trials.push(event); + }, + isCancelled: () => false, + }); + + expect(evaluated).toHaveLength(30); + for (const values of evaluated) { + expect(Object.keys(values).sort()).toEqual([ + "count", + "enabled", + "offset", + "rate", + ]); + expect(asNumber(values.rate)).toBeGreaterThanOrEqual(0.1); + expect(asNumber(values.rate)).toBeLessThanOrEqual(10); + expect(asNumber(values.offset)).toBeGreaterThanOrEqual(-5); + expect(asNumber(values.offset)).toBeLessThanOrEqual(5); + expect(asNumber(values.count) % 2).toBe(0); + expect(asNumber(values.count)).toBeGreaterThanOrEqual(2); + expect(asNumber(values.count)).toBeLessThanOrEqual(20); + expect(typeof values.enabled).toBe("boolean"); + } + + expect(trials.map((event) => event.trial)).toEqual( + Array.from({ length: 30 }, (_, index) => index), + ); + let bestSoFar = Number.POSITIVE_INFINITY; + for (const [index, event] of trials.entries()) { + expect(event.parameters).toEqual(evaluated[index]); + if (index === 3) { + expect(event.state).toBe("pruned"); + expect(event.objective).toBeNull(); + } else { + expect(event.state).toBe("complete"); + expect(event.objective).toBeCloseTo(objectiveOf(event.parameters), 9); + bestSoFar = Math.min(bestSoFar, event.objective ?? Infinity); + } + if (index === 0) { + expect(event.best).toEqual({ + trial: 0, + parameters: event.parameters, + objective: event.objective, + }); + } + expect(event.best?.objective).toBeCloseTo(bestSoFar, 9); + } + + expect(summary).toEqual({ + requestedTrials: 30, + completedTrials: 29, + prunedTrials: 1, + failedTrials: 0, + best: trials.at(-1)?.best, + }); + }, + loadTimeout, + ); + + test( + "stops early once the host reports cancellation", + async ({ skip }) => { + skipWhenOffline(skip); + let evaluations = 0; + const trials: OptimizerTrialPayload[] = []; + + const summary = await runner.run({ + description, + evaluate: async (_trial, suggestedValues) => { + evaluations += 1; + return { kind: "objective", objective: objectiveOf(suggestedValues) }; + }, + onTrial: (event) => { + trials.push(event); + }, + isCancelled: () => evaluations >= 5, + }); + + expect(evaluations).toBe(5); + expect(trials).toHaveLength(4); + expect(summary).toMatchObject({ + requestedTrials: 30, + completedTrials: 4, + cancelled: true, + }); + }, + loadTimeout, + ); + + test( + "samples deterministically for a seed and rejects a non-finite objective", + async ({ skip }) => { + skipWhenOffline(skip); + const sequences: Record[][] = []; + for (let repeat = 0; repeat < 2; repeat++) { + const sequence: Record[] = []; + await runner.run({ + description: { + ...description, + study: { ...description.study, trials: 8 }, + }, + evaluate: async (_trial, suggestedValues) => { + sequence.push(suggestedValues); + return { + kind: "objective", + objective: objectiveOf(suggestedValues), + }; + }, + onTrial: () => {}, + isCancelled: () => false, + }); + sequences.push(sequence); + } + expect(sequences[0]).toEqual(sequences[1]); + + await expect( + runner.run({ + description, + evaluate: async () => ({ kind: "objective", objective: Number.NaN }), + onTrial: () => {}, + isCancelled: () => false, + }), + ).rejects.toThrow("trial objective must be a finite number"); + }, + loadTimeout, + ); +}); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts new file mode 100644 index 00000000000..f5696e60911 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts @@ -0,0 +1,222 @@ +/** + * @talksTo optimizer-core via Python sources loaded into Pyodide + */ +import { micropipRequirements } from "./pyodide-config"; +import { isPyProxyLike } from "./pyodide-like"; + +import type { + OptimizationScalar, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationTrialOutcome, +} from "../optimization"; +import type { + OptimizerBestTrial, + OptimizerStudySummary, + OptimizerTrialPayload, +} from "./messages"; +import type { OptimizerPyodideConfig } from "./pyodide-config"; +import type { LoadPyodide, PyodideLike } from "./pyodide-like"; + +export type OptimizerStudyRunInput = { + description: PetrinautOptimizationDescribeResult; + evaluate( + trial: number, + suggestedValues: Record, + ): Promise; + onTrial(event: OptimizerTrialPayload): void; + isCancelled(): boolean; +}; + +export type OptimizerStudyRunner = { + /** Settles once Pyodide, the packages and the optimizer sources are loaded. */ + readonly ready: Promise; + /** Runs one study; concurrent calls run one after another. */ + run(input: OptimizerStudyRunInput): Promise; +}; + +/** The outcome shape `ask_tell.run_study` expects from its evaluate callback. */ +type PythonTrialOutcome = { objective: number } | { pruned: string }; + +type PyodideEntryModule = { + run_browser_study( + descriptionJson: string, + evaluate: (values: unknown) => Promise, + onTrial: (payload: unknown) => void, + isCancelled: () => boolean, + ): Promise; +}; + +const pythonSourceRoot = "/home/pyodide"; + +const toJsValue = (value: unknown): unknown => + isPyProxyLike(value) + ? value.toJs({ dict_converter: Object.fromEntries }) + : value; + +const asRecord = (value: unknown, what: string): Record => { + if (typeof value !== "object" || value === null || Array.isArray(value)) { + throw new Error(`The optimizer returned a malformed ${what}`); + } + return value as Record; +}; + +const asNumber = (value: unknown, what: string): number => { + if (typeof value !== "number") { + throw new Error(`The optimizer returned a non-numeric ${what}`); + } + return value; +}; + +const asParameters = (value: unknown): Record => { + const parameters: Record = {}; + for (const [identifier, scalar] of Object.entries( + asRecord(value, "parameter set"), + )) { + if (typeof scalar !== "number" && typeof scalar !== "boolean") { + throw new Error( + `The optimizer suggested a non-scalar value for "${identifier}"`, + ); + } + parameters[identifier] = scalar; + } + return parameters; +}; + +const asBest = (value: unknown): OptimizerBestTrial | null => { + if (value === null || value === undefined) { + return null; + } + const record = asRecord(value, "best trial"); + return { + trial: asNumber(record.trial, "best trial number"), + parameters: asParameters(record.parameters), + objective: asNumber(record.objective, "best objective"), + }; +}; + +const asTrialState = (value: unknown): OptimizerTrialPayload["state"] => { + if (value === "complete" || value === "pruned" || value === "failed") { + return value; + } + throw new Error(`The optimizer reported an unknown trial state`); +}; + +const normalizeTrialPayload = (value: unknown): OptimizerTrialPayload => { + const record = asRecord(value, "trial payload"); + const objective = record.objective; + return { + trial: asNumber(record.trial, "trial number"), + parameters: asParameters(record.parameters), + objective: + objective === null || objective === undefined + ? null + : asNumber(objective, "trial objective"), + state: asTrialState(record.state), + best: asBest(record.best), + }; +}; + +const normalizeSummary = (value: unknown): OptimizerStudySummary => { + const record = asRecord(value, "study summary"); + return { + requestedTrials: asNumber(record.requestedTrials, "requested trial count"), + completedTrials: asNumber(record.completedTrials, "completed trial count"), + prunedTrials: asNumber(record.prunedTrials, "pruned trial count"), + failedTrials: asNumber(record.failedTrials, "failed trial count"), + best: asBest(record.best), + ...(record.cancelled === true ? { cancelled: true } : {}), + }; +}; + +const toPythonOutcome = ( + outcome: PetrinautOptimizationTrialOutcome, +): PythonTrialOutcome => + outcome.kind === "objective" + ? { objective: outcome.objective } + : { pruned: outcome.reason }; + +const writePythonSources = ( + pyodide: PyodideLike, + sources: Readonly>, +): void => { + for (const [path, source] of Object.entries(sources)) { + const absolutePath = `${pythonSourceRoot}/${path}`; + pyodide.FS.mkdirTree(absolutePath.slice(0, absolutePath.lastIndexOf("/"))); + pyodide.FS.writeFile(absolutePath, source); + } +}; + +const loadOptimizerEntry = async (options: { + loadPyodide: LoadPyodide; + config: OptimizerPyodideConfig; + pythonSources: Readonly>; +}): Promise => { + const { config } = options; + const pyodide = await options.loadPyodide({ indexURL: config.indexURL }); + await pyodide.loadPackage([...config.distributionPackages, "micropip"]); + const requirements = JSON.stringify(micropipRequirements(config)); + await pyodide.runPythonAsync( + `import micropip\nawait micropip.install(${requirements}, deps=False)`, + ); + writePythonSources(pyodide, options.pythonSources); + const root = JSON.stringify(pythonSourceRoot); + await pyodide.runPythonAsync( + `import sys\nif ${root} not in sys.path:\n sys.path.insert(0, ${root})`, + ); + return pyodide.pyimport( + "petrinaut_optimizer_core.pyodide_entry", + ) as PyodideEntryModule; +}; + +export const createOptimizerStudyRunner = (options: { + loadPyodide: LoadPyodide; + config: OptimizerPyodideConfig; + pythonSources: Readonly>; +}): OptimizerStudyRunner => { + const entry = loadOptimizerEntry(options); + const ready = entry.then(() => undefined); + let queue: Promise = ready; + + const runStudy = async ( + input: OptimizerStudyRunInput, + ): Promise => { + const module = await entry; + // Optuna numbers an in-memory study's trials densely from 0 in ask order, + // so the evaluate call count is the trial number. + let nextTrial = 0; + const evaluate = (values: unknown): Promise => { + // Argument proxies are destroyed when this call returns: convert before + // the first await. + const suggestedValues = asParameters(toJsValue(values)); + const trial = nextTrial; + nextTrial += 1; + return input.evaluate(trial, suggestedValues).then(toPythonOutcome); + }; + const onTrial = (payload: unknown): void => { + input.onTrial(normalizeTrialPayload(toJsValue(payload))); + }; + const result = await module.run_browser_study( + JSON.stringify(input.description), + evaluate, + onTrial, + () => input.isCancelled(), + ); + const summary = normalizeSummary(toJsValue(result)); + if (isPyProxyLike(result)) { + result.destroy(); + } + return summary; + }; + + return { + ready, + run(input) { + const result = queue.then(() => runStudy(input)); + queue = result.then( + () => undefined, + () => undefined, + ); + return result; + }, + }; +}; diff --git a/libs/@hashintel/petrinaut-core/src/environment.ts b/libs/@hashintel/petrinaut-core/src/environment.ts index 3231eb13378..f6a2772729d 100644 --- a/libs/@hashintel/petrinaut-core/src/environment.ts +++ b/libs/@hashintel/petrinaut-core/src/environment.ts @@ -92,3 +92,14 @@ export function createWorkerThreadRuntime< }, }; } + +export interface AbortControllerLike { + readonly signal: AbortSignalLike; + abort(): void; +} + +declare const AbortController: new () => AbortControllerLike; + +export function createAbortController(): AbortControllerLike { + return new AbortController(); +} diff --git a/libs/@hashintel/petrinaut-core/src/optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization.ts index ed6c8d6356b..7addb8cd669 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.ts @@ -13,6 +13,9 @@ export const PETRINAUT_OPTIMIZATION_MAX_SEEDS_PER_TRIAL = 100; const optimizationScalarSchema = z.union([z.number(), z.boolean()]); +/** A value Optuna may suggest or a fixed binding may hold. */ +export type OptimizationScalar = z.infer; + export const petrinautContinuousOptimizationDomainSchema = z .strictObject({ kind: z.literal("continuous"), @@ -700,3 +703,93 @@ export type PetrinautOptimization = { /** Idempotently stop a detached run server-side. */ cancelOptimizationRun(runId: string): Promise; }; + +/** + * One trial's computation, as the optimizer hands it to whoever runs + * simulations for it. + * + * The optimizer never simulates. It proposes values and asks its channel for + * the objective, so the host decides where and how a trial's runs happen and + * can show them as they compute. + */ +export type PetrinautOptimizationTrialRequest = { + readonly runId: string; + /** Optuna's trial number, from 0. */ + readonly trial: number; + /** The frozen study the trial belongs to. */ + readonly manifest: PetrinautOptimizationManifest; + /** The optimizer's suggestions for the optimized parameters only. */ + readonly suggestedValues: Readonly>; + /** + * Every scenario parameter's value for this trial: fixed bindings merged + * with the suggestions, booleans as 0/1 as the scenario compiler expects. + */ + readonly scenarioParameterValues: Readonly>; + /** + * The seeds the trial's simulations run with. The same sequence for every + * trial (common random numbers), derived as the CLI derives them. + */ + readonly seeds: readonly number[]; + /** Aborted when the run is cancelled; the host should stop the trial's runs. */ + readonly signal: AbortSignalLike; +}; + +export type PetrinautOptimizationTrialOutcome = + | { + readonly kind: "objective"; + /** The mean of the per-seed objectives; finite. */ + readonly objective: number; + readonly replicates?: readonly { seed: number; objective: number }[]; + } + | { + /** The host could not run the trial; Optuna records it as pruned. */ + readonly kind: "pruned"; + readonly reason: string; + }; + +/** + * A way to communicate with the running optimization. + * + * The host implements it; the optimizer calls it once per trial. Everything + * the optimizer needs computed goes through here, so the host can stream those + * runs into its own metrics views instead of receiving only a number. + */ +export type PetrinautOptimizationChannel = { + evaluateTrial( + this: void, + request: PetrinautOptimizationTrialRequest, + ): Promise; +}; + +/** + * An optimization capability that runs where the host runs and needs the + * host's compute: connect it to a channel to obtain the capability. + * + * The remote capability (`PetrinautOptimization`) is self-contained because + * the service owns its simulations; this one is not, by design. + */ +export type PetrinautConnectedOptimization = { + readonly kind: "connected"; + connect( + this: void, + channel: PetrinautOptimizationChannel, + ): PetrinautOptimization & { dispose(this: void): void }; +}; + +/** What a host supplies: a remote capability, or one to connect locally. */ +export type PetrinautOptimizationSource = + | PetrinautOptimization + | PetrinautConnectedOptimization; + +export const isConnectedOptimization = ( + source: PetrinautOptimizationSource, +): source is PetrinautConnectedOptimization => + (source as Partial).kind === "connected"; + +export { + deriveOptimizationTrialSeeds, + describeOptimization, + describeOptimizationParameter, + resolveTrialScenarioParameterValues, + validateSuggestedOptimizationValue, +} from "./optimization/describe"; diff --git a/libs/@hashintel/petrinaut-core/src/optimization/describe.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/describe.test.ts new file mode 100644 index 00000000000..1b471951b74 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/describe.test.ts @@ -0,0 +1,140 @@ +import { describe, expect, it } from "vitest"; + +import { createOptimizationManifest } from "../shared/optimization-manifest.fixtures"; +import { + deriveOptimizationTrialSeeds, + describeOptimization, + resolveTrialScenarioParameterValues, +} from "./describe"; + +describe("describeOptimization", () => { + it("reports the direction, the seeded study and only the optimized parameters", () => { + const manifest = createOptimizationManifest(); + + expect(describeOptimization(manifest)).toEqual({ + direction: "maximize", + study: { trials: 20, sampler: "tpe", seed: 42, seedsPerTrial: 1 }, + parameters: [ + { + identifier: "rate", + type: "float", + default: 0.5, + minimum: 0.1, + maximum: 2, + scale: "log", + }, + { + identifier: "count", + type: "int", + default: 10, + minimum: 2, + maximum: 10, + step: 2, + scale: "linear", + }, + { identifier: "enabled", type: "boolean", default: true }, + ], + }); + }); + + it("passes the configured seeds per trial through", () => { + const manifest = createOptimizationManifest({ + execution: { seed: 7, dt: 0.1, maxTime: 100, seedsPerTrial: 3 }, + }); + + expect(describeOptimization(manifest).study).toEqual({ + trials: 20, + sampler: "tpe", + seed: 7, + seedsPerTrial: 3, + }); + }); +}); + +describe("deriveOptimizationTrialSeeds", () => { + it("keeps the base seed first and derives the documented sequence", () => { + expect(deriveOptimizationTrialSeeds(42, 1)).toEqual([42]); + expect(deriveOptimizationTrialSeeds(42, 2)).toEqual([42, 1_013_904_268]); + + const seeds = deriveOptimizationTrialSeeds(42, 100); + expect(new Set(seeds).size).toBe(100); + for (const seed of seeds) { + expect(Number.isInteger(seed)).toBe(true); + expect(seed).toBeGreaterThanOrEqual(0); + expect(seed).toBeLessThanOrEqual(2_147_483_647); + } + }); +}); + +describe("resolveTrialScenarioParameterValues", () => { + const manifest = createOptimizationManifest(); + + it("merges fixed bindings with the suggestions and encodes booleans as 0/1", () => { + expect( + resolveTrialScenarioParameterValues(manifest, { + rate: 1.5, + count: 6, + enabled: false, + }), + ).toEqual({ rate: 1.5, count: 6, enabled: 0, share: 0.25 }); + expect( + resolveTrialScenarioParameterValues(manifest, { + rate: 0.1, + count: 10, + enabled: true, + }).enabled, + ).toBe(1); + }); + + it("requires every and only optimized value", () => { + expect(() => + resolveTrialScenarioParameterValues(manifest, { rate: 1, count: 2 }), + ).toThrow('Missing optimized parameter "enabled"'); + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: 1, + count: 2, + enabled: true, + share: 0.5, + }), + ).toThrow('Unexpected optimization parameter "share"'); + }); + + it("validates each suggestion against its domain", () => { + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: 3, + count: 2, + enabled: true, + }), + ).toThrow('Optimization parameter "rate" must be between 0.1 and 2'); + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: 1, + count: 5, + enabled: true, + }), + ).toThrow('Optimization parameter "count" must align with step 2 from 2'); + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: 1, + count: 2.5, + enabled: true, + }), + ).toThrow('Optimization parameter "count" must be an integer'); + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: 1, + count: 2, + enabled: 1, + }), + ).toThrow('Optimization parameter "enabled" must be boolean'); + expect(() => + resolveTrialScenarioParameterValues(manifest, { + rate: true, + count: 2, + enabled: true, + }), + ).toThrow('Optimization parameter "rate" must be numeric'); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/optimization/describe.ts b/libs/@hashintel/petrinaut-core/src/optimization/describe.ts new file mode 100644 index 00000000000..b1a5937766b --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/describe.ts @@ -0,0 +1,192 @@ +/** + * @layerRoot core.optimization + * @role Derives an Optuna study description, trial seeds and per-trial scenario parameter values from an optimization manifest, shared by the CLI and the browser runtime + */ +import { deriveRunSeed } from "../simulation/monte-carlo/run-state"; + +import type { + OptimizationScalar, + PetrinautOptimizationDescribeParameter, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationDomain, + PetrinautOptimizationManifest, +} from "../optimization"; +import type { Scenario } from "../types/sdcpn"; + +type ScenarioParameter = Scenario["scenarioParameters"][number]; + +type OptimizedParameter = { + parameter: ScenarioParameter; + domain: PetrinautOptimizationDomain; +}; + +export const describeOptimizationParameter = ( + parameter: ScenarioParameter, + domain: PetrinautOptimizationDomain, +): PetrinautOptimizationDescribeParameter => { + switch (domain.kind) { + case "continuous": + return { + identifier: parameter.identifier, + type: "float", + default: parameter.default, + minimum: domain.minimum, + maximum: domain.maximum, + scale: domain.scale, + }; + case "integer": + return { + identifier: parameter.identifier, + type: "int", + default: parameter.default, + minimum: domain.minimum, + maximum: domain.maximum, + step: domain.step, + scale: domain.scale, + }; + case "boolean": + return { + identifier: parameter.identifier, + type: "boolean", + default: parameter.default !== 0, + }; + } +}; + +export const validateSuggestedOptimizationValue = ( + parameter: ScenarioParameter, + domain: PetrinautOptimizationDomain, + value: OptimizationScalar, +): void => { + if (domain.kind === "boolean") { + if (typeof value !== "boolean") { + throw new Error( + `Optimization parameter "${parameter.identifier}" must be boolean`, + ); + } + return; + } + if (typeof value !== "number") { + throw new Error( + `Optimization parameter "${parameter.identifier}" must be numeric`, + ); + } + if (value < domain.minimum || value > domain.maximum) { + throw new Error( + `Optimization parameter "${parameter.identifier}" must be between ${domain.minimum} and ${domain.maximum}`, + ); + } + if (domain.kind === "integer") { + if (!Number.isInteger(value)) { + throw new Error( + `Optimization parameter "${parameter.identifier}" must be an integer`, + ); + } + if ((value - domain.minimum) % domain.step !== 0) { + throw new Error( + `Optimization parameter "${parameter.identifier}" must align with step ${domain.step} from ${domain.minimum}`, + ); + } + } +}; + +/** + * Derives one trial's run seeds. Run 0 keeps the base seed, so a single-seed + * trial matches a plain seeded run; later runs use the Monte Carlo derivation. + * Every trial gets the same sequence: common random numbers. + */ +export const deriveOptimizationTrialSeeds = ( + baseSeed: number, + seedsPerTrial: number, +): number[] => + Array.from({ length: seedsPerTrial }, (_, index) => + index === 0 ? baseSeed : deriveRunSeed(baseSeed, index), + ); + +const getOptimizationScenario = ( + manifest: PetrinautOptimizationManifest, +): Scenario => { + const scenario = manifest.model.definition.scenarios?.[0]; + const metric = manifest.model.definition.metrics?.[0]; + if (!scenario || !metric) { + throw new Error( + "An optimization manifest requires exactly one scenario and one metric", + ); + } + return scenario; +}; + +const listOptimizedParameters = ( + manifest: PetrinautOptimizationManifest, + scenario: Scenario, +): OptimizedParameter[] => + scenario.scenarioParameters.flatMap((parameter) => { + const binding = manifest.scenario.parameterBindings[parameter.identifier]; + return binding?.kind === "optimize" + ? [{ parameter, domain: binding.domain }] + : []; + }); + +export const describeOptimization = ( + manifest: PetrinautOptimizationManifest, +): PetrinautOptimizationDescribeResult => { + const scenario = getOptimizationScenario(manifest); + return { + direction: manifest.objective.direction, + study: { + ...manifest.study, + seed: manifest.execution.seed, + seedsPerTrial: manifest.execution.seedsPerTrial ?? 1, + }, + parameters: listOptimizedParameters(manifest, scenario).map( + ({ parameter, domain }) => + describeOptimizationParameter(parameter, domain), + ), + }; +}; + +/** + * Every scenario parameter's value for one trial: the fixed bindings merged + * with the validated suggestions, booleans as 0/1 as the scenario compiler + * expects. Throws when a suggestion is missing, unexpected or off-domain. + */ +export const resolveTrialScenarioParameterValues = ( + manifest: PetrinautOptimizationManifest, + suggestedValues: Readonly>, +): Record => { + const scenario = getOptimizationScenario(manifest); + const optimizedParameters = listOptimizedParameters(manifest, scenario); + const optimizedIdentifiers = new Set( + optimizedParameters.map(({ parameter }) => parameter.identifier), + ); + for (const { parameter } of optimizedParameters) { + if (!Object.hasOwn(suggestedValues, parameter.identifier)) { + throw new Error(`Missing optimized parameter "${parameter.identifier}"`); + } + } + for (const identifier of Object.keys(suggestedValues)) { + if (!optimizedIdentifiers.has(identifier)) { + throw new Error(`Unexpected optimization parameter "${identifier}"`); + } + } + + const scenarioParameterValues: Record = {}; + for (const parameter of scenario.scenarioParameters) { + const binding = manifest.scenario.parameterBindings[parameter.identifier]; + if (!binding) { + throw new Error( + `Scenario parameter "${parameter.identifier}" has no binding`, + ); + } + const value = + binding.kind === "fixed" + ? binding.value + : suggestedValues[parameter.identifier]!; + if (binding.kind === "optimize") { + validateSuggestedOptimizationValue(parameter, binding.domain, value); + } + scenarioParameterValues[parameter.identifier] = + typeof value === "boolean" ? (value ? 1 : 0) : value; + } + return scenarioParameterValues; +}; diff --git a/libs/@hashintel/petrinaut-core/src/shared/optimization-manifest.fixtures.ts b/libs/@hashintel/petrinaut-core/src/shared/optimization-manifest.fixtures.ts new file mode 100644 index 00000000000..f19052c4508 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/shared/optimization-manifest.fixtures.ts @@ -0,0 +1,80 @@ +import { petrinautOptimizationManifestSchema } from "../optimization"; + +import type { PetrinautOptimizationManifest } from "../optimization"; + +const scenario = { + id: "baseline", + name: "Baseline", + scenarioParameters: [ + { identifier: "rate", type: "real" as const, default: 0.5 }, + { identifier: "count", type: "integer" as const, default: 10 }, + { identifier: "enabled", type: "boolean" as const, default: 1 }, + { identifier: "share", type: "ratio" as const, default: 0.25 }, + ], + parameterOverrides: {}, + initialState: { type: "per_place" as const, content: {} }, +}; + +/** + * A study over two optimized numeric parameters and one optimized boolean, + * with one fixed ratio: enough to exercise every describe and resolve branch. + */ +export const createOptimizationManifestInput = () => ({ + kind: "petrinaut-optimization" as const, + version: 1 as const, + name: "Find the best rate", + model: { + title: "Example", + definition: { + places: [], + transitions: [], + types: [], + differentialEquations: [], + parameters: [], + subnets: [], + componentInstances: [], + scenarios: [scenario], + metrics: [{ id: "profit", name: "Profit", code: "return 1;" }], + }, + }, + scenario: { + id: "baseline", + parameterBindings: { + rate: { + kind: "optimize" as const, + domain: { + kind: "continuous" as const, + minimum: 0.1, + maximum: 2, + scale: "log" as const, + }, + }, + count: { + kind: "optimize" as const, + domain: { + kind: "integer" as const, + minimum: 2, + maximum: 10, + step: 2, + scale: "linear" as const, + }, + }, + enabled: { + kind: "optimize" as const, + domain: { kind: "boolean" as const }, + }, + share: { kind: "fixed" as const, value: 0.25 }, + }, + }, + objective: { metricId: "profit", direction: "maximize" as const }, + execution: { seed: 42, dt: 0.1, maxTime: 100 }, + study: { trials: 20, sampler: "tpe" as const }, +}); + +export const createOptimizationManifest = ( + overrides: Record = {}, +): PetrinautOptimizationManifest => + petrinautOptimizationManifestSchema.parse({ + ...createOptimizationManifestInput(), + ...overrides, + }); diff --git a/libs/@hashintel/petrinaut-core/src/vite-types.d.ts b/libs/@hashintel/petrinaut-core/src/vite-types.d.ts index b58971305e1..371b53306a1 100644 --- a/libs/@hashintel/petrinaut-core/src/vite-types.d.ts +++ b/libs/@hashintel/petrinaut-core/src/vite-types.d.ts @@ -13,3 +13,8 @@ declare module "*?worker&inline" { export default WorkerConstructor; } + +declare module "*?worker&url" { + const workerUrl: string; + export default workerUrl; +} diff --git a/libs/@hashintel/petrinaut-core/src/workers/README.md b/libs/@hashintel/petrinaut-core/src/workers/README.md index 9e01e0793f2..4087cdb5f52 100644 --- a/libs/@hashintel/petrinaut-core/src/workers/README.md +++ b/libs/@hashintel/petrinaut-core/src/workers/README.md @@ -14,6 +14,7 @@ the runtimes stay testable on the main thread. | `lsp.ts` | the language server | | `simulation.ts` | frame computation for a single run | | `monte-carlo.ts` | batched runs reporting only metrics | +| `optimizer.ts` | the Optuna study, in Pyodide | Separate export subpaths rather than one worker, so a host pays only for the threads it uses — an editor with no experiments open never loads the Monte diff --git a/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts b/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts new file mode 100644 index 00000000000..9f5bef00c34 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts @@ -0,0 +1,23 @@ +export { + createOptimizerWorker, + type OptimizerWorkerErrorEvent, + type OptimizerWorkerLike, +} from "../browser-optimization/create-optimizer-worker"; +export type { + OptimizerBestTrial, + OptimizerCancelledMessage, + OptimizerCancelMessage, + OptimizerCompleteMessage, + OptimizerErrorMessage, + OptimizerEvaluatedMessage, + OptimizerEvaluateMessage, + OptimizerInitErrorMessage, + OptimizerInitMessage, + OptimizerReadyMessage, + OptimizerStartMessage, + OptimizerStudySummary, + OptimizerToMainMessage, + OptimizerToWorkerMessage, + OptimizerTrialMessage, + OptimizerTrialPayload, +} from "../browser-optimization/messages"; diff --git a/libs/@hashintel/petrinaut-core/vite.config.ts b/libs/@hashintel/petrinaut-core/vite.config.ts index 045cc84c3e6..6f0f9d8edfc 100644 --- a/libs/@hashintel/petrinaut-core/vite.config.ts +++ b/libs/@hashintel/petrinaut-core/vite.config.ts @@ -23,6 +23,11 @@ export default defineConfig(({ command }) => ({ // Dependency-free instantiation of compiled HIR artifacts. "hir-runtime": resolve(packageRoot, "src/hir-runtime.ts"), optimization: resolve(packageRoot, "src/optimization.ts"), + // Runs the Optuna study in a Pyodide worker; inlines the Python sources. + "browser-optimization": resolve( + packageRoot, + "src/browser-optimization.ts", + ), // Dependency-free entry: the selection vocabulary alone, for hosts that // validate selection in a route or a server function. selection: resolve(packageRoot, "src/selection.ts"), @@ -40,6 +45,7 @@ export default defineConfig(({ command }) => ({ packageRoot, "src/workers/monte-carlo.ts", ), + "workers/optimizer": resolve(packageRoot, "src/workers/optimizer.ts"), "workers/simulation": resolve(packageRoot, "src/workers/simulation.ts"), }, fileName: (_format, entryName) => `${entryName}.js`, @@ -94,9 +100,12 @@ export default defineConfig(({ command }) => ({ ], experimental: { + // A worker URL resolved against the importing module survives bundling by + // a host: it becomes `new URL(, import.meta.url)`, which the host's + // bundler copies as an asset, where a page-relative path would 404. renderBuiltUrl: (filename) => { if (filename.includes(".worker")) { - return `./${filename}`; + return { relative: true }; } return filename; }, diff --git a/yarn.lock b/yarn.lock index 8968ded828a..d5615575c70 100644 --- a/yarn.lock +++ b/yarn.lock @@ -7679,6 +7679,7 @@ __metadata: version: 0.0.0-use.local resolution: "@hashintel/petrinaut-core@workspace:libs/@hashintel/petrinaut-core" dependencies: + "@local/petrinaut-optimizer-core": "workspace:*" "@types/js-yaml": "npm:^4" "@types/node": "npm:22.18.13" "@typescript/native-preview": "npm:7.0.0-dev.20260511.1" @@ -7688,6 +7689,7 @@ __metadata: js-yaml: "npm:4.3.1" oxlint: "npm:1.63.0" oxlint-tsgolint: "npm:0.22.1" + pyodide: "npm:314.0.6" rolldown: "npm:1.2.6" rolldown-plugin-dts: "npm:0.28.3" typescript: "npm:5.9.3" @@ -18594,6 +18596,13 @@ __metadata: languageName: node linkType: hard +"@types/emscripten@npm:^1.41.4": + version: 1.41.5 + resolution: "@types/emscripten@npm:1.41.5" + checksum: 10c0/ae816da716f896434e59df7a71b67c71ae7e85ca067a32aef1616572fc4757459515d42ade6f5b8fd8d69733a9dbd0cf23010fec5b2f41ce52c09501aa350e45 + languageName: node + linkType: hard + "@types/eslint-scope@npm:^3.7.7": version: 3.7.7 resolution: "@types/eslint-scope@npm:3.7.7" @@ -39309,6 +39318,16 @@ __metadata: languageName: node linkType: hard +"pyodide@npm:314.0.6": + version: 314.0.6 + resolution: "pyodide@npm:314.0.6" + dependencies: + "@types/emscripten": "npm:^1.41.4" + ws: "npm:^8.5.0" + checksum: 10c0/0991c470ec771330dac202380bb875487625abf816dda4ab282ae87d276845a32bd863e5e5666dd2e4695c81ebcfe416dbeefc85841b0aa242555d5ac885245a + languageName: node + linkType: hard + "qs@npm:6.15.2": version: 6.15.2 resolution: "qs@npm:6.15.2" @@ -47017,7 +47036,7 @@ __metadata: languageName: node linkType: hard -"ws@npm:^8.17.1, ws@npm:^8.18.0, ws@npm:^8.18.2, ws@npm:^8.18.3, ws@npm:^8.19.0, ws@npm:^8.21.1": +"ws@npm:^8.17.1, ws@npm:^8.18.0, ws@npm:^8.18.2, ws@npm:^8.18.3, ws@npm:^8.19.0, ws@npm:^8.21.1, ws@npm:^8.5.0": version: 8.21.3 resolution: "ws@npm:8.21.3" peerDependencies: From 5e25a00b517aa157323ef72c693f22479bceca54 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Fri, 4 Sep 2026 05:26:27 +0200 Subject: [PATCH 2/4] Extend, release and parallelise browser optimization runs --- .changeset/browser-optimization-runtime.md | 2 +- .../src/browser-optimization.ts | 2 + .../browser-optimization.test.ts | 407 ++++++++++++++++++ .../browser-optimization.ts | 282 +++++++++--- .../src/browser-optimization/messages.ts | 41 +- .../browser-optimization/optimizer.worker.ts | 159 +++++-- .../src/browser-optimization/run-log.test.ts | 58 ++- .../src/browser-optimization/run-log.ts | 39 +- .../study-runner.pyodide.test.ts | 304 ++++++++++--- .../src/browser-optimization/study-runner.ts | 185 ++++++-- .../petrinaut-core/src/optimization.ts | 51 ++- .../petrinaut-core/src/workers/optimizer.ts | 4 + 12 files changed, 1310 insertions(+), 224 deletions(-) diff --git a/.changeset/browser-optimization-runtime.md b/.changeset/browser-optimization-runtime.md index 6bd95875612..4f2596606df 100644 --- a/.changeset/browser-optimization-runtime.md +++ b/.changeset/browser-optimization-runtime.md @@ -2,4 +2,4 @@ "@hashintel/petrinaut-core": patch --- -Adds an in-browser optimization capability that runs the Optuna study in a Pyodide worker and evaluates trials through a host channel. +Adds an in-browser optimization capability that runs the Optuna study in a Pyodide worker and evaluates trials through a host channel. A study that completed or was cancelled stays in the worker until it is released, so the connected capability can extend it with more trials on the same sampler history, and a run may keep up to four trials in flight at once. diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization.ts index a6a02805749..a5eb9f2a39f 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization.ts @@ -13,6 +13,8 @@ export { export type { OptimizationScalar, PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, PetrinautOptimizationChannel, PetrinautOptimizationSource, PetrinautOptimizationTrialOutcome, diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts index 6482ffe9ad0..69be9093068 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, vi } from "vitest"; +import { createAbortController } from "../environment"; import { PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE } from "../optimization"; import { createOptimizationManifestInput } from "../shared/optimization-manifest.fixtures"; import { createBrowserOptimization } from "./browser-optimization"; @@ -208,6 +209,7 @@ describe("createBrowserOptimization", () => { direction: "maximize", study: { trials: 20, sampler: "tpe", seed: 42, seedsPerTrial: 1 }, }, + parallelism: 1, }); context.worker.emit({ @@ -356,6 +358,88 @@ describe("createBrowserOptimization", () => { }); }); + it("answers an evaluate posted after cancel with a cancelled outcome without running it", async () => { + const context = setUp(); + const runId = await startRun(context); + await context.capability.cancelOptimizationRun(runId); + + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + + expect(context.evaluateTrial).not.toHaveBeenCalled(); + expect(context.worker.sentOfType("evaluated")).toEqual([ + { + type: "evaluated", + requestId: 1, + outcome: { kind: "pruned", reason: "cancelled" }, + }, + ]); + }); + + it("ignores an evaluation of a stopped segment that settles after the next segment started", async () => { + let rejectStale: (error: Error) => void = () => {}; + const context = setUp({ + evaluateTrial: async (request) => { + if (request.trial === 0) { + return new Promise((_resolve, reject) => { + rejectStale = reject; + }); + } + return { kind: "objective", objective: 1 }; + }, + }); + const runId = await startRun(context); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + await context.capability.cancelOptimizationRun(runId); + context.worker.emit({ type: "cancelled", runId }); + await context.capability.extendOptimizationRun(runId, 1); + await flush(); + expect(context.worker.sentOfType("extend")).toHaveLength(1); + + rejectStale(new Error("late failure")); + await flush(); + + expect(context.worker.sentOfType("evaluated")).toHaveLength(0); + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + expect(context.worker.sentOfType("release")).toHaveLength(0); + + context.worker.emit({ + type: "evaluate", + runId, + requestId: 2, + trial: 1, + suggestedValues: { rate: 1, count: 4, enabled: false }, + }); + await flush(); + expect(context.worker.sentOfType("evaluated")).toEqual([ + { + type: "evaluated", + requestId: 2, + outcome: { kind: "objective", objective: 1 }, + }, + ]); + context.worker.emit({ type: "complete", runId, summary }); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId, { cursor: 2 }), + ); + expect(events.map((event) => event.type)).toEqual(["started", "complete"]); + }); + it("cancels a run that is still waiting for the runtime without starting it", async () => { const context = setUp(); const { runId } = await context.capability.createOptimizationRun( @@ -678,6 +762,329 @@ describe("createBrowserOptimization", () => { ).rejects.toThrow("disposed"); }); + it("extends a completed study with trials that continue the numbering", async () => { + const context = setUp(); + const runId = await startRun(context); + context.worker.emit({ type: "trial", runId, event: completedTrial }); + context.worker.emit({ type: "complete", runId, summary }); + expect( + ( + await collectEvents(context.capability.attachOptimizationRun(runId)) + ).map((event) => event.type), + ).toEqual(["started", "trial", "complete"]); + + await context.capability.extendOptimizationRun(runId, 5); + const tail = collectEvents( + context.capability.attachOptimizationRun(runId, { cursor: 3 }), + ); + await flush(); + + expect(context.worker.sentOfType("extend")).toEqual([ + { type: "extend", runId, trials: 5, parallelism: 1 }, + ]); + context.worker.emit({ type: "started", runId, requestedTrials: 6 }); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 2, + trial: 1, + suggestedValues: { rate: 1, count: 4, enabled: false }, + }); + await flush(); + const request = context.evaluateTrial.mock.calls.at(-1)?.[0]; + expect(request).toMatchObject({ runId, trial: 1 }); + expect(request?.signal.aborted).toBe(false); + expect(context.worker.sentOfType("evaluated").at(-1)).toEqual({ + type: "evaluated", + requestId: 2, + outcome: { kind: "objective", objective: 2 }, + }); + + const secondTrial: OptimizerTrialPayload = { + trial: 1, + parameters: { rate: 1, count: 4, enabled: false }, + objective: 2, + state: "complete", + best: { + trial: 1, + parameters: { rate: 1, count: 4, enabled: false }, + objective: 2, + }, + }; + context.worker.emit({ type: "trial", runId, event: secondTrial }); + context.worker.emit({ + type: "complete", + runId, + summary: { + ...summary, + requestedTrials: 6, + completedTrials: 2, + best: secondTrial.best, + }, + }); + + expect(await tail).toEqual([ + { type: "started", requestedTrials: 6, seq: 4 }, + { + type: "trial", + trial: 1, + parameters: secondTrial.parameters, + objective: 2, + state: "complete", + best: secondTrial.best, + seq: 5, + }, + { + type: "complete", + requestedTrials: 6, + completedTrials: 2, + prunedTrials: 0, + failedTrials: 0, + best: secondTrial.best, + seq: 6, + }, + ]); + expect( + ( + await collectEvents(context.capability.attachOptimizationRun(runId)) + ).map((event) => event.seq), + ).toEqual([1, 2, 3]); + }); + + it("extends a stopped study from the trials it was told, with a fresh signal", async () => { + const context = setUp(); + const runId = await startRun(context); + context.worker.emit({ type: "trial", runId, event: completedTrial }); + await context.capability.cancelOptimizationRun(runId); + context.worker.emit({ type: "cancelled", runId }); + + await context.capability.extendOptimizationRun(runId, 2); + await flush(); + + expect(context.worker.sentOfType("extend")).toEqual([ + { type: "extend", runId, trials: 2, parallelism: 1 }, + ]); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 2, + trial: 2, + suggestedValues: { rate: 1, count: 4, enabled: false }, + }); + await flush(); + expect(context.evaluateTrial.mock.calls.at(-1)?.[0]?.signal.aborted).toBe( + false, + ); + expect(context.worker.sentOfType("evaluated")).toHaveLength(1); + + context.worker.emit({ type: "complete", runId, summary }); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId, { cursor: 3 }), + ); + expect(events.map((event) => [event.type, event.seq])).toEqual([ + ["started", 4], + ["complete", 5], + ]); + expect(events[0]).toMatchObject({ requestedTrials: 3 }); + }); + + it("rejects extending a run that is running, unknown, released or failed, or past the caps", async () => { + const context = setUp(); + const runId = await startRun(context); + + await expect( + context.capability.extendOptimizationRun(runId, 1), + ).rejects.toThrow("is still running"); + await expect( + context.capability.extendOptimizationRun("missing", 1), + ).rejects.toThrow( + expect.objectContaining({ category: "http", httpStatus: 404 }), + ); + + context.worker.emit({ type: "trial", runId, event: completedTrial }); + context.worker.emit({ type: "complete", runId, summary }); + await expect( + context.capability.extendOptimizationRun(runId, 0), + ).rejects.toThrow("positive whole number of trials"); + await expect( + context.capability.extendOptimizationRun(runId, 1000), + ).rejects.toThrow(/at most .* trials in total; 1 already ran/); + await expect( + context.capability.extendOptimizationRun(runId, 2, { parallelism: 5 }), + ).rejects.toThrow("between 1 and 4"); + expect(context.worker.sentOfType("extend")).toHaveLength(0); + + await context.capability.extendOptimizationRun(runId, 999); + await flush(); + expect(context.worker.sentOfType("extend")).toEqual([ + { type: "extend", runId, trials: 999, parallelism: 1 }, + ]); + context.worker.emit({ type: "complete", runId, summary }); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId, { cursor: 3 }), + ); + expect(events[0]).toMatchObject({ + type: "started", + requestedTrials: 1000, + }); + + await context.capability.releaseOptimizationRun(runId); + await context.capability.releaseOptimizationRun(runId); + expect(context.worker.sentOfType("release")).toEqual([ + { type: "release", runId }, + ]); + await expect( + context.capability.extendOptimizationRun(runId, 1), + ).rejects.toThrow("released or failed"); + + const failed = await startRun(context); + context.worker.emit({ type: "error", runId: failed, message: "boom" }); + await expect( + context.capability.extendOptimizationRun(failed, 1), + ).rejects.toThrow("released or failed"); + }); + + it("releasing a running study stops it and ends its stream", async () => { + const context = setUp(); + const runId = await startRun(context); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + + await context.capability.releaseOptimizationRun(runId); + await flush(); + + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + expect(context.worker.sentOfType("release")).toEqual([ + { type: "release", runId }, + ]); + expect(context.worker.sentOfType("evaluated")).toHaveLength(0); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toMatchObject({ + type: "error", + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + seq: 2, + }); + }); + + it("stopping a queued extension keeps the study and lets the queue move on", async () => { + const context = setUp(); + const first = await startRun(context); + context.worker.emit({ type: "trial", runId: first, event: completedTrial }); + context.worker.emit({ type: "complete", runId: first, summary }); + const { runId: second } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + await flush(); + + await context.capability.extendOptimizationRun(first, 3); + expect(context.worker.sentOfType("extend")).toHaveLength(0); + await context.capability.cancelOptimizationRun(first); + expect( + ( + await collectEvents( + context.capability.attachOptimizationRun(first, { cursor: 3 }), + ) + ).map((event) => [event.type, event.seq]), + ).toEqual([ + ["started", 4], + ["error", 5], + ]); + + await context.capability.extendOptimizationRun(first, 2); + context.worker.emit({ type: "complete", runId: second, summary }); + await flush(); + + expect(context.worker.sentOfType("extend")).toEqual([ + { type: "extend", runId: first, trials: 2, parallelism: 1 }, + ]); + context.worker.emit({ type: "complete", runId: first, summary }); + const events = await collectEvents( + context.capability.attachOptimizationRun(first, { cursor: 5 }), + ); + expect(events.map((event) => [event.type, event.seq])).toEqual([ + ["started", 6], + ["complete", 7], + ]); + expect(events[0]).toMatchObject({ requestedTrials: 3 }); + }); + + it("passes the parallelism to the worker and evaluates trials in flight together", async () => { + const context = setUp(); + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + { parallelism: 3 }, + ); + context.worker.emit({ type: "ready" }); + await flush(); + expect(context.worker.sentOfType("start")[0]?.parallelism).toBe(3); + + for (const [requestId, trial] of [ + [1, 0], + [2, 1], + ] as const) { + context.worker.emit({ + type: "evaluate", + runId, + requestId, + trial, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + } + await flush(); + expect(context.evaluateTrial).toHaveBeenCalledTimes(2); + expect( + context.worker.sentOfType("evaluated").map(({ requestId }) => requestId), + ).toEqual([1, 2]); + + context.worker.emit({ type: "complete", runId, summary }); + await context.capability.extendOptimizationRun(runId, 2); + await flush(); + expect(context.worker.sentOfType("extend").at(-1)).toMatchObject({ + trials: 2, + parallelism: 3, + }); + + context.worker.emit({ type: "complete", runId, summary }); + await context.capability.extendOptimizationRun(runId, 1, { + parallelism: 1, + }); + await flush(); + expect(context.worker.sentOfType("extend").at(-1)).toMatchObject({ + trials: 1, + parallelism: 1, + }); + + await expect( + context.capability.createOptimizationRun( + createOptimizationManifestInput(), + { parallelism: 0 }, + ), + ).rejects.toThrow("between 1 and 4"); + }); + + it("refuses to create a run for an already-aborted signal", async () => { + const context = setUp(); + const controller = createAbortController(); + controller.abort(); + + await expect( + context.capability.createOptimizationRun( + createOptimizationManifestInput(), + { signal: controller.signal }, + ), + ).rejects.toThrow(expect.objectContaining({ name: "AbortError" })); + expect(context.workers).toHaveLength(0); + }); + it("rejects an invalid manifest before allocating a run", async () => { const context = setUp(); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts index c6b2c11cfe0..b39d2ea49de 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/browser-optimization.ts @@ -9,6 +9,8 @@ import { deriveOptimizationTrialSeeds, describeOptimization, PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + PETRINAUT_OPTIMIZATION_MAX_PARALLELISM, + PETRINAUT_OPTIMIZATION_MAX_TRIALS, petrinautOptimizationManifestSchema, resolveTrialScenarioParameterValues, } from "../optimization"; @@ -28,10 +30,10 @@ import { type OptimizationRunLogEvent, } from "./run-log"; -import type { AbortSignalLike } from "../environment"; +import type { AbortControllerLike, AbortSignalLike } from "../environment"; import type { PetrinautConnectedOptimization, - PetrinautOptimization, + PetrinautConnectedOptimizationCapability, PetrinautOptimizationChannel, PetrinautOptimizationDescribeResult, PetrinautOptimizationEvent, @@ -41,7 +43,10 @@ import type { } from "../optimization"; import type { OptimizerEvaluateMessage, + OptimizerExtendMessage, + OptimizerStartMessage, OptimizerToMainMessage, + OptimizerToWorkerMessage, } from "./messages"; export type CreateBrowserOptimizationOptions = { @@ -49,7 +54,17 @@ export type CreateBrowserOptimizationOptions = { createWorker?: () => OptimizerWorkerLike; }; -type RunStatus = "queued" | "starting" | "running" | "finished"; +/** + * `queued` and `starting` wait for the worker to take the run's segment, + * `running` has it posted, `finished-resumable` ended a segment with the + * study kept in the worker, and `finished` has no study to return to. + */ +type RunStatus = + | "queued" + | "starting" + | "running" + | "finished-resumable" + | "finished"; type RunRecord = { readonly runId: string; @@ -57,9 +72,11 @@ type RunRecord = { readonly description: PetrinautOptimizationDescribeResult; readonly seeds: readonly number[]; readonly log: OptimizationRunLog; - readonly signal: AbortSignalLike; - readonly abort: () => void; status: RunStatus; + /** The segment the worker runs when it takes this run. */ + command: OptimizerStartMessage | OptimizerExtendMessage; + /** Aborted when the segment is cancelled; the next segment gets a fresh one. */ + controller: AbortControllerLike; }; type WorkerSession = { @@ -97,6 +114,13 @@ const isAbortError = (error: unknown): boolean => "name" in error && error.name === "AbortError"; +const isSettled = (run: RunRecord): boolean => + run.status === "finished" || run.status === "finished-resumable"; + +/** Trials the study was told an outcome for, which is where the next segment counts from. */ +const toldTrials = (log: OptimizationRunLog): number => + log.events.filter((event) => event.type === "trial").length; + const invalidManifestError = ( issues: readonly { path: PropertyKey[]; message: string }[], ): Error => @@ -116,6 +140,53 @@ const unknownRunError = (runId: string): Error => httpStatus: 404, }); +const disposedError = (): Error => + new Error("The in-browser optimizer was disposed"); + +const creationAbortedError = (): Error => { + const error = new Error("optimization run creation aborted"); + error.name = "AbortError"; + return error; +}; + +const notResumableError = (run: RunRecord): Error => + new Error( + run.status === "finished" + ? `Optimization run "${run.runId}" has no study to extend: it was released or failed` + : `Optimization run "${run.runId}" is still running`, + ); + +const validParallelism = (value: number | undefined): number => { + if (value === undefined) { + return 1; + } + if ( + !Number.isInteger(value) || + value < 1 || + value > PETRINAUT_OPTIMIZATION_MAX_PARALLELISM + ) { + throw new Error( + `Optimization parallelism must be an integer between 1 and ${PETRINAUT_OPTIMIZATION_MAX_PARALLELISM}`, + ); + } + return value; +}; + +const requestedTotal = (told: number, trials: number): number => { + if (!Number.isInteger(trials) || trials < 1) { + throw new Error( + "An optimization extends by a positive whole number of trials", + ); + } + const total = told + trials; + if (total > PETRINAUT_OPTIMIZATION_MAX_TRIALS) { + throw new Error( + `An optimization may run at most ${PETRINAUT_OPTIMIZATION_MAX_TRIALS.toLocaleString()} trials in total; ${told.toLocaleString()} already ran`, + ); + } + return total; +}; + async function* attachToLog( log: OptimizationRunLog, options: { @@ -132,7 +203,7 @@ const connectBrowserOptimization = (options: { channel: PetrinautOptimizationChannel; pyodide: OptimizerPyodideConfig; createWorker: () => OptimizerWorkerLike; -}): PetrinautOptimization & { dispose(this: void): void } => { +}): PetrinautConnectedOptimizationCapability => { const { channel } = options; const runs = new Map(); const queue: RunRecord[] = []; @@ -145,6 +216,10 @@ const connectBrowserOptimization = (options: { return run && run.status === "running" ? run : null; }; + const post = (message: OptimizerToWorkerMessage): void => { + session?.worker.postMessage(message); + }; + const resetSession = (stale: WorkerSession): void => { if (session === stale) { session = null; @@ -152,12 +227,17 @@ const connectBrowserOptimization = (options: { } }; - const finish = (run: RunRecord, event: OptimizationRunLogEvent): void => { - if (run.status === "finished") { + /** Ends the run's segment with `event` and lets the queue move on. */ + const finish = ( + run: RunRecord, + event: OptimizationRunLogEvent, + status: "finished-resumable" | "finished", + ): void => { + if (isSettled(run)) { return; } // eslint-disable-next-line no-param-reassign -- the record's status is the session state this helper advances - run.status = "finished"; + run.status = status; run.log.append(event); const queuedAt = queue.indexOf(run); if (queuedAt !== -1) { @@ -174,18 +254,30 @@ const connectBrowserOptimization = (options: { requestId: number, outcome: PetrinautOptimizationTrialOutcome, ): void => { - session?.worker.postMessage({ type: "evaluated", requestId, outcome }); + post({ type: "evaluated", requestId, outcome }); + }; + + /** Stops the segment on the worker (when one is posted) and drops the study. */ + const discardStudy = (run: RunRecord): void => { + run.controller.abort(); + if (run.status === "running") { + post({ type: "cancel", runId: run.runId }); + } + post({ type: "release", runId: run.runId }); }; const failTrialEvaluation = (run: RunRecord, error: unknown): void => { - run.abort(); - session?.worker.postMessage({ type: "cancel", runId: run.runId }); - finish(run, { - type: "error", - code: "trial_evaluation_failed", - message: errorMessage(error), - retryable: false, - }); + discardStudy(run); + finish( + run, + { + type: "error", + code: "trial_evaluation_failed", + message: errorMessage(error), + retryable: false, + }, + "finished", + ); }; const handleEvaluate = (message: OptimizerEvaluateMessage): void => { @@ -193,6 +285,13 @@ const connectBrowserOptimization = (options: { if (!run) { return; } + const { controller } = run; + if (controller.signal.aborted) { + reply(message.requestId, { kind: "pruned", reason: "cancelled" }); + return; + } + const segmentIsCurrent = (): boolean => + run.controller === controller && run.status === "running"; let request: PetrinautOptimizationTrialRequest; try { request = { @@ -205,19 +304,19 @@ const connectBrowserOptimization = (options: { message.suggestedValues, ), seeds: run.seeds, - signal: run.signal, + signal: controller.signal, }; } catch (error) { failTrialEvaluation(run, error); return; } const evaluated = (outcome: PetrinautOptimizationTrialOutcome): void => { - if (run.status === "running") { + if (segmentIsCurrent()) { reply(message.requestId, outcome); } }; const evaluationFailed = (error: unknown): void => { - if (run.status !== "running") { + if (!segmentIsCurrent()) { return; } if (isAbortError(error)) { @@ -238,8 +337,12 @@ const connectBrowserOptimization = (options: { const handleWorkerMessage = (message: OptimizerToMainMessage): void => { switch (message.type) { + // The session promise settles on `ready` and `init-error`; `started` and + // `released` acknowledge segments the log and run status already record. case "ready": case "init-error": + case "started": + case "released": return; case "evaluate": handleEvaluate(message); @@ -261,33 +364,41 @@ const connectBrowserOptimization = (options: { const run = activeRunFor(message.runId); const { summary } = message; if (run) { - finish(run, { - type: "complete", - requestedTrials: summary.requestedTrials, - completedTrials: summary.completedTrials, - prunedTrials: summary.prunedTrials, - failedTrials: summary.failedTrials, - best: summary.best, - }); + finish( + run, + { + type: "complete", + requestedTrials: summary.requestedTrials, + completedTrials: summary.completedTrials, + prunedTrials: summary.prunedTrials, + failedTrials: summary.failedTrials, + best: summary.best, + }, + "finished-resumable", + ); } return; } case "cancelled": { const run = activeRunFor(message.runId); if (run) { - finish(run, cancelledEvent); + finish(run, cancelledEvent, "finished-resumable"); } return; } case "error": { const run = activeRunFor(message.runId); if (run) { - finish(run, { - type: "error", - code: "study_failed", - message: message.message, - retryable: false, - }); + finish( + run, + { + type: "error", + code: "study_failed", + message: message.message, + retryable: false, + }, + "finished", + ); } } } @@ -335,59 +446,82 @@ const connectBrowserOptimization = (options: { try { current = ensureSession(); } catch (error) { - finish(run, unavailableEvent(error)); + finish(run, unavailableEvent(error), "finished"); return; } current.ready.then( () => { if (run.status === "starting" && session === current) { run.status = "running"; - current.worker.postMessage({ - type: "start", - runId: run.runId, - description: run.description, - }); + current.worker.postMessage(run.command); } }, (error: unknown) => { resetSession(current); - finish(run, unavailableEvent(error)); + finish(run, unavailableEvent(error), "finished"); }, ); }; + const enqueue = (run: RunRecord, requestedTrials: number): void => { + run.log.append({ type: "started", requestedTrials }); + queue.push(run); + startNext(); + }; + return { - async createOptimizationRun(input) { + async createOptimizationRun(input, runOptions = {}) { if (disposed) { - throw new Error("The in-browser optimizer was disposed"); + throw disposedError(); + } + if (runOptions.signal?.aborted) { + throw creationAbortedError(); } + const parallelism = validParallelism(runOptions.parallelism); const parsed = petrinautOptimizationManifestSchema.safeParse(input); if (!parsed.success) { throw invalidManifestError(parsed.error.issues); } const manifest = parsed.data; - const controller = createAbortController(); + const runId = generateUuid(); + const description = describeOptimization(manifest); const run: RunRecord = { - runId: generateUuid(), + runId, manifest, - description: describeOptimization(manifest), + description, seeds: deriveOptimizationTrialSeeds( manifest.execution.seed, manifest.execution.seedsPerTrial ?? 1, ), log: createOptimizationRunLog(), - signal: controller.signal, - abort: () => controller.abort(), status: "queued", + command: { type: "start", runId, description, parallelism }, + controller: createAbortController(), }; - runs.set(run.runId, run); - run.log.append({ - type: "started", - requestedTrials: manifest.study.trials, - }); - queue.push(run); - startNext(); - return { runId: run.runId }; + runs.set(runId, run); + enqueue(run, manifest.study.trials); + return { runId }; + }, + async extendOptimizationRun(runId, trials, extendOptions = {}) { + if (disposed) { + throw disposedError(); + } + const run = runs.get(runId); + if (!run) { + throw unknownRunError(runId); + } + if (run.status !== "finished-resumable") { + throw notResumableError(run); + } + const parallelism = + extendOptions.parallelism === undefined + ? run.command.parallelism + : validParallelism(extendOptions.parallelism); + const total = requestedTotal(toldTrials(run.log), trials); + run.command = { type: "extend", runId, trials, parallelism }; + run.controller = createAbortController(); + run.status = "queued"; + enqueue(run, total); }, attachOptimizationRun(runId, attachOptions = {}) { const run = runs.get(runId); @@ -398,24 +532,42 @@ const connectBrowserOptimization = (options: { }, async cancelOptimizationRun(runId) { const run = runs.get(runId); - if (!run || run.status === "finished") { + if (!run || isSettled(run)) { return; } - run.abort(); + run.controller.abort(); if (run.status === "running") { - session?.worker.postMessage({ type: "cancel", runId }); + post({ type: "cancel", runId }); return; } - finish(run, cancelledEvent); + // The segment never reached the worker: an extension leaves its study + // kept, a first run has no study at all. + finish( + run, + cancelledEvent, + run.command.type === "extend" ? "finished-resumable" : "finished", + ); + }, + async releaseOptimizationRun(runId) { + const run = runs.get(runId); + if (!run || run.status === "finished") { + return; + } + discardStudy(run); + if (isSettled(run)) { + run.status = "finished"; + } else { + finish(run, cancelledEvent, "finished"); + } }, dispose() { disposed = true; for (const run of runs.values()) { - if (run.status !== "finished") { - run.abort(); - run.status = "finished"; + if (!isSettled(run)) { + run.controller.abort(); run.log.append(cancelledEvent); } + run.status = "finished"; } queue.length = 0; active = null; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts index 9aaa7f11197..4d2a3e759ac 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/messages.ts @@ -20,13 +20,14 @@ export type OptimizerTrialPayload = { best: OptimizerBestTrial | null; }; +/** Counts over the whole study so far, across every segment it ran. */ export type OptimizerStudySummary = { requestedTrials: number; completedTrials: number; prunedTrials: number; failedTrials: number; best: OptimizerBestTrial | null; - /** Set when the study stopped early because the run was cancelled. */ + /** Set when the segment stopped early because the run was cancelled. */ cancelled?: boolean; }; @@ -36,10 +37,21 @@ export type OptimizerInitMessage = { pythonSources: Readonly>; }; +/** Creates the study and runs its first `description.study.trials` trials. */ export type OptimizerStartMessage = { type: "start"; runId: string; description: PetrinautOptimizationDescribeResult; + /** Trials kept in flight at once. */ + parallelism: number; +}; + +/** Runs `trials` more on the kept study; trial numbers continue. */ +export type OptimizerExtendMessage = { + type: "extend"; + runId: string; + trials: number; + parallelism: number; }; export type OptimizerEvaluatedMessage = { @@ -48,16 +60,25 @@ export type OptimizerEvaluatedMessage = { outcome: PetrinautOptimizationTrialOutcome; }; +/** Stops the running segment; the study stays in memory. */ export type OptimizerCancelMessage = { type: "cancel"; runId: string; }; +/** Drops the kept study. */ +export type OptimizerReleaseMessage = { + type: "release"; + runId: string; +}; + export type OptimizerToWorkerMessage = | OptimizerInitMessage | OptimizerStartMessage + | OptimizerExtendMessage | OptimizerEvaluatedMessage - | OptimizerCancelMessage; + | OptimizerCancelMessage + | OptimizerReleaseMessage; export type OptimizerReadyMessage = { type: "ready"; @@ -68,6 +89,13 @@ export type OptimizerInitErrorMessage = { message: string; }; +/** A segment began; `requestedTrials` is the study's cumulative total. */ +export type OptimizerStartedMessage = { + type: "started"; + runId: string; + requestedTrials: number; +}; + export type OptimizerEvaluateMessage = { type: "evaluate"; runId: string; @@ -99,11 +127,18 @@ export type OptimizerErrorMessage = { message: string; }; +export type OptimizerReleasedMessage = { + type: "released"; + runId: string; +}; + export type OptimizerToMainMessage = | OptimizerReadyMessage | OptimizerInitErrorMessage + | OptimizerStartedMessage | OptimizerEvaluateMessage | OptimizerTrialMessage | OptimizerCompleteMessage | OptimizerCancelledMessage - | OptimizerErrorMessage; + | OptimizerErrorMessage + | OptimizerReleasedMessage; diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts index bc768c2a95f..c935bc28428 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/optimizer.worker.ts @@ -5,15 +5,22 @@ import type { PetrinautOptimizationTrialOutcome } from "../optimization"; import type { OptimizerCancelMessage, OptimizerEvaluatedMessage, + OptimizerExtendMessage, OptimizerInitMessage, + OptimizerReleaseMessage, OptimizerStartMessage, + OptimizerStudySummary, OptimizerToMainMessage, OptimizerToWorkerMessage, } from "./messages"; import type { LoadPyodide } from "./pyodide-like"; -import type { OptimizerStudyRunner } from "./study-runner"; +import type { + OptimizerStudyCallbacks, + OptimizerStudyRunner, +} from "./study-runner"; -type ActiveRun = { +/** A run's segment, from the message that posted it until its outcome is reported. */ +type ActiveSegment = { cancelled: boolean; pendingRequestIds: Set; }; @@ -24,7 +31,7 @@ const workerRuntime = createWorkerThreadRuntime< >(); let runner: OptimizerStudyRunner | null = null; -const runs = new Map(); +const segments = new Map(); const pendingEvaluations = new Map< number, (outcome: PetrinautOptimizationTrialOutcome) => void @@ -70,45 +77,57 @@ const handleInit = (message: OptimizerInitMessage): void => { ); }; -const handleStart = (message: OptimizerStartMessage): void => { - const { runId } = message; +const runnerFor = (runId: string): OptimizerStudyRunner | null => { if (!runner) { workerRuntime.postMessage({ type: "error", runId, message: "The optimizer worker received a study before its runtime", }); - return; } - const run: ActiveRun = { cancelled: false, pendingRequestIds: new Set() }; - runs.set(runId, run); - runner - .run({ - description: message.description, - evaluate: (trial, suggestedValues) => - new Promise((resolve) => { - const requestId = nextRequestId; - nextRequestId += 1; - pendingEvaluations.set(requestId, resolve); - run.pendingRequestIds.add(requestId); - workerRuntime.postMessage({ - type: "evaluate", - runId, - requestId, - trial, - suggestedValues, - }); - }), - onTrial: (event) => - workerRuntime.postMessage({ type: "trial", runId, event }), - isCancelled: () => run.cancelled, - }) + return runner; +}; + +const beginSegment = (runId: string): OptimizerStudyCallbacks => { + const segment: ActiveSegment = { + cancelled: false, + pendingRequestIds: new Set(), + }; + segments.set(runId, segment); + return { + onStarted: (requestedTrials) => + workerRuntime.postMessage({ type: "started", runId, requestedTrials }), + evaluate: (trial, suggestedValues) => + new Promise((resolve) => { + const requestId = nextRequestId; + nextRequestId += 1; + pendingEvaluations.set(requestId, resolve); + segment.pendingRequestIds.add(requestId); + workerRuntime.postMessage({ + type: "evaluate", + runId, + requestId, + trial, + suggestedValues, + }); + }), + onTrial: (event) => + workerRuntime.postMessage({ type: "trial", runId, event }), + isCancelled: () => segment.cancelled, + }; +}; + +const reportSegment = ( + runId: string, + summary: Promise, +): void => { + summary .then( - (summary) => + (result) => workerRuntime.postMessage( - summary.cancelled === true + result.cancelled === true ? { type: "cancelled", runId } - : { type: "complete", runId, summary }, + : { type: "complete", runId, summary: result }, ), (error: unknown) => workerRuntime.postMessage({ @@ -118,34 +137,86 @@ const handleStart = (message: OptimizerStartMessage): void => { }), ) .finally(() => { - runs.delete(runId); + segments.delete(runId); }); }; +const handleStart = (message: OptimizerStartMessage): void => { + const { runId } = message; + const current = runnerFor(runId); + if (!current) { + return; + } + reportSegment( + runId, + current.start({ + runId, + description: message.description, + parallelism: message.parallelism, + callbacks: beginSegment(runId), + }), + ); +}; + +const handleExtend = (message: OptimizerExtendMessage): void => { + const { runId } = message; + const current = runnerFor(runId); + if (!current) { + return; + } + reportSegment( + runId, + current.extend({ + runId, + trials: message.trials, + parallelism: message.parallelism, + callbacks: beginSegment(runId), + }), + ); +}; + const handleEvaluated = (message: OptimizerEvaluatedMessage): void => { const resolve = pendingEvaluations.get(message.requestId); if (!resolve) { return; } pendingEvaluations.delete(message.requestId); - for (const run of runs.values()) { - run.pendingRequestIds.delete(message.requestId); + for (const segment of segments.values()) { + segment.pendingRequestIds.delete(message.requestId); } resolve(message.outcome); }; -const handleCancel = (message: OptimizerCancelMessage): void => { - const run = runs.get(message.runId); - if (!run) { +const cancelSegment = (runId: string): void => { + const segment = segments.get(runId); + if (!segment) { return; } - run.cancelled = true; - for (const requestId of run.pendingRequestIds) { + segment.cancelled = true; + for (const requestId of segment.pendingRequestIds) { const resolve = pendingEvaluations.get(requestId); pendingEvaluations.delete(requestId); resolve?.({ kind: "pruned", reason: "cancelled" }); } - run.pendingRequestIds.clear(); + segment.pendingRequestIds.clear(); +}; + +const handleCancel = (message: OptimizerCancelMessage): void => { + cancelSegment(message.runId); +}; + +const handleRelease = (message: OptimizerReleaseMessage): void => { + const { runId } = message; + cancelSegment(runId); + runner?.release(runId).then( + () => workerRuntime.postMessage({ type: "released", runId }), + (error: unknown) => + workerRuntime.postMessage({ + type: "error", + runId, + message: errorMessage(error), + }), + ); }; workerRuntime.onMessage((message) => { @@ -156,11 +227,17 @@ workerRuntime.onMessage((message) => { case "start": handleStart(message); break; + case "extend": + handleExtend(message); + break; case "evaluated": handleEvaluated(message); break; case "cancel": handleCancel(message); break; + case "release": + handleRelease(message); + break; } }); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts index df6fa014170..205539abee6 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.test.ts @@ -36,20 +36,30 @@ const complete = { } as const; describe("createOptimizationRunLog", () => { - it("stamps dense sequence numbers from 1 and closes on a terminal event", () => { + it("stamps dense sequence numbers from 1 and settles on a terminal event", () => { const log = createOptimizationRunLog(); expect(log.append({ type: "started", requestedTrials: 2 }).seq).toBe(1); expect(log.append(trial(0)).seq).toBe(2); - expect(log.closed).toBe(false); + expect(log.settled).toBe(false); expect(log.append(complete).seq).toBe(3); - expect(log.closed).toBe(true); + expect(log.settled).toBe(true); expect(log.events.map((event) => event.seq)).toEqual([1, 2, 3]); expect(() => log.append(trial(1))).toThrow( - "optimization run log is closed", + "a settled optimization run log accepts only a started event", ); }); + it("begins a new segment with a started event after a terminal one", () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 2 }); + log.append(complete); + + expect(log.append({ type: "started", requestedTrials: 4 }).seq).toBe(3); + expect(log.settled).toBe(false); + expect(log.append(trial(2)).seq).toBe(4); + }); + it("replays the stored events past the cursor, then tails until the terminal event", async () => { const log = createOptimizationRunLog(); log.append({ type: "started", requestedTrials: 2 }); @@ -67,7 +77,45 @@ describe("createOptimizationRunLog", () => { ]); }); - it("ends at once when attached past the end of a closed log", async () => { + it("ends a replay at the first terminal event after the cursor", async () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 1 }); + log.append(trial(0)); + log.append(complete); + log.append({ type: "started", requestedTrials: 2 }); + log.append(trial(1)); + log.append({ ...complete, requestedTrials: 2 }); + + expect((await collect(log.replay())).map((event) => event.seq)).toEqual([ + 1, 2, 3, + ]); + expect( + (await collect(log.replay({ cursor: 3 }))).map((event) => event.seq), + ).toEqual([4, 5, 6]); + expect( + (await collect(log.replay({ cursor: 2 }))).map((event) => event.seq), + ).toEqual([3]); + }); + + it("tails a segment begun after the cursor until its terminal event", async () => { + const log = createOptimizationRunLog(); + log.append({ type: "started", requestedTrials: 1 }); + log.append(complete); + log.append({ type: "started", requestedTrials: 2 }); + + const replay = collect(log.replay({ cursor: 2 })); + await Promise.resolve(); + log.append(trial(1)); + log.append({ ...complete, requestedTrials: 2 }); + + expect((await replay).map((event) => [event.type, event.seq])).toEqual([ + ["started", 3], + ["trial", 4], + ["complete", 5], + ]); + }); + + it("ends at once when attached past the end of a settled log", async () => { const log = createOptimizationRunLog(); log.append({ type: "started", requestedTrials: 2 }); log.append(complete); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts index 6bc67e50a6c..9522571bc48 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/run-log.ts @@ -6,17 +6,25 @@ type WithoutSeq = TEvent extends unknown ? Omit : never; /** An optimization event before the log stamps its sequence number. */ export type OptimizationRunLogEvent = WithoutSeq; +/** + * One run's events, in segments: each segment begins with `started` and ends + * with a terminal `complete`/`error`, and a study kept in memory may begin + * another segment when it is extended. + */ export type OptimizationRunLog = { readonly events: readonly PetrinautOptimizationEvent[]; - /** True once a terminal `complete`/`error` event was appended. */ - readonly closed: boolean; - /** Stamps the next dense `seq` (from 1) and stores the event. Throws once closed. */ + /** True while the latest event is terminal, so a replay past it ends at once. */ + readonly settled: boolean; + /** + * Stamps the next dense `seq` (from 1) and stores the event. Once settled, + * only a `started` event may follow. + */ append(event: OptimizationRunLogEvent): PetrinautOptimizationEvent; subscribe(listener: (event: PetrinautOptimizationEvent) => void): () => void; /** * Yields the stored events with `seq` greater than `cursor`, then tails live - * events until the terminal one. Aborting the signal ends the iteration with - * an `AbortError`. + * events, and ends at the first terminal event after the cursor. Aborting + * the signal ends the iteration with an `AbortError`. */ replay(options?: { cursor?: number; @@ -36,7 +44,11 @@ const createAbortError = (): Error => { export const createOptimizationRunLog = (): OptimizationRunLog => { const events: PetrinautOptimizationEvent[] = []; const listeners = new Set<(event: PetrinautOptimizationEvent) => void>(); - let closed = false; + + const isSettled = (): boolean => { + const latest = events.at(-1); + return latest !== undefined && isTerminalEvent(latest); + }; const subscribe = ( listener: (event: PetrinautOptimizationEvent) => void, @@ -66,18 +78,17 @@ export const createOptimizationRunLog = (): OptimizationRunLog => { get events() { return events; }, - get closed() { - return closed; + get settled() { + return isSettled(); }, append(event) { - if (closed) { - throw new Error("optimization run log is closed"); + if (isSettled() && event.type !== "started") { + throw new Error( + "a settled optimization run log accepts only a started event", + ); } const stamped = { ...event, seq: events.length + 1 }; events.push(stamped); - if (isTerminalEvent(stamped)) { - closed = true; - } for (const listener of listeners) { listener(stamped); } @@ -102,7 +113,7 @@ export const createOptimizationRunLog = (): OptimizationRunLog => { } continue; } - if (closed) { + if (isSettled()) { return; } await waitForNextEvent(signal); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts index 1b2be8c3e7e..9e75e63a59a 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts @@ -8,24 +8,27 @@ * import error included, fails the suite. */ import { loadPyodide } from "pyodide"; -import { beforeAll, describe, expect, test } from "vitest"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; import { defaultOptimizerPyodideConfig } from "./pyodide-config"; import { optimizerPythonSources } from "./python-sources"; import { createOptimizerStudyRunner, + type OptimizerStudyCallbacks, type OptimizerStudyRunner, } from "./study-runner"; import type { OptimizationScalar, PetrinautOptimizationDescribeResult, + PetrinautOptimizationTrialOutcome, } from "../optimization"; import type { OptimizerTrialPayload } from "./messages"; declare const process: { readonly env: Readonly>; }; +declare const setTimeout: (handler: () => void, timeout?: number) => unknown; const loadTimeout = 180_000; const runningInCi = (process.env.CI ?? "") !== ""; @@ -72,6 +75,11 @@ const description: PetrinautOptimizationDescribeResult = { ], }; +const withTrials = (trials: number): PetrinautOptimizationDescribeResult => ({ + ...description, + study: { ...description.study, trials }, +}); + const asNumber = (value: OptimizationScalar | undefined): number => { if (typeof value !== "number") { throw new Error(`expected a number, received ${String(value)}`); @@ -85,6 +93,37 @@ const objectiveOf = (values: Record): number => (asNumber(values.count) - 8) ** 2 / 16 + (values.enabled === true ? 0 : 1); +/** Callbacks that record what the study asked and reported. */ +const recorder = (options?: { + evaluate?: ( + trial: number, + suggestedValues: Record, + ) => Promise; + isCancelled?: () => boolean; +}) => { + const evaluated: Record[] = []; + const trialNumbers: number[] = []; + const trials: OptimizerTrialPayload[] = []; + const started: number[] = []; + const callbacks: OptimizerStudyCallbacks = { + onStarted: (requestedTrials) => { + started.push(requestedTrials); + }, + evaluate: async (trial, suggestedValues) => { + trialNumbers.push(trial); + evaluated.push(suggestedValues); + return options?.evaluate + ? options.evaluate(trial, suggestedValues) + : { kind: "objective", objective: objectiveOf(suggestedValues) }; + }, + onTrial: (event) => { + trials.push(event); + }, + isCancelled: options?.isCancelled ?? (() => false), + }; + return { evaluated, trialNumbers, trials, started, callbacks }; +}; + let runner: OptimizerStudyRunner; let loadFailure: string | null = null; @@ -105,6 +144,12 @@ beforeAll(async () => { } }, loadTimeout); +afterAll(async () => { + if (loadFailure === null) { + await runner.dispose(); + } +}); + const skipWhenOffline = (skip: (note?: string) => never): void => { if (loadFailure !== null) { skip(`Pyodide packages could not be downloaded in CI: ${loadFailure}`); @@ -116,26 +161,25 @@ describe("createOptimizerStudyRunner", () => { "drives a seeded TPE study through the host evaluate callback", async ({ skip }) => { skipWhenOffline(skip); - const trials: OptimizerTrialPayload[] = []; - const evaluated: Record[] = []; + const run = recorder({ + evaluate: async (trial, suggestedValues) => + trial === 3 + ? { kind: "pruned", reason: "no frames" } + : { kind: "objective", objective: objectiveOf(suggestedValues) }, + }); - const summary = await runner.run({ + const summary = await runner.start({ + runId: "seeded", description, - evaluate: async (trial, suggestedValues) => { - expect(trial).toBe(evaluated.length); - evaluated.push(suggestedValues); - return trial === 3 - ? { kind: "pruned", reason: "no frames" } - : { kind: "objective", objective: objectiveOf(suggestedValues) }; - }, - onTrial: (event) => { - trials.push(event); - }, - isCancelled: () => false, + parallelism: 1, + callbacks: run.callbacks, }); - expect(evaluated).toHaveLength(30); - for (const values of evaluated) { + expect(run.started).toEqual([30]); + expect(run.trialNumbers).toEqual( + Array.from({ length: 30 }, (_, index) => index), + ); + for (const values of run.evaluated) { expect(Object.keys(values).sort()).toEqual([ "count", "enabled", @@ -152,12 +196,12 @@ describe("createOptimizerStudyRunner", () => { expect(typeof values.enabled).toBe("boolean"); } - expect(trials.map((event) => event.trial)).toEqual( + expect(run.trials.map((event) => event.trial)).toEqual( Array.from({ length: 30 }, (_, index) => index), ); let bestSoFar = Number.POSITIVE_INFINITY; - for (const [index, event] of trials.entries()) { - expect(event.parameters).toEqual(evaluated[index]); + for (const [index, event] of run.trials.entries()) { + expect(event.parameters).toEqual(run.evaluated[index]); if (index === 3) { expect(event.state).toBe("pruned"); expect(event.objective).toBeNull(); @@ -181,76 +225,214 @@ describe("createOptimizerStudyRunner", () => { completedTrials: 29, prunedTrials: 1, failedTrials: 0, - best: trials.at(-1)?.best, + best: run.trials.at(-1)?.best, }); }, loadTimeout, ); test( - "stops early once the host reports cancellation", + "continues a study on the same sampler history, numbering onwards", async ({ skip }) => { skipWhenOffline(skip); - let evaluations = 0; - const trials: OptimizerTrialPayload[] = []; + const whole = recorder(); + await runner.start({ + runId: "whole", + description: withTrials(16), + parallelism: 1, + callbacks: whole.callbacks, + }); - const summary = await runner.run({ - description, - evaluate: async (_trial, suggestedValues) => { - evaluations += 1; - return { kind: "objective", objective: objectiveOf(suggestedValues) }; - }, - onTrial: (event) => { - trials.push(event); - }, - isCancelled: () => evaluations >= 5, + const first = recorder(); + await runner.start({ + runId: "split", + description: withTrials(8), + parallelism: 1, + callbacks: first.callbacks, + }); + const second = recorder(); + const summary = await runner.extend({ + runId: "split", + trials: 8, + parallelism: 1, + callbacks: second.callbacks, }); - expect(evaluations).toBe(5); - expect(trials).toHaveLength(4); + // TPE leaves its random start-up after 10 trials, so equal sequences + // prove the extension sampled from the first segment's history. + expect([...first.evaluated, ...second.evaluated]).toEqual( + whole.evaluated, + ); + expect(first.started).toEqual([8]); + expect(second.started).toEqual([16]); + expect(second.trialNumbers).toEqual([8, 9, 10, 11, 12, 13, 14, 15]); + expect(second.trials.map((event) => event.trial)).toEqual( + second.trialNumbers, + ); expect(summary).toMatchObject({ + requestedTrials: 16, + completedTrials: 16, + prunedTrials: 0, + }); + expect(summary.best).toEqual(whole.trials.at(-1)?.best); + }, + loadTimeout, + ); + + test( + "stops early once the host reports cancellation and resumes from the trials told", + async ({ skip }) => { + skipWhenOffline(skip); + const run = recorder({ isCancelled: () => run.trialNumbers.length >= 5 }); + + const stopped = await runner.start({ + runId: "stopped", + description, + parallelism: 1, + callbacks: run.callbacks, + }); + + expect(run.trialNumbers).toHaveLength(5); + expect(run.trials).toHaveLength(4); + // The trial in flight at the stop is told as failed so Optuna keeps no + // running trial behind; the count carries into the resumed summary. + expect(stopped).toMatchObject({ requestedTrials: 30, completedTrials: 4, + failedTrials: 1, cancelled: true, }); + + const resumed = recorder(); + const summary = await runner.extend({ + runId: "stopped", + trials: 2, + parallelism: 1, + callbacks: resumed.callbacks, + }); + + expect(resumed.started).toEqual([6]); + expect(resumed.trialNumbers).toEqual([5, 6]); + expect(resumed.trials.map((event) => event.trial)).toEqual([5, 6]); + expect(summary).toEqual({ + requestedTrials: 6, + completedTrials: 6, + prunedTrials: 0, + failedTrials: 1, + best: resumed.trials.at(-1)?.best, + }); }, loadTimeout, ); test( - "samples deterministically for a seed and rejects a non-finite objective", + "keeps up to the parallelism in flight and tells outcomes as they settle", async ({ skip }) => { skipWhenOffline(skip); - const sequences: Record[][] = []; - for (let repeat = 0; repeat < 2; repeat++) { - const sequence: Record[] = []; - await runner.run({ - description: { - ...description, - study: { ...description.study, trials: 8 }, - }, - evaluate: async (_trial, suggestedValues) => { - sequence.push(suggestedValues); - return { - kind: "objective", - objective: objectiveOf(suggestedValues), - }; - }, - onTrial: () => {}, - isCancelled: () => false, - }); - sequences.push(sequence); + let inFlight = 0; + let mostInFlight = 0; + const run = recorder({ + evaluate: async (trial, suggestedValues) => { + inFlight += 1; + mostInFlight = Math.max(mostInFlight, inFlight); + await new Promise((resolve) => { + setTimeout(resolve, trial % 2 === 0 ? 30 : 1); + }); + inFlight -= 1; + return { kind: "objective", objective: objectiveOf(suggestedValues) }; + }, + }); + + const summary = await runner.start({ + runId: "parallel", + description: withTrials(6), + parallelism: 2, + callbacks: run.callbacks, + }); + + expect(mostInFlight).toBe(2); + expect(run.trialNumbers).toEqual([0, 1, 2, 3, 4, 5]); + expect( + run.trials + .map((event) => event.trial) + .sort((left, right) => left - right), + ).toEqual([0, 1, 2, 3, 4, 5]); + for (const event of run.trials) { + expect(event.parameters).toEqual(run.evaluated[event.trial]); + expect(event.objective).toBeCloseTo(objectiveOf(event.parameters), 9); } - expect(sequences[0]).toEqual(sequences[1]); + expect(summary).toMatchObject({ + requestedTrials: 6, + completedTrials: 6, + prunedTrials: 0, + }); + }, + loadTimeout, + ); + + test( + "drops a released study, and one whose segment failed", + async ({ skip }) => { + skipWhenOffline(skip); + await runner.start({ + runId: "released", + description: withTrials(2), + parallelism: 1, + callbacks: recorder().callbacks, + }); + await runner.release("released"); + await runner.release("never-started"); + + await expect( + runner.extend({ + runId: "released", + trials: 1, + parallelism: 1, + callbacks: recorder().callbacks, + }), + ).rejects.toThrow('Optimization study "released" is not kept'); await expect( - runner.run({ + runner.start({ + runId: "failing", description, - evaluate: async () => ({ kind: "objective", objective: Number.NaN }), - onTrial: () => {}, - isCancelled: () => false, + parallelism: 1, + callbacks: recorder({ + evaluate: async () => ({ + kind: "objective", + objective: Number.NaN, + }), + }).callbacks, }), ).rejects.toThrow("trial objective must be a finite number"); + await expect( + runner.extend({ + runId: "failing", + trials: 1, + parallelism: 1, + callbacks: recorder().callbacks, + }), + ).rejects.toThrow('Optimization study "failing" is not kept'); + }, + loadTimeout, + ); + + test( + "samples deterministically for a seed", + async ({ skip }) => { + skipWhenOffline(skip); + const sequences: Record[][] = []; + for (let repeat = 0; repeat < 2; repeat++) { + const run = recorder(); + await runner.start({ + runId: `repeat-${repeat}`, + description: withTrials(8), + parallelism: 1, + callbacks: run.callbacks, + }); + sequences.push(run.evaluated); + } + expect(sequences[0]).toEqual(sequences[1]); }, loadTimeout, ); diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts index f5696e60911..c16df423731 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.ts @@ -15,10 +15,11 @@ import type { OptimizerTrialPayload, } from "./messages"; import type { OptimizerPyodideConfig } from "./pyodide-config"; -import type { LoadPyodide, PyodideLike } from "./pyodide-like"; +import type { LoadPyodide, PyodideLike, PyProxyLike } from "./pyodide-like"; -export type OptimizerStudyRunInput = { - description: PetrinautOptimizationDescribeResult; +export type OptimizerStudyCallbacks = { + /** The segment began; `requestedTrials` is the study's cumulative total. */ + onStarted(requestedTrials: number): void; evaluate( trial: number, suggestedValues: Record, @@ -27,23 +28,67 @@ export type OptimizerStudyRunInput = { isCancelled(): boolean; }; +export type OptimizerStudyStartInput = { + runId: string; + description: PetrinautOptimizationDescribeResult; + /** Trials kept in flight at once. */ + parallelism: number; + callbacks: OptimizerStudyCallbacks; +}; + +export type OptimizerStudyExtendInput = { + runId: string; + trials: number; + parallelism: number; + callbacks: OptimizerStudyCallbacks; +}; + export type OptimizerStudyRunner = { /** Settles once Pyodide, the packages and the optimizer sources are loaded. */ readonly ready: Promise; - /** Runs one study; concurrent calls run one after another. */ - run(input: OptimizerStudyRunInput): Promise; + /** + * Creates the study for `runId` and runs its first `description.study.trials` + * trials. The segments of every study run one after another, in call order. + */ + start(input: OptimizerStudyStartInput): Promise; + /** Runs `trials` more on the kept study; the trial numbers continue. */ + extend(input: OptimizerStudyExtendInput): Promise; + /** Drops the kept study once the segments queued before it have run. */ + release(runId: string): Promise; + /** Drops every kept study. */ + dispose(): Promise; }; /** The outcome shape `ask_tell.run_study` expects from its evaluate callback. */ type PythonTrialOutcome = { objective: number } | { pruned: string }; +/** The Python `StudyHandle`; `requested` is the study's cumulative trial total. */ +type StudyHandleProxy = PyProxyLike & { readonly requested: number }; + type PyodideEntryModule = { - run_browser_study( + create_browser_study( descriptionJson: string, + parallelism: number, + ): StudyHandleProxy; + run_browser_study( + handle: StudyHandleProxy, + trials: number, evaluate: (values: unknown) => Promise, onTrial: (payload: unknown) => void, isCancelled: () => boolean, + parallelism: number, ): Promise; + release_browser_study(handle: StudyHandleProxy): void; +}; + +type KeptStudy = { + readonly handle: StudyHandleProxy; + /** + * Optuna numbers a study's trials densely in ask order and every ask leads + * to one evaluate call, so the count of evaluate calls is the next trial's + * number, across segments and across trials a stop left untold. + */ + nextTrial: number; }; const pythonSourceRoot = "/home/pyodide"; @@ -175,48 +220,122 @@ export const createOptimizerStudyRunner = (options: { }): OptimizerStudyRunner => { const entry = loadOptimizerEntry(options); const ready = entry.then(() => undefined); + const studies = new Map(); let queue: Promise = ready; - const runStudy = async ( - input: OptimizerStudyRunInput, + const enqueue = ( + task: (module: PyodideEntryModule) => Promise, + ): Promise => { + const result = queue.then(() => entry).then(task); + queue = result.then( + () => undefined, + () => undefined, + ); + return result; + }; + + const dropStudy = (module: PyodideEntryModule, runId: string): void => { + const study = studies.get(runId); + if (!study) { + return; + } + studies.delete(runId); + module.release_browser_study(study.handle); + study.handle.destroy(); + }; + + const runSegment = async ( + module: PyodideEntryModule, + runId: string, + trials: number, + parallelism: number, + callbacks: OptimizerStudyCallbacks, ): Promise => { - const module = await entry; - // Optuna numbers an in-memory study's trials densely from 0 in ask order, - // so the evaluate call count is the trial number. - let nextTrial = 0; + const study = studies.get(runId); + if (!study) { + throw new Error( + `Optimization study "${runId}" is not kept: it was released, failed or never started`, + ); + } const evaluate = (values: unknown): Promise => { // Argument proxies are destroyed when this call returns: convert before // the first await. const suggestedValues = asParameters(toJsValue(values)); - const trial = nextTrial; - nextTrial += 1; - return input.evaluate(trial, suggestedValues).then(toPythonOutcome); + const trial = study.nextTrial; + study.nextTrial += 1; + return callbacks.evaluate(trial, suggestedValues).then(toPythonOutcome); }; const onTrial = (payload: unknown): void => { - input.onTrial(normalizeTrialPayload(toJsValue(payload))); + callbacks.onTrial(normalizeTrialPayload(toJsValue(payload))); }; - const result = await module.run_browser_study( - JSON.stringify(input.description), - evaluate, - onTrial, - () => input.isCancelled(), - ); - const summary = normalizeSummary(toJsValue(result)); - if (isPyProxyLike(result)) { - result.destroy(); + try { + const pending = module.run_browser_study( + study.handle, + trials, + evaluate, + onTrial, + () => callbacks.isCancelled(), + parallelism, + ); + callbacks.onStarted(study.handle.requested); + const result = await pending; + const summary = normalizeSummary(toJsValue(result)); + if (isPyProxyLike(result)) { + result.destroy(); + } + return summary; + } catch (error) { + // The run that owns a failed study ends as failed and never releases it. + dropStudy(module, runId); + throw error; } - return summary; }; return { ready, - run(input) { - const result = queue.then(() => runStudy(input)); - queue = result.then( - () => undefined, - () => undefined, - ); - return result; + start(input) { + return enqueue(async (module) => { + if (studies.has(input.runId)) { + throw new Error(`Optimization study "${input.runId}" already exists`); + } + studies.set(input.runId, { + handle: module.create_browser_study( + JSON.stringify(input.description), + input.parallelism, + ), + nextTrial: 0, + }); + return runSegment( + module, + input.runId, + input.description.study.trials, + input.parallelism, + input.callbacks, + ); + }); + }, + extend(input) { + return enqueue(async (module) => { + return runSegment( + module, + input.runId, + input.trials, + input.parallelism, + input.callbacks, + ); + }); + }, + release(runId) { + return enqueue(async (module) => { + dropStudy(module, runId); + }); + }, + dispose() { + return enqueue(async (module) => { + for (const runId of studies.keys()) { + dropStudy(module, runId); + } + }); }, }; }; diff --git a/libs/@hashintel/petrinaut-core/src/optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization.ts index 7addb8cd669..db102ee09f1 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization.ts @@ -761,6 +761,55 @@ export type PetrinautOptimizationChannel = { ): Promise; }; +export const PETRINAUT_OPTIMIZATION_MAX_PARALLELISM = 4; + +/** Options of a run on a connected capability; the remote one takes `signal` only. */ +export type PetrinautConnectedRunOptions = { + /** + * Creation rejects with an `AbortError` when this is already aborted; a + * run that exists is stopped through `cancelOptimizationRun`. + */ + signal?: AbortSignalLike; + /** + * How many trials the study keeps in flight, 1 to + * `PETRINAUT_OPTIMIZATION_MAX_PARALLELISM`. Defaults to 1, which samples + * exactly as a sequential study does. + */ + parallelism?: number; +}; + +/** + * The capability a connected source yields. A study that completed or was + * cancelled stays in memory until it is released, so more trials can be run + * on it with the sampler's history intact. + */ +export type PetrinautConnectedOptimizationCapability = Omit< + PetrinautOptimization, + "createOptimizationRun" +> & { + createOptimizationRun( + input: PetrinautOptimizationInput, + options?: PetrinautConnectedRunOptions, + ): Promise<{ runId: string }>; + /** + * Run `trials` more on a run that completed or was cancelled. Its event + * stream gains a `started` event carrying the cumulative `requestedTrials`, + * the new trials continue the numbering, and the next `complete` reports + * counts over the whole study. Rejects for a run that is running, was + * released or failed, and when the total would exceed + * `PETRINAUT_OPTIMIZATION_MAX_TRIALS`. + */ + extendOptimizationRun( + runId: string, + trials: number, + options?: Pick, + ): Promise; + /** Drop the study behind a run, which can then no longer be extended. Idempotent. */ + releaseOptimizationRun(runId: string): Promise; + /** Cancel every run, drop every study and free the runtime. */ + dispose(this: void): void; +}; + /** * An optimization capability that runs where the host runs and needs the * host's compute: connect it to a channel to obtain the capability. @@ -773,7 +822,7 @@ export type PetrinautConnectedOptimization = { connect( this: void, channel: PetrinautOptimizationChannel, - ): PetrinautOptimization & { dispose(this: void): void }; + ): PetrinautConnectedOptimizationCapability; }; /** What a host supplies: a remote capability, or one to connect locally. */ diff --git a/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts b/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts index 9f5bef00c34..15a4dd5f184 100644 --- a/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts +++ b/libs/@hashintel/petrinaut-core/src/workers/optimizer.ts @@ -11,9 +11,13 @@ export type { OptimizerErrorMessage, OptimizerEvaluatedMessage, OptimizerEvaluateMessage, + OptimizerExtendMessage, OptimizerInitErrorMessage, OptimizerInitMessage, OptimizerReadyMessage, + OptimizerReleasedMessage, + OptimizerReleaseMessage, + OptimizerStartedMessage, OptimizerStartMessage, OptimizerStudySummary, OptimizerToMainMessage, From 6e677b1d88b01b8cb78527727cbed10d08b3a611 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Fri, 4 Sep 2026 06:53:51 +0200 Subject: [PATCH 3/4] Record petrinaut-core's task dependency on the optimizer core library --- libs/@hashintel/petrinaut-core/docs/task-dependencies.json | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/@hashintel/petrinaut-core/docs/task-dependencies.json b/libs/@hashintel/petrinaut-core/docs/task-dependencies.json index 0d31f2ec9b1..d8af6dec7da 100644 --- a/libs/@hashintel/petrinaut-core/docs/task-dependencies.json +++ b/libs/@hashintel/petrinaut-core/docs/task-dependencies.json @@ -1,6 +1,8 @@ { "package": "@hashintel/petrinaut-core", - "dependencies": [], + "dependencies": [ + "@local/petrinaut-optimizer-core" + ], "tasks": { "build": [], "fix:eslint": [ From b597b24440eb9e2428a3cd597e6458ec41b0c0a4 Mon Sep 17 00:00:00 2001 From: Chris Feijoo Date: Fri, 4 Sep 2026 14:15:27 +0200 Subject: [PATCH 4/4] State the continuation test as a sampler-history claim --- .../study-runner.pyodide.test.ts | 68 +++++++++---------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts index 9e75e63a59a..30249835888 100644 --- a/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization/study-runner.pyodide.test.ts @@ -235,46 +235,44 @@ describe("createOptimizerStudyRunner", () => { "continues a study on the same sampler history, numbering onwards", async ({ skip }) => { skipWhenOffline(skip); - const whole = recorder(); - await runner.start({ - runId: "whole", - description: withTrials(16), - parallelism: 1, - callbacks: whole.callbacks, - }); - - const first = recorder(); - await runner.start({ - runId: "split", - description: withTrials(8), - parallelism: 1, - callbacks: first.callbacks, - }); - const second = recorder(); - const summary = await runner.extend({ - runId: "split", - trials: 8, - parallelism: 1, - callbacks: second.callbacks, - }); - - // TPE leaves its random start-up after 10 trials, so equal sequences - // prove the extension sampled from the first segment's history. - expect([...first.evaluated, ...second.evaluated]).toEqual( - whole.evaluated, - ); - expect(first.started).toEqual([8]); - expect(second.started).toEqual([16]); - expect(second.trialNumbers).toEqual([8, 9, 10, 11, 12, 13, 14, 15]); - expect(second.trials.map((event) => event.trial)).toEqual( - second.trialNumbers, + const split = async (runId: string) => { + const first = recorder(); + await runner.start({ + runId, + description: withTrials(8), + parallelism: 1, + callbacks: first.callbacks, + }); + const second = recorder(); + const summary = await runner.extend({ + runId, + trials: 8, + parallelism: 1, + callbacks: second.callbacks, + }); + return { first, second, summary }; + }; + const study = await split("split"); + const again = await split("split-again"); + + // A seeded restart repeats the first segment exactly, so a second + // segment that differs from the first proves the extension sampled + // from the history the first segment left instead of restarting. + expect(again.first.evaluated).toEqual(study.first.evaluated); + expect(again.second.evaluated).toEqual(study.second.evaluated); + expect(study.second.evaluated).not.toEqual(study.first.evaluated); + expect(study.first.started).toEqual([8]); + expect(study.second.started).toEqual([16]); + expect(study.second.trialNumbers).toEqual([8, 9, 10, 11, 12, 13, 14, 15]); + expect(study.second.trials.map((event) => event.trial)).toEqual( + study.second.trialNumbers, ); - expect(summary).toMatchObject({ + expect(study.summary).toMatchObject({ requestedTrials: 16, completedTrials: 16, prunedTrials: 0, }); - expect(summary.best).toEqual(whole.trials.at(-1)?.best); + expect(study.summary.best).toEqual(study.second.trials.at(-1)?.best); }, loadTimeout, );