9ca3682Stan remote sampling: compile via stan-wasm-wasi, sample via wasm-execJeremy Magland 1import { useCallback, useRef, useState } from "react";
2import compileStanProgram, {
3 DEFAULT_WASI_SERVER_URL,
4} from "./compileStanProgram";
5import runRemoteChains from "./runRemoteChains";
6import {
7 ChainStatusMap,
8 defaultSamplingOpts,
9 MultiChainResult,
10 SamplingOpts,
11} from "./types";
13const DEFAULT_STAN_PROGRAM = `data {
14 int<lower=0> N;
15 array[N] int<lower=0, upper=1> y;
16}
17parameters {
18 real<lower=0, upper=1> theta;
19}
20model {
21 theta ~ beta(1, 1);
22 y ~ bernoulli(theta);
23}
24`;
26const DEFAULT_DATA_JSON = `{
27 "N": 10,
28 "y": [0, 1, 0, 0, 0, 0, 0, 0, 0, 1]
29}
30`;
32const mean = (xs: number[]) => xs.reduce((a, b) => a + b, 0) / xs.length;
33const std = (xs: number[]) => {
34 const m = mean(xs);
35 return Math.sqrt(mean(xs.map((x) => (x - m) ** 2)));
36};
38function App() {
39 const [stanProgram, setStanProgram] = useState(DEFAULT_STAN_PROGRAM);
40 const [dataJson, setDataJson] = useState(DEFAULT_DATA_JSON);
42 const [wasiServerUrl, setWasiServerUrl] = useState(
43 () => localStorage.getItem("wasiServerUrl") ?? DEFAULT_WASI_SERVER_URL,
44 );
45 const [wasmExecUrl, setWasmExecUrl] = useState(
46 () => localStorage.getItem("wasmExecUrl") ?? "https://wasm-exec.fly.dev",
47 );
48 const [wasmExecKey, setWasmExecKey] = useState(
49 () => localStorage.getItem("wasmExecKey") ?? "",
50 );
52 const [compileStatus, setCompileStatus] = useState("");
53 const [wasiWasmUrl, setWasiWasmUrl] = useState<string | undefined>(undefined);
54 const [compiling, setCompiling] = useState(false);
56 const [samplingOpts, setSamplingOpts] =
57 useState<SamplingOpts>(defaultSamplingOpts);
58 const [sampling, setSampling] = useState(false);
59 const [chainStatus, setChainStatus] = useState<ChainStatusMap>({});
60 const [samplingError, setSamplingError] = useState("");
61 const [result, setResult] = useState<MultiChainResult | undefined>(undefined);
62 const cancelRef = useRef<(() => void) | undefined>(undefined);
64 const handleCompile = useCallback(async () => {
65 setCompiling(true);
66 setResult(undefined);
67 setWasiWasmUrl(undefined);
68 const { artifactUrl } = await compileStanProgram(
69 wasiServerUrl,
70 stanProgram,
71 setCompileStatus,
72 );
73 setWasiWasmUrl(artifactUrl);
74 setCompiling(false);
75 }, [wasiServerUrl, stanProgram]);
77 const handleSample = useCallback(async () => {
78 if (!wasiWasmUrl) return;
79 setSampling(true);
80 setSamplingError("");
81 setResult(undefined);
82 setChainStatus({});
84 try {
85 // fetch the compiled WASI module from the compile server (browser
86 // HTTP cache makes repeat samples cheap)
87 const resp = await fetch(wasiWasmUrl);
88 if (!resp.ok)
89 throw new Error(`failed to fetch ${wasiWasmUrl}: ${resp.statusText}`);
90 const moduleBytes = new Uint8Array(await resp.arrayBuffer());
92 const run = runRemoteChains(
93 wasmExecUrl,
94 wasmExecKey,
95 moduleBytes,
96 dataJson,
97 samplingOpts,
98 (chainId, status) =>
99 setChainStatus((prev) => ({ ...prev, [chainId]: status })),
100 );
101 cancelRef.current = run.cancel;
102 setResult(await run.result);
103 } catch (e) {
104 setSamplingError(`${e}`);
105 } finally {
106 cancelRef.current = undefined;
107 setSampling(false);
108 }
109 }, [wasiWasmUrl, wasmExecUrl, wasmExecKey, dataJson, samplingOpts]);
111 const handleCancel = useCallback(() => {
112 cancelRef.current?.();
113 cancelRef.current = undefined;
114 setSampling(false);
115 setSamplingError("canceled");
116 }, []);
118 const setOpt = (key: keyof SamplingOpts, value: number) =>
119 setSamplingOpts((prev) => ({ ...prev, [key]: value }));
121 return (
122 <div style={{ fontFamily: "sans-serif", margin: 20, maxWidth: 1100 }}>
123 <h2>Stan Remote Sampling</h2>
124 <p style={{ color: "#555", fontSize: 14 }}>
125 The Stan program is compiled to a pure-WASI module by{" "}
126 <code>stan-wasm-wasi</code>; sampling runs one chain per parallel job
127 on <code>wasm-exec</code> workers.
128 </p>
130 <div style={{ display: "flex", gap: 20 }}>
131 <div style={{ flex: 2 }}>
132 <h3>Stan program</h3>
133 <textarea
134 value={stanProgram}
135 onChange={(e) => setStanProgram(e.target.value)}
136 rows={16}
137 style={{ width: "100%", fontFamily: "monospace" }}
138 spellCheck={false}
139 />
140 </div>
141 <div style={{ flex: 1 }}>
142 <h3>data.json</h3>
143 <textarea
144 value={dataJson}
145 onChange={(e) => setDataJson(e.target.value)}
146 rows={16}
147 style={{ width: "100%", fontFamily: "monospace" }}
148 spellCheck={false}
149 />
150 </div>
151 </div>
153 <h3>Compile</h3>
154 <div style={{ display: "flex", gap: 10, alignItems: "center", flexWrap: "wrap" }}>
155 <label>
156 Compile server:{" "}
157 <input
158 value={wasiServerUrl}
159 onChange={(e) => {
160 setWasiServerUrl(e.target.value);
161 localStorage.setItem("wasiServerUrl", e.target.value);
162 }}
163 style={{ width: 280 }}
164 />
165 </label>
166 <button onClick={handleCompile} disabled={compiling || sampling}>
167 Compile
168 </button>
169 <span>{compileStatus}</span>
170 </div>
172 <h3>Sample</h3>
173 <div
174 style={{
175 display: "flex",
176 gap: 15,
177 alignItems: "center",
178 flexWrap: "wrap",
179 fontSize: 14,
180 marginBottom: 10,
181 }}
182 >
183 <label>
184 wasm-exec server:{" "}
185 <input
186 value={wasmExecUrl}
187 onChange={(e) => {
188 setWasmExecUrl(e.target.value);
189 localStorage.setItem("wasmExecUrl", e.target.value);
190 }}
191 style={{ width: 200 }}
192 />
193 </label>
194 <label>
195 Client key (ck_...):{" "}
196 <input
197 value={wasmExecKey}
198 onChange={(e) => {
199 setWasmExecKey(e.target.value);
200 localStorage.setItem("wasmExecKey", e.target.value);
201 }}
202 style={{ width: 220 }}
203 />
204 </label>
205 </div>
207 <div style={{ display: "flex", gap: 15, alignItems: "center", flexWrap: "wrap" }}>
208 <label>
209 Chains (parallel jobs):{" "}
210 <input
211 type="number"
212 min={1}
213 max={16}
214 value={samplingOpts.num_chains}
215 onChange={(e) => setOpt("num_chains", parseInt(e.target.value) || 1)}
216 style={{ width: 60 }}
217 />
218 </label>
219 <label>
220 Warmup:{" "}
221 <input
222 type="number"
223 min={0}
224 value={samplingOpts.num_warmup}
225 onChange={(e) => setOpt("num_warmup", parseInt(e.target.value) || 0)}
226 style={{ width: 80 }}
227 />
228 </label>
229 <label>
230 Samples:{" "}
231 <input
232 type="number"
233 min={1}
234 value={samplingOpts.num_samples}
235 onChange={(e) => setOpt("num_samples", parseInt(e.target.value) || 1)}
236 style={{ width: 80 }}
237 />
238 </label>
239 <button
240 onClick={handleSample}
241 disabled={sampling || compiling || !wasiWasmUrl || !wasmExecKey}
242 >
243 Sample
244 </button>
245 {sampling && <button onClick={handleCancel}>Cancel</button>}
246 </div>
248 {Object.keys(chainStatus).length > 0 && (
249 <div style={{ marginTop: 15 }}>
250 {Array.from({ length: samplingOpts.num_chains }, (_, i) => i + 1).map(
251 (chainId) => {
252 const s = chainStatus[chainId];
253 let text = "submitted / queued...";
254 let pct = 0;
255 if (s === "done") {
256 text = "done";
257 pct = 100;
258 } else if (s === "error") {
259 text = "error";
260 } else if (s) {
261 pct = Math.round((s.iteration / s.totalIterations) * 100);
262 text = `${s.iteration} / ${s.totalIterations} ${s.warmup ? "(warmup)" : "(sampling)"}`;
263 }
264 return (
265 <div
266 key={chainId}
267 style={{ display: "flex", alignItems: "center", gap: 10, marginBottom: 4 }}
268 >
269 <span style={{ width: 70 }}>Chain {chainId}</span>
270 <div style={{ width: 300, height: 14, background: "#eee", border: "1px solid #ccc" }}>
271 <div
272 style={{
273 width: `${pct}%`,
274 height: "100%",
275 background: s === "error" ? "#c33" : "#4a90d9",
276 }}
277 />
278 </div>
279 <span style={{ fontSize: 13, color: "#555" }}>{text}</span>
280 </div>
281 );
282 },
283 )}
284 </div>
285 )}
287 {samplingError && (
288 <pre style={{ color: "#c33", whiteSpace: "pre-wrap" }}>{samplingError}</pre>
289 )}
291 {result && (
292 <div style={{ marginTop: 15 }}>
293 <h3>Results</h3>
294 <p style={{ fontSize: 14 }}>
295 {result.numChains} chains, {result.draws[0]?.length ?? 0} total draws,{" "}
296 {result.computeTimeSec.toFixed(2)} sec
297 </p>
298 <table style={{ borderCollapse: "collapse" }}>
299 <thead>
300 <tr>
301 {["parameter", "mean", "std"].map((h) => (
302 <th
303 key={h}
304 style={{ border: "1px solid #ccc", padding: "4px 12px", textAlign: "left" }}
305 >
306 {h}
307 </th>
308 ))}
309 </tr>
310 </thead>
311 <tbody>
312 {result.paramNames.map((name, p) => (
313 <tr key={name}>
314 <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>{name}</td>
315 <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>
316 {mean(result.draws[p]).toFixed(4)}
317 </td>
318 <td style={{ border: "1px solid #ccc", padding: "4px 12px" }}>
319 {std(result.draws[p]).toFixed(4)}
320 </td>
321 </tr>
322 ))}
323 </tbody>
324 </table>
325 {result.consoleText.trim() && (
326 <details style={{ marginTop: 10 }}>
327 <summary>Console output</summary>
328 <pre style={{ fontSize: 12, background: "#f6f6f6", padding: 10 }}>
329 {result.consoleText}
330 </pre>
331 </details>
332 )}
333 </div>
334 )}
335 </div>
336 );
337}
339export default App;