/ concept-collection / turing-surface-cache
Sign in
concept-collection / turing-surface-cache
turing-surface-cache / src / cache / runSpec.ts
317 lines · 11.9 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;
104 /** Read at every upload point, so a key entered mid-run still contributes. */
105 apiKey(): string;
106 events?: RunEvents;
109/**
110 * Nothing that is not a number gets uploaded. A combination whose timestep is
111 * too large for its reaction blows up rather than failing, and an unattended
112 * walk would happily publish the wreckage under a hash someone later trusts.
113 */
114const stateIsFinite = (state: Record<string, Float32Array>): boolean =>
115 Object.values(state).every((a) => a.every(Number.isFinite));
117/** Run the solver to the spec's end time, contributing everything it passes. */
118export async function runSpec(opts: RunOptions): Promise<RunOutcome> {
119 const { solver, spec, adapter, apiKey } = opts;
120 const ev = opts.events ?? {};
121 const cancelled = (): boolean => ev.cancelled?.() ?? false;
122 const session = solver.live;
123 const steps = stepsFor(spec);
124 const dt = spec.params.dt;
126 // Warm start: the state is Markovian in (U, V), so a cached run of the same
127 // spec at a smaller listed end time is an exact prefix of this one. Take the
128 // longest one there is and continue from its final state rather than
129 // recomputing it.
130 let warm: { tEnd: number; decoded: DecodedCacheFile } | null = null;
131 const earlier = T_END_CHOICE.values.filter((T) => T < spec.tEnd).sort((a, b) => b - a);
132 if (earlier.length) ev.onPhase?.({ kind: 'warm-search' });
133 for (const T of earlier) {
134 const lookup = await lookupFor({ ...spec, tEnd: T });
135 let bytes: Uint8Array | null = null;
136 try {
137 bytes = await fetchCached(lookup);
138 } catch {
139 break; // cache unreachable: no point probing further down the ladder
140 }
141 if (cancelled()) return { kind: 'abandoned' };
142 if (!bytes) continue;
143 try {
144 warm = { tEnd: T, decoded: await decodeCacheFile(bytes, lookup.specJson, solver.model.state) };
145 break;
146 } catch {
147 continue; // an unreadable candidate is skipped, not fatal
148 }
149 }
150 if (cancelled()) return { kind: 'abandoned' };
152 let initial: Record<string, Float32Array>;
153 if (warm) {
154 session.loadState(warm.decoded.final);
155 // loadState resets the clock; put it at the cached run's end so the loop
156 // below computes only the remainder.
157 session.steps = Math.round(warm.tEnd / dt);
158 session.t = warm.tEnd;
159 // The t = 0 state travels with every file of the chain, so files written
160 // from this continuation carry the same initial state as the one resumed.
161 initial = warm.decoded.initial;
162 } else {
163 ev.onPhase?.({ kind: 'seeding' });
164 await session.seed(spec.seed);
165 if (cancelled()) return { kind: 'abandoned' };
166 initial = await session.readState();
167 if (cancelled()) return { kind: 'abandoned' };
168 }
169 const startSteps = session.steps;
171 // Snapshot points: every listed end time strictly between the starting point
172 // and this run's end. The run passes through each exactly (all are whole
173 // multiples of every dt choice).
174 const snapshotAt = new Map<number, number>(); // step index -> tEnd value
175 for (const T of T_END_CHOICE.values) {
176 if (T < spec.tEnd && T > (warm?.tEnd ?? 0)) snapshotAt.set(Math.round(T / dt), T);
177 }
178 const snapshots: { tEnd: number; state: Record<string, Float32Array> }[] = [];
180 // Everything a cache file needs exists before the run starts, so a snapshot
181 // is encoded and uploaded the moment it is captured, overlapping the network
182 // with the GPU still stepping, rather than queued for the end.
183 const geometryCoeffs = {
184 X: session.geometry.X,
185 Y: session.geometry.Y,
186 Z: session.geometry.Z,
187 };
188 const encode = (t: number, state: Record<string, Float32Array>) =>
189 encodeCacheFile({
190 spec: { ...spec, tEnd: t },
191 grid: session.cfg,
192 species: solver.model.state,
193 geometry: geometryCoeffs,
194 initial,
195 final: state,
196 adapter,
197 });
198 const uploadedTimes: number[] = [];
199 const uploadErrors: string[] = [];
200 let uploadsStarted = 0;
201 const pendingUploads: Promise<void>[] = [];
202 /** Encode + upload without the stepping loop waiting. A captured snapshot is
203 * a complete solution of its own spec, so this stays valid even if the run
204 * is stopped afterwards. */
205 const uploadInBackground = (
206 t: number,
207 state: Record<string, Float32Array>,
208 key: string,
209 preEncoded?: Uint8Array,
210 ): void => {
211 uploadsStarted++;
212 pendingUploads.push(
213 (async () => {
214 const bytes = preEncoded ?? (await encode(t, state));
215 const lookup = await lookupFor({ ...spec, tEnd: t });
216 await uploadCacheFile(key, lookup.fileName, bytes);
217 uploadedTimes.push(t);
218 ev.onUploaded?.(t);
219 })().catch((e) => {
220 uploadErrors.push(`t = ${t}: ${e instanceof Error ? e.message : e}`);
221 }),
222 );
223 };
225 ev.onStepping?.();
226 const t0 = performance.now();
227 while (session.steps < steps) {
228 if (cancelled()) return { kind: 'abandoned' };
229 if (ev.stopRequested?.()) {
230 return { kind: 'stopped', t: session.steps * dt, uploaded: [...uploadedTimes] };
231 }
232 // One chunk: up to CHUNK_STEPS steps submitted back to back (each
233 // submission stays under the dispatch budget), then a single sync and at
234 // most one render. Reading back and drawing after every submission is what
235 // made the run advance at a fraction of the solver's rate — a readback
236 // costs several times the 3-4 steps it fenced. The chunk stops exactly at
237 // snapshot points so those states are still captured exactly.
238 let target = Math.min(steps, session.steps + CHUNK_STEPS);
239 for (const s of snapshotAt.keys()) {
240 if (s > session.steps && s < target) target = s;
241 }
242 while (session.steps < target) {
243 session.step(Math.min(solver.stepsPerSubmit, target - session.steps));
244 }
245 // The sync bounds how far the CPU runs ahead of the GPU, and (being a
246 // promise) yields to the event loop, which is what keeps a Stop button
247 // clickable.
248 await session.sync();
249 if (cancelled()) return { kind: 'abandoned' };
250 const hit = snapshotAt.get(session.steps);
251 if (hit !== undefined) {
252 const state = await session.readState();
253 if (cancelled()) return { kind: 'abandoned' };
254 if (!stateIsFinite(state)) return { kind: 'diverged', t: session.steps * dt };
255 // With a key on hand the snapshot goes straight to the cache; without one
256 // it is kept, in case a key is entered before the run ends.
257 const key = apiKey();
258 if (key) uploadInBackground(hit, state, key);
259 else snapshots.push({ tEnd: hit, state });
260 }
261 ev.onProgress?.({
262 t: session.steps * dt,
263 tEnd: spec.tEnd,
264 steps: session.steps,
265 totalSteps: steps,
266 fraction: (session.steps - startSteps) / (steps - startSteps),
267 rate: (session.steps - startSteps) / ((performance.now() - t0) / 1000),
268 warmFrom: warm?.tEnd ?? null,
269 uploadsStarted,
270 uploadsDone: uploadedTimes.length,
271 });
272 await ev.onTick?.();
273 }
274 if (cancelled()) return { kind: 'abandoned' };
276 const final = await session.readState();
277 if (cancelled()) return { kind: 'abandoned' };
278 if (!stateIsFinite(final)) return { kind: 'diverged', t: spec.tEnd };
279 await ev.onFinal?.(spec.tEnd);
280 const summary: RunSummary = {
281 tEnd: spec.tEnd,
282 seconds: (performance.now() - t0) / 1000,
283 warmFrom: warm?.tEnd ?? null,
284 };
286 ev.onPhase?.({ kind: 'encoding', run: summary });
287 const finalBytes = await encode(spec.tEnd, final);
288 if (cancelled()) return { kind: 'abandoned' };
289 const fileName = (await lookupFor(spec)).fileName.split('/').pop()!;
290 ev.onFile?.(finalBytes, fileName);
292 // The final solution, plus any snapshots captured before a key was entered.
293 const key = apiKey();
294 if (key) {
295 uploadInBackground(spec.tEnd, final, key, finalBytes);
296 for (const snap of snapshots) uploadInBackground(snap.tEnd, snap.state, key);
297 }
298 if (uploadsStarted > 0) {
299 ev.onPhase?.({
300 kind: 'uploading',
301 run: summary,
302 started: uploadsStarted,
303 uploaded: uploadedTimes.length,
304 });
305 await Promise.all(pendingUploads);
306 if (cancelled()) return { kind: 'abandoned' };
307 }
308 return {
309 kind: 'done',
310 ...summary,
311 fileName,
312 bytes: finalBytes,
313 uploaded: [...uploadedTimes],
314 uploadsStarted,
315 uploadErrors,
316 };
moveopenescclose