1// Web worker that runs ONE MCMC chain: instantiates the compiled Stan model
2// (a pure-WASI command module from the compile server) and runs it like a
3// CLI. Stan's console output arrives on stderr — progress lines are parsed
4// into structured reports, the rest streams back as console messages — and
5// the draws arrive on stdout as CSV (param-name header + one row per draw).
7import type { ChainRunConfig, Progress, WorkerRequest, WorkerResponse } from './protocol';
8import { runWasiModule } from './wasi';
10function post(message: WorkerResponse): void {
11 self.postMessage(message);
12}
14// Stan progress lines look like (spacing varies):
15// Iteration: 800 / 2000 [ 40%] (Warmup)
16// There is no "Chain [n]" prefix — each module run is a single chain; the
17// chain id comes from this worker's config.
18function parseProgress(line: string, chainId: number): Progress | undefined {
19 const match = line.match(/^Iteration:\s*(\d+)\s*\/\s*(\d+)\s*\[\s*(\d+)%\]\s*\((Warmup|Sampling)\)/);
20 if (!match) {
21 return undefined;
22 }
23 return {
24 chain: chainId,
25 iteration: parseInt(match[1], 10),
26 totalIterations: parseInt(match[2], 10),
27 percent: parseInt(match[3], 10),
28 warmup: match[4] === 'Warmup',
29 };
30}
32/** Splits a byte stream into decoded lines (UTF-8-safe across chunks). */
33function lineSplitter(onLine: (line: string) => void): { push(bytes: Uint8Array): void; flush(): void } {
34 const decoder = new TextDecoder();
35 let pending = '';
36 const drain = () => {
37 let index;
38 while ((index = pending.indexOf('\n')) >= 0) {
39 onLine(pending.slice(0, index));
40 pending = pending.slice(index + 1);
41 }
42 };
43 return {
44 push(bytes) {
45 pending += decoder.decode(bytes, { stream: true });
46 drain();
47 },
48 flush() {
49 pending += decoder.decode();
50 drain();
51 if (pending) {
52 onLine(pending);
53 pending = '';
54 }
55 },
56 };
57}
59function concatBytes(chunks: Uint8Array[]): Uint8Array {
60 const result = new Uint8Array(chunks.reduce((size, chunk) => size + chunk.length, 0));
61 let offset = 0;
62 for (const chunk of chunks) {
63 result.set(chunk, offset);
64 offset += chunk.length;
65 }
66 return result;
67}
69/** The driver's stdout: a param-name header, then one CSV row per draw.
70 * Returns draws[param][draw]. */
71function parseDrawsCsv(text: string): { paramNames: string[]; draws: number[][] } {
72 const lines = text.split('\n').filter((line) => line.length > 0);
73 if (lines.length < 2) {
74 throw new Error('the model produced no draws');
75 }
76 const paramNames = lines[0].split(',');
77 const draws: number[][] = paramNames.map(() => new Array<number>(lines.length - 1));
78 for (let row = 1; row < lines.length; row++) {
79 const values = lines[row].split(',');
80 for (let p = 0; p < paramNames.length; p++) {
81 draws[p][row - 1] = Number(values[p]);
82 }
83 }
84 return { paramNames, draws };
85}
87async function runChain(module: WebAssembly.Module, config: ChainRunConfig): Promise<void> {
88 const stdoutChunks: Uint8Array[] = [];
89 const stderrTail: string[] = [];
90 const stderr = lineSplitter((line) => {
91 const report = parseProgress(line, config.chainId);
92 if (report) {
93 post({ type: 'progress', report });
94 return;
95 }
96 if (line.trim()) {
97 stderrTail.push(line);
98 if (stderrTail.length > 5) {
99 stderrTail.shift();
100 }
101 post({ type: 'console', text: line, level: line.startsWith('error:') ? 'error' : 'log' });
102 }
103 });
105 const exitCode = await runWasiModule({
106 module,
107 args: [
108 config.data,
109 String(config.seed),
110 String(config.chainId),
111 String(config.numWarmup),
112 String(config.numSamples),
113 String(config.initRadius),
114 String(config.refresh),
115 ],
116 onStdout: (bytes) => stdoutChunks.push(bytes),
117 onStderr: (bytes) => stderr.push(bytes),
118 });
119 stderr.flush();
121 if (exitCode !== 0) {
122 const detail = stderrTail.join('\n');
123 post({ type: 'error', message: `chain ${config.chainId} failed (exit code ${exitCode})${detail ? `:\n${detail}` : ''}` });
124 return;
125 }
127 const { paramNames, draws } = parseDrawsCsv(new TextDecoder().decode(concatBytes(stdoutChunks)));
128 post({ type: 'done', paramNames, draws });
129}
131self.onmessage = (event: MessageEvent<WorkerRequest>) => {
132 const { module, config } = event.data;
133 runChain(module, config).catch((error) => {
134 post({ type: 'error', message: `chain ${config.chainId} failed: ${error}` });
135 });
136};