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';
4import type { StanSampleConfig, WorkerResponse } from './protocol';
5import { setRunState, updateChainProgress } from './runEvents';
6import { dirnameOf, parseSampleFile, resolveProjectPath, type SampleFileConfig } from './sampleConfig';
7import { getServerUrl } from './settings';
9// The .sample runner: compile the referenced Stan program on the compile
10// server, run NUTS-HMC sampling in a web worker (tinystan), stream progress
11// to the output channel (and to the .sample view's progress bars via
12// runEvents), then write draws + summary into the output directory.
14interface ActiveRun {
15 uriKey: string;
16 worker: Worker;
17 /** Resolves the run() promise; the worker is terminated afterwards. */
18 finish: () => void;
19 stopped: boolean;
20}
22export interface StanRunner {
23 runner: FileRunner;
24 /** Stops the in-flight run, if any. */
25 stop(): void;
26 dispose(): void;
27}
29export function createStanRunner(fs: WorkspaceFileSystem): StanRunner {
30 let active: ActiveRun | undefined;
32 const run = async ({ uri, getText, output }: RunContext): Promise<void> => {
33 const uriKey = uri.toString();
34 const fail = (message: string): void => {
35 output.error(message);
36 setRunState(uriKey, { phase: 'failed', message });
37 };
39 output.appendLine('');
40 output.info(`run ${uri.path}`);
42 // 1. the .sample config
43 const { config, errors, warnings } = parseSampleFile(await getText());
44 for (const warning of warnings) {
45 output.warn(warning);
46 }
47 if (errors.length > 0) {
48 for (const error of errors) {
49 output.error(error);
50 }
51 setRunState(uriKey, { phase: 'failed', message: errors[0] });
52 return;
53 }
55 // 2. referenced files
56 const sampleDir = dirnameOf(uri.path);
57 const stanPath = resolveProjectPath(sampleDir, config.stan!);
58 const dataPath = resolveProjectPath(sampleDir, config.data!);
59 const outputDir = resolveProjectPath(sampleDir, config.output_dir!);
60 if (outputDir === '/') {
61 return fail("'output_dir' must not be the project root (its contents are replaced on each run)");
62 }
63 for (const [name, path] of [['.sample file', uri.path], ['stan file', stanPath], ['data file', dataPath]] as const) {
64 if (path === outputDir || path.startsWith(`${outputDir}/`)) {
65 return fail(`'output_dir' (${outputDir}) would overwrite the ${name} (${path})`);
66 }
67 }
69 const stanText = await readProjectText(fs, stanPath);
70 if (stanText === undefined) {
71 return fail(`Stan program not found: ${stanPath}`);
72 }
73 const dataText = await readProjectText(fs, dataPath);
74 if (dataText === undefined) {
75 return fail(`data file not found: ${dataPath}`);
76 }
77 try {
78 JSON.parse(dataText);
79 } catch (error) {
80 return fail(`data file ${dataPath} is not valid JSON: ${error instanceof Error ? error.message : error}`);
81 }
83 // 3. compile (server-side, cached by source hash)
84 setRunState(uriKey, { phase: 'compiling', message: 'compiling...' });
85 const serverUrl = getServerUrl();
86 output.info(`compiling ${stanPath} (server: ${serverUrl})`);
87 const compiled = await compileStanProgram(serverUrl, stanText, (status) => {
88 output.info(`[compile] ${status}`);
89 setRunState(uriKey, { phase: 'compiling', message: status });
90 });
91 if (!compiled.mainJsUrl) {
92 return fail(compiled.error ?? 'compilation failed');
93 }
95 // 4. sample in a fresh worker
96 const seed = config.seed ?? Math.floor(Math.random() * Math.pow(2, 32));
97 const sampleConfig: StanSampleConfig = {
98 data: dataText,
99 num_chains: config.num_chains,
100 num_warmup: config.num_warmup,
101 num_samples: config.num_samples,
102 init_radius: config.init_radius,
103 seed,
104 refresh: reasonableRefreshRate(config),
105 // one thread per chain: chains run in parallel (issue mirrors
106 // stan-playground's setting)
107 num_threads: config.num_chains,
108 };
110 setRunState(uriKey, { phase: 'loading', message: 'loading model...' });
111 const worker = new Worker(new URL('./samplerWorker.ts', import.meta.url), { type: 'module' });
112 const consoleLines: string[] = [];
113 let samplingStarted = 0;
114 let computeTimeSec = 0;
116 await new Promise<void>((resolve) => {
117 const current: ActiveRun = { uriKey, worker, finish: resolve, stopped: false };
118 active = current;
120 worker.onmessage = async (event: MessageEvent<WorkerResponse>) => {
121 if (current.stopped) {
122 return;
123 }
124 const message = event.data;
125 switch (message.type) {
126 case 'loaded': {
127 output.info(`model loaded (Stan v${message.stanVersion}); sampling: ${config.num_chains} chains × (${config.num_warmup} warmup + ${config.num_samples} samples), seed ${seed}`);
128 setRunState(uriKey, { phase: 'sampling', message: 'sampling...' });
129 samplingStarted = performance.now();
130 worker.postMessage({ type: 'sample', config: sampleConfig });
131 break;
132 }
133 case 'progress': {
134 const r = message.report;
135 updateChainProgress(uriKey, config.num_chains, r);
136 const line = `Chain ${r.chain} Iteration: ${r.iteration} / ${r.totalIterations} [${String(r.percent).padStart(3)}%] (${r.warmup ? 'Warmup' : 'Sampling'})`;
137 consoleLines.push(line);
138 output.appendLine(line);
139 break;
140 }
141 case 'console': {
142 consoleLines.push(message.text);
143 output.appendLine(message.text);
144 break;
145 }
146 case 'done': {
147 computeTimeSec = (performance.now() - samplingStarted) / 1000;
148 setRunState(uriKey, { phase: 'writing', message: 'writing outputs...' });
149 try {
150 const written = await writeRunOutputs(fs, outputDir, {
151 draws: message.draws,
152 paramNames: message.paramNames,
153 numChains: config.num_chains,
154 consoleText: consoleLines.join('\n') + '\n',
155 samplingOpts: {
156 stan: stanPath,
157 data: dataPath,
158 output_dir: outputDir,
159 num_chains: config.num_chains,
160 num_warmup: config.num_warmup,
161 num_samples: config.num_samples,
162 init_radius: config.init_radius,
163 seed,
164 compute_time_sec: Number(computeTimeSec.toFixed(3)),
165 },
166 computeTimeSec,
167 });
168 output.info(`sampling completed in ${computeTimeSec.toFixed(2)}s — wrote ${written.length} files to ${outputDir}`);
169 setRunState(uriKey, { phase: 'done', message: `completed in ${computeTimeSec.toFixed(2)}s → ${outputDir}`, computeTimeSec });
170 } catch (error) {
171 fail(`failed to write outputs: ${error}`);
172 }
173 resolve();
174 break;
175 }
176 case 'error': {
177 fail(message.message);
178 resolve();
179 break;
180 }
181 }
182 };
183 worker.onerror = (event) => {
184 fail(`worker error: ${event.message ?? 'failed to load'}`);
185 resolve();
186 };
187 worker.postMessage({ type: 'load', mainJsUrl: compiled.mainJsUrl });
188 }).finally(() => {
189 worker.terminate();
190 if (active?.worker === worker) {
191 active = undefined;
192 }
193 });
194 };
196 const stop = (): void => {
197 if (active) {
198 active.stopped = true;
199 setRunState(active.uriKey, { phase: 'failed', message: 'stopped' });
200 active.finish();
201 }
202 };
204 return {
205 runner: {
206 id: 'stan.sample',
207 displayName: 'Run sampling',
208 selector: [{ filenamePattern: '*.sample' }],
209 run,
210 stop,
211 },
212 stop,
213 dispose(): void {
214 active?.finish();
215 },
216 };
217}
219/** Progress lines roughly every 2.5% of total iterations (min every 15). */
220function reasonableRefreshRate(config: SampleFileConfig): number {
221 const total = (config.num_samples + config.num_warmup) * config.num_chains;
222 const nearestTen = Math.round(Math.floor(total / 40) / 10) * 10;
223 return Math.max(15, nearestTen);
224}
226/** Reads a project file as text, preferring an open editor's contents. */
227async function readProjectText(fs: WorkspaceFileSystem, path: string): Promise<string | undefined> {
228 const uri = fs.root.with({ path });
229 const model = monaco.editor.getModel(uri);
230 if (model) {
231 return model.getValue();
232 }
233 if (!(await fs.fileService.exists(uri))) {
234 return undefined;
235 }
236 return (await fs.fileService.readFile(uri)).value.toString();
237}