21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 1import { monaco, type FileRunner, type RunContext, type WorkspaceFileSystem } from 'minwebide';
2import { compileStanProgram } from './compile';
3import { writeRunOutputs } from './outputs';
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 4import type { ChainRunConfig, WorkerResponse } from './protocol';
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 5import { setRunState, updateChainProgress } from './runEvents';
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 6import { dirnameOf, outputDirFor, parseSampleFile, resolveProjectPath, type SampleFileConfig } from './sampleConfig';
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 7import { getServerUrl } from './settings';
9// The .sample runner: compile the referenced Stan program on the compile
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 10// server (to a pure-WASI module), run NUTS-HMC sampling locally — one web
11// worker per chain, each invoking the module CLI-style — stream progress to
12// the output channel (and to the .sample view's progress bars via
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 13// runEvents), then write draws + summary into the output directory.
15interface ActiveRun {
16 uriKey: string;
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 17 workers: Worker[];
18 /** Resolves the run() promise; the workers are terminated afterwards. */
20 stopped: boolean;
21}
23export interface StanRunner {
24 runner: FileRunner;
25 /** Stops the in-flight run, if any. */
26 stop(): void;
27 dispose(): void;
28}
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 30export function createStanRunner(fs: WorkspaceFileSystem, openFile: (path: string) => Promise<unknown>): StanRunner {
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 31 let active: ActiveRun | undefined;
33 const run = async ({ uri, getText, output }: RunContext): Promise<void> => {
34 const uriKey = uri.toString();
35 const fail = (message: string): void => {
36 output.error(message);
37 setRunState(uriKey, { phase: 'failed', message });
38 };
40 output.appendLine('');
41 output.info(`run ${uri.path}`);
43 // 1. the .sample config
44 const { config, errors, warnings } = parseSampleFile(await getText());
45 for (const warning of warnings) {
46 output.warn(warning);
47 }
48 if (errors.length > 0) {
49 for (const error of errors) {
50 output.error(error);
51 }
52 setRunState(uriKey, { phase: 'failed', message: errors[0] });
53 return;
54 }
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 56 // 2. referenced files; the output directory is derived from the
57 // .sample file's name (fit.sample → fit.out next to it)
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 58 const sampleDir = dirnameOf(uri.path);
59 const stanPath = resolveProjectPath(sampleDir, config.stan!);
60 const dataPath = resolveProjectPath(sampleDir, config.data!);
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 61 const outputDir = outputDirFor(uri.path);
62 for (const [name, path] of [['stan file', stanPath], ['data file', dataPath]] as const) {
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 63 if (path === outputDir || path.startsWith(`${outputDir}/`)) {
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 64 return fail(`the output directory (${outputDir}) would overwrite the ${name} (${path}) — move it out of ${outputDir}`);
66 }
68 const stanText = await readProjectText(fs, stanPath);
69 if (stanText === undefined) {
70 return fail(`Stan program not found: ${stanPath}`);
71 }
72 const dataText = await readProjectText(fs, dataPath);
73 if (dataText === undefined) {
74 return fail(`data file not found: ${dataPath}`);
75 }
76 try {
77 JSON.parse(dataText);
78 } catch (error) {
79 return fail(`data file ${dataPath} is not valid JSON: ${error instanceof Error ? error.message : error}`);
80 }
82 // 3. compile (server-side, cached by source hash)
83 setRunState(uriKey, { phase: 'compiling', message: 'compiling...' });
84 const serverUrl = getServerUrl();
85 output.info(`compiling ${stanPath} (server: ${serverUrl})`);
86 const compiled = await compileStanProgram(serverUrl, stanText, (status) => {
87 output.info(`[compile] ${status}`);
88 setRunState(uriKey, { phase: 'compiling', message: status });
89 });
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 90 if (!compiled.mainWasmUrl) {
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 91 return fail(compiled.error ?? 'compilation failed');
92 }
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 94 // 4. download + compile the wasm module once; WebAssembly.Module is
95 // structured-cloneable, so the chain workers share the compiled code
96 setRunState(uriKey, { phase: 'loading', message: 'loading model...' });
97 let module: WebAssembly.Module;
98 let moduleBytes = 0;
99 try {
100 const response = await fetch(compiled.mainWasmUrl);
101 if (!response.ok) {
102 return fail(`failed to download compiled model: ${response.status} ${response.statusText}`);
103 }
104 const buffer = await response.arrayBuffer();
105 moduleBytes = buffer.byteLength;
106 module = await WebAssembly.compile(buffer);
107 } catch (error) {
108 return fail(`failed to load compiled model: ${error}`);
109 }
111 // 5. sample: one worker per chain (CmdStan convention — same seed,
112 // chain ids 1..n differentiate the streams)
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 113 const seed = config.seed ?? Math.floor(Math.random() * Math.pow(2, 32));
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 114 output.info(`model loaded (${(moduleBytes / 1024).toFixed(0)} kB wasm); sampling: ${config.num_chains} chains × (${config.num_warmup} warmup + ${config.num_samples} samples), seed ${seed}`);
115 setRunState(uriKey, { phase: 'sampling', message: 'sampling...' });
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 116
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 117 const workers = Array.from({ length: config.num_chains }, () =>
118 new Worker(new URL('./samplerWorker.ts', import.meta.url), { type: 'module' }));
119 const chainResults: ({ paramNames: string[]; draws: number[][] } | undefined)[] = new Array(config.num_chains);
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 120 const consoleLines: string[] = [];
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 121 let failed = false;
122 const samplingStarted = performance.now();
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 123
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 124 const current: ActiveRun = { uriKey, workers, finish: () => {}, stopped: false };
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 125 await new Promise<void>((resolve) => {
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 126 current.finish = resolve;
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 128 let remaining = config.num_chains;
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 129
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 130 workers.forEach((worker, index) => {
131 const chainId = index + 1;
132 worker.onmessage = (event: MessageEvent<WorkerResponse>) => {
133 if (current.stopped || failed) {
134 return;
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 136 const message = event.data;
137 switch (message.type) {
138 case 'progress': {
139 const r = message.report;
140 updateChainProgress(uriKey, config.num_chains, r);
141 const line = `Chain ${r.chain} Iteration: ${r.iteration} / ${r.totalIterations} [${String(r.percent).padStart(3)}%] (${r.warmup ? 'Warmup' : 'Sampling'})`;
142 consoleLines.push(line);
143 output.appendLine(line);
144 break;
145 }
146 case 'console': {
147 const line = config.num_chains > 1 ? `[chain ${chainId}] ${message.text}` : message.text;
148 consoleLines.push(line);
149 output.appendLine(line);
150 break;
151 }
152 case 'done': {
153 chainResults[index] = { paramNames: message.paramNames, draws: message.draws };
154 remaining -= 1;
155 if (remaining === 0) {
156 resolve();
157 }
158 break;
159 }
160 case 'error': {
161 failed = true;
162 fail(message.message);
163 resolve();
164 break;
166 }
168 worker.onerror = (event) => {
169 if (!current.stopped && !failed) {
170 failed = true;
171 fail(`worker error: ${event.message ?? 'failed to load'}`);
173 }
175 const chainConfig: ChainRunConfig = {
176 data: dataText,
177 seed,
178 chainId,
179 numWarmup: config.num_warmup,
180 numSamples: config.num_samples,
181 initRadius: config.init_radius,
182 refresh: reasonableRefreshRate(config),
183 };
184 worker.postMessage({ type: 'run', module, config: chainConfig });
185 });
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 187 for (const worker of workers) {
188 worker.terminate();
189 }
190 if (active === current) {
192 }
193 });
195 if (current.stopped || failed) {
196 return; // already reported
197 }
198 if (chainResults.some((result) => !result)) {
199 return; // finished early without all chains (e.g. disposed mid-run)
200 }
202 // 6. merge chains and write outputs: draws[param][draw], chains
203 // concatenated along the draw axis (the layout outputs.ts expects)
204 const computeTimeSec = (performance.now() - samplingStarted) / 1000;
205 const results = chainResults as { paramNames: string[]; draws: number[][] }[];
206 const paramNames = results[0].paramNames;
207 const draws = paramNames.map((_, p) => {
208 const merged: number[] = [];
209 for (const chain of results) {
210 merged.push(...chain.draws[p]);
211 }
212 return merged;
213 });
215 setRunState(uriKey, { phase: 'writing', message: 'writing outputs...' });
216 try {
217 const written = await writeRunOutputs(fs, outputDir, {
218 draws,
219 paramNames,
220 numChains: config.num_chains,
221 consoleText: consoleLines.join('\n') + '\n',
222 samplingOpts: {
223 stan: stanPath,
224 data: dataPath,
225 output_dir: outputDir,
226 num_chains: config.num_chains,
227 num_warmup: config.num_warmup,
228 num_samples: config.num_samples,
229 init_radius: config.init_radius,
230 seed,
231 compute_time_sec: Number(computeTimeSec.toFixed(3)),
232 },
233 computeTimeSec,
234 });
235 output.info(`sampling completed in ${computeTimeSec.toFixed(2)}s — wrote ${written.length} files to ${outputDir}`);
236 setRunState(uriKey, { phase: 'done', message: `completed in ${computeTimeSec.toFixed(2)}s → ${outputDir}`, computeTimeSec });
237 // open (or refresh) the results dashboard
238 openFile(`${outputDir}/run.json`).catch(() => {});
239 } catch (error) {
240 fail(`failed to write outputs: ${error}`);
241 }
244 const stop = (): void => {
245 if (active) {
246 active.stopped = true;
247 setRunState(active.uriKey, { phase: 'failed', message: 'stopped' });
248 active.finish();
249 }
250 };
252 return {
253 runner: {
254 id: 'stan.sample',
255 displayName: 'Run sampling',
256 selector: [{ filenamePattern: '*.sample' }],
257 run,
258 stop,
259 },
260 stop,
261 dispose(): void {
262 active?.finish();
263 },
264 };
265}
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 267/** Progress lines roughly every 2.5% of a chain's iterations (min every 15). */
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 268function reasonableRefreshRate(config: SampleFileConfig): number {
b1dbd02Switch sampling to pure-WASI web workers; add results dashboardJeremy Magland 269 const total = config.num_samples + config.num_warmup;
21efcb6stan web IDE: run Stan sampling in the browserJeremy Magland 270 const nearestTen = Math.round(Math.floor(total / 40) / 10) * 10;
271 return Math.max(15, nearestTen);
272}
274/** Reads a project file as text, preferring an open editor's contents. */
275async function readProjectText(fs: WorkspaceFileSystem, path: string): Promise<string | undefined> {
276 const uri = fs.root.with({ path });
277 const model = monaco.editor.getModel(uri);
278 if (model) {
279 return model.getValue();
280 }
281 if (!(await fs.fileService.exists(uri))) {
282 return undefined;
283 }
284 return (await fs.fileService.readFile(uri)).value.toString();
285}