1module Checkpoint.Envelope
2
3import Data.Bytes
4import Data.SHA256Digest
5import Std.Foundation
6import Std.Natural
7
8-- THE CHECKPOINT ENVELOPE: what a native host writes around a system's raw
9-- training state, and refuses to resume from unless it holds.
10--
11-- A checkpoint file is a header of `headerBytes` (a whole number of pages)
12-- and then the payload, unmodified (for Bob and Coppelius: the P/M/V
13-- planes). The header, little-endian:
14-- 0 magic "ALPHACKP"
15-- 8 version (1)
16-- 16 header bytes
17-- 24 payload bytes
18-- 32 chunk bytes
19-- 40 full chunks (payload bytes / chunk bytes)
20-- 48 tail bytes (payload bytes mod chunk bytes)
21-- 56 completed updates (the learner's step count)
22-- 64 completed invocations (the sampler's position: batches consumed)
23-- 72 SHA-256 of the state schema's identity (model, state layout)
24-- 104 SHA-256 of the learner's identity (optimizer, hyperparameters)
25-- 136 SHA-256 of the data contract's identity (tokenizer, batch format)
26-- 168 SHA-256 of the initialization's identity (the RNG seed)
27-- 200 the SHA-256 of each chunk of the payload, in order (the tail last)
28-- then zeros to `headerBytes`. The file is complete when its size is
29-- exactly header + payload bytes.
30--
31-- A system states its contract (below); the first 56 bytes and the four
32-- identity digests of any checkpoint it resumes from must be the ones the
33-- contract gives, and every chunk must hash to its digest. The first
34-- failure is named (`CheckpointEnvelopeVerdict`): a short or long file is
35-- truncated, another version or layout is refused, another system's,
36-- learner's, data's or seed's state is foreign, and a chunk whose digest
37-- differs -- a swapped or corrupted file -- is refused by its index.
38
39-- The native persistence owner may commit the complete header with SHA-256
40-- after writing/verifying all chunk digests. That compact manifest commits to
41-- the payload through its chunk hashes and includes the continuation counters.
42-- A manifest digest is distinct from the digest of the complete file. AOS
43-- retains the full manifest hash beside its truncated typed state/checkpoint
44-- reference; see docs/observability/STATE.md.
45
46family CheckpointEnvelopeContract : Type 0
47constructor CheckpointEnvelopeContractValue
48field unrestricted checkpointEnvelopeSchemaIdentity : Bytes
49field unrestricted checkpointEnvelopeLearnerIdentity : Bytes
50field unrestricted checkpointEnvelopeDataIdentity : Bytes
51field unrestricted checkpointEnvelopeSeedIdentity : Bytes
52field unrestricted checkpointEnvelopePayloadBytes : Nat
53field unrestricted checkpointEnvelopeChunkBytes : Nat
54field unrestricted checkpointEnvelopeUpdatesPerInvocation : Nat
55end-family
56
57family CheckpointEnvelopeVerdict : Type 0
58constructor CheckpointEnvelopeAccepted
59constructor CheckpointEnvelopeTruncated
60constructor CheckpointEnvelopeLayout
61constructor CheckpointEnvelopeForeignSchema
62constructor CheckpointEnvelopeForeignLearner
63constructor CheckpointEnvelopeForeignData
64constructor CheckpointEnvelopeForeignSeed
65constructor CheckpointEnvelopeChunkMismatch
66field unrestricted checkpointEnvelopeMismatchedChunk : Nat
67end-family
68
69def checkpointEnvelopeMagic : Bytes = b"ALPHACKP"
70def checkpointEnvelopeVersion : Nat = 1
71def checkpointEnvelopePageBytes : Nat = 4096
72def checkpointEnvelopeDigestBytes : Nat = 32
73
74-- the header's fields, as byte offsets
75def checkpointEnvelopeVersionAt : Nat = 8
76def checkpointEnvelopeHeaderBytesAt : Nat = 16
77def checkpointEnvelopePayloadBytesAt : Nat = 24
78def checkpointEnvelopeChunkBytesAt : Nat = 32
79def checkpointEnvelopeFullChunksAt : Nat = 40
80def checkpointEnvelopeTailBytesAt : Nat = 48
81def checkpointEnvelopeUpdatesAt : Nat = 56
82def checkpointEnvelopeInvocationsAt : Nat = 64
83def checkpointEnvelopeSchemaAt : Nat = 72
84def checkpointEnvelopeLearnerAt : Nat = 104
85def checkpointEnvelopeDataAt : Nat = 136
86def checkpointEnvelopeSeedAt : Nat = 168
87def checkpointEnvelopeDigestsAt : Nat = 200
88-- the bytes every checkpoint of a contract begins with: magic to tail
89def checkpointEnvelopeFixedBytes : Nat = 56
90
91def checkpointEnvelopeContractPayload =
92 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
93 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
94 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . payload)))
95
96def checkpointEnvelopeContractChunk =
97 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
98 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
99 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . chunk)))
100
101def checkpointEnvelopeContractUpdates =
102 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
103 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Nat) contract
104 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates . updates)))
105
106-- the four identities' digests, in header order
107def checkpointEnvelopeIdentityDigests =
108 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
109 (eliminate CheckpointEnvelopeContract (lambda unrestricted current : (family CheckpointEnvelopeContract) . Bytes) contract
110 (branch CheckpointEnvelopeContractValue schema learner data seed payload chunk updates .
111 (bytes-append (sha256RawDigestOrEmpty schema)
112 (bytes-append (sha256RawDigestOrEmpty learner)
113 (bytes-append (sha256RawDigestOrEmpty data) (sha256RawDigestOrEmpty seed)))))))
114
115def checkpointEnvelopeFullChunks =
116 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
117 (naturalDivideUnchecked (checkpointEnvelopeContractPayload contract) (checkpointEnvelopeContractChunk contract)))
118
119def checkpointEnvelopeTailBytes =
120 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
121 (naturalModuloUnchecked (checkpointEnvelopeContractPayload contract) (checkpointEnvelopeContractChunk contract)))
122
123-- the chunks a payload has: the full ones and, when there is one, the tail
124def checkpointEnvelopeChunkCount =
125 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
126 (naturalAdd (checkpointEnvelopeFullChunks contract) (naturalNonzero (checkpointEnvelopeTailBytes contract))))
127
128-- The final digest covers only bytes present in the payload. The generic
129-- byte-prefix operation zero-pads, so passing the full chunk extent for the
130-- tail would disagree with native file hashing while still passing a codec
131-- round trip that made the same mistake on both sides.
132def checkpointEnvelopeChunkExtent =
133 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
134 (lambda unrestricted index : Nat .
135 (naturalSelect (naturalLess index (checkpointEnvelopeFullChunks contract))
136 (checkpointEnvelopeContractChunk contract)
137 (naturalSelect (naturalEqual index (checkpointEnvelopeFullChunks contract))
138 (checkpointEnvelopeTailBytes contract) 0))))
139
140def checkpointEnvelopeHeaderBytes =
141 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
142 (naturalMultiply
143 (naturalDivideUnchecked
144 (naturalAdd
145 (naturalAdd checkpointEnvelopeDigestsAt (naturalMultiply checkpointEnvelopeDigestBytes (checkpointEnvelopeChunkCount contract)))
146 (naturalSaturatingSubtract checkpointEnvelopePageBytes 1))
147 checkpointEnvelopePageBytes)
148 checkpointEnvelopePageBytes))
149
150def checkpointEnvelopeFileBytes =
151 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
152 (naturalAdd (checkpointEnvelopeHeaderBytes contract) (checkpointEnvelopeContractPayload contract)))
153
154-- a natural as eight little-endian bytes (it is below 2^64)
155def checkpointEnvelopeWord =
156 (lambda unrestricted value : Nat .
157 (bytes-builder-build
158 (app (nat-eliminate
159 (lambda unrestricted current : Nat . (pi unrestricted rest : Nat . BytesBuilder))
160 (lambda unrestricted rest : Nat . (bytes-builder-chunk b""))
161 (lambda unrestricted p : Nat .
162 (lambda unrestricted induction : (pi unrestricted rest : Nat . BytesBuilder) .
163 (lambda unrestricted rest : Nat .
164 (bytes-builder-append
165 (bytes-builder-chunk (bytes (nat-to-byte (naturalModuloUnchecked rest 256))))
166 (induction (naturalDivideUnchecked rest 256))))))
167 8)
168 value)))
169
170-- the fixed bytes: magic, version, header, payload, chunk, full, tail
171def checkpointEnvelopeFixed =
172 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
173 (bytes-append checkpointEnvelopeMagic
174 (bytes-append (checkpointEnvelopeWord checkpointEnvelopeVersion)
175 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeHeaderBytes contract))
176 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeContractPayload contract))
177 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeContractChunk contract))
178 (bytes-append (checkpointEnvelopeWord (checkpointEnvelopeFullChunks contract))
179 (checkpointEnvelopeWord (checkpointEnvelopeTailBytes contract)))))))))
180
181-- a choice between verdicts on a 0/1 natural
182def checkpointEnvelopeIf =
183 (lambda unrestricted condition : Nat .
184 (lambda unrestricted whenTrue : (family CheckpointEnvelopeVerdict) .
185 (lambda unrestricted whenFalse : (family CheckpointEnvelopeVerdict) .
186 (nat-eliminate
187 (lambda unrestricted current : Nat . (family CheckpointEnvelopeVerdict))
188 whenFalse
189 (lambda unrestricted p : Nat . (lambda unrestricted ignored : (family CheckpointEnvelopeVerdict) . whenTrue))
190 condition))))
191
192def checkpointEnvelopeZeros =
193 (lambda unrestricted count : Nat .
194 (bytes-builder-build
195 (nat-eliminate
196 (lambda unrestricted current : Nat . BytesBuilder)
197 (bytes-builder-chunk b"")
198 (lambda unrestricted p : Nat . (lambda unrestricted induction : BytesBuilder . (bytes-builder-append (bytes-builder-chunk (bytes 0)) induction)))
199 count)))
200
201-- ---- the model's reading of a checkpoint file ----
202def checkpointEnvelopeSlice =
203 (lambda unrestricted file : Bytes .
204 (lambda unrestricted offset : Nat .
205 (lambda unrestricted length : Nat .
206 (dataBytesTakeValidated length (dataBytesDropValidated offset file)))))
207
208-- the index of the first chunk whose digest differs, or the chunk count
209def checkpointEnvelopeFirstBadChunk =
210 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
211 (lambda unrestricted file : Bytes .
212 (let unrestricted count = (checkpointEnvelopeChunkCount contract) in
213 (let unrestricted header = (checkpointEnvelopeHeaderBytes contract) in
214 (let unrestricted chunk = (checkpointEnvelopeContractChunk contract) in
215 (app
216 (nat-eliminate
217 (lambda unrestricted remaining : Nat . (pi unrestricted index : Nat . Nat))
218 (lambda unrestricted index : Nat . index)
219 (lambda unrestricted p : Nat .
220 (lambda unrestricted induction : (pi unrestricted index : Nat . Nat) .
221 (lambda unrestricted index : Nat .
222 (nat-eliminate
223 (lambda unrestricted same : Nat . Nat)
224 index
225 (lambda unrestricted q : Nat . (lambda unrestricted ignored : Nat . (induction (succ index))))
226 (bytes-equal
227 (checkpointEnvelopeSlice file (naturalAdd checkpointEnvelopeDigestsAt (naturalMultiply checkpointEnvelopeDigestBytes index)) checkpointEnvelopeDigestBytes)
228 (sha256RawDigestOrEmpty (checkpointEnvelopeSlice file (naturalAdd header (naturalMultiply chunk index)) (checkpointEnvelopeChunkExtent contract index))))))))
229 count)
230 0))))))
231
232def checkpointEnvelopeVerdictOf =
233 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
234 (lambda unrestricted file : Bytes .
235 (let unrestricted digests = (checkpointEnvelopeIdentityDigests contract) in
236 (let unrestricted identity = (lambda unrestricted index : Nat . (lambda unrestricted at : Nat .
237 (bytes-equal (checkpointEnvelopeSlice file at checkpointEnvelopeDigestBytes)
238 (checkpointEnvelopeSlice digests (naturalMultiply index checkpointEnvelopeDigestBytes) checkpointEnvelopeDigestBytes)))) in
239 (let unrestricted bad = (checkpointEnvelopeFirstBadChunk contract file) in
240 (checkpointEnvelopeIf
241 (naturalEqual (bytes-length file) (checkpointEnvelopeFileBytes contract))
242 (checkpointEnvelopeIf
243 (bytes-equal (checkpointEnvelopeSlice file 0 checkpointEnvelopeFixedBytes) (checkpointEnvelopeFixed contract))
244 (checkpointEnvelopeIf (identity 0 checkpointEnvelopeSchemaAt)
245 (checkpointEnvelopeIf (identity 1 checkpointEnvelopeLearnerAt)
246 (checkpointEnvelopeIf (identity 2 checkpointEnvelopeDataAt)
247 (checkpointEnvelopeIf (identity 3 checkpointEnvelopeSeedAt)
248 (checkpointEnvelopeIf (naturalEqual bad (checkpointEnvelopeChunkCount contract))
249 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeAccepted)
250 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeChunkMismatch bad))
251 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignSeed))
252 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignData))
253 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignLearner))
254 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeForeignSchema))
255 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeLayout))
256 (constructor CheckpointEnvelopeVerdict CheckpointEnvelopeTruncated)))))))
257
258-- the checkpoint a contract's host writes: the header -- fixed bytes,
259-- `updates` and `invocations`, the identity digests, each chunk's digest
260-- -- and the payload
261def checkpointEnvelopeWrite =
262 (lambda unrestricted contract : (family CheckpointEnvelopeContract) .
263 (lambda unrestricted updates : Nat .
264 (lambda unrestricted invocations : Nat .
265 (lambda unrestricted payload : Bytes .
266 (let unrestricted chunk = (checkpointEnvelopeContractChunk contract) in
267 (let unrestricted digests =
268 (bytes-builder-build
269 (app (nat-eliminate
270 (lambda unrestricted remaining : Nat . (pi unrestricted index : Nat . BytesBuilder))
271 (lambda unrestricted index : Nat . (bytes-builder-chunk b""))
272 (lambda unrestricted p : Nat .
273 (lambda unrestricted induction : (pi unrestricted index : Nat . BytesBuilder) .
274 (lambda unrestricted index : Nat .
275 (bytes-builder-append
276 (bytes-builder-chunk (sha256RawDigestOrEmpty (checkpointEnvelopeSlice payload (naturalMultiply chunk index) (checkpointEnvelopeChunkExtent contract index))))
277 (induction (succ index))))))
278 (checkpointEnvelopeChunkCount contract))
279 0)) in
280 (let unrestricted body =
281 (bytes-append (checkpointEnvelopeFixed contract)
282 (bytes-append (checkpointEnvelopeWord updates)
283 (bytes-append (checkpointEnvelopeWord invocations)
284 (bytes-append (checkpointEnvelopeIdentityDigests contract) digests)))) in
285 (bytes-append body
286 (bytes-append (checkpointEnvelopeZeros (naturalSaturatingSubtract (checkpointEnvelopeHeaderBytes contract) (bytes-length body)))
287 payload)))))))))The compiler supplied declaration spans and resolved links from this source snapshot. This page does not assert that this file belongs to a checked closure.