1use std::collections::{BTreeMap, BTreeSet};
2
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5
6use super::{
7 invalid_plan, AllocationLifetime, BufferUsage, ContractVersion, DynamicResourceDemand,
8 DynamicResourceDescriptor, ElementType, ExecutionPlan, NodeId, PlanHash, ProgramValueId,
9 ProviderId, ResolvedValueBinding, ResolvedValueRole, ResourceId, StateId, TensorAccess,
10 TokenSpanWork, VNextError, WeightId,
11};
12use crate::vnext::{
13 BatchParticipantTokenRange, ProviderExecutionContractFingerprint, ProviderReplayEquivalence,
14};
15
16pub const EXECUTION_DETERMINISM_WITNESS_VERSION: ContractVersion = ContractVersion::new(4, 0);
17const MAX_EXECUTION_DETERMINISM_WITNESS_WIRE_BYTES: usize = 64 * 1024 * 1024;
18const MAX_EXECUTION_DETERMINISM_WITNESS_NODES: usize = 65_536;
19const MAX_EXECUTION_DETERMINISM_INITIALIZATIONS: usize = 262_144;
20const MAX_EXECUTION_DETERMINISM_WITNESSES: usize = 1_048_576;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
28#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
29pub enum ExecutionDeterminismValueExtent {
30 Fixed,
31 ImmediateTokenSpan {
32 bytes_per_token: u64,
33 maximum_tokens: u64,
34 },
35 ActiveTokenPrefix {
36 bytes_per_token: u64,
37 maximum_tokens: u64,
38 maximum_storage_bytes: u64,
39 },
40}
41
42#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
48#[serde(deny_unknown_fields)]
49pub struct ExecutionDeterminismValueLocation {
50 node_id: NodeId,
51 value_id: ProgramValueId,
52 role: ResolvedValueRole,
53 ordinal: u32,
54 usage: BufferUsage,
55 storage_component_ordinal: u32,
56 storage_component_id: Option<WeightId>,
57 resource_id: ResourceId,
58 logical_offset_bytes: u64,
59 declared_length_bytes: u64,
60 element_type: ElementType,
61 extent: ExecutionDeterminismValueExtent,
62}
63
64impl ExecutionDeterminismValueLocation {
65 fn from_binding(
66 node: &super::PlanNode,
67 binding: &ResolvedValueBinding,
68 dynamic_descriptors: &[DynamicResourceDescriptor],
69 ) -> Result<Vec<Self>, VNextError> {
70 let components = binding.storage().components();
71 let single_component_logical_length = (components.len() == 1)
72 .then(|| binding.tensor().minimum_storage_bytes())
73 .transpose()?;
74 components
75 .iter()
76 .enumerate()
77 .map(|(component_ordinal, component)| {
78 let declared_length_bytes =
82 single_component_logical_length.unwrap_or(component.length_bytes());
83 if declared_length_bytes == 0
84 || declared_length_bytes > component.length_bytes()
85 || component
86 .offset_bytes()
87 .checked_add(declared_length_bytes)
88 .is_none()
89 {
90 return Err(invalid_plan(format!(
91 "node `{}` determinism value `{}` component {component_ordinal} has an invalid physical range",
92 node.id(),
93 binding.value_id()
94 )));
95 }
96 let immediate_projection = node
97 .work()
98 .token_projection(binding.role(), binding.ordinal())
99 .map(|projection| {
100 let canonical_extent = projection.canonical_extent();
101 if canonical_extent == 0
102 || declared_length_bytes % canonical_extent != 0
103 {
104 return Err(invalid_plan(format!(
105 "node `{}` determinism value `{}` component {component_ordinal} has a non-integral token projection",
106 node.id(),
107 binding.value_id()
108 )));
109 }
110 declared_length_bytes
111 .checked_div(canonical_extent)
112 .filter(|bytes| *bytes > 0)
113 .map(|bytes_per_token| (bytes_per_token, canonical_extent))
114 .ok_or_else(|| {
115 invalid_plan(format!(
116 "node `{}` determinism value `{}` component {component_ordinal} has an invalid token projection",
117 node.id(),
118 binding.value_id()
119 ))
120 })
121 })
122 .transpose()?;
123 let dynamic_descriptor = dynamic_descriptors
124 .iter()
125 .find(|descriptor| descriptor.base_resource_id() == component.resource_id());
126 let descriptor_tokens = dynamic_descriptor
127 .map(|descriptor| {
128 if descriptor.base_resource_id() != component.resource_id()
129 || descriptor.usage() != binding.usage()
130 || descriptor.element_type() != component.element_type()
131 {
132 return Err(invalid_plan(format!(
133 "node `{}` determinism value `{}` differs from its dynamic resource descriptor",
134 node.id(),
135 binding.value_id()
136 )));
137 }
138 Ok(match descriptor.demand() {
139 DynamicResourceDemand::Tokens {
140 bytes_per_token,
141 maximum_tokens,
142 } => Some((*bytes_per_token, *maximum_tokens)),
143 _ => None,
144 })
145 })
146 .transpose()?
147 .flatten();
148 let extent = match (binding.usage(), immediate_projection, descriptor_tokens) {
149 (
150 BufferUsage::State,
151 None,
152 Some((bytes_per_token, maximum_tokens)),
153 ) => {
154 let maximum_storage_bytes =
155 dynamic_descriptor
156 .expect("token demand came from this descriptor")
157 .theoretical_maximum_request_bytes()?;
158 if components.len() != 1
159 || component.offset_bytes() != 0
160 || bytes_per_token < declared_length_bytes
161 || bytes_per_token % component.element_type().size_bytes() != 0
162 || maximum_storage_bytes < bytes_per_token
163 {
164 return Err(invalid_plan(format!(
165 "node `{}` token-scaled state `{}` has no lossless active-prefix projection",
166 node.id(),
167 binding.value_id()
168 )));
169 }
170 ExecutionDeterminismValueExtent::ActiveTokenPrefix {
171 bytes_per_token,
172 maximum_tokens,
173 maximum_storage_bytes,
174 }
175 }
176 (
177 BufferUsage::State,
178 Some(_),
179 Some(_) | None,
180 ) => {
181 return Err(invalid_plan(format!(
182 "node `{}` state `{}` cannot use a transient token projection",
183 node.id(),
184 binding.value_id()
185 )));
186 }
187 (
188 _,
189 Some((projected_bytes, projected_tokens)),
190 descriptor_tokens,
191 ) => {
192 let (bytes_per_token, maximum_tokens) =
193 descriptor_tokens.unwrap_or((projected_bytes, projected_tokens));
194 if bytes_per_token != projected_bytes
195 || bytes_per_token % component.element_type().size_bytes() != 0
196 {
197 return Err(invalid_plan(format!(
198 "node `{}` determinism value `{}` has inconsistent token extent evidence: projected_bytes={projected_bytes}, projected_tokens={projected_tokens}, descriptor_bytes={bytes_per_token}, descriptor_maximum_tokens={maximum_tokens}",
199 node.id(),
200 binding.value_id()
201 )));
202 }
203 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
204 bytes_per_token,
205 maximum_tokens,
206 }
207 }
208 (_, None, Some(_)) => {
209 return Err(invalid_plan(format!(
210 "node `{}` token-scaled value `{}` lacks a typed work projection",
211 node.id(),
212 binding.value_id()
213 )));
214 }
215 (_, None, None) => ExecutionDeterminismValueExtent::Fixed,
216 };
217 Ok(Self {
218 node_id: node.id().clone(),
219 value_id: binding.value_id().clone(),
220 role: binding.role(),
221 ordinal: binding.ordinal(),
222 usage: binding.usage(),
223 storage_component_ordinal: u32::try_from(component_ordinal).map_err(|_| {
224 invalid_plan("determinism value component ordinal exceeds u32")
225 })?,
226 storage_component_id: component.component_id().cloned(),
227 resource_id: component.resource_id().clone(),
228 logical_offset_bytes: component.offset_bytes(),
229 declared_length_bytes,
230 element_type: component.element_type(),
231 extent,
232 })
233 })
234 .collect()
235 }
236
237 pub fn node_id(&self) -> &NodeId {
238 &self.node_id
239 }
240
241 pub fn value_id(&self) -> &ProgramValueId {
242 &self.value_id
243 }
244
245 pub const fn role(&self) -> ResolvedValueRole {
246 self.role
247 }
248
249 pub const fn ordinal(&self) -> u32 {
250 self.ordinal
251 }
252
253 pub const fn usage(&self) -> BufferUsage {
254 self.usage
255 }
256
257 pub const fn storage_component_ordinal(&self) -> u32 {
258 self.storage_component_ordinal
259 }
260
261 pub fn storage_component_id(&self) -> Option<&WeightId> {
262 self.storage_component_id.as_ref()
263 }
264
265 pub fn resource_id(&self) -> &ResourceId {
266 &self.resource_id
267 }
268
269 pub const fn logical_offset_bytes(&self) -> u64 {
270 self.logical_offset_bytes
271 }
272
273 pub const fn declared_length_bytes(&self) -> u64 {
274 self.declared_length_bytes
275 }
276
277 pub const fn element_type(&self) -> ElementType {
278 self.element_type
279 }
280
281 pub const fn extent(&self) -> ExecutionDeterminismValueExtent {
282 self.extent
283 }
284
285 pub fn maximum_bound_length_bytes(&self) -> Result<u64, VNextError> {
286 match self.extent {
287 ExecutionDeterminismValueExtent::Fixed => Ok(self.declared_length_bytes),
288 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
289 bytes_per_token,
290 maximum_tokens,
291 } => bytes_per_token
292 .checked_mul(maximum_tokens)
293 .ok_or_else(|| invalid_plan("determinism value maximum byte extent overflows")),
294 ExecutionDeterminismValueExtent::ActiveTokenPrefix {
295 maximum_storage_bytes,
296 ..
297 } => Ok(maximum_storage_bytes),
298 }
299 }
300
301 pub fn bound_length_bytes(
302 &self,
303 token_span: &TokenSpanWork,
304 token_range: &BatchParticipantTokenRange,
305 ) -> Result<u64, VNextError> {
306 if token_range.immediate_tokens() != token_span.immediate_tokens()
307 || token_range.source_token_range() != token_span.immediate_token_range()
308 {
309 return Err(invalid_plan(
310 "determinism value work span differs from its participant token range",
311 ));
312 }
313 self.bound_length_bytes_for_source_end(token_span, token_range.source_token_range().end)
314 }
315
316 fn bound_length_bytes_for_source_end(
317 &self,
318 token_span: &TokenSpanWork,
319 source_end_tokens: u64,
320 ) -> Result<u64, VNextError> {
321 let bytes = match self.extent {
322 ExecutionDeterminismValueExtent::Fixed => self.declared_length_bytes,
323 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
324 bytes_per_token, ..
325 } => bytes_per_token
326 .checked_mul(token_span.immediate_tokens())
327 .ok_or_else(|| invalid_plan("determinism value active byte extent overflows"))?,
328 ExecutionDeterminismValueExtent::ActiveTokenPrefix {
329 bytes_per_token, ..
330 } => bytes_per_token
331 .checked_mul(source_end_tokens)
332 .ok_or_else(|| invalid_plan("determinism state prefix byte extent overflows"))?,
333 };
334 if bytes == 0 || bytes > self.maximum_bound_length_bytes()? {
335 return Err(invalid_plan(
336 "determinism value active byte extent exceeds its immutable bound",
337 ));
338 }
339 Ok(bytes)
340 }
341}
342
343#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
344#[serde(tag = "kind", rename_all = "snake_case")]
345pub enum ExecutionDeterminismWitnessKind {
346 Output {
347 value_id: ProgramValueId,
348 output_ordinal: u32,
349 },
350 StateEffect {
351 state_id: StateId,
352 state_value_id: ProgramValueId,
353 lifetime: AllocationLifetime,
354 access: TensorAccess,
355 },
356}
357
358#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
359#[serde(deny_unknown_fields)]
360pub struct ExecutionDeterminismWitnessSpec {
361 provider_id: ProviderId,
362 provider_implementation_fingerprint: String,
363 provider_execution_contract_fingerprint: ProviderExecutionContractFingerprint,
364 kind: ExecutionDeterminismWitnessKind,
365 location: ExecutionDeterminismValueLocation,
366}
367
368impl ExecutionDeterminismWitnessSpec {
369 fn from_binding(
370 node: &super::PlanNode,
371 kind: ExecutionDeterminismWitnessKind,
372 binding: &ResolvedValueBinding,
373 dynamic_descriptors: &[DynamicResourceDescriptor],
374 ) -> Result<Vec<Self>, VNextError> {
375 ExecutionDeterminismValueLocation::from_binding(node, binding, dynamic_descriptors)?
376 .into_iter()
377 .map(|location| {
378 Ok(Self {
379 provider_id: node.selection().selected_provider().clone(),
380 provider_implementation_fingerprint: node
381 .provider_implementation_fingerprint()
382 .to_owned(),
383 provider_execution_contract_fingerprint: node
384 .provider_execution_semantics()
385 .contract_fingerprint(),
386 kind: kind.clone(),
387 location,
388 })
389 })
390 .collect()
391 }
392
393 pub fn node_id(&self) -> &NodeId {
394 self.location.node_id()
395 }
396
397 pub fn provider_id(&self) -> &ProviderId {
398 &self.provider_id
399 }
400
401 pub fn provider_implementation_fingerprint(&self) -> &str {
402 &self.provider_implementation_fingerprint
403 }
404
405 pub const fn provider_execution_contract_fingerprint(
406 &self,
407 ) -> ProviderExecutionContractFingerprint {
408 self.provider_execution_contract_fingerprint
409 }
410
411 pub fn kind(&self) -> &ExecutionDeterminismWitnessKind {
412 &self.kind
413 }
414
415 pub fn location(&self) -> &ExecutionDeterminismValueLocation {
416 &self.location
417 }
418
419 pub const fn storage_component_ordinal(&self) -> u32 {
420 self.location.storage_component_ordinal()
421 }
422
423 pub fn storage_component_id(&self) -> Option<&WeightId> {
424 self.location.storage_component_id()
425 }
426
427 pub fn resource_id(&self) -> &ResourceId {
428 self.location.resource_id()
429 }
430
431 pub const fn logical_offset_bytes(&self) -> u64 {
432 self.location.logical_offset_bytes()
433 }
434
435 pub const fn declared_length_bytes(&self) -> u64 {
436 self.location.declared_length_bytes()
437 }
438
439 pub const fn element_type(&self) -> ElementType {
440 self.location.element_type()
441 }
442
443 pub const fn extent(&self) -> ExecutionDeterminismValueExtent {
444 self.location.extent()
445 }
446
447 pub fn maximum_bound_length_bytes(&self) -> Result<u64, VNextError> {
448 self.location.maximum_bound_length_bytes()
449 }
450
451 pub fn bound_length_bytes(
452 &self,
453 token_span: &TokenSpanWork,
454 token_range: &BatchParticipantTokenRange,
455 ) -> Result<u64, VNextError> {
456 self.location.bound_length_bytes(token_span, token_range)
457 }
458}
459
460#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
461#[serde(tag = "kind", rename_all = "snake_case")]
462pub enum ExecutionDeterminismInitializationKind {
463 ExternalInput {
464 value_id: ProgramValueId,
465 },
466 State {
467 state_id: StateId,
468 state_value_id: ProgramValueId,
469 lifetime: AllocationLifetime,
470 access: TensorAccess,
471 },
472}
473
474#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
477#[serde(deny_unknown_fields)]
478pub struct ExecutionDeterminismInitializationSpec {
479 kind: ExecutionDeterminismInitializationKind,
480 location: ExecutionDeterminismValueLocation,
481 consumer_node_ids: Vec<NodeId>,
482}
483
484impl ExecutionDeterminismInitializationSpec {
485 pub fn kind(&self) -> &ExecutionDeterminismInitializationKind {
486 &self.kind
487 }
488
489 pub fn location(&self) -> &ExecutionDeterminismValueLocation {
490 &self.location
491 }
492
493 pub fn consumer_node_ids(&self) -> &[NodeId] {
494 &self.consumer_node_ids
495 }
496}
497
498#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
499#[serde(deny_unknown_fields)]
500pub struct ProviderDeterminismCoverageRequirement {
501 provider_id: ProviderId,
502 provider_implementation_fingerprint: String,
503 provider_execution_contract_fingerprint: ProviderExecutionContractFingerprint,
504 node_ids: Vec<NodeId>,
505}
506
507impl ProviderDeterminismCoverageRequirement {
508 pub fn provider_id(&self) -> &ProviderId {
509 &self.provider_id
510 }
511
512 pub fn provider_implementation_fingerprint(&self) -> &str {
513 &self.provider_implementation_fingerprint
514 }
515
516 pub const fn provider_execution_contract_fingerprint(
517 &self,
518 ) -> ProviderExecutionContractFingerprint {
519 self.provider_execution_contract_fingerprint
520 }
521
522 pub fn node_ids(&self) -> &[NodeId] {
523 &self.node_ids
524 }
525}
526
527#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
528#[serde(deny_unknown_fields)]
529pub struct ExecutionDeterminismWitnessPlan {
530 schema_version: ContractVersion,
531 plan_hash: PlanHash,
532 node_ids: Vec<NodeId>,
533 replay_provider_requirements: Vec<ProviderDeterminismCoverageRequirement>,
534 initializations: Vec<ExecutionDeterminismInitializationSpec>,
535 witnesses: Vec<ExecutionDeterminismWitnessSpec>,
536}
537
538impl ExecutionDeterminismWitnessPlan {
539 pub const fn schema_version(&self) -> ContractVersion {
540 self.schema_version
541 }
542
543 pub fn plan_hash(&self) -> &PlanHash {
544 &self.plan_hash
545 }
546
547 pub fn node_ids(&self) -> &[NodeId] {
548 &self.node_ids
549 }
550
551 pub fn replay_provider_requirements(&self) -> &[ProviderDeterminismCoverageRequirement] {
552 &self.replay_provider_requirements
553 }
554
555 pub fn initializations(&self) -> &[ExecutionDeterminismInitializationSpec] {
556 &self.initializations
557 }
558
559 pub fn witnesses(&self) -> &[ExecutionDeterminismWitnessSpec] {
560 &self.witnesses
561 }
562
563 pub fn to_json(&self) -> Result<Vec<u8>, VNextError> {
564 self.validate_shape()?;
565 serde_json::to_vec_pretty(self).map_err(|error| VNextError::Serialization {
566 context: "serialize execution determinism witness plan",
567 message: error.to_string(),
568 })
569 }
570
571 pub fn fingerprint(&self) -> Result<String, VNextError> {
572 self.validate_shape()?;
573 let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
574 context: "fingerprint execution determinism witness plan",
575 message: error.to_string(),
576 })?;
577 Ok(format!("{:x}", Sha256::digest(bytes)))
578 }
579
580 pub fn decode_untrusted(bytes: &[u8]) -> Result<Self, VNextError> {
581 if bytes.len() > MAX_EXECUTION_DETERMINISM_WITNESS_WIRE_BYTES {
582 return Err(invalid_plan(
583 "execution determinism witness plan exceeds its wire bound",
584 ));
585 }
586 let plan =
587 serde_json::from_slice::<Self>(bytes).map_err(|error| VNextError::Serialization {
588 context: "decode execution determinism witness plan",
589 message: error.to_string(),
590 })?;
591 plan.validate_shape()?;
592 Ok(plan)
593 }
594
595 pub(super) fn validate_shape(&self) -> Result<(), VNextError> {
596 if self.schema_version != EXECUTION_DETERMINISM_WITNESS_VERSION
597 || self.node_ids.is_empty()
598 || self.node_ids.len() > MAX_EXECUTION_DETERMINISM_WITNESS_NODES
599 || self.initializations.len() > MAX_EXECUTION_DETERMINISM_INITIALIZATIONS
600 || self.witnesses.is_empty()
601 || self.witnesses.len() > MAX_EXECUTION_DETERMINISM_WITNESSES
602 {
603 return Err(invalid_plan(
604 "execution determinism witness plan identity or cardinality is invalid",
605 ));
606 }
607 let node_ids = self.node_ids.iter().collect::<BTreeSet<_>>();
608 if node_ids.len() != self.node_ids.len() {
609 return Err(invalid_plan(
610 "execution determinism witness plan node scope is not unique",
611 ));
612 }
613
614 let mut replay_provider_ids = BTreeSet::new();
615 let mut replay_nodes = BTreeSet::new();
616 for requirement in &self.replay_provider_requirements {
617 if !replay_provider_ids.insert(&requirement.provider_id)
618 || !super::is_canonical_sha256(&requirement.provider_implementation_fingerprint)
619 || requirement.node_ids.is_empty()
620 || requirement.node_ids.iter().collect::<BTreeSet<_>>().len()
621 != requirement.node_ids.len()
622 || requirement
623 .node_ids
624 .iter()
625 .any(|node_id| !node_ids.contains(node_id))
626 {
627 return Err(invalid_plan(
628 "execution determinism replay provider denominator is invalid",
629 ));
630 }
631 for node_id in &requirement.node_ids {
632 if !replay_nodes.insert(node_id) {
633 return Err(invalid_plan(
634 "execution determinism replay node belongs to multiple providers",
635 ));
636 }
637 }
638 }
639
640 let mut initialization_rows = BTreeSet::new();
641 for initialization in &self.initializations {
642 let location = initialization.location();
643 if initialization.consumer_node_ids.is_empty()
644 || initialization
645 .consumer_node_ids
646 .iter()
647 .collect::<BTreeSet<_>>()
648 .len()
649 != initialization.consumer_node_ids.len()
650 || initialization
651 .consumer_node_ids
652 .iter()
653 .any(|node_id| !node_ids.contains(node_id))
654 || !node_ids.contains(location.node_id())
655 {
656 return Err(invalid_plan(
657 "execution determinism initialization scope is invalid",
658 ));
659 }
660 validate_determinism_location(location)?;
661 let row = serde_json::to_string(initialization).map_err(|error| {
662 VNextError::Serialization {
663 context: "canonicalize execution determinism initialization",
664 message: error.to_string(),
665 }
666 })?;
667 if !initialization_rows.insert(row) {
668 return Err(invalid_plan(
669 "execution determinism initialization denominator contains duplicates",
670 ));
671 }
672 }
673
674 let mut witness_rows = BTreeSet::new();
675 let mut output_nodes = BTreeSet::new();
676 let mut replay_witness_nodes = BTreeSet::new();
677 for witness in &self.witnesses {
678 if !node_ids.contains(witness.node_id())
679 || !super::is_canonical_sha256(&witness.provider_implementation_fingerprint)
680 {
681 return Err(invalid_plan(
682 "execution determinism witness identity is invalid",
683 ));
684 }
685 validate_determinism_location(witness.location())?;
686 if matches!(
687 witness.kind(),
688 ExecutionDeterminismWitnessKind::Output { .. }
689 ) {
690 output_nodes.insert(witness.node_id());
691 }
692 if replay_provider_ids.contains(witness.provider_id()) {
693 replay_witness_nodes.insert(witness.node_id());
694 }
695 let row =
696 serde_json::to_string(witness).map_err(|error| VNextError::Serialization {
697 context: "canonicalize execution determinism witness",
698 message: error.to_string(),
699 })?;
700 if !witness_rows.insert(row) {
701 return Err(invalid_plan(
702 "execution determinism witness denominator contains duplicates",
703 ));
704 }
705 }
706 if output_nodes != node_ids || replay_witness_nodes != replay_nodes {
707 return Err(invalid_plan(
708 "execution determinism witness denominator does not cover its node scope",
709 ));
710 }
711 Ok(())
712 }
713}
714
715fn validate_determinism_location(
716 location: &ExecutionDeterminismValueLocation,
717) -> Result<(), VNextError> {
718 let element_bytes = location.element_type().size_bytes();
719 if location.declared_length_bytes == 0
720 || location.logical_offset_bytes % element_bytes != 0
721 || location.declared_length_bytes % element_bytes != 0
722 || location
723 .logical_offset_bytes
724 .checked_add(location.declared_length_bytes)
725 .is_none()
726 {
727 return Err(invalid_plan(
728 "execution determinism value location has an invalid byte range",
729 ));
730 }
731 match location.extent {
732 ExecutionDeterminismValueExtent::Fixed => {}
733 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
734 bytes_per_token,
735 maximum_tokens,
736 } => {
737 if bytes_per_token == 0
738 || maximum_tokens == 0
739 || bytes_per_token % element_bytes != 0
740 || location.declared_length_bytes % bytes_per_token != 0
741 || location.maximum_bound_length_bytes()? == 0
742 {
743 return Err(invalid_plan(
744 "execution determinism immediate-token location has an invalid dynamic extent",
745 ));
746 }
747 }
748 ExecutionDeterminismValueExtent::ActiveTokenPrefix {
749 bytes_per_token,
750 maximum_tokens,
751 maximum_storage_bytes,
752 } => {
753 if bytes_per_token == 0
754 || maximum_tokens == 0
755 || bytes_per_token % element_bytes != 0
756 || maximum_storage_bytes < bytes_per_token
757 || maximum_storage_bytes < location.declared_length_bytes
758 {
759 return Err(invalid_plan(
760 "execution determinism active-prefix location has an invalid dynamic extent",
761 ));
762 }
763 }
764 }
765 Ok(())
766}
767
768impl ExecutionPlan {
769 pub fn determinism_witness_plan(&self) -> Result<ExecutionDeterminismWitnessPlan, VNextError> {
773 let node_ids = self
774 .payload()
775 .nodes()
776 .iter()
777 .map(|node| node.id().clone())
778 .collect::<Vec<_>>();
779 self.determinism_witness_plan_for_nodes(&node_ids)
780 }
781
782 pub fn determinism_witness_plan_for_nodes(
787 &self,
788 node_ids: &[NodeId],
789 ) -> Result<ExecutionDeterminismWitnessPlan, VNextError> {
790 if node_ids.is_empty() || node_ids.iter().collect::<BTreeSet<_>>().len() != node_ids.len() {
791 return Err(invalid_plan(
792 "execution determinism node scope must be non-empty and unique",
793 ));
794 }
795 let requested = node_ids.iter().collect::<BTreeSet<_>>();
796 let nodes = self
797 .payload()
798 .nodes()
799 .iter()
800 .filter(|node| requested.contains(node.id()))
801 .collect::<Vec<_>>();
802 let canonical_node_ids = nodes
803 .iter()
804 .map(|node| node.id().clone())
805 .collect::<Vec<_>>();
806 if canonical_node_ids != node_ids {
807 return Err(invalid_plan(
808 "execution determinism node scope is unknown or not in canonical plan order",
809 ));
810 }
811 let produced_values = nodes
812 .iter()
813 .flat_map(|node| {
814 node.values().iter().filter_map(|binding| {
815 (binding.role() == ResolvedValueRole::Output)
816 .then(|| binding.value_id().clone())
817 })
818 })
819 .collect::<BTreeSet<_>>();
820 let dynamic_descriptors = self.payload().memory().dynamic_descriptors();
821 let mut external_inputs = BTreeMap::<
822 (
823 ProgramValueId,
824 ResourceId,
825 u64,
826 u64,
827 ElementType,
828 u32,
829 Option<WeightId>,
830 ExecutionDeterminismValueExtent,
831 ),
832 (ExecutionDeterminismValueLocation, BTreeSet<NodeId>),
833 >::new();
834 let mut initial_state = BTreeMap::<
835 (
836 StateId,
837 ProgramValueId,
838 AllocationLifetime,
839 ResourceId,
840 u64,
841 u64,
842 ElementType,
843 u32,
844 Option<WeightId>,
845 ExecutionDeterminismValueExtent,
846 ),
847 (
848 ExecutionDeterminismValueLocation,
849 TensorAccess,
850 BTreeSet<NodeId>,
851 ),
852 >::new();
853
854 for node in &nodes {
855 for binding in node.values().iter().filter(|binding| {
856 binding.role() == ResolvedValueRole::Input
857 && binding.usage() == BufferUsage::Activations
858 && matches!(
859 binding.access(),
860 TensorAccess::Read | TensorAccess::ReadWrite
861 )
862 && !produced_values.contains(binding.value_id())
863 }) {
864 for location in ExecutionDeterminismValueLocation::from_binding(
865 node,
866 binding,
867 dynamic_descriptors,
868 )? {
869 let key = (
870 binding.value_id().clone(),
871 location.resource_id().clone(),
872 location.logical_offset_bytes(),
873 location.declared_length_bytes(),
874 location.element_type(),
875 location.storage_component_ordinal(),
876 location.storage_component_id().cloned(),
877 location.extent(),
878 );
879 let (_, consumers) = external_inputs
880 .entry(key)
881 .or_insert_with(|| (location, BTreeSet::new()));
882 consumers.insert(node.id().clone());
883 }
884 }
885
886 for effect in node.state_effects().iter().filter(|effect| {
887 matches!(
888 effect.access(),
889 TensorAccess::Read | TensorAccess::ReadWrite
890 )
891 }) {
892 let mut matched_read_binding = false;
893 for binding in node.values().iter().filter(|binding| {
894 binding.value_id() == effect.state_value_id()
895 && binding.usage() == BufferUsage::State
896 && matches!(
897 binding.access(),
898 TensorAccess::Read | TensorAccess::ReadWrite
899 )
900 }) {
901 for location in ExecutionDeterminismValueLocation::from_binding(
902 node,
903 binding,
904 dynamic_descriptors,
905 )? {
906 matched_read_binding = true;
907 let key = (
908 effect.state_id().clone(),
909 effect.state_value_id().clone(),
910 effect.lifetime(),
911 location.resource_id().clone(),
912 location.logical_offset_bytes(),
913 location.declared_length_bytes(),
914 location.element_type(),
915 location.storage_component_ordinal(),
916 location.storage_component_id().cloned(),
917 location.extent(),
918 );
919 let (_, access, consumers) = initial_state
920 .entry(key)
921 .or_insert_with(|| (location, effect.access(), BTreeSet::new()));
922 if effect.access() == TensorAccess::ReadWrite {
923 *access = TensorAccess::ReadWrite;
924 }
925 consumers.insert(node.id().clone());
926 }
927 }
928 if !matched_read_binding {
929 return Err(invalid_plan(format!(
930 "node `{}` readable state `{}` has no exact determinism initialization closure",
931 node.id(),
932 effect.state_id()
933 )));
934 }
935 }
936 }
937
938 let mut initializations =
939 Vec::with_capacity(external_inputs.len().saturating_add(initial_state.len()));
940 initializations.extend(external_inputs.into_iter().map(
941 |((value_id, _, _, _, _, _, _, _), (location, consumer_node_ids))| {
942 ExecutionDeterminismInitializationSpec {
943 kind: ExecutionDeterminismInitializationKind::ExternalInput { value_id },
944 location,
945 consumer_node_ids: consumer_node_ids.into_iter().collect(),
946 }
947 },
948 ));
949 initializations.extend(initial_state.into_iter().map(
950 |(
951 (state_id, state_value_id, lifetime, _, _, _, _, _, _, _),
952 (location, access, consumer_node_ids),
953 )| {
954 ExecutionDeterminismInitializationSpec {
955 kind: ExecutionDeterminismInitializationKind::State {
956 state_id,
957 state_value_id,
958 lifetime,
959 access,
960 },
961 location,
962 consumer_node_ids: consumer_node_ids.into_iter().collect(),
963 }
964 },
965 ));
966
967 let mut witnesses = Vec::new();
968 let mut replay_providers = BTreeMap::<
969 ProviderId,
970 (
971 String,
972 ProviderExecutionContractFingerprint,
973 BTreeSet<NodeId>,
974 ),
975 >::new();
976
977 for node in &nodes {
978 let semantics = node.provider_execution_semantics();
979 if semantics.replay_equivalence() == ProviderReplayEquivalence::BitwiseEagerEquivalent {
980 match replay_providers.entry(node.selection().selected_provider().clone()) {
981 std::collections::btree_map::Entry::Vacant(entry) => {
982 entry.insert((
983 node.provider_implementation_fingerprint().to_owned(),
984 semantics.contract_fingerprint(),
985 BTreeSet::from([node.id().clone()]),
986 ));
987 }
988 std::collections::btree_map::Entry::Occupied(mut entry) => {
989 let (implementation, contract, nodes) = entry.get_mut();
990 if implementation != node.provider_implementation_fingerprint()
991 || *contract != semantics.contract_fingerprint()
992 {
993 return Err(invalid_plan(format!(
994 "provider `{}` has inconsistent determinism identity in one plan",
995 node.selection().selected_provider()
996 )));
997 }
998 nodes.insert(node.id().clone());
999 }
1000 }
1001 }
1002
1003 for binding in node
1004 .values()
1005 .iter()
1006 .filter(|binding| binding.role() == super::ResolvedValueRole::Output)
1007 {
1008 witnesses.extend(ExecutionDeterminismWitnessSpec::from_binding(
1009 node,
1010 ExecutionDeterminismWitnessKind::Output {
1011 value_id: binding.value_id().clone(),
1012 output_ordinal: binding.ordinal(),
1013 },
1014 binding,
1015 dynamic_descriptors,
1016 )?);
1017 }
1018
1019 for effect in node.state_effects().iter().filter(|effect| {
1020 matches!(
1021 effect.access(),
1022 TensorAccess::Write | TensorAccess::ReadWrite
1023 )
1024 }) {
1025 let mut matched_resources = BTreeSet::new();
1026 for binding in node.values().iter().filter(|binding| {
1027 binding.value_id() == effect.state_value_id()
1028 && matches!(
1029 binding.access(),
1030 TensorAccess::Write | TensorAccess::ReadWrite
1031 )
1032 }) {
1033 let specs = ExecutionDeterminismWitnessSpec::from_binding(
1034 node,
1035 ExecutionDeterminismWitnessKind::StateEffect {
1036 state_id: effect.state_id().clone(),
1037 state_value_id: effect.state_value_id().clone(),
1038 lifetime: effect.lifetime(),
1039 access: effect.access(),
1040 },
1041 binding,
1042 dynamic_descriptors,
1043 )?;
1044 matched_resources.extend(specs.iter().map(|spec| spec.resource_id().clone()));
1045 witnesses.extend(specs);
1046 }
1047 let expected_resources = effect
1048 .resource_ids()
1049 .iter()
1050 .cloned()
1051 .collect::<BTreeSet<_>>();
1052 if matched_resources.is_empty() || matched_resources != expected_resources {
1053 return Err(invalid_plan(format!(
1054 "node `{}` writable state `{}` has no exact determinism witness closure",
1055 node.id(),
1056 effect.state_id()
1057 )));
1058 }
1059 }
1060 }
1061
1062 if witnesses.is_empty() {
1063 return Err(invalid_plan(
1064 "execution determinism witness plan has no declared outputs or writable state",
1065 ));
1066 }
1067
1068 let replay_provider_requirements = replay_providers
1069 .into_iter()
1070 .map(
1071 |(provider_id, (provider_implementation_fingerprint, contract, node_ids))| {
1072 ProviderDeterminismCoverageRequirement {
1073 provider_id,
1074 provider_implementation_fingerprint,
1075 provider_execution_contract_fingerprint: contract,
1076 node_ids: node_ids.into_iter().collect(),
1077 }
1078 },
1079 )
1080 .collect();
1081
1082 let witness_plan = ExecutionDeterminismWitnessPlan {
1083 schema_version: EXECUTION_DETERMINISM_WITNESS_VERSION,
1084 plan_hash: self.plan_hash().clone(),
1085 node_ids: canonical_node_ids,
1086 replay_provider_requirements,
1087 initializations,
1088 witnesses,
1089 };
1090 witness_plan.validate_shape()?;
1091 Ok(witness_plan)
1092 }
1093}
1094
1095#[cfg(test)]
1096mod tests {
1097 use super::validate_determinism_location;
1098 use crate::vnext::{
1099 AliasPolicy, AllocationKind, AllocationLifetime, BufferUsage, DynamicResourceDemand,
1100 DynamicResourceDescriptor, DynamicStorageAllocator, DynamicStorageContract,
1101 DynamicStorageProfile, DynamicStorageView, ElementType, ExecutionDeterminismValueExtent,
1102 ExecutionDeterminismWitnessKind, ExecutionDeterminismWitnessSpec, NodeId,
1103 NodeTokenBindingProjection, NodeWorkContract, PlanNode, ProgramValueId,
1104 ResolvedTensorLayout, ResolvedTensorSpec, ResolvedValueBinding, ResolvedValueRole,
1105 ResolvedValueStorage, ResourceId, StateInitialization, TensorAccess, TokenSpanWork,
1106 };
1107
1108 fn dynamic_activation_descriptor(
1109 resource_id: ResourceId,
1110 bytes_per_token: u64,
1111 maximum_tokens: u64,
1112 ) -> DynamicResourceDescriptor {
1113 let profile = DynamicStorageProfile::new(
1114 DynamicStorageAllocator::LinearArena,
1115 DynamicStorageView::Contiguous,
1116 )
1117 .unwrap();
1118 let storage =
1119 DynamicStorageContract::resource_test_contract(profile, "1".repeat(64)).unwrap();
1120 DynamicResourceDescriptor::new(
1121 resource_id,
1122 DynamicResourceDemand::tokens(bytes_per_token, maximum_tokens).unwrap(),
1123 16,
1124 BufferUsage::Activations,
1125 ElementType::F16,
1126 AllocationLifetime::Step,
1127 AllocationKind::Value,
1128 storage,
1129 StateInitialization::None,
1130 32,
1131 )
1132 .unwrap()
1133 }
1134
1135 #[test]
1136 fn token_witness_uses_immediate_span_and_rejects_capacity_overrun() {
1137 let value_id = ProgramValueId::new("value/output").unwrap();
1138 let projection = NodeTokenBindingProjection {
1139 value_id: value_id.clone(),
1140 role: ResolvedValueRole::Output,
1141 ordinal: 0,
1142 axis: 0,
1143 rank: 2,
1144 canonical_extent: 8,
1145 };
1146 let mut node = PlanNode::resource_test_node(NodeId::new("node/token-output").unwrap());
1147 node.work = NodeWorkContract::Tokens {
1148 source: projection.clone(),
1149 projections: vec![projection],
1150 };
1151 let binding = ResolvedValueBinding::new(
1152 value_id.clone(),
1153 ResolvedValueRole::Output,
1154 0,
1155 ResolvedTensorSpec::new(
1156 vec![8, 8],
1157 ElementType::F16,
1158 ResolvedTensorLayout::Contiguous,
1159 )
1160 .unwrap(),
1161 TensorAccess::Write,
1162 AliasPolicy::NoAlias,
1163 BufferUsage::Activations,
1164 None,
1165 ResolvedValueStorage::single(
1166 ResourceId::new("resource/token-output").unwrap(),
1167 32,
1168 128,
1169 ElementType::F16,
1170 )
1171 .unwrap(),
1172 )
1173 .unwrap();
1174
1175 let witnesses = ExecutionDeterminismWitnessSpec::from_binding(
1176 &node,
1177 ExecutionDeterminismWitnessKind::Output {
1178 value_id,
1179 output_ordinal: 0,
1180 },
1181 &binding,
1182 &[],
1183 )
1184 .unwrap();
1185 assert_eq!(witnesses.len(), 1);
1186 let witness = &witnesses[0];
1187 assert_eq!(witness.declared_length_bytes(), 128);
1188 assert_eq!(
1189 witness.extent(),
1190 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
1191 bytes_per_token: 16,
1192 maximum_tokens: 8,
1193 }
1194 );
1195 assert_eq!(
1196 witness
1197 .location()
1198 .bound_length_bytes_for_source_end(
1199 &TokenSpanWork::from_token_ids(&[1, 2, 3, 4, 5, 6, 7, 8], 2..5).unwrap(),
1200 5,
1201 )
1202 .unwrap(),
1203 48
1204 );
1205 assert!(witness
1206 .location()
1207 .bound_length_bytes_for_source_end(
1208 &TokenSpanWork::from_token_ids(&[1, 2, 3, 4, 5, 6, 7, 8, 9], 0..9).unwrap(),
1209 9,
1210 )
1211 .is_err());
1212 }
1213
1214 #[test]
1215 fn token_witness_accepts_scheduler_capacity_below_canonical_extent() {
1216 let value_id = ProgramValueId::new("value/scheduled-output").unwrap();
1217 let resource_id = ResourceId::new("resource/scheduled-output").unwrap();
1218 let projection = NodeTokenBindingProjection {
1219 value_id: value_id.clone(),
1220 role: ResolvedValueRole::Output,
1221 ordinal: 0,
1222 axis: 0,
1223 rank: 2,
1224 canonical_extent: 8,
1225 };
1226 let mut node = PlanNode::resource_test_node(NodeId::new("node/scheduled-output").unwrap());
1227 node.work = NodeWorkContract::Tokens {
1228 source: projection.clone(),
1229 projections: vec![projection],
1230 };
1231 let binding = ResolvedValueBinding::new(
1232 value_id.clone(),
1233 ResolvedValueRole::Output,
1234 0,
1235 ResolvedTensorSpec::new(
1236 vec![8, 8],
1237 ElementType::F16,
1238 ResolvedTensorLayout::Contiguous,
1239 )
1240 .unwrap(),
1241 TensorAccess::Write,
1242 AliasPolicy::NoAlias,
1243 BufferUsage::Activations,
1244 None,
1245 ResolvedValueStorage::single(resource_id.clone(), 0, 128, ElementType::F16).unwrap(),
1246 )
1247 .unwrap();
1248 let descriptor = dynamic_activation_descriptor(resource_id, 16, 4);
1249
1250 let witnesses = ExecutionDeterminismWitnessSpec::from_binding(
1251 &node,
1252 ExecutionDeterminismWitnessKind::Output {
1253 value_id,
1254 output_ordinal: 0,
1255 },
1256 &binding,
1257 &[descriptor],
1258 )
1259 .unwrap();
1260 let witness = &witnesses[0];
1261 validate_determinism_location(witness.location()).unwrap();
1262 assert_eq!(
1263 witness.extent(),
1264 ExecutionDeterminismValueExtent::ImmediateTokenSpan {
1265 bytes_per_token: 16,
1266 maximum_tokens: 4,
1267 }
1268 );
1269 assert_eq!(witness.maximum_bound_length_bytes().unwrap(), 64);
1270 assert_eq!(
1271 witness
1272 .location()
1273 .bound_length_bytes_for_source_end(
1274 &TokenSpanWork::from_token_ids(&[1, 2, 3, 4], 0..4).unwrap(),
1275 4,
1276 )
1277 .unwrap(),
1278 64
1279 );
1280 assert!(witness
1281 .location()
1282 .bound_length_bytes_for_source_end(
1283 &TokenSpanWork::from_token_ids(&[1, 2, 3, 4, 5], 0..5).unwrap(),
1284 5,
1285 )
1286 .is_err());
1287 }
1288
1289 #[test]
1290 fn non_contiguous_witness_retains_the_complete_typed_span() {
1291 let mut node = PlanNode::resource_test_node(NodeId::new("node/strided-output").unwrap());
1292 node.work = NodeWorkContract::Fixed;
1293 let value_id = ProgramValueId::new("value/strided-output").unwrap();
1294 let binding = ResolvedValueBinding::new(
1295 value_id.clone(),
1296 ResolvedValueRole::Output,
1297 0,
1298 ResolvedTensorSpec::new(
1299 vec![2, 2],
1300 ElementType::F16,
1301 ResolvedTensorLayout::Strided {
1302 byte_strides: vec![8, 2],
1303 },
1304 )
1305 .unwrap(),
1306 TensorAccess::Write,
1307 AliasPolicy::NoAlias,
1308 BufferUsage::Activations,
1309 None,
1310 ResolvedValueStorage::single(
1311 ResourceId::new("resource/strided-output").unwrap(),
1312 16,
1313 12,
1314 ElementType::F16,
1315 )
1316 .unwrap(),
1317 )
1318 .unwrap();
1319
1320 let witnesses = ExecutionDeterminismWitnessSpec::from_binding(
1321 &node,
1322 ExecutionDeterminismWitnessKind::Output {
1323 value_id,
1324 output_ordinal: 0,
1325 },
1326 &binding,
1327 &[],
1328 )
1329 .unwrap();
1330 assert_eq!(witnesses.len(), 1);
1331 assert_eq!(witnesses[0].logical_offset_bytes(), 16);
1332 assert_eq!(witnesses[0].declared_length_bytes(), 12);
1333 assert_eq!(witnesses[0].element_type(), ElementType::F16);
1334 }
1335}