1use serde::Serialize;
2
3use crate::vnext::{
4 BufferUsage, DeviceCommandPhase, DeviceComputePathRequirement, DeviceExecutionPath,
5 ElementType, ExecutionDeterminismWitnessKind, TensorAccess, VNextError,
6};
7
8use super::foundation::invalid_operation;
9use super::SubmissionWaveDeterminismEvidence;
10
11const MAX_ARTIFACT_EXECUTION_ID_BYTES: usize = 128;
12
13#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
14pub struct SubmissionWaveDeterminismArtifactInitializationIdentity {
15 input_sha256: String,
16 rng_sha256: String,
17 initial_state_sha256: String,
18}
19
20impl SubmissionWaveDeterminismArtifactInitializationIdentity {
21 pub fn input_sha256(&self) -> &str {
22 &self.input_sha256
23 }
24
25 pub fn rng_sha256(&self) -> &str {
26 &self.rng_sha256
27 }
28
29 pub fn initial_state_sha256(&self) -> &str {
30 &self.initial_state_sha256
31 }
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
35pub struct SubmissionWaveDeterminismArtifactPhysicalCommand {
36 command_index: u32,
37 node_id: Option<String>,
38 command_phase: String,
39 native_op_id: String,
40 execution_path: String,
41 batching_form: String,
42 participant_count: u32,
43 token_count: u64,
44 compute_dispatch_count: u64,
45 transfer_command_count: u64,
46 reusable_graph_node_count: Option<u64>,
47}
48
49#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
50pub struct SubmissionWaveDeterminismArtifactLogicalCommand {
51 logical_command_ordinal: u32,
52 node_id: String,
53 native_op_id: String,
54 batching_form: String,
55 participant_count: u32,
56 token_count: u64,
57 compute_dispatch_count: u64,
58 transfer_command_count: u64,
59 reusable_graph_node_count: u64,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
63pub struct SubmissionWaveDeterminismArtifactReplayedSegment {
64 physical_command_index: u32,
65 reusable_program_fingerprint: String,
66 reusable_executable_fingerprint: String,
67 logical_commands: Vec<SubmissionWaveDeterminismArtifactLogicalCommand>,
68}
69
70#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
71pub struct SubmissionWaveDeterminismArtifactAttribution {
72 batch_identity_fingerprint: String,
73 submission_fingerprint: String,
74 physical_commands: Vec<SubmissionWaveDeterminismArtifactPhysicalCommand>,
75 replayed_segments: Vec<SubmissionWaveDeterminismArtifactReplayedSegment>,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
79pub struct SubmissionWaveDeterminismArtifactWitness {
80 kind: String,
81 semantic_id: String,
82 node_id: String,
83 resource_id: String,
84 access: String,
85 participant_index: u32,
86 logical_offset_bytes: u64,
87 length_bytes: u64,
88 element_type: String,
89 raw_sha256: String,
90}
91
92impl SubmissionWaveDeterminismArtifactWitness {
93 pub fn kind(&self) -> &str {
94 &self.kind
95 }
96
97 pub fn semantic_id(&self) -> &str {
98 &self.semantic_id
99 }
100
101 pub fn node_id(&self) -> &str {
102 &self.node_id
103 }
104
105 pub fn resource_id(&self) -> &str {
106 &self.resource_id
107 }
108
109 pub fn access(&self) -> &str {
110 &self.access
111 }
112
113 pub const fn participant_index(&self) -> u32 {
114 self.participant_index
115 }
116
117 pub const fn logical_offset_bytes(&self) -> u64 {
118 self.logical_offset_bytes
119 }
120
121 pub const fn length_bytes(&self) -> u64 {
122 self.length_bytes
123 }
124
125 pub fn element_type(&self) -> &str {
126 &self.element_type
127 }
128
129 pub fn raw_sha256(&self) -> &str {
130 &self.raw_sha256
131 }
132}
133
134#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
135pub struct SubmissionWaveDeterminismArtifactExecution {
136 execution_id: String,
137 mode: String,
138 compute_path_requirement: String,
139 reusable_program_fingerprint: Option<String>,
140 declared_eager_boundary_node_ids: Vec<String>,
141 restore_sha256: String,
142 initialization_identity: SubmissionWaveDeterminismArtifactInitializationIdentity,
143 submission_fingerprint: String,
144 receipt_fingerprint: String,
145 attribution: SubmissionWaveDeterminismArtifactAttribution,
146 witnesses: Vec<SubmissionWaveDeterminismArtifactWitness>,
147}
148
149impl SubmissionWaveDeterminismArtifactExecution {
150 pub fn execution_id(&self) -> &str {
151 &self.execution_id
152 }
153
154 pub fn mode(&self) -> &str {
155 &self.mode
156 }
157
158 pub fn compute_path_requirement(&self) -> &str {
159 &self.compute_path_requirement
160 }
161
162 pub fn reusable_program_fingerprint(&self) -> Option<&str> {
163 self.reusable_program_fingerprint.as_deref()
164 }
165
166 pub fn declared_eager_boundary_node_ids(&self) -> &[String] {
167 &self.declared_eager_boundary_node_ids
168 }
169
170 pub fn restore_sha256(&self) -> &str {
171 &self.restore_sha256
172 }
173
174 pub const fn initialization_identity(
175 &self,
176 ) -> &SubmissionWaveDeterminismArtifactInitializationIdentity {
177 &self.initialization_identity
178 }
179
180 pub fn witnesses(&self) -> &[SubmissionWaveDeterminismArtifactWitness] {
181 &self.witnesses
182 }
183
184 pub fn replayed_segments(&self) -> &[SubmissionWaveDeterminismArtifactReplayedSegment] {
185 &self.attribution.replayed_segments
186 }
187}
188
189impl SubmissionWaveDeterminismEvidence {
190 pub fn into_artifact_execution(
191 self,
192 execution_id: impl Into<String>,
193 ) -> Result<SubmissionWaveDeterminismArtifactExecution, VNextError> {
194 let execution_id = execution_id.into();
195 if execution_id.is_empty()
196 || execution_id.len() > MAX_ARTIFACT_EXECUTION_ID_BYTES
197 || !execution_id.bytes().all(|byte| {
198 byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b':' | b'/' | b'-')
199 })
200 {
201 return Err(invalid_operation(
202 "determinism artifact execution id is empty, too long, or non-portable",
203 ));
204 }
205
206 let attribution = self.attribution();
207 let batch_identity = attribution.batch_identity();
208 let physical_commands = attribution
209 .device()
210 .commands()
211 .iter()
212 .map(|command| {
213 let node_id = command
214 .node_index()
215 .map(|index| {
216 usize::try_from(index)
217 .ok()
218 .and_then(|index| batch_identity.node_id_at(index))
219 .map(ToString::to_string)
220 .ok_or_else(|| {
221 invalid_operation(
222 "determinism physical command has an unknown plan node index",
223 )
224 })
225 })
226 .transpose()?;
227 Ok(SubmissionWaveDeterminismArtifactPhysicalCommand {
228 command_index: command.command_index(),
229 node_id,
230 command_phase: command_phase_label(command.command_phase()).to_owned(),
231 native_op_id: command.native_op_id().to_owned(),
232 execution_path: command.execution_path().as_str().to_owned(),
233 batching_form: command.batching_form().as_str().to_owned(),
234 participant_count: command.participant_count(),
235 token_count: command.token_count(),
236 compute_dispatch_count: command.compute_dispatch_count(),
237 transfer_command_count: command.transfer_command_count(),
238 reusable_graph_node_count: command.reusable_graph_node_count(),
239 })
240 })
241 .collect::<Result<Vec<_>, VNextError>>()?;
242 let replayed_segments = attribution
243 .device()
244 .replayed_segments()
245 .iter()
246 .map(|segment| {
247 let logical_commands = segment
248 .logical_commands()
249 .iter()
250 .map(|command| {
251 let node_index = usize::try_from(command.node_index()).map_err(|_| {
252 invalid_operation("determinism replay logical node index exceeds usize")
253 })?;
254 let node_id = batch_identity.node_id_at(node_index).ok_or_else(|| {
255 invalid_operation(
256 "determinism replay logical command has an unknown plan node index",
257 )
258 })?;
259 Ok(SubmissionWaveDeterminismArtifactLogicalCommand {
260 logical_command_ordinal: command.logical_command_ordinal(),
261 node_id: node_id.to_string(),
262 native_op_id: command.native_op_id().to_owned(),
263 batching_form: command.batching_form().as_str().to_owned(),
264 participant_count: command.participant_count(),
265 token_count: command.token_count(),
266 compute_dispatch_count: command.compute_dispatch_count(),
267 transfer_command_count: command.transfer_command_count(),
268 reusable_graph_node_count: command.reusable_graph_node_count(),
269 })
270 })
271 .collect::<Result<Vec<_>, VNextError>>()?;
272 Ok(SubmissionWaveDeterminismArtifactReplayedSegment {
273 physical_command_index: segment.physical_command_index(),
274 reusable_program_fingerprint: segment.program_id().fingerprint(),
275 reusable_executable_fingerprint: segment
276 .reusable_executable_fingerprint()
277 .to_owned(),
278 logical_commands,
279 })
280 })
281 .collect::<Result<Vec<_>, VNextError>>()?;
282
283 let mut witnesses = self
284 .witnesses()
285 .iter()
286 .map(|witness| {
287 let physical_index =
288 usize::try_from(witness.physical_readback_index()).map_err(|_| {
289 invalid_operation(
290 "determinism witness physical readback index exceeds usize",
291 )
292 })?;
293 let physical = self
294 .physical_readbacks()
295 .get(physical_index)
296 .ok_or_else(|| {
297 invalid_operation(
298 "determinism witness references an absent physical readback",
299 )
300 })?;
301 let request = physical.request();
302 let spec = witness.witness();
303 if request.node_id() != spec.node_id()
304 || request.resource_id() != spec.resource_id()
305 || request.participant_index() != witness.physical_participant_index()
306 || request.expected_usage() != spec.location().usage()
307 {
308 return Err(invalid_operation(
309 "determinism witness differs from its exact physical readback request",
310 ));
311 }
312 let (kind, semantic_id, access) = match spec.kind() {
313 ExecutionDeterminismWitnessKind::Output { value_id, .. } => {
314 ("declared_output", value_id.to_string(), TensorAccess::Write)
315 }
316 ExecutionDeterminismWitnessKind::StateEffect {
317 state_id, access, ..
318 } => ("state_effect", state_id.to_string(), *access),
319 };
320 if kind == "state_effect" && request.expected_usage() != BufferUsage::State {
321 return Err(invalid_operation(
322 "determinism state witness readback does not use state backing",
323 ));
324 }
325 let logical_range = witness.logical_range();
326 if logical_range.length_bytes() != request.output_layout().byte_len()? {
327 return Err(invalid_operation(
328 "determinism semantic witness length differs from its physical readback",
329 ));
330 }
331 Ok(SubmissionWaveDeterminismArtifactWitness {
332 kind: kind.to_owned(),
333 semantic_id,
334 node_id: spec.node_id().to_string(),
335 resource_id: spec.resource_id().to_string(),
336 access: tensor_access_label(access).to_owned(),
337 participant_index: witness.participant_index(),
338 logical_offset_bytes: logical_range.logical_offset_bytes(),
339 length_bytes: logical_range.length_bytes(),
340 element_type: element_type_label(spec.element_type()).to_owned(),
341 raw_sha256: physical.raw_sha256().to_owned(),
342 })
343 })
344 .collect::<Result<Vec<_>, VNextError>>()?;
345 witnesses.sort_by(|left, right| {
346 (
347 &left.kind,
348 &left.semantic_id,
349 &left.node_id,
350 &left.resource_id,
351 &left.access,
352 left.participant_index,
353 left.logical_offset_bytes,
354 left.length_bytes,
355 &left.element_type,
356 )
357 .cmp(&(
358 &right.kind,
359 &right.semantic_id,
360 &right.node_id,
361 &right.resource_id,
362 &right.access,
363 right.participant_index,
364 right.logical_offset_bytes,
365 right.length_bytes,
366 &right.element_type,
367 ))
368 });
369
370 let initialization = self.initialization_identity();
371 let reusable_program_fingerprint = self.reusable_program_fingerprint();
372 let declared_eager_boundary_node_ids = self
373 .declared_eager_boundary_node_ids()
374 .iter()
375 .map(ToString::to_string)
376 .collect();
377 Ok(SubmissionWaveDeterminismArtifactExecution {
378 execution_id,
379 mode: match self.expected_execution_path() {
380 DeviceExecutionPath::Eager => "eager",
381 DeviceExecutionPath::Replayed => "replay",
382 }
383 .to_owned(),
384 compute_path_requirement: compute_path_requirement_label(
385 self.expected_compute_path_requirement(),
386 )
387 .to_owned(),
388 reusable_program_fingerprint,
389 declared_eager_boundary_node_ids,
390 restore_sha256: self.restore_fingerprint().to_owned(),
391 initialization_identity: SubmissionWaveDeterminismArtifactInitializationIdentity {
392 input_sha256: initialization.input_sha256().to_owned(),
393 rng_sha256: initialization.rng_sha256().to_owned(),
394 initial_state_sha256: initialization.initial_state_sha256().to_owned(),
395 },
396 submission_fingerprint: self.submission_receipt_fingerprint().to_owned(),
397 receipt_fingerprint: self.terminal_receipt_fingerprint().to_owned(),
398 attribution: SubmissionWaveDeterminismArtifactAttribution {
399 batch_identity_fingerprint: batch_identity.fingerprint().to_owned(),
400 submission_fingerprint: attribution.submission_fingerprint().to_owned(),
401 physical_commands,
402 replayed_segments,
403 },
404 witnesses,
405 })
406 }
407}
408
409const fn compute_path_requirement_label(value: DeviceComputePathRequirement) -> &'static str {
410 match value {
411 DeviceComputePathRequirement::Adaptive => "adaptive",
412 DeviceComputePathRequirement::EagerOnly => "eager_only",
413 DeviceComputePathRequirement::ReplayedOnly => "replayed_only",
414 DeviceComputePathRequirement::ReplayedWithDeclaredEagerBoundaries => {
415 "replayed_with_declared_eager_boundaries"
416 }
417 }
418}
419
420const fn command_phase_label(value: DeviceCommandPhase) -> &'static str {
421 match value {
422 DeviceCommandPhase::Initialization => "initialization",
423 DeviceCommandPhase::DynamicBinding => "dynamic_binding",
424 DeviceCommandPhase::Compute => "compute",
425 DeviceCommandPhase::ResultBinding => "result_binding",
426 }
427}
428
429const fn tensor_access_label(value: TensorAccess) -> &'static str {
430 match value {
431 TensorAccess::Read => "read",
432 TensorAccess::Write => "write",
433 TensorAccess::ReadWrite => "read_write",
434 }
435}
436
437const fn element_type_label(value: ElementType) -> &'static str {
438 match value {
439 ElementType::Bool => "bool",
440 ElementType::U8 => "u8",
441 ElementType::U32 => "u32",
442 ElementType::I8 => "i8",
443 ElementType::I32 => "i32",
444 ElementType::F16 => "f16",
445 ElementType::Bf16 => "bf16",
446 ElementType::F32 => "f32",
447 }
448}