/ concept-collection / turing-surface-cache
Sign in
concept-collection / turing-surface-cache
turing-surface-cache / src / cache / runSpec.ts
320 lines · 12.1 KBCodeBlameHistory
2 * Computing one cached solution: the run behind both the page's Compute
3 * solution button and the command line's fill walk.
4 *
5 * What the run does besides stepping to the end time is the part worth having
6 * in one place. It starts from the longest cached shorter run of the same
7 * spec instead of from t = 0, since the state is Markovian in the spectral
8 * coefficients; it captures the state at every listed end time it passes, and
9 * encodes and uploads each one while the solver keeps stepping; and it
10 * refuses to publish a state that has gone non-finite. All of that is worth
11 * exactly one implementation.
12 *
13 * Everything a caller wants to say about a run in progress — status text,
14 * rendering, when to stop — arrives through RunEvents, so nothing here knows
15 * whether it is driving a page or a terminal.
16 */
17import { T_END_CHOICE } from './options.ts';
18import { stepsFor, type CacheSpec } from './spec.ts';
19import { lookupFor, fetchCached, uploadCacheFile } from './client.ts';
20import { encodeCacheFile, decodeCacheFile, type DecodedCacheFile } from './h5file.ts';
21import type { SolverSession } from './solver.ts';
23/**
24 * Steps between syncs: many small submissions queued back to back, one wait.
25 * The readbacks and renders that pace a live view happen per chunk, not per
26 * submission — that is what lets a run advance at close to the solver's own
27 * rate.
28 */
29const CHUNK_STEPS = 32;
31/** What a finished run amounts to, and all a caller needs to describe it. */
32export interface RunSummary {
33 tEnd: number;
34 seconds: number;
35 /** The end time this run resumed from, if it resumed from one. */
36 warmFrom: number | null;
39export type RunPhase =
40 | { kind: 'warm-search' }
41 | { kind: 'seeding' }
42 | { kind: 'encoding'; run: RunSummary }
43 | { kind: 'uploading'; run: RunSummary; started: number; uploaded: number };
45export interface RunProgress {
46 /** Simulation time reached, and where the run ends. */
47 t: number;
48 tEnd: number;
49 steps: number;
50 totalSteps: number;
51 /** Fraction of *this* run's work done: a warm start begins at 0 here. */
52 fraction: number;
53 /** Steps per second over the run so far. */
54 rate: number;
55 /** The end time this run resumed from, if it resumed from one. */
56 warmFrom: number | null;
57 uploadsStarted: number;
58 uploadsDone: number;
61export interface RunEvents {
62 onPhase?(phase: RunPhase): void;
63 /** The warm start or the seeding is done and the stepping is about to
64 * begin: whatever is on display now belongs to the previous run. */
65 onStepping?(): void;
66 /** Once per chunk. Callers throttle their own display. */
67 onProgress?(p: RunProgress): void;
68 /** After each chunk: the caller's chance to draw, and to yield. */
69 onTick?(): Promise<void> | void;
70 /** The final state is computed and finite, before the file is written. */
71 onFinal?(tEnd: number): Promise<void> | void;
72 onUploaded?(tEnd: number): void;
73 /** The finished file, named as it is in the cache. */
74 onFile?(bytes: Uint8Array, fileName: string): void;
75 /**
76 * Abandon the run at the next safe point, reporting nothing: something else
77 * has taken the session over. Distinct from stopRequested, which is a
78 * deliberate stop whose partial results still count.
79 */
80 cancelled?(): boolean;
81 /** Stop cleanly at the next chunk boundary, keeping what was uploaded. */
82 stopRequested?(): boolean;
85export type RunOutcome =
86 | (RunSummary & {
87 kind: 'done';
88 fileName: string;
89 bytes: Uint8Array;
90 /** End times uploaded, in the order they completed. */
91 uploaded: number[];
92 uploadsStarted: number;
93 uploadErrors: string[];
94 })
95 | { kind: 'stopped'; t: number; uploaded: number[] }
96 | { kind: 'diverged'; t: number }
97 | { kind: 'abandoned' };
99export interface RunOptions {
100 solver: SolverSession;
101 spec: CacheSpec;
102 /** Which GPU computed it; recorded in every file it writes. */
103 adapter: string;
b0546a9Fix the command line's file writing, and how it gets updatedJeremy Magland 104 /** What is driving that GPU: 'browser-webgpu', or the command line's Dawn. */
105 runtime: string;
795cdccMove the run and the walk out of the pageJeremy Magland 106 /** Read at every upload point, so a key entered mid-run still contributes. */
107 apiKey(): string;
108 events?: RunEvents;
111/**
112 * Nothing that is not a number gets uploaded. A combination whose timestep is
113 * too large for its reaction blows up rather than failing, and an unattended
114 * walk would happily publish the wreckage under a hash someone later trusts.
115 */
116const stateIsFinite = (state: Record<string, Float32Array>): boolean =>
117 Object.values(state).every((a) => a.every(Number.isFinite));
119/** Run the solver to the spec's end time, contributing everything it passes. */
120export async function runSpec(opts: RunOptions): Promise<RunOutcome> {
b0546a9Fix the command line's file writing, and how it gets updatedJeremy Magland 121 const { solver, spec, adapter, runtime, apiKey } = opts;
795cdccMove the run and the walk out of the pageJeremy Magland 122 const ev = opts.events ?? {};
123 const cancelled = (): boolean => ev.cancelled?.() ?? false;
124 const session = solver.live;
125 const steps = stepsFor(spec);
126 const dt = spec.params.dt;
128 // Warm start: the state is Markovian in (U, V), so a cached run of the same
129 // spec at a smaller listed end time is an exact prefix of this one. Take the
130 // longest one there is and continue from its final state rather than
131 // recomputing it.
132 let warm: { tEnd: number; decoded: DecodedCacheFile } | null = null;
133 const earlier = T_END_CHOICE.values.filter((T) => T < spec.tEnd).sort((a, b) => b - a);
134 if (earlier.length) ev.onPhase?.({ kind: 'warm-search' });
135 for (const T of earlier) {
136 const lookup = await lookupFor({ ...spec, tEnd: T });
137 let bytes: Uint8Array | null = null;
138 try {
139 bytes = await fetchCached(lookup);
140 } catch {
141 break; // cache unreachable: no point probing further down the ladder
142 }
143 if (cancelled()) return { kind: 'abandoned' };
144 if (!bytes) continue;
145 try {
146 warm = { tEnd: T, decoded: await decodeCacheFile(bytes, lookup.specJson, solver.model.state) };
147 break;
148 } catch {
149 continue; // an unreadable candidate is skipped, not fatal
150 }
151 }
152 if (cancelled()) return { kind: 'abandoned' };
154 let initial: Record<string, Float32Array>;
155 if (warm) {
156 session.loadState(warm.decoded.final);
157 // loadState resets the clock; put it at the cached run's end so the loop
158 // below computes only the remainder.
159 session.steps = Math.round(warm.tEnd / dt);
160 session.t = warm.tEnd;
161 // The t = 0 state travels with every file of the chain, so files written
162 // from this continuation carry the same initial state as the one resumed.
163 initial = warm.decoded.initial;
164 } else {
165 ev.onPhase?.({ kind: 'seeding' });
166 await session.seed(spec.seed);
167 if (cancelled()) return { kind: 'abandoned' };
168 initial = await session.readState();
169 if (cancelled()) return { kind: 'abandoned' };
170 }
171 const startSteps = session.steps;
173 // Snapshot points: every listed end time strictly between the starting point
174 // and this run's end. The run passes through each exactly (all are whole
175 // multiples of every dt choice).
176 const snapshotAt = new Map<number, number>(); // step index -> tEnd value
177 for (const T of T_END_CHOICE.values) {
178 if (T < spec.tEnd && T > (warm?.tEnd ?? 0)) snapshotAt.set(Math.round(T / dt), T);
179 }
180 const snapshots: { tEnd: number; state: Record<string, Float32Array> }[] = [];
182 // Everything a cache file needs exists before the run starts, so a snapshot
183 // is encoded and uploaded the moment it is captured, overlapping the network
184 // with the GPU still stepping, rather than queued for the end.
185 const geometryCoeffs = {
186 X: session.geometry.X,
187 Y: session.geometry.Y,
188 Z: session.geometry.Z,
189 };
190 const encode = (t: number, state: Record<string, Float32Array>) =>
191 encodeCacheFile({
192 spec: { ...spec, tEnd: t },
193 grid: session.cfg,
194 species: solver.model.state,
195 geometry: geometryCoeffs,
196 initial,
197 final: state,
198 adapter,
201 const uploadedTimes: number[] = [];
202 const uploadErrors: string[] = [];
203 let uploadsStarted = 0;
204 const pendingUploads: Promise<void>[] = [];
205 /** Encode + upload without the stepping loop waiting. A captured snapshot is
206 * a complete solution of its own spec, so this stays valid even if the run
207 * is stopped afterwards. */
208 const uploadInBackground = (
209 t: number,
210 state: Record<string, Float32Array>,
211 key: string,
212 preEncoded?: Uint8Array,
213 ): void => {
214 uploadsStarted++;
215 pendingUploads.push(
216 (async () => {
217 const bytes = preEncoded ?? (await encode(t, state));
218 const lookup = await lookupFor({ ...spec, tEnd: t });
219 await uploadCacheFile(key, lookup.fileName, bytes);
220 uploadedTimes.push(t);
221 ev.onUploaded?.(t);
222 })().catch((e) => {
223 uploadErrors.push(`t = ${t}: ${e instanceof Error ? e.message : e}`);
224 }),
225 );
226 };
228 ev.onStepping?.();
229 const t0 = performance.now();
230 while (session.steps < steps) {
231 if (cancelled()) return { kind: 'abandoned' };
232 if (ev.stopRequested?.()) {
233 return { kind: 'stopped', t: session.steps * dt, uploaded: [...uploadedTimes] };
234 }
235 // One chunk: up to CHUNK_STEPS steps submitted back to back (each
236 // submission stays under the dispatch budget), then a single sync and at
237 // most one render. Reading back and drawing after every submission is what
238 // made the run advance at a fraction of the solver's rate — a readback
239 // costs several times the 3-4 steps it fenced. The chunk stops exactly at
240 // snapshot points so those states are still captured exactly.
241 let target = Math.min(steps, session.steps + CHUNK_STEPS);
242 for (const s of snapshotAt.keys()) {
243 if (s > session.steps && s < target) target = s;
244 }
245 while (session.steps < target) {
246 session.step(Math.min(solver.stepsPerSubmit, target - session.steps));
247 }
248 // The sync bounds how far the CPU runs ahead of the GPU, and (being a
249 // promise) yields to the event loop, which is what keeps a Stop button
250 // clickable.
251 await session.sync();
252 if (cancelled()) return { kind: 'abandoned' };
253 const hit = snapshotAt.get(session.steps);
254 if (hit !== undefined) {
255 const state = await session.readState();
256 if (cancelled()) return { kind: 'abandoned' };
257 if (!stateIsFinite(state)) return { kind: 'diverged', t: session.steps * dt };
258 // With a key on hand the snapshot goes straight to the cache; without one
259 // it is kept, in case a key is entered before the run ends.
260 const key = apiKey();
261 if (key) uploadInBackground(hit, state, key);
262 else snapshots.push({ tEnd: hit, state });
263 }
264 ev.onProgress?.({
265 t: session.steps * dt,
266 tEnd: spec.tEnd,
267 steps: session.steps,
268 totalSteps: steps,
269 fraction: (session.steps - startSteps) / (steps - startSteps),
270 rate: (session.steps - startSteps) / ((performance.now() - t0) / 1000),
271 warmFrom: warm?.tEnd ?? null,
272 uploadsStarted,
273 uploadsDone: uploadedTimes.length,
274 });
275 await ev.onTick?.();
276 }
277 if (cancelled()) return { kind: 'abandoned' };
279 const final = await session.readState();
280 if (cancelled()) return { kind: 'abandoned' };
281 if (!stateIsFinite(final)) return { kind: 'diverged', t: spec.tEnd };
282 await ev.onFinal?.(spec.tEnd);
283 const summary: RunSummary = {
284 tEnd: spec.tEnd,
285 seconds: (performance.now() - t0) / 1000,
286 warmFrom: warm?.tEnd ?? null,
287 };
289 ev.onPhase?.({ kind: 'encoding', run: summary });
290 const finalBytes = await encode(spec.tEnd, final);
291 if (cancelled()) return { kind: 'abandoned' };
292 const fileName = (await lookupFor(spec)).fileName.split('/').pop()!;
293 ev.onFile?.(finalBytes, fileName);
295 // The final solution, plus any snapshots captured before a key was entered.
296 const key = apiKey();
297 if (key) {
298 uploadInBackground(spec.tEnd, final, key, finalBytes);
299 for (const snap of snapshots) uploadInBackground(snap.tEnd, snap.state, key);
300 }
301 if (uploadsStarted > 0) {
302 ev.onPhase?.({
303 kind: 'uploading',
304 run: summary,
305 started: uploadsStarted,
306 uploaded: uploadedTimes.length,
307 });
308 await Promise.all(pendingUploads);
309 if (cancelled()) return { kind: 'abandoned' };
310 }
311 return {
312 kind: 'done',
313 ...summary,
314 fileName,
315 bytes: finalBytes,
316 uploaded: [...uploadedTimes],
317 uploadsStarted,
318 uploadErrors,
319 };
moveopenescclose