diff --git a/.changeset/browser-optimization-runtime.md b/.changeset/browser-optimization-runtime.md new file mode 100644 index 00000000000..d5a6833f7d3 --- /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. 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. Pyodide and Optuna download from jsDelivr and PyPI by default; `pyodide.indexURL` overrides where the runtime loads from. The `./optimization` entry gains `parseOptimizationManifest`, `isUnknownOptimizationRunError` and `PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE`, and the main entry gains `createReadableStore`. diff --git a/libs/@hashintel/petrinaut-cli/src/commands/built-cli.test.ts b/libs/@hashintel/petrinaut-cli/src/commands/built-cli.test.ts index a2addeb9f0b..228626a0536 100644 --- a/libs/@hashintel/petrinaut-cli/src/commands/built-cli.test.ts +++ b/libs/@hashintel/petrinaut-cli/src/commands/built-cli.test.ts @@ -5,7 +5,8 @@ import { fileURLToPath } from "node:url"; import { describe, expect, it } from "vitest"; -import { deriveTrialSeeds } from "../runtime/optimization"; +import { deriveOptimizationTrialSeeds } from "@hashintel/petrinaut-core/optimization"; + import { createOptimizationManifest } from "./optimization-manifest.fixtures"; const distCliPath = fileURLToPath( @@ -57,7 +58,7 @@ describe.skipIf(!existsSync(distCliPath))("built CLI", () => { .split("\n") .filter(Boolean) .map((line) => JSON.parse(line) as unknown); - const seeds = deriveTrialSeeds(42, 2); + const seeds = deriveOptimizationTrialSeeds(42, 2); expect(responses).toEqual([ { id: 1, diff --git a/libs/@hashintel/petrinaut-cli/src/commands/stdio.ts b/libs/@hashintel/petrinaut-cli/src/commands/stdio.ts index 6f12b673ba7..89a3839a6a6 100644 --- a/libs/@hashintel/petrinaut-cli/src/commands/stdio.ts +++ b/libs/@hashintel/petrinaut-cli/src/commands/stdio.ts @@ -2,13 +2,13 @@ import { resolve } from "node:path"; import { createInterface } from "node:readline"; import { compilePetrinautModel } from "@hashintel/petrinaut-core/compiled-model"; +import { parseOptimizationManifest } from "@hashintel/petrinaut-core/optimization"; import { loadSdcpnModel, parseSdcpnModel } from "../runtime/load-model"; import { createNodeSimulationWorkerFactory } from "../runtime/node-simulation-worker"; import { createOptimizationProtocol, loadOptimizationManifest, - parseOptimizationManifest, } from "../runtime/optimization"; import { handleProtocolLine, diff --git a/libs/@hashintel/petrinaut-cli/src/commands/transports.test.ts b/libs/@hashintel/petrinaut-cli/src/commands/transports.test.ts index 9f864635d76..be52431bd60 100644 --- a/libs/@hashintel/petrinaut-cli/src/commands/transports.test.ts +++ b/libs/@hashintel/petrinaut-cli/src/commands/transports.test.ts @@ -8,7 +8,8 @@ import { fileURLToPath } from "node:url"; import { afterEach, describe, expect, it, vi } from "vitest"; -import { deriveTrialSeeds } from "../runtime/optimization"; +import { deriveOptimizationTrialSeeds } from "@hashintel/petrinaut-core/optimization"; + import { MAX_REQUEST_LINE_BYTES } from "../runtime/protocol"; import { createOptimizationManifest } from "./optimization-manifest.fixtures"; import { serve } from "./serve"; @@ -231,7 +232,7 @@ describe("CLI transports", () => { ); await serving; - const seeds = deriveTrialSeeds(42, 3); + const seeds = deriveOptimizationTrialSeeds(42, 3); expect(seeds[0]).toBe(42); expect(parseResponses(stdout)).toEqual([ { diff --git a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.test.ts b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.test.ts index abb0d2d4561..97a48b645a6 100644 --- a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.test.ts +++ b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.test.ts @@ -7,12 +7,14 @@ import { describe, expect, it } from "vitest"; import { serializeDocument } from "@hashintel/petrinaut-core"; import { compilePetrinautModel } from "@hashintel/petrinaut-core/compiled-model"; +import { + deriveOptimizationTrialSeeds, + parseOptimizationManifest, +} from "@hashintel/petrinaut-core/optimization"; import { createOptimizationProtocol, - deriveTrialSeeds, loadOptimizationManifest, - parseOptimizationManifest, } from "./optimization"; import type { @@ -451,7 +453,7 @@ describe("createOptimizationProtocol", () => { it("runs every trial seed and aggregates the objectives by mean", async () => { const manifest = await createSeededManifest(3); - const seeds = deriveTrialSeeds(42, 3); + const seeds = deriveOptimizationTrialSeeds(42, 3); // Objectives 1, 2 and 3 in seed order, so the mean and the per-seed // echoes are both observable. const { factory, calls } = createFakeExperimentFactory( @@ -551,7 +553,7 @@ describe("createOptimizationProtocol", () => { parameterValues: { infected_ratio: 0.1 }, }); expect(first.replicates?.map((replicate) => replicate.seed)).toEqual( - deriveTrialSeeds(42, 2), + deriveOptimizationTrialSeeds(42, 2), ); for (const replicate of first.replicates ?? []) { expect(Number.isFinite(replicate.objective)).toBe(true); @@ -583,22 +585,3 @@ describe("createOptimizationProtocol", () => { ).rejects.toThrow("1 of 1 optimization replicates failed"); }); }); - -describe("deriveTrialSeeds", () => { - it("keeps the base seed first and derives a stable, in-range sequence", () => { - expect(deriveTrialSeeds(42, 1)).toEqual([42]); - // Pins the documented derivation |seed + (i + 1) x 2654435761| mod 2^31, - // which the other tests' expected seed sequences depend on. - expect(deriveTrialSeeds(42, 2)).toEqual([42, 1_013_904_268]); - - const seeds = deriveTrialSeeds(42, 100); - expect(seeds[0]).toBe(42); - expect(seeds).toEqual(deriveTrialSeeds(42, 100)); - expect(new Set(seeds).size).toBe(seeds.length); - for (const seed of seeds) { - expect(Number.isInteger(seed)).toBe(true); - expect(seed).toBeGreaterThanOrEqual(0); - expect(seed).toBeLessThanOrEqual(2_147_483_647); - } - }); -}); diff --git a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts index 3eb73674aa0..2af9fccc6ed 100644 --- a/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts +++ b/libs/@hashintel/petrinaut-cli/src/runtime/optimization.ts @@ -3,57 +3,37 @@ 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, + parseOptimizationManifest, + 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"]; - -function formatManifestIssues( - prefix: string, +function invalidParamsError( issues: readonly { path: PropertyKey[]; message: string }[], ): Error { const details = issues .map( ({ path, message }) => - `${path.length > 0 ? path.join(".") : "manifest"}: ${message}`, + `${path.length > 0 ? path.join(".") : "params"}: ${message}`, ) .join("; "); - return new Error(`${prefix}: ${details}`); -} - -export function parseOptimizationManifest( - data: unknown, -): PetrinautOptimizationManifest { - const parsed = petrinautOptimizationManifestSchema.safeParse(data); - if (!parsed.success) { - throw formatManifestIssues( - "Invalid optimization manifest", - parsed.error.issues, - ); - } - return parsed.data; + return new Error(`Invalid optimization.evaluate params: ${details}`); } export async function loadOptimizationManifest( @@ -67,95 +47,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 +100,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 +111,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,54 +125,18 @@ 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 = petrinautOptimizationEvaluateParamsSchema.safeParse(params); if (!parsed.success) { - throw formatManifestIssues( - "Invalid optimization.evaluate params", - 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; + throw invalidParamsError(parsed.error.issues); } + 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/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": [ diff --git a/libs/@hashintel/petrinaut-core/package.json b/libs/@hashintel/petrinaut-core/package.json index 32b4826b368..42b01169803 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" @@ -98,12 +102,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..c62aee28cd3 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/browser-optimization.ts @@ -0,0 +1,3 @@ +// The architecture graph resolves a package subpath from the root source file +// of the same name; the browser runtime's entry lives with its layer. +export * from "./optimization/browser/index"; 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/index.ts b/libs/@hashintel/petrinaut-core/src/index.ts index 5196310c0ec..70d3097a2da 100644 --- a/libs/@hashintel/petrinaut-core/src/index.ts +++ b/libs/@hashintel/petrinaut-core/src/index.ts @@ -2,7 +2,8 @@ * Public surface for `@hashintel/petrinaut-core` — the headless engine. * * No React, no DOM, no Monaco. Stateful handles, streams, and pure logic for - * SDCPN documents, simulation, LSP, and playback. + * SDCPN documents, simulation, LSP, and playback, and an in-browser + * optimization runtime (Optuna under Pyodide in a worker). * * @layerRoot core * @role SDCPN document model, compiler, simulation runtimes and LSP, with no UI framework @@ -52,7 +53,7 @@ export { type PetrinautHistory, type PetrinautPatch, } from "./handle"; -export type { ReadableStore } from "./store"; +export { createReadableStore, type ReadableStore } from "./store"; export { DEFAULT_PETRINAUT_EXTENSIONS, PETRINAUT_EXTENSION_NAMES, diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/README.md b/libs/@hashintel/petrinaut-core/src/optimization/browser/README.md new file mode 100644 index 00000000000..ff7e12cc2db --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/README.md @@ -0,0 +1,71 @@ +--- +layer: core.optimization.browser +role: Runs the Optuna study in a Pyodide worker and evaluates trials through the host channel +--- + +# Browser optimization runtime + +`createBrowserOptimization` is a `PetrinautConnectedOptimization`: a host +connects it to a `PetrinautOptimizationChannel` and receives the same +capability a remote optimization service offers, plus `extendOptimizationRun` +and `releaseOptimizationRun` for the study the worker keeps between segments. + +## Three tiers + +| Tier | Where it runs | Files | +| ------------ | ------------------ | ------------------------------------------------------------------ | +| Capability | the main thread | `browser-optimization.ts`, `run-log.ts`, `messages.ts` | +| Worker | a module worker | `worker/optimizer.worker.ts`, `worker/attach-optimizer-worker.ts` | +| Python study | Pyodide, in-worker | `worker/study-runner.ts` driving `@local/petrinaut-optimizer-core` | + +The capability queues runs, keeps one event log per run and answers each +`evaluate` message by calling `channel.evaluateTrial`. The worker loads Pyodide +from `pyodide.indexURL` (jsDelivr by default), installs the packages pinned in +`runtime-lock.json` with micropip, writes the Python sources from +`python-sources.ts` into Pyodide's filesystem and pairs each evaluate request +with the study loop. The Python study asks Optuna for values, awaits the +worker's evaluate callback and tells the outcome back. + +## Segments + +A run advances in segments. `start` runs the manifest's trial count on a new +study; `extend` runs more trials on the kept study, numbering onwards. Each +segment begins with a `started` event in the run log and ends with a terminal +`complete` or `error` event. + +```text +queued ──worker ready──▶ running ──complete / cancelled──▶ finished-resumable + │ │ │ + │ cancel (first run) │ trial evaluation failed, │ extend + │ │ study error, worker error ▼ + ▼ ▼ queued +finished ◀──────────── finished (again) + ▲ + └── release, from any status +``` + +`queued` waits for the worker to take the 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. Runs execute one at a time on a shared +worker. + +## Who owns what + +- **Cancellation** is the capability's. `cancelOptimizationRun` aborts the + segment's signal so the channel stops the trial's runs, and posts `cancel`; + the worker resolves the segment's pending evaluations as pruned and the + Python loop tells the trials in flight as failed without reporting them, then + returns early. The capability appends the cancelled error event when the + worker confirms with `cancelled`. +- **Trial numbering** is the study runner's. Optuna numbers trials densely in + ask order and every ask leads to one evaluate call, so the count of evaluate + calls made for a study is the next trial's number, across segments and across + trials a stop left untold. +- **The trial cap** is the capability's, checked when a run is created or + extended against `PETRINAUT_OPTIMIZATION_MAX_TRIALS` and the trials the log + already holds. Python checks its own cap for `handle.requested`. +- **`requestedTrials`** is the capability's. It appends `started` with the + cumulative total it computed; the Python summary reports its own count in + `complete`. +- **Parallelism** is fixed when the study is created and every extension + inherits it. diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/browser-optimization.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/browser-optimization.test.ts new file mode 100644 index 00000000000..e0caaa15740 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/browser-optimization.test.ts @@ -0,0 +1,1181 @@ +import { describe, expect, it, vi } from "vitest"; + +import { createAbortController } from "../../environment"; +import { createOptimizationManifestInput } from "../../shared/optimization-manifest.fixtures"; +import { + isUnknownOptimizationRunError, + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE, +} from "../index"; +import { createBrowserOptimization } from "./browser-optimization"; + +import type { WorkerMessageHandler } from "../../environment"; +import type { + PetrinautOptimizationChannel, + PetrinautOptimizationEvent, + PetrinautOptimizationTrialRequest, +} from "../index"; +import type { + OptimizerStudySummary, + OptimizerToMainMessage, + OptimizerToWorkerMessage, + OptimizerTrialPayload, +} from "./messages"; +import type { + OptimizerWorkerErrorEvent, + OptimizerWorkerLike, +} from "./worker/create-optimizer-worker"; + +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 }, + }, + parallelism: 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, + resumable: true, + 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); + context.worker.emit({ + type: "evaluate", + runId, + requestId: 2, + trial: 1, + suggestedValues: { rate: 0.5, count: 6, enabled: true }, + }); + await flush(); + expect(context.worker.sentOfType("cancel")).toEqual([ + { type: "cancel", runId }, + ]); + // The worker prunes the segment's pending evaluations itself, so the + // aborted trial gets no reply and the late evaluate is never run. + expect(context.worker.sentOfType("evaluated")).toHaveLength(0); + expect(context.evaluateTrial).toHaveBeenCalledTimes(1); + + 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, + resumable: true, + seq: 2, + }); + }); + + 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( + 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"]); + // No study was ever created for it, so the host offers no continuation. + expect(events[1]).toMatchObject({ + code: PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + resumable: false, + }); + }); + + it("rejects attaching to an unknown run with the unknown-run code", () => { + const context = setUp(); + + let thrown: unknown; + try { + context.capability.attachOptimizationRun("missing"); + } catch (error) { + thrown = error; + } + expect(thrown).toMatchObject({ + code: PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE, + }); + expect(isUnknownOptimizationRunError(thrown)).toBe(true); + expect(isUnknownOptimizationRunError(new Error("other"))).toBe(false); + }); + + 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, + resumable: 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, + resumable: 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.each([ + { + failure: "the runtime does not load", + createWorker: undefined, + fail: (context: ReturnType) => { + context.worker.emit({ type: "init-error", message: "offline" }); + }, + message: "The in-browser optimizer could not start: offline", + }, + { + failure: "the worker script does not load", + createWorker: undefined, + fail: (context: ReturnType) => { + context.worker.emitError({}); + }, + message: + "The in-browser optimizer could not start: The optimizer worker failed", + }, + { + failure: "the worker cannot be created", + createWorker: (attempt: number) => { + if (attempt === 1) { + throw new Error("SecurityError: cross-origin worker script"); + } + return createFakeWorker(); + }, + fail: () => {}, + message: + "The in-browser optimizer could not start: SecurityError: cross-origin worker script", + }, + ])( + "fails a run as retryable when $failure and retries with a fresh worker", + async ({ createWorker, fail, message }) => { + const context = setUp({ createWorker }); + const { runId } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + const failedWorkers = [...context.workers]; + + fail(context); + await flush(); + + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "optimizer_unavailable", + message, + retryable: true, + resumable: false, + seq: 2, + }); + for (const worker of failedWorkers) { + expect(worker.terminated).toBe(true); + } + + const second = await startRun(context); + expect(context.workers).toHaveLength(failedWorkers.length + 1); + expect(context.worker.sentOfType("init")).toHaveLength(1); + expect( + context.worker.sentOfType("start").map(({ runId: started }) => started), + ).toEqual([second]); + }, + ); + + it("fails the running study as retryable when the worker crashes and respawns for the queued run", async () => { + const context = setUp(); + const running = await startRun(context); + const { runId: queued } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + const crashed = context.worker; + + crashed.emitError({ message: "RangeError: out of memory" }); + await flush(); + + expect(crashed.terminated).toBe(true); + const events = await collectEvents( + context.capability.attachOptimizationRun(running), + ); + expect(events.at(-1)).toEqual({ + type: "error", + code: "optimizer_unavailable", + message: + "The in-browser optimizer could not start: RangeError: out of memory", + retryable: true, + resumable: false, + seq: 2, + }); + await expect( + context.capability.extendOptimizationRun(running, 1), + ).rejects.toThrow("released or failed"); + + expect(context.workers).toHaveLength(2); + expect(context.worker.sentOfType("init")).toHaveLength(1); + context.worker.emit({ type: "ready" }); + await flush(); + expect( + context.worker.sentOfType("start").map(({ runId }) => runId), + ).toEqual([queued]); + }); + + it("a worker crash marks the studies it kept as gone, so a continuation is refused", async () => { + const context = setUp(); + const completed = await startRun(context); + context.worker.emit({ + type: "trial", + runId: completed, + event: completedTrial, + }); + context.worker.emit({ type: "complete", runId: completed, summary }); + await flush(); + const { runId: running } = await context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ); + await flush(); + expect( + context.worker.sentOfType("start").map(({ runId }) => runId), + ).toEqual([completed, running]); + + context.worker.emit({ type: "error", runId: running, message: "crashed" }); + context.worker.emitError({ message: "RangeError: out of memory" }); + await flush(); + + await expect( + context.capability.extendOptimizationRun(completed, 1), + ).rejects.toThrow("optimizer worker was lost"); + expect(context.workers).toHaveLength(1); + const events = await collectEvents( + context.capability.attachOptimizationRun(completed), + ); + expect(events.map((event) => event.type)).toEqual([ + "started", + "trial", + "complete", + ]); + }); + + it("a worker crash aborts the trial evaluation in flight", async () => { + const context = setUp({ evaluateTrial: () => new Promise(() => {}) }); + 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(); + const request = context.evaluateTrial.mock + .calls[0]?.[0] as PetrinautOptimizationTrialRequest; + expect(request.signal.aborted).toBe(false); + + context.worker.emitError({ message: "RangeError: out of memory" }); + await flush(); + + expect(request.signal.aborted).toBe(true); + }); + + it("a study error aborts the trial evaluation in flight", async () => { + const context = setUp({ evaluateTrial: () => new Promise(() => {}) }); + 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(); + const request = context.evaluateTrial.mock + .calls[0]?.[0] as PetrinautOptimizationTrialRequest; + + context.worker.emit({ type: "error", runId, message: "ValueError: nope" }); + await flush(); + + expect(request.signal.aborted).toBe(true); + const events = await collectEvents( + context.capability.attachOptimizationRun(runId), + ); + expect(events.at(-1)).toMatchObject({ + type: "error", + code: "study_failed", + }); + }); + + 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, + resumable: false, + }); + } + await expect( + context.capability.createOptimizationRun( + createOptimizationManifestInput(), + ), + ).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 }, + ]); + 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, + resumable: true, + 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 }, + ]); + 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({ + code: PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE, + }), + ); + + 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/); + 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 }, + ]); + 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, + resumable: false, + 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 }, + ]); + 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 once 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")).toEqual([ + { type: "extend", runId, trials: 2 }, + ]); + + 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(); + + 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/optimization/browser/browser-optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/browser-optimization.ts new file mode 100644 index 00000000000..786a8915280 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/browser-optimization.ts @@ -0,0 +1,548 @@ +import { v4 as generateUuid } from "uuid"; + +import { createAbortController } from "../../environment"; +import { + deriveOptimizationTrialSeeds, + describeOptimization, + PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE, + PETRINAUT_OPTIMIZATION_MAX_PARALLELISM, + PETRINAUT_OPTIMIZATION_MAX_TRIALS, + PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE, + parseOptimizationManifest, + resolveTrialScenarioParameterValues, +} from "../index"; +import { + defaultOptimizerPyodideConfig, + type OptimizerPyodideConfig, +} from "./pyodide-config"; +import { optimizerPythonSources } from "./python-sources"; +import { + createOptimizationRunLog, + type OptimizationRunLog, + type OptimizationRunLogEvent, +} from "./run-log"; +import { + createOptimizerWorker, + type OptimizerWorkerErrorEvent, + type OptimizerWorkerLike, +} from "./worker/create-optimizer-worker"; + +import type { AbortControllerLike, AbortSignalLike } from "../../environment"; +import type { + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautOptimizationChannel, + PetrinautOptimizationEvent, + PetrinautOptimizationManifest, +} from "../index"; +import type { + OptimizerEvaluateMessage, + OptimizerExtendMessage, + OptimizerStartMessage, + OptimizerToMainMessage, + OptimizerToWorkerMessage, +} from "./messages"; + +export type CreateBrowserOptimizationOptions = { + pyodide?: Partial; + createWorker?: () => OptimizerWorkerLike; +}; + +/** + * `queued` waits 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" | "running" | "finished-resumable" | "finished"; + +type RunRecord = { + readonly runId: string; + readonly manifest: PetrinautOptimizationManifest; + readonly seeds: readonly number[]; + readonly log: OptimizationRunLog; + 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 = { + readonly worker: OptimizerWorkerLike; + /** Settles on `ready`; a session that fails before or after is torn down instead. */ + readonly ready: Promise; + readonly markReady: () => void; +}; + +/** The events that end a segment: what `finish` appends, with `resumable`. */ +type TerminalRunLogEvent = Extract< + OptimizationRunLogEvent, + { type: "complete" | "error" } +>; + +const cancelledEvent: TerminalRunLogEvent = { + 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" + : event.message, + ); + +const unavailableEvent = (error: unknown): TerminalRunLogEvent => ({ + type: "error", + code: "optimizer_unavailable", + message: `The in-browser optimizer could not start: ${errorMessage(error)}`, + retryable: true, +}); + +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 unknownRunError = (runId: string): Error => + Object.assign(new Error(`Unknown optimization run "${runId}"`), { + code: PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE, + }); + +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, or its optimizer worker was lost` + : `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: { + 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; +}): PetrinautConnectedOptimizationCapability => { + const { channel } = options; + const runs = new Map(); + const queue: RunRecord[] = []; + /** The run waiting for the worker or running on it. */ + 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 post = (message: OptimizerToWorkerMessage): void => { + session?.worker.postMessage(message); + }; + + /** Ends the run's segment with `event` and lets the queue move on. */ + const finish = ( + run: RunRecord, + event: TerminalRunLogEvent, + 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 = status; + // The host learns from the event whether Continue has a study to return + // to: a first segment stopped before it reached the worker has none. + run.log.append({ ...event, resumable: status === "finished-resumable" }); + // A failure can leave the host evaluating trials whose outcomes no one + // will read; a completed or stopped segment has nothing left in flight. + run.controller.abort(); + 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(); + } + }; + + /** + * Drops a worker that failed to load or crashed. The run on it, if any, + * fails as retryable, the studies it kept for finished runs are gone with + * it, and the queue moves on to a fresh worker. + */ + const failSession = (stale: WorkerSession, error: unknown): void => { + if (session !== stale) { + return; + } + session = null; + stale.worker.terminate(); + for (const run of runs.values()) { + if (run.status === "finished-resumable") { + run.status = "finished"; + } + } + if (active) { + finish(active, unavailableEvent(error), "finished"); + } + }; + + /** 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 => { + discardStudy(run); + finish( + run, + { + type: "error", + code: "trial_evaluation_failed", + message: errorMessage(error), + retryable: false, + }, + "finished", + ); + }; + + const handleEvaluate = async ( + message: OptimizerEvaluateMessage, + ): Promise => { + const run = activeRunFor(message.runId); + if (!run) { + return; + } + const { controller } = run; + // The worker answers the evaluations of a cancelled segment itself. + const segmentIsCurrent = (): boolean => + run.controller === controller && + run.status === "running" && + !controller.signal.aborted; + if (!segmentIsCurrent()) { + return; + } + try { + const outcome = await channel.evaluateTrial({ + runId: run.runId, + trial: message.trial, + manifest: run.manifest, + suggestedValues: message.suggestedValues, + scenarioParameterValues: resolveTrialScenarioParameterValues( + run.manifest, + message.suggestedValues, + ), + seeds: run.seeds, + signal: controller.signal, + }); + if (segmentIsCurrent()) { + post({ type: "evaluated", requestId: message.requestId, outcome }); + } + } catch (error) { + if (segmentIsCurrent()) { + failTrialEvaluation(run, error); + } + } + }; + + const handleWorkerMessage = ( + current: WorkerSession, + message: OptimizerToMainMessage, + ): void => { + switch (message.type) { + case "ready": + current.markReady(); + return; + case "init-error": + failSession(current, new Error(message.message)); + return; + case "evaluate": + void 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, + }, + "finished-resumable", + ); + } + return; + } + case "cancelled": { + const run = activeRunFor(message.runId); + if (run) { + 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, + }, + "finished", + ); + } + } + } + }; + + const ensureSession = (): WorkerSession => { + if (session) { + return session; + } + const worker = options.createWorker(); + const { promise: ready, resolve: markReady } = + Promise.withResolvers(); + const current: WorkerSession = { worker, ready, markReady }; + worker.addEventListener("message", ({ data }) => { + handleWorkerMessage(current, data); + }); + worker.addEventListener("error", (event) => { + failSession(current, workerLoadError(event)); + }); + worker.postMessage({ + type: "init", + pyodide: options.pyodide, + pythonSources: optimizerPythonSources, + }); + session = current; + return current; + }; + + const startNext = (): void => { + if (disposed || active) { + return; + } + const run = queue.shift(); + if (!run) { + return; + } + active = run; + let current: WorkerSession; + try { + current = ensureSession(); + } catch (error) { + finish(run, unavailableEvent(error), "finished"); + return; + } + void current.ready.then(() => { + if (active === run && run.status === "queued" && session === current) { + run.status = "running"; + current.worker.postMessage(run.command); + } + }); + }; + + const enqueue = (run: RunRecord, requestedTrials: number): void => { + run.log.append({ type: "started", requestedTrials }); + queue.push(run); + startNext(); + }; + + return { + async createOptimizationRun(input, runOptions = {}) { + if (disposed) { + throw disposedError(); + } + if (runOptions.signal?.aborted) { + throw creationAbortedError(); + } + const parallelism = validParallelism(runOptions.parallelism); + const manifest = parseOptimizationManifest(input); + const runId = generateUuid(); + const run: RunRecord = { + runId, + manifest, + seeds: deriveOptimizationTrialSeeds( + manifest.execution.seed, + manifest.execution.seedsPerTrial ?? 1, + ), + log: createOptimizationRunLog(), + status: "queued", + command: { + type: "start", + runId, + description: describeOptimization(manifest), + parallelism, + }, + controller: createAbortController(), + }; + runs.set(runId, run); + enqueue(run, manifest.study.trials); + return { runId }; + }, + async extendOptimizationRun(runId, trials) { + if (disposed) { + throw disposedError(); + } + const run = runs.get(runId); + if (!run) { + throw unknownRunError(runId); + } + if (run.status !== "finished-resumable") { + throw notResumableError(run); + } + const total = requestedTotal(toldTrials(run.log), trials); + run.command = { type: "extend", runId, trials }; + run.controller = createAbortController(); + run.status = "queued"; + enqueue(run, total); + }, + 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 || isSettled(run)) { + return; + } + run.controller.abort(); + if (run.status === "running") { + post({ type: "cancel", runId }); + return; + } + // 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 (!isSettled(run)) { + run.controller.abort(); + run.log.append({ ...cancelledEvent, resumable: false }); + } + run.status = "finished"; + } + 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/optimization/browser/index.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/index.ts new file mode 100644 index 00000000000..d5ff8af5e80 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/index.ts @@ -0,0 +1,22 @@ +export { + createBrowserOptimization, + type CreateBrowserOptimizationOptions, +} from "./browser-optimization"; +export type { + OptimizerWorkerErrorEvent, + OptimizerWorkerLike, +} from "./worker/create-optimizer-worker"; +export { + defaultOptimizerPyodideConfig, + type OptimizerPyodideConfig, +} from "./pyodide-config"; +export type { + OptimizationScalar, + PetrinautConnectedOptimization, + PetrinautConnectedOptimizationCapability, + PetrinautConnectedRunOptions, + PetrinautOptimizationChannel, + PetrinautOptimizationSource, + PetrinautOptimizationTrialOutcome, + PetrinautOptimizationTrialRequest, +} from "../index"; diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/messages.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/messages.ts new file mode 100644 index 00000000000..d5b9add4fe7 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/messages.ts @@ -0,0 +1,129 @@ +import type { + OptimizationScalar, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationTrialOutcome, +} from "../index"; +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; +}; + +/** 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 segment stopped early because the run was cancelled. */ + cancelled?: boolean; +}; + +export type OptimizerInitMessage = { + type: "init"; + pyodide: OptimizerPyodideConfig; + 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, at its parallelism; trial numbers continue. */ +export type OptimizerExtendMessage = { + type: "extend"; + runId: string; + trials: number; +}; + +export type OptimizerEvaluatedMessage = { + type: "evaluated"; + requestId: number; + outcome: PetrinautOptimizationTrialOutcome; +}; + +/** Stops the running segment; the study stays in memory. */ +export type OptimizerCancelMessage = { + type: "cancel"; + runId: string; +}; + +/** Drops the kept study; nothing is posted back. */ +export type OptimizerReleaseMessage = { + type: "release"; + runId: string; +}; + +export type OptimizerToWorkerMessage = + | OptimizerInitMessage + | OptimizerStartMessage + | OptimizerExtendMessage + | OptimizerEvaluatedMessage + | OptimizerCancelMessage + | OptimizerReleaseMessage; + +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/optimization/browser/pyodide-config.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/pyodide-config.test.ts new file mode 100644 index 00000000000..a5e96b5de54 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/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/optimization/browser/pyodide-config.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/pyodide-config.ts new file mode 100644 index 00000000000..40f5acc1bb7 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/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/optimization/browser/python-sources.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/python-sources.ts new file mode 100644 index 00000000000..45e1cc98a0c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/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/optimization/browser/run-log.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/run-log.test.ts new file mode 100644 index 00000000000..fc296b23a95 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/run-log.test.ts @@ -0,0 +1,135 @@ +import { describe, expect, it } from "vitest"; + +import { createAbortController } from "../../environment"; +import { createOptimizationRunLog } from "./run-log"; + +import type { PetrinautOptimizationEvent } from "../index"; + +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 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.append(complete).seq).toBe(3); + expect(log.events.map((event) => event.seq)).toEqual([1, 2, 3]); + expect(() => log.append(trial(1))).toThrow( + "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.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 }); + 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 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); + + expect(await collect(log.replay({ cursor: 2 }))).toEqual([]); + expect(await collect(log.replay())).toHaveLength(2); + }); + + 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/optimization/browser/run-log.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/run-log.ts new file mode 100644 index 00000000000..a9d854e69b2 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/run-log.ts @@ -0,0 +1,116 @@ +import type { AbortSignalLike } from "../../environment"; +import type { PetrinautOptimizationEvent } from "../index"; + +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[]; + /** + * Stamps the next dense `seq` (from 1) and stores the event. After a + * terminal event, only a `started` event may follow. + */ + append(event: OptimizationRunLogEvent): PetrinautOptimizationEvent; + /** + * Yields the stored events with `seq` greater than `cursor`, then tails live + * events, and ends at the first terminal event after the cursor. 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>(); + + const isSettled = (): boolean => { + const latest = events.at(-1); + return latest !== undefined && isTerminalEvent(latest); + }; + + 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; + }, + append(event) { + 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); + for (const listener of listeners) { + listener(stamped); + } + return stamped; + }, + 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 (isSettled()) { + return; + } + await waitForNextEvent(signal); + } + }, + }; +}; diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.test.ts new file mode 100644 index 00000000000..2dd54f3f65d --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.test.ts @@ -0,0 +1,367 @@ +import { describe, expect, it, vi } from "vitest"; + +import { attachOptimizerWorker } from "./attach-optimizer-worker"; + +import type { WorkerThreadRuntime } from "../../../environment"; +import type { PetrinautOptimizationDescribeResult } from "../../index"; +import type { + OptimizerStudySummary, + OptimizerToMainMessage, + OptimizerToWorkerMessage, + OptimizerTrialPayload, +} from "../messages"; +import type { + OptimizerStudyCallbacks, + OptimizerStudyRunner, +} from "./study-runner"; + +type FakeRuntime = WorkerThreadRuntime< + OptimizerToWorkerMessage, + OptimizerToMainMessage +> & { + readonly posted: OptimizerToMainMessage[]; + /** Deliver a message as the main thread would post it. */ + receive(message: OptimizerToWorkerMessage): void; + postedOfType( + type: TType, + ): Extract[]; +}; + +const createFakeRuntime = (): FakeRuntime => { + const posted: OptimizerToMainMessage[] = []; + let listener: ((message: OptimizerToWorkerMessage) => void) | null = null; + return { + posted, + postMessage(message) { + posted.push(message); + }, + onMessage(next) { + listener = next; + }, + delay: () => Promise.resolve(), + receive(message) { + if (!listener) { + throw new Error("the protocol registered no message listener"); + } + listener(message); + }, + postedOfType(type) { + return posted.filter( + (message): message is Extract => + message.type === type, + ); + }, + }; +}; + +/** One segment the fake runner was asked to run, settled by the test. */ +type Segment = { + readonly runId: string; + readonly trials: number; + readonly callbacks: OptimizerStudyCallbacks; + readonly settle: (summary: OptimizerStudySummary) => void; + readonly fail: (error: unknown) => void; +}; + +const createFakeRunner = (ready: Promise) => { + const segments: Segment[] = []; + const released: string[] = []; + const begin = ( + runId: string, + trials: number, + callbacks: OptimizerStudyCallbacks, + ): Promise => { + const { promise, resolve, reject } = + Promise.withResolvers(); + segments.push({ runId, trials, callbacks, settle: resolve, fail: reject }); + return promise; + }; + const start = vi.fn((input) => + begin(input.runId, input.description.study.trials, input.callbacks), + ); + const extend = vi.fn((input) => + begin(input.runId, input.trials, input.callbacks), + ); + const runner: OptimizerStudyRunner = { + ready, + start, + extend, + release: async (runId) => { + released.push(runId); + }, + }; + return { runner, start, extend, segments, released }; +}; + +const description: PetrinautOptimizationDescribeResult = { + direction: "minimize", + study: { trials: 3, sampler: "tpe", seed: 1, seedsPerTrial: 1 }, + parameters: [ + { + identifier: "rate", + type: "float", + default: 1, + minimum: 0, + maximum: 2, + scale: "linear", + }, + ], +}; + +const summary: OptimizerStudySummary = { + requestedTrials: 3, + completedTrials: 3, + prunedTrials: 0, + failedTrials: 0, + best: { trial: 1, parameters: { rate: 0.5 }, objective: 1 }, +}; + +const trialPayload: OptimizerTrialPayload = { + trial: 0, + parameters: { rate: 0.5 }, + objective: 1, + state: "complete", + best: summary.best, +}; + +const flush = async (): Promise => { + for (let index = 0; index < 5; index++) { + await Promise.resolve(); + } +}; + +const setUp = (options?: { ready?: Promise }) => { + const runtime = createFakeRuntime(); + const fake = createFakeRunner(options?.ready ?? Promise.resolve()); + const createRunner = vi.fn(() => fake.runner); + attachOptimizerWorker(runtime, createRunner); + return { runtime, createRunner, ...fake }; +}; + +const init = (context: ReturnType): void => { + context.runtime.receive({ + type: "init", + pyodide: { + indexURL: "https://example.test/pyodide/", + packages: { optuna: "4.9.0" }, + distributionPackages: ["numpy"], + }, + pythonSources: { "petrinaut_optimizer_core/__init__.py": "" }, + }); +}; + +const segmentOf = ( + context: ReturnType, + index: number, +): Segment => { + const segment = context.segments[index]; + if (!segment) { + throw new Error(`no segment ${index} was started`); + } + return segment; +}; + +describe("attachOptimizerWorker", () => { + it("creates the runner from the init message and reports readiness", async () => { + const context = setUp(); + + init(context); + await flush(); + + expect(context.createRunner).toHaveBeenCalledWith({ + pyodide: { + indexURL: "https://example.test/pyodide/", + packages: { optuna: "4.9.0" }, + distributionPackages: ["numpy"], + }, + pythonSources: { "petrinaut_optimizer_core/__init__.py": "" }, + }); + expect(context.runtime.posted).toEqual([{ type: "ready" }]); + }); + + it("reports a runtime that fails to load as init-error", async () => { + const context = setUp({ ready: Promise.reject(new Error("offline")) }); + + init(context); + await flush(); + + expect(context.runtime.posted).toEqual([ + { type: "init-error", message: "offline" }, + ]); + }); + + it("answers a study posted before init with an error", () => { + const context = setUp(); + + context.runtime.receive({ + type: "start", + runId: "early", + description, + parallelism: 1, + }); + + expect(context.runtime.posted).toEqual([ + { + type: "error", + runId: "early", + message: "The optimizer worker received a study before its runtime", + }, + ]); + expect(context.start).not.toHaveBeenCalled(); + }); + + it("starts the study and completes an evaluate round trip through the main thread", async () => { + const context = setUp(); + init(context); + context.runtime.receive({ + type: "start", + runId: "study", + description, + parallelism: 2, + }); + + expect(context.start).toHaveBeenCalledWith( + expect.objectContaining({ runId: "study", description, parallelism: 2 }), + ); + const segment = segmentOf(context, 0); + expect(segment.callbacks.isCancelled()).toBe(false); + + const outcome = segment.callbacks.evaluate(0, { rate: 0.5 }); + expect(context.runtime.postedOfType("evaluate")).toEqual([ + { + type: "evaluate", + runId: "study", + requestId: 1, + trial: 0, + suggestedValues: { rate: 0.5 }, + }, + ]); + context.runtime.receive({ + type: "evaluated", + requestId: 1, + outcome: { kind: "objective", objective: 1 }, + }); + expect(await outcome).toEqual({ kind: "objective", objective: 1 }); + + segment.callbacks.onTrial(trialPayload); + segment.settle(summary); + await flush(); + expect(context.runtime.posted.slice(-2)).toEqual([ + { type: "trial", runId: "study", event: trialPayload }, + { type: "complete", runId: "study", summary }, + ]); + }); + + it("ignores an evaluated message for a request it no longer holds", () => { + const context = setUp(); + init(context); + + expect(() => + context.runtime.receive({ + type: "evaluated", + requestId: 7, + outcome: { kind: "objective", objective: 1 }, + }), + ).not.toThrow(); + }); + + it("cancel prunes the segment's pending evaluations, flags the loop and posts cancelled when it stops", async () => { + const context = setUp(); + init(context); + context.runtime.receive({ + type: "start", + runId: "stopped", + description, + parallelism: 2, + }); + context.runtime.receive({ + type: "start", + runId: "other", + description, + parallelism: 1, + }); + const stopped = segmentOf(context, 0); + const other = segmentOf(context, 1); + const pending = [ + stopped.callbacks.evaluate(0, { rate: 0.1 }), + stopped.callbacks.evaluate(1, { rate: 0.2 }), + ]; + const otherPending = other.callbacks.evaluate(0, { rate: 0.3 }); + + context.runtime.receive({ type: "cancel", runId: "stopped" }); + + expect(stopped.callbacks.isCancelled()).toBe(true); + expect(other.callbacks.isCancelled()).toBe(false); + expect(await Promise.all(pending)).toEqual([ + { kind: "pruned", reason: "cancelled" }, + { kind: "pruned", reason: "cancelled" }, + ]); + context.runtime.receive({ + type: "evaluated", + requestId: 3, + outcome: { kind: "objective", objective: 3 }, + }); + expect(await otherPending).toEqual({ kind: "objective", objective: 3 }); + + stopped.settle({ ...summary, cancelled: true }); + await flush(); + expect(context.runtime.posted.at(-1)).toEqual({ + type: "cancelled", + runId: "stopped", + }); + + // The next segment of the same study starts uncancelled. + context.runtime.receive({ type: "extend", runId: "stopped", trials: 2 }); + expect(context.extend).toHaveBeenCalledWith( + expect.objectContaining({ runId: "stopped", trials: 2 }), + ); + expect(segmentOf(context, 2).callbacks.isCancelled()).toBe(false); + }); + + it("release prunes pending evaluations and drops the study without posting back", async () => { + const context = setUp(); + init(context); + await flush(); + context.runtime.receive({ + type: "start", + runId: "released", + description, + parallelism: 1, + }); + const segment = segmentOf(context, 0); + const pending = segment.callbacks.evaluate(0, { rate: 0.1 }); + const postedBefore = context.runtime.posted.length; + + context.runtime.receive({ type: "release", runId: "released" }); + + // The loop sees the cancellation until the release, queued behind the + // segment in the real runner, settles. + expect(segment.callbacks.isCancelled()).toBe(true); + expect(await pending).toEqual({ kind: "pruned", reason: "cancelled" }); + await flush(); + expect(context.released).toEqual(["released"]); + expect(context.runtime.posted).toHaveLength(postedBefore); + }); + + it("reports a segment the runner rejects as an error for its run", async () => { + const context = setUp(); + init(context); + context.runtime.receive({ + type: "start", + runId: "failing", + description, + parallelism: 1, + }); + + segmentOf(context, 0).fail( + new Error("ValueError: trial objective must be a finite number"), + ); + await flush(); + + expect(context.runtime.posted.at(-1)).toEqual({ + type: "error", + runId: "failing", + message: "ValueError: trial objective must be a finite number", + }); + }); +}); diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.ts new file mode 100644 index 00000000000..20412f93a1c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/attach-optimizer-worker.ts @@ -0,0 +1,191 @@ +/** + * The optimizer worker protocol, detached from any thread host: a + * {@link WorkerThreadRuntime} posts and receives the messages, and + * `createRunner` supplies the study runner, so a test drives the protocol with + * a fake of each. + */ +import type { WorkerThreadRuntime } from "../../../environment"; +import type { PetrinautOptimizationTrialOutcome } from "../../index"; +import type { + OptimizerInitMessage, + OptimizerStudySummary, + OptimizerToMainMessage, + OptimizerToWorkerMessage, +} from "../messages"; +import type { + OptimizerStudyCallbacks, + OptimizerStudyRunner, +} from "./study-runner"; + +export type OptimizerRunnerFactory = ( + init: Pick, +) => OptimizerStudyRunner; + +/** An evaluate request posted to the main thread and not yet answered. */ +type PendingEvaluation = { + readonly runId: string; + readonly resolve: (outcome: PetrinautOptimizationTrialOutcome) => void; +}; + +const errorMessage = (error: unknown): string => + error instanceof Error ? error.message : String(error); + +/** + * Runs the optimizer worker protocol against `runtime`. + * + * Handles `init`, `start`, `extend`, `evaluated`, `cancel` and `release`; + * posts `ready` or `init-error` once, then one `evaluate` per trial and one + * `complete`, `cancelled` or `error` per segment. + */ +export const attachOptimizerWorker = ( + runtime: WorkerThreadRuntime< + OptimizerToWorkerMessage, + OptimizerToMainMessage + >, + createRunner: OptimizerRunnerFactory, +): void => { + let runner: OptimizerStudyRunner | null = null; + const pending = new Map(); + /** Runs whose current segment was cancelled; cleared when the segment ends. */ + const cancelled = new Set(); + let nextRequestId = 1; + + const postError = (runId: string, error: unknown): void => { + runtime.postMessage({ type: "error", runId, message: errorMessage(error) }); + }; + + const initialize = (message: OptimizerInitMessage): void => { + const current = createRunner({ + pyodide: message.pyodide, + pythonSources: message.pythonSources, + }); + runner = current; + current.ready.then( + () => runtime.postMessage({ type: "ready" }), + (error: unknown) => + runtime.postMessage({ + type: "init-error", + message: errorMessage(error), + }), + ); + }; + + const runnerFor = (runId: string): OptimizerStudyRunner | null => { + if (!runner) { + postError( + runId, + new Error("The optimizer worker received a study before its runtime"), + ); + } + return runner; + }; + + const beginSegment = (runId: string): OptimizerStudyCallbacks => { + cancelled.delete(runId); + return { + evaluate: (trial, suggestedValues) => + new Promise((resolve) => { + const requestId = nextRequestId; + nextRequestId += 1; + pending.set(requestId, { runId, resolve }); + runtime.postMessage({ + type: "evaluate", + runId, + requestId, + trial, + suggestedValues, + }); + }), + onTrial: (event) => runtime.postMessage({ type: "trial", runId, event }), + isCancelled: () => cancelled.has(runId), + }; + }; + + const reportSegment = ( + runId: string, + summary: Promise, + ): void => { + summary + .then( + (result) => + runtime.postMessage( + result.cancelled === true + ? { type: "cancelled", runId } + : { type: "complete", runId, summary: result }, + ), + (error: unknown) => postError(runId, error), + ) + .finally(() => cancelled.delete(runId)); + }; + + /** Ends the run's segment early: its loop stops at the next poll, and its trials in flight are pruned. */ + const cancelSegment = (runId: string): void => { + cancelled.add(runId); + for (const [requestId, evaluation] of pending) { + if (evaluation.runId === runId) { + pending.delete(requestId); + evaluation.resolve({ kind: "pruned", reason: "cancelled" }); + } + } + }; + + runtime.onMessage((message) => { + switch (message.type) { + case "init": + initialize(message); + return; + case "start": { + const current = runnerFor(message.runId); + if (current) { + reportSegment( + message.runId, + current.start({ + runId: message.runId, + description: message.description, + parallelism: message.parallelism, + callbacks: beginSegment(message.runId), + }), + ); + } + return; + } + case "extend": { + const current = runnerFor(message.runId); + if (current) { + reportSegment( + message.runId, + current.extend({ + runId: message.runId, + trials: message.trials, + callbacks: beginSegment(message.runId), + }), + ); + } + return; + } + case "evaluated": { + const evaluation = pending.get(message.requestId); + if (evaluation) { + pending.delete(message.requestId); + evaluation.resolve(message.outcome); + } + return; + } + case "cancel": + cancelSegment(message.runId); + return; + case "release": { + if (!runner) { + return; + } + cancelSegment(message.runId); + // The release runs after the segments queued before it, so once it + // settles no segment of the run remains to observe the cancellation. + runner.release(message.runId).then( + () => cancelled.delete(message.runId), + (error: unknown) => postError(message.runId, error), + ); + } + } + }); +}; diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/create-optimizer-worker.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/create-optimizer-worker.ts new file mode 100644 index 00000000000..2e985dc441c --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/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/optimization/browser/worker/optimizer.worker.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/optimizer.worker.ts new file mode 100644 index 00000000000..2a8524e6dd0 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/optimizer.worker.ts @@ -0,0 +1,40 @@ +/** + * @layerRoot core.optimization.browser.worker + * @role Hosts Optuna in Pyodide off the main thread and pairs each evaluate request with the study loop + */ +import { createWorkerThreadRuntime } from "../../../environment"; +import { attachOptimizerWorker } from "./attach-optimizer-worker"; +import { createOptimizerStudyRunner } from "./study-runner"; + +import type { LoadPyodide } from "./pyodide-like"; + +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; +}; + +attachOptimizerWorker( + createWorkerThreadRuntime(), + ({ pyodide, pythonSources }) => { + const indexURL = ensureTrailingSlash(pyodide.indexURL); + return createOptimizerStudyRunner({ + loadPyodide: async (options) => + (await importLoadPyodide(indexURL))(options), + config: { ...pyodide, indexURL }, + pythonSources, + }); + }, +); diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/pyodide-like.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/pyodide-like.ts new file mode 100644 index 00000000000..3e0d6456c88 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/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/optimization/browser/worker/study-runner.pyodide.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/study-runner.pyodide.test.ts new file mode 100644 index 00000000000..dfed595b1bf --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/study-runner.pyodide.test.ts @@ -0,0 +1,413 @@ +/** + * 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 { 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 "../../index"; +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 ?? "") !== ""; + +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 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)}`); + } + 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); + +/** 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 callbacks: OptimizerStudyCallbacks = { + 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, callbacks }; +}; + +let runner: OptimizerStudyRunner; +let loadFailure: string | null = null; +const startedStudies: string[] = []; + +/** Starts a study and remembers it, so the suite releases every study it kept. */ +const startStudy = ( + input: Parameters[0], +): ReturnType => { + startedStudies.push(input.runId); + return runner.start(input); +}; + +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); + +afterAll(async () => { + if (loadFailure === null) { + await Promise.all(startedStudies.map((runId) => runner.release(runId))); + } +}); + +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 run = recorder({ + evaluate: async (trial, suggestedValues) => + trial === 3 + ? { kind: "pruned", reason: "no frames" } + : { kind: "objective", objective: objectiveOf(suggestedValues) }, + }); + + const summary = await startStudy({ + runId: "seeded", + description, + parallelism: 1, + callbacks: run.callbacks, + }); + + expect(run.trialNumbers).toEqual( + Array.from({ length: 30 }, (_, index) => index), + ); + for (const values of run.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(run.trials.map((event) => event.trial)).toEqual( + Array.from({ length: 30 }, (_, index) => index), + ); + let bestSoFar = Number.POSITIVE_INFINITY; + 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(); + } 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: run.trials.at(-1)?.best, + }); + }, + loadTimeout, + ); + + test( + "continues a study on the same sampler history, numbering onwards", + async ({ skip }) => { + skipWhenOffline(skip); + const split = async (runId: string) => { + const first = recorder(); + await startStudy({ + runId, + description: withTrials(8), + parallelism: 1, + callbacks: first.callbacks, + }); + const second = recorder(); + const summary = await runner.extend({ + runId, + trials: 8, + 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.second.trialNumbers).toEqual([8, 9, 10, 11, 12, 13, 14, 15]); + expect(study.second.trials.map((event) => event.trial)).toEqual( + study.second.trialNumbers, + ); + expect(study.summary).toMatchObject({ + requestedTrials: 16, + completedTrials: 16, + prunedTrials: 0, + }); + expect(study.summary.best).toEqual(study.second.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 startStudy({ + 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 failed and marked unreported, + // so Optuna keeps no running trial behind and the counters skip it. + expect(stopped).toMatchObject({ + completedTrials: 4, + failedTrials: 0, + cancelled: true, + }); + + const resumed = recorder(); + const summary = await runner.extend({ + runId: "stopped", + trials: 2, + callbacks: resumed.callbacks, + }); + + 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: 0, + best: resumed.trials.at(-1)?.best, + }); + }, + loadTimeout, + ); + + test( + "keeps up to the parallelism in flight and tells outcomes as they settle", + async ({ skip }) => { + skipWhenOffline(skip); + 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 startStudy({ + 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(summary).toMatchObject({ + requestedTrials: 6, + completedTrials: 6, + prunedTrials: 0, + }); + }, + loadTimeout, + ); + + test( + "drops a released study, and one whose segment failed", + async ({ skip }) => { + skipWhenOffline(skip); + await startStudy({ + 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, + callbacks: recorder().callbacks, + }), + ).rejects.toThrow('Optimization study "released" is not kept'); + + await expect( + startStudy({ + runId: "failing", + description, + 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, + callbacks: recorder().callbacks, + }), + ).rejects.toThrow('Optimization study "failing" is not kept'); + }, + loadTimeout, + ); +}); diff --git a/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/study-runner.ts b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/study-runner.ts new file mode 100644 index 00000000000..8287fb9c298 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/browser/worker/study-runner.ts @@ -0,0 +1,320 @@ +/** + * @talksTo optimizer-core via Python sources written into Pyodide's filesystem + */ +import { micropipRequirements } from "../pyodide-config"; +import { isPyProxyLike } from "./pyodide-like"; + +import type { + OptimizationScalar, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationTrialOutcome, +} from "../../index"; +import type { + OptimizerBestTrial, + OptimizerStudySummary, + OptimizerTrialPayload, +} from "../messages"; +import type { OptimizerPyodideConfig } from "../pyodide-config"; +import type { LoadPyodide, PyodideLike, PyProxyLike } from "./pyodide-like"; + +export type OptimizerStudyCallbacks = { + evaluate( + trial: number, + suggestedValues: Record, + ): Promise; + onTrial(event: OptimizerTrialPayload): void; + 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; + callbacks: OptimizerStudyCallbacks; +}; + +export type OptimizerStudyRunner = { + /** Settles once Pyodide, the packages and the optimizer sources are loaded. */ + readonly ready: 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, at the parallelism it was created + * with; the trial numbers continue. + */ + extend(input: OptimizerStudyExtendInput): Promise; + /** Drops the kept study once the segments queued before it have run. */ + release(runId: string): Promise; +}; + +/** The outcome shape `ask_tell.run_study` expects from its evaluate callback. */ +type PythonTrialOutcome = { objective: number } | { pruned: string }; + +/** The Python `StudyHandle`, opaque on this side. */ +type StudyHandleProxy = PyProxyLike; + +type PyodideEntryModule = { + 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, + ): 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"; + +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); + const studies = new Map(); + let queue: Promise = ready; + + 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, + callbacks: OptimizerStudyCallbacks, + ): Promise => { + 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 = study.nextTrial; + study.nextTrial += 1; + return callbacks.evaluate(trial, suggestedValues).then(toPythonOutcome); + }; + const onTrial = (payload: unknown): void => { + callbacks.onTrial(normalizeTrialPayload(toJsValue(payload))); + }; + try { + const result = await module.run_browser_study( + study.handle, + trials, + evaluate, + onTrial, + () => callbacks.isCancelled(), + ); + 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 { + ready, + 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.callbacks, + ); + }); + }, + extend(input) { + return enqueue(async (module) => + runSegment(module, input.runId, input.trials, input.callbacks), + ); + }, + release(runId) { + return enqueue(async (module) => { + dropStudy(module, runId); + }); + }, + }; +}; 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..24ebbc356f6 --- /dev/null +++ b/libs/@hashintel/petrinaut-core/src/optimization/describe.ts @@ -0,0 +1,188 @@ +import { deriveRunSeed } from "../simulation/monte-carlo/run-state"; + +import type { Scenario } from "../types/sdcpn"; +import type { + OptimizationScalar, + PetrinautOptimizationDescribeParameter, + PetrinautOptimizationDescribeResult, + PetrinautOptimizationDomain, + PetrinautOptimizationManifest, +} from "./index"; + +type ScenarioParameter = Scenario["scenarioParameters"][number]; + +type OptimizedParameter = { + parameter: ScenarioParameter; + domain: PetrinautOptimizationDomain; +}; + +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, + }; + } +}; + +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/optimization.ts b/libs/@hashintel/petrinaut-core/src/optimization/index.ts similarity index 76% rename from libs/@hashintel/petrinaut-core/src/optimization.ts rename to libs/@hashintel/petrinaut-core/src/optimization/index.ts index ed6c8d6356b..2efa57b1ead 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization/index.ts @@ -1,9 +1,13 @@ +/** + * @layerRoot core.optimization + * @role The optimization contract shared by the CLI, the service client and the browser runtime: manifest and event schemas, capability, channel and connected-source types, study description and seed derivation + */ import { z } from "zod"; -import { parseSDCPNFile } from "./file-format/parse-sdcpn-file"; -import { sdcpnSchema } from "./file-format/types"; +import { parseSDCPNFile } from "../file-format/parse-sdcpn-file"; +import { sdcpnSchema } from "../file-format/types"; -import type { AbortSignalLike } from "./environment"; +import type { AbortSignalLike } from "../environment"; export const PETRINAUT_OPTIMIZATION_MAX_SEED = 2_147_483_647; export const PETRINAUT_OPTIMIZATION_MAX_TRIALS = 1_000; @@ -13,6 +17,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"), @@ -451,6 +458,23 @@ export const petrinautOptimizationManifestSchema = z "A versioned, self-contained study over a flat set of scenario parameters.", }); +/** Parses a manifest, naming every invalid field in the error it throws. */ +export const parseOptimizationManifest = ( + data: unknown, +): PetrinautOptimizationManifest => { + const parsed = petrinautOptimizationManifestSchema.safeParse(data); + if (!parsed.success) { + const details = parsed.error.issues + .map( + ({ path, message }) => + `${path.length > 0 ? path.join(".") : "manifest"}: ${message}`, + ) + .join("; "); + throw new Error(`Invalid optimization manifest: ${details}`); + } + return parsed.data; +}; + /** The application optimization request is the immutable CLI manifest. */ export const petrinautOptimizationInputSchema = petrinautOptimizationManifestSchema; @@ -591,6 +615,14 @@ export const petrinautOptimizationTrialEventSchema = z }) .meta({ description: "One completed Optuna trial and the running best." }); +/** + * Whether the study behind the run stays available to `extendOptimizationRun` + * after this terminal event. A connected capability sets it on every terminal + * event: `false` for a segment that never reached the worker or a run that + * failed. A remote service keeps no study and omits it. + */ +const optimizationResumableSchema = z.boolean().optional(); + export const petrinautOptimizationCompleteEventSchema = z .strictObject({ type: z.literal("complete"), @@ -599,6 +631,7 @@ export const petrinautOptimizationCompleteEventSchema = z prunedTrials: z.number().int().nonnegative(), failedTrials: z.number().int().nonnegative(), best: optimizationBestSchema.nullable(), + resumable: optimizationResumableSchema, seq: optimizationEventSeqSchema, }) .meta({ description: "The final optimization summary." }); @@ -614,12 +647,26 @@ export const petrinautOptimizationCompleteEventSchema = z export const PETRINAUT_OPTIMIZATION_CANCELLED_ERROR_CODE = "optimization_cancelled"; +/** + * The `code` on the error a connected capability throws for a run id it does + * not hold: one it never created, or one it dropped when the tab reloaded. + * A stored run that meets it is stale and can be forgotten. + */ +export const PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE = + "optimization_unknown_run"; + +export const isUnknownOptimizationRunError = (error: unknown): boolean => + error instanceof Error && + "code" in error && + error.code === PETRINAUT_OPTIMIZATION_UNKNOWN_RUN_ERROR_CODE; + export const petrinautOptimizationErrorEventSchema = z .strictObject({ type: z.literal("error"), code: z.string(), message: z.string(), retryable: z.boolean(), + resumable: optimizationResumableSchema, seq: optimizationEventSeqSchema, }) .meta({ description: "A terminal optimizer error." }); @@ -669,7 +716,9 @@ export type PetrinautOptimizationTrialEvent = z.infer< >; /** - * Host-provided optimization capability for Petrinaut. + * Host-provided optimization capability for Petrinaut: the self-contained + * variant, backed by a remote service that owns its simulations. The + * in-tab variant is {@link PetrinautConnectedOptimization}. * * A run is detached from any one connection: it is created by id, its event * stream can be (re-)attached with a `seq` cursor, and it is cancelled @@ -700,3 +749,134 @@ 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; + } + | { + /** 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; +}; + +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`, for the run and every + * extension of it. 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): 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, in the tab, and + * needs the host's compute: connect it to a channel to obtain the + * capability. The self-contained variant is {@link PetrinautOptimization}. + */ +export type PetrinautConnectedOptimization = { + readonly kind: "connected"; + connect( + this: void, + channel: PetrinautOptimizationChannel, + ): PetrinautConnectedOptimizationCapability; +}; + +/** 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, + resolveTrialScenarioParameterValues, +} from "./describe"; diff --git a/libs/@hashintel/petrinaut-core/src/optimization.test.ts b/libs/@hashintel/petrinaut-core/src/optimization/optimization.test.ts similarity index 99% rename from libs/@hashintel/petrinaut-core/src/optimization.test.ts rename to libs/@hashintel/petrinaut-core/src/optimization/optimization.test.ts index 0a30083787d..9ea6651483c 100644 --- a/libs/@hashintel/petrinaut-core/src/optimization.test.ts +++ b/libs/@hashintel/petrinaut-core/src/optimization/optimization.test.ts @@ -3,7 +3,7 @@ import { describe, expect, it } from "vitest"; import { petrinautOptimizationEventSchema, petrinautOptimizationManifestSchema, -} from "./optimization"; +} from "./index"; const scenario = { id: "baseline", 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/vite.config.ts b/libs/@hashintel/petrinaut-core/vite.config.ts index 045cc84c3e6..020dca7bb0b 100644 --- a/libs/@hashintel/petrinaut-core/vite.config.ts +++ b/libs/@hashintel/petrinaut-core/vite.config.ts @@ -22,7 +22,12 @@ export default defineConfig(({ command }) => ({ hir: resolve(packageRoot, "src/hir.ts"), // Dependency-free instantiation of compiled HIR artifacts. "hir-runtime": resolve(packageRoot, "src/hir-runtime.ts"), - optimization: resolve(packageRoot, "src/optimization.ts"), + optimization: resolve(packageRoot, "src/optimization/index.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"), @@ -94,9 +99,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/libs/@local/petrinaut-arch-docs/dependency-cruiser.tsconfig.json b/libs/@local/petrinaut-arch-docs/dependency-cruiser.tsconfig.json index afe5f0c20e9..743e923c6ea 100644 --- a/libs/@local/petrinaut-arch-docs/dependency-cruiser.tsconfig.json +++ b/libs/@local/petrinaut-arch-docs/dependency-cruiser.tsconfig.json @@ -22,7 +22,7 @@ "libs/@hashintel/petrinaut-core/src/compiled-model.ts" ], "@hashintel/petrinaut-core/optimization": [ - "libs/@hashintel/petrinaut-core/src/optimization.ts" + "libs/@hashintel/petrinaut-core/src/optimization/index.ts" ], "@hashintel/petrinaut-core/workers/lsp": [ "libs/@hashintel/petrinaut-core/src/workers/lsp.ts" diff --git a/yarn.lock b/yarn.lock index 271a7c611b2..0c24057700f 100644 --- a/yarn.lock +++ b/yarn.lock @@ -7686,6 +7686,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" @@ -7695,6 +7696,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" @@ -18499,6 +18501,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" @@ -39198,6 +39207,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.16.0": version: 6.16.0 resolution: "qs@npm:6.16.0" @@ -46737,7 +46756,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: