ferrum_interfaces/vnext/execution/
checkpoint.rs

1use super::{
2    invalid_plan, BTreeSet, BufferUsage, Deserialize, NodeId, PlanNode, PreparedModelFamily,
3    ProgramValueId, ResolvedTensorLayout, ResolvedTensorSpec, ResolvedValueBinding,
4    ResolvedValueRole, ResourceId, Serialize, TensorAccess, VNextError,
5};
6use crate::vnext::{CompletionReadbackRequest, HostTransferLayout};
7
8/// Explicit semantic activations that must remain readable at the terminal
9/// completion fence. The empty default preserves the normal execution plan.
10#[derive(Debug, Clone, Default, PartialEq, Eq)]
11pub struct CompletionRetentionSpec {
12    values: BTreeSet<ProgramValueId>,
13}
14
15impl CompletionRetentionSpec {
16    pub fn new(values: BTreeSet<ProgramValueId>) -> Self {
17        Self { values }
18    }
19
20    /// Retains every semantic operation output for a terminal determinism
21    /// readback. Duplicate output identities are rejected because they cannot
22    /// be attributed to exactly one producer.
23    pub fn for_determinism_outputs(family: &PreparedModelFamily) -> Result<Self, VNextError> {
24        let mut values = BTreeSet::new();
25        for node in family
26            .program()
27            .blocks()
28            .iter()
29            .flat_map(|block| &block.nodes)
30        {
31            for value_id in &node.outputs {
32                if !values.insert(value_id.clone()) {
33                    return Err(invalid_plan(format!(
34                        "determinism output `{value_id}` has more than one producer"
35                    )));
36                }
37            }
38        }
39        if values.is_empty() {
40            return Err(invalid_plan(
41                "determinism output retention requires at least one operation output",
42            ));
43        }
44        Ok(Self { values })
45    }
46
47    pub fn insert(&mut self, value_id: ProgramValueId) -> bool {
48        self.values.insert(value_id)
49    }
50
51    pub fn values(&self) -> &BTreeSet<ProgramValueId> {
52        &self.values
53    }
54
55    pub fn is_empty(&self) -> bool {
56        self.values.is_empty()
57    }
58}
59
60/// Immutable plan evidence for one semantic activation retained until the
61/// terminal completion fence. Callers never reconstruct this binding from raw
62/// node and resource strings.
63#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
64pub struct RetainedCompletionValue {
65    value_id: ProgramValueId,
66    producer_node_id: NodeId,
67    output_ordinal: u32,
68    resource_id: ResourceId,
69    logical_offset_bytes: u64,
70    tensor: ResolvedTensorSpec,
71}
72
73#[derive(Deserialize)]
74#[serde(deny_unknown_fields)]
75struct RetainedCompletionValueWire {
76    value_id: ProgramValueId,
77    producer_node_id: NodeId,
78    output_ordinal: u32,
79    resource_id: ResourceId,
80    logical_offset_bytes: u64,
81    tensor: ResolvedTensorSpec,
82}
83
84impl<'de> Deserialize<'de> for RetainedCompletionValue {
85    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
86    where
87        D: serde::Deserializer<'de>,
88    {
89        let wire = RetainedCompletionValueWire::deserialize(deserializer)?;
90        Self::new(
91            wire.value_id,
92            wire.producer_node_id,
93            wire.output_ordinal,
94            wire.resource_id,
95            wire.logical_offset_bytes,
96            wire.tensor,
97        )
98        .map_err(serde::de::Error::custom)
99    }
100}
101
102impl RetainedCompletionValue {
103    fn new(
104        value_id: ProgramValueId,
105        producer_node_id: NodeId,
106        output_ordinal: u32,
107        resource_id: ResourceId,
108        logical_offset_bytes: u64,
109        tensor: ResolvedTensorSpec,
110    ) -> Result<Self, VNextError> {
111        let byte_len = tensor.minimum_storage_bytes()?;
112        if logical_offset_bytes.checked_add(byte_len).is_none() {
113            return Err(invalid_plan(
114                "retained completion value range overflows u64",
115            ));
116        }
117        Ok(Self {
118            value_id,
119            producer_node_id,
120            output_ordinal,
121            resource_id,
122            logical_offset_bytes,
123            tensor,
124        })
125    }
126
127    pub fn value_id(&self) -> &ProgramValueId {
128        &self.value_id
129    }
130
131    pub fn producer_node_id(&self) -> &NodeId {
132        &self.producer_node_id
133    }
134
135    pub const fn output_ordinal(&self) -> u32 {
136        self.output_ordinal
137    }
138
139    pub fn resource_id(&self) -> &ResourceId {
140        &self.resource_id
141    }
142
143    pub const fn logical_offset_bytes(&self) -> u64 {
144        self.logical_offset_bytes
145    }
146
147    pub fn tensor(&self) -> &ResolvedTensorSpec {
148        &self.tensor
149    }
150
151    pub fn readback_request(
152        &self,
153        participant_index: u32,
154        output_layout: HostTransferLayout,
155    ) -> Result<CompletionReadbackRequest, VNextError> {
156        if !matches!(self.tensor.layout(), ResolvedTensorLayout::Contiguous) {
157            return Err(invalid_plan(
158                "completion readback currently requires contiguous retained storage",
159            ));
160        }
161        if output_layout.element_type() != self.tensor.element_type() {
162            return Err(invalid_plan(
163                "completion readback element type differs from retained activation",
164            ));
165        }
166        if output_layout.byte_len()? > self.tensor.minimum_storage_bytes()? {
167            return Err(invalid_plan(
168                "completion readback exceeds retained activation capacity",
169            ));
170        }
171        CompletionReadbackRequest::new(
172            self.producer_node_id.clone(),
173            participant_index,
174            self.resource_id.clone(),
175            self.logical_offset_bytes,
176            output_layout,
177        )
178    }
179}
180
181pub(super) fn resolve_retained_completion_values(
182    nodes: &[PlanNode],
183    spec: &CompletionRetentionSpec,
184) -> Result<Vec<RetainedCompletionValue>, VNextError> {
185    let mut retained = Vec::with_capacity(spec.values.len());
186    for value_id in &spec.values {
187        let matches = nodes
188            .iter()
189            .enumerate()
190            .flat_map(|(node_index, node)| {
191                node.values()
192                    .iter()
193                    .map(move |binding| (node_index, node, binding))
194            })
195            .filter(|(_, _, binding)| {
196                binding.value_id() == value_id && binding.role() == ResolvedValueRole::Output
197            })
198            .collect::<Vec<_>>();
199        let [(producer_index, producer, binding)] = matches.as_slice() else {
200            return Err(invalid_plan(format!(
201                "completion retention value `{value_id}` must identify exactly one plan-node output"
202            )));
203        };
204        if binding.usage() != BufferUsage::Activations {
205            return Err(invalid_plan(format!(
206                "completion retention value `{value_id}` is not an activation"
207            )));
208        }
209        let [component] = binding.storage().components() else {
210            return Err(invalid_plan(format!(
211                "completion retention value `{value_id}` does not use one physical component"
212            )));
213        };
214        let retained_end = component
215            .offset_bytes()
216            .checked_add(binding.tensor().minimum_storage_bytes()?)
217            .ok_or_else(|| invalid_plan("retained completion range overflows u64"))?;
218        reject_later_overlapping_write(
219            nodes,
220            *producer_index,
221            binding,
222            component.resource_id(),
223            component.offset_bytes(),
224            retained_end,
225        )?;
226        retained.push(RetainedCompletionValue::new(
227            value_id.clone(),
228            producer.id().clone(),
229            binding.ordinal(),
230            component.resource_id().clone(),
231            component.offset_bytes(),
232            binding.tensor().clone(),
233        )?);
234    }
235    Ok(retained)
236}
237
238fn reject_later_overlapping_write(
239    nodes: &[PlanNode],
240    producer_index: usize,
241    retained_binding: &ResolvedValueBinding,
242    retained_resource: &ResourceId,
243    retained_start: u64,
244    retained_end: u64,
245) -> Result<(), VNextError> {
246    for (node_index, node) in nodes.iter().enumerate().skip(producer_index) {
247        for binding in node.values() {
248            let is_target = node_index == producer_index
249                && binding.role() == retained_binding.role()
250                && binding.ordinal() == retained_binding.ordinal()
251                && binding.value_id() == retained_binding.value_id();
252            if is_target
253                || !matches!(
254                    binding.access(),
255                    TensorAccess::Write | TensorAccess::ReadWrite
256                )
257            {
258                continue;
259            }
260            for component in binding.storage().components() {
261                if component.resource_id() != retained_resource {
262                    continue;
263                }
264                let write_end = component
265                    .offset_bytes()
266                    .checked_add(component.length_bytes())
267                    .ok_or_else(|| invalid_plan("plan value write range overflows u64"))?;
268                if component.offset_bytes() < retained_end && retained_start < write_end {
269                    return Err(invalid_plan(format!(
270                        "completion retention value `{}` is overwritten by node `{}` after its producer",
271                        retained_binding.value_id(),
272                        node.id()
273                    )));
274                }
275            }
276        }
277    }
278    Ok(())
279}
ferrum_interfaces/vnext/execution/checkpoint.rs

ferrum_interfaces/vnext/execution/
checkpoint.rs
