9ca3682Stan remote sampling: compile via stan-wasm-wasi, sample via wasm-execJeremy Magland 1import { mergeChainResults } from "./mergeChains";
2import { parseProgress } from "./parseProgress";
3import type {
4 ChainResult,
5 ChainStatus,
6 MultiChainResult,
7 SamplingOpts,
8} from "./types";
10// Runs chains remotely on a wasm-exec service: one POST /v1/execute job per
11// chain, module = the model compiled as a pure-WASI command module (see
12// wasi-prototype/). The driver's argv: data_json seed chain_id num_warmup
13// num_samples init_radius refresh. Progress arrives as SSE stderr events,
14// draws as CSV on SSE stdout events.
16type SseEvent = { event: string; data: any };
18async function* sseEvents(body: ReadableStream<Uint8Array>): AsyncGenerator<SseEvent> {
19 const reader = body.getReader();
20 const decoder = new TextDecoder();
21 let buf = "";
22 while (true) {
23 const { done, value } = await reader.read();
24 if (done) break;
25 buf += decoder.decode(value, { stream: true });
26 let sep;
27 while ((sep = buf.indexOf("\n\n")) !== -1) {
28 const block = buf.slice(0, sep);
29 buf = buf.slice(sep + 2);
30 let event = "message";
31 let data = "";
32 for (const line of block.split("\n")) {
33 if (line.startsWith("event:")) event = line.slice(6).trim();
34 else if (line.startsWith("data:")) data += line.slice(5).trim();
35 }
36 if (data) yield { event, data: JSON.parse(data) };
37 }
38 }
39}
41const runOneRemoteChain = async (
42 serverUrl: string,
43 clientKey: string,
44 moduleBytes: Uint8Array,
45 args: string[],
46 chainId: number,
47 signal: AbortSignal,
48 onStatus: (chainId: number, status: ChainStatus) => void,
49): Promise<ChainResult> => {
50 const form = new FormData();
51 // POC: send the module bytes with every job. The server caches by sha256,
52 // so this could be optimized to module_sha256 with a 404 fallback.
53 form.append("module", new Blob([moduleBytes as BlobPart]), "main-wasi.wasm");
54 form.append("args", JSON.stringify(args));
55 form.append("timeout_ms", "300000");
57 const resp = await fetch(`${serverUrl}/v1/execute`, {
58 method: "POST",
59 headers: { Authorization: `Bearer ${clientKey}` },
60 body: form,
61 signal,
62 });
63 if (!resp.ok || !resp.body) {
64 let message = resp.statusText;
65 try {
66 message = (await resp.json()).message ?? message;
67 } catch {
68 // keep statusText
69 }
70 throw new Error(`chain ${chainId}: submit failed: ${message}`);
71 }
73 let stdoutText = "";
74 let stderrText = "";
75 let stderrLineBuf = "";
76 let exitCode: number | undefined;
78 for await (const { event, data } of sseEvents(resp.body)) {
79 if (event === "stdout") {
80 stdoutText += atob(data.b64);
81 } else if (event === "stderr") {
82 const text = atob(data.b64);
83 stderrText += text;
84 stderrLineBuf += text;
85 let nl;
86 while ((nl = stderrLineBuf.indexOf("\n")) !== -1) {
87 const line = stderrLineBuf.slice(0, nl).trim();
88 stderrLineBuf = stderrLineBuf.slice(nl + 1);
89 const progress = parseProgress(line);
90 if (progress) onStatus(chainId, progress);
91 }
92 } else if (event === "exit") {
93 exitCode = data.exitCode;
94 } else if (event === "error") {
95 throw new Error(`chain ${chainId}: ${data.code}: ${data.message}`);
96 }
97 }
99 if (exitCode === undefined)
100 throw new Error(`chain ${chainId}: stream ended without exit event`);
101 if (exitCode !== 0)
102 throw new Error(
103 `chain ${chainId}: module exited with code ${exitCode}\n${stderrText.slice(-1000)}`,
104 );
106 // stdout: header line of param names, then one CSV row per draw
107 const lines = stdoutText.trim().split("\n");
108 const paramNames = lines[0].split(",");
109 const rows = lines.slice(1).map((l) => l.split(",").map(Number));
110 const draws = paramNames.map((_, p) => rows.map((r) => r[p]));
111 return { chainId, paramNames, draws, consoleText: stderrText };
112};
114const runRemoteChains = (
115 serverUrl: string,
116 clientKey: string,
117 moduleBytes: Uint8Array,
118 data: string,
119 opts: SamplingOpts,
120 onChainStatus: (chainId: number, status: ChainStatus) => void,
121): { result: Promise<MultiChainResult>; cancel: () => void } => {
122 const seed = opts.seed ?? Math.floor(Math.random() * Math.pow(2, 32));
123 const refresh = Math.max(
124 10,
125 Math.floor((opts.num_warmup + opts.num_samples) / 100) * 10,
126 );
127 const startTimeSec = performance.now() / 1000;
128 const abort = new AbortController();
130 const chainPromises = Array.from({ length: opts.num_chains }, (_, i) => {
131 const chainId = i + 1;
132 const args = [
133 data,
134 String(seed),
135 String(chainId),
136 String(opts.num_warmup),
137 String(opts.num_samples),
138 String(opts.init_radius),
139 String(refresh),
140 ];
141 return runOneRemoteChain(
142 serverUrl.replace(/\/$/, ""),
143 clientKey,
144 moduleBytes,
145 args,
146 chainId,
147 abort.signal,
148 onChainStatus,
149 ).then(
150 (r) => {
151 onChainStatus(chainId, "done");
152 return r;
153 },
154 (e) => {
155 onChainStatus(chainId, "error");
156 throw e;
157 },
158 );
159 });
161 const result = Promise.all(chainPromises).then((chainResults) =>
162 mergeChainResults(chainResults, performance.now() / 1000 - startTimeSec),
163 );
165 return { result, cancel: () => abort.abort() };
166};
168export default runRemoteChains;