Skip to main content

ferrum_interfaces/vnext/execution/sequence_checkpoint/
inputs.rs

1use std::collections::BTreeMap;
2
3use super::*;
4use crate::vnext::ResolvedTensorLayout;
5
6/// Full canonical content of one non-token input. Tensor identity and all bytes
7/// participate in equality; a hash alone is never the final matching proof.
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct CheckpointCanonicalInput {
10    tensor: ResolvedTensorSpec,
11    bytes: Vec<u8>,
12}
13
14impl CheckpointCanonicalInput {
15    pub fn new(tensor: ResolvedTensorSpec, bytes: Vec<u8>) -> Result<Self, VNextError> {
16        if !matches!(tensor.layout(), ResolvedTensorLayout::Contiguous)
17            || u64::try_from(bytes.len()).ok() != Some(tensor.minimum_storage_bytes()?)
18        {
19            return Err(invalid_plan(
20                "checkpoint input requires exact contiguous canonical tensor bytes",
21            ));
22        }
23        Ok(Self { tensor, bytes })
24    }
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct SequenceCheckpointInputIdentity {
29    layout_fingerprint: String,
30    token_input: ProgramValueId,
31    tokens: Vec<u32>,
32    conditioning: BTreeMap<ProgramValueId, CheckpointCanonicalInput>,
33    requires_entire_input: bool,
34    completed_input_capture: CheckpointCompletedInputCapture,
35}
36
37impl SequenceCheckpointInputIdentity {
38    pub fn tokens(&self) -> &[u32] {
39        &self.tokens
40    }
41
42    /// Checks exact contents after an indexed lookup. This is only the input
43    /// portion of matching: plan/loading instance, saved N and numerical
44    /// partition identity remain mandatory independent checks.
45    pub fn matches_at(&self, target: &Self, boundary: usize) -> bool {
46        boundary > 0
47            && boundary < target.tokens.len()
48            && (boundary < self.tokens.len()
49                || (boundary == self.tokens.len()
50                    && self.completed_input_capture == CheckpointCompletedInputCapture::Supported))
51            && self.layout_fingerprint == target.layout_fingerprint
52            && self.token_input == target.token_input
53            && self.conditioning == target.conditioning
54            && self.requires_entire_input == target.requires_entire_input
55            && self.tokens[..boundary] == target.tokens[..boundary]
56            && (!self.requires_entire_input || self.tokens == target.tokens)
57    }
58}
59
60impl SequenceCheckpointLayout {
61    pub fn bind_inputs(
62        &self,
63        token_input: &ProgramValueId,
64        tokens: &[u32],
65        conditioning: &BTreeMap<ProgramValueId, CheckpointCanonicalInput>,
66    ) -> Result<SequenceCheckpointInputIdentity, Vec<SequenceCheckpointUnsupportedReason>> {
67        use SequenceCheckpointUnsupportedReason as Reason;
68        let declared = self.inputs();
69        let mut reasons = Vec::new();
70        if token_input != declared.token_input() || tokens.is_empty() {
71            reasons.push(Reason::InvalidInput {
72                value_id: declared.token_input().clone(),
73            });
74        }
75        for value_id in declared.conditioning_inputs() {
76            if !conditioning.contains_key(value_id) {
77                reasons.push(Reason::MissingInput {
78                    value_id: value_id.clone(),
79                });
80            }
81        }
82        for value_id in conditioning.keys() {
83            if !declared.conditioning_inputs().contains(value_id) {
84                reasons.push(Reason::InvalidInput {
85                    value_id: value_id.clone(),
86                });
87            }
88        }
89        if !reasons.is_empty() {
90            return Err(reasons);
91        }
92        let layout_fingerprint = self
93            .fingerprint()
94            .map_err(|_| vec![Reason::InputCoverage])?;
95        Ok(SequenceCheckpointInputIdentity {
96            layout_fingerprint,
97            token_input: token_input.clone(),
98            tokens: tokens.to_vec(),
99            conditioning: conditioning.clone(),
100            requires_entire_input: self.input_dependency()
101                == CheckpointInputDependency::EntireTokenInput,
102            completed_input_capture: self.completed_input_capture(),
103        })
104    }
105}