Skip to main content

eredu_runtime/
component.rs

1//! Typed multimodal component graphs and residency classes.
2
3use std::collections::{BTreeMap, BTreeSet, VecDeque};
4
5/// Logical value domain crossing a component boundary.
6#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
7pub enum ComponentDomain {
8    /// Integer token identities.
9    TokenIds,
10    /// Prepacked image or video patches.
11    PatchMatrix,
12    /// Prepared audio features or codebook identities.
13    AudioFeatures,
14    /// Decoder-width hidden activations.
15    HiddenStates,
16    /// Vocabulary logits.
17    Logits,
18    /// Ordered target block states consumed by a draft model.
19    TargetStates,
20}
21
22/// Semantic execution-unit kind without family equations.
23#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
24pub enum ComponentKind {
25    /// Token embedding, normalization, or another immutable text root.
26    StaticText,
27    /// Vision encoder or projector unit.
28    Vision,
29    /// Audio encoder or projector unit.
30    Audio,
31    /// Ordered text/media assembly.
32    Assembly,
33    /// One decoder layer.
34    Decoder,
35    /// Embedded multi-token prediction unit.
36    Prediction,
37    /// External assistant unit.
38    Assistant,
39    /// Final vocabulary projection.
40    OutputProjection,
41}
42
43/// Immutable-weight residency accounting class.
44#[derive(Debug, Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
45pub enum ComponentResidencyClass {
46    /// Small static text modules.
47    Static,
48    /// Modality-specific tower and prepared-media workspace.
49    Media,
50    /// Independently resident or streamed decoder unit.
51    Decoder,
52    /// Independently leased routed experts.
53    Experts,
54    /// Embedded or external draft modules.
55    Draft,
56}
57
58/// One typed unit in a composite execution graph.
59#[derive(Debug, Clone, Eq, PartialEq)]
60pub struct ComponentSpec {
61    /// Stable architecture-owned identity.
62    pub id: String,
63    /// Semantic component kind.
64    pub kind: ComponentKind,
65    /// External input domains in declared order.
66    pub external_inputs: Vec<ComponentDomain>,
67    /// Dependency unit identities in declared order.
68    pub dependencies: Vec<String>,
69    /// Required output domain for each dependency.
70    pub dependency_inputs: Vec<ComponentDomain>,
71    /// Unit output domain.
72    pub output: ComponentDomain,
73    /// Weight-residency accounting class.
74    pub residency: ComponentResidencyClass,
75}
76
77/// Validated typed component graph with one or more observable outputs.
78#[derive(Debug, Clone, Eq, PartialEq)]
79pub struct ComponentGraph {
80    units: Vec<ComponentSpec>,
81    execution_order: Vec<usize>,
82    outputs: Vec<usize>,
83}
84
85impl ComponentGraph {
86    /// Validates identities, dependency domains, acyclicity, and outputs.
87    pub fn new(
88        units: Vec<ComponentSpec>,
89        outputs: impl IntoIterator<Item = impl AsRef<str>>,
90    ) -> Result<Self, ComponentGraphError> {
91        if units.is_empty() {
92            return Err(ComponentGraphError::Empty);
93        }
94        let mut by_id = BTreeMap::new();
95        for (index, unit) in units.iter().enumerate() {
96            if unit.id.trim().is_empty() {
97                return Err(ComponentGraphError::EmptyId);
98            }
99            if unit.dependencies.len() != unit.dependency_inputs.len() {
100                return Err(ComponentGraphError::DependencyArity {
101                    unit: unit.id.clone(),
102                    dependencies: unit.dependencies.len(),
103                    domains: unit.dependency_inputs.len(),
104                });
105            }
106            if by_id.insert(unit.id.clone(), index).is_some() {
107                return Err(ComponentGraphError::DuplicateId(unit.id.clone()));
108            }
109        }
110        let mut edges = vec![Vec::new(); units.len()];
111        let mut indegree = vec![0_usize; units.len()];
112        for (unit_index, unit) in units.iter().enumerate() {
113            let mut seen = BTreeSet::new();
114            for (slot, dependency) in unit.dependencies.iter().enumerate() {
115                let dependency_index = by_id.get(dependency).copied().ok_or_else(|| {
116                    ComponentGraphError::UnknownDependency {
117                        unit: unit.id.clone(),
118                        dependency: dependency.clone(),
119                    }
120                })?;
121                if !seen.insert(dependency_index) {
122                    return Err(ComponentGraphError::DuplicateDependency {
123                        unit: unit.id.clone(),
124                        dependency: dependency.clone(),
125                    });
126                }
127                let actual = units[dependency_index].output;
128                let expected = unit.dependency_inputs[slot];
129                if actual != expected {
130                    return Err(ComponentGraphError::DomainMismatch {
131                        unit: unit.id.clone(),
132                        dependency: dependency.clone(),
133                        expected,
134                        actual,
135                    });
136                }
137                edges[dependency_index].push(unit_index);
138                indegree[unit_index] += 1;
139            }
140        }
141        let mut ready = indegree
142            .iter()
143            .enumerate()
144            .filter_map(|(index, degree)| (*degree == 0).then_some(index))
145            .collect::<VecDeque<_>>();
146        let mut execution_order = Vec::with_capacity(units.len());
147        while let Some(index) = ready.pop_front() {
148            execution_order.push(index);
149            for dependent in &edges[index] {
150                indegree[*dependent] -= 1;
151                if indegree[*dependent] == 0 {
152                    ready.push_back(*dependent);
153                }
154            }
155        }
156        if execution_order.len() != units.len() {
157            return Err(ComponentGraphError::Cycle);
158        }
159        let outputs = outputs
160            .into_iter()
161            .map(|output| {
162                let output = output.as_ref();
163                by_id
164                    .get(output)
165                    .copied()
166                    .ok_or_else(|| ComponentGraphError::UnknownOutput(output.to_owned()))
167            })
168            .collect::<Result<Vec<_>, _>>()?;
169        if outputs.is_empty() {
170            return Err(ComponentGraphError::NoOutputs);
171        }
172        if outputs.iter().copied().collect::<BTreeSet<_>>().len() != outputs.len() {
173            return Err(ComponentGraphError::DuplicateOutput);
174        }
175        Ok(Self {
176            units,
177            execution_order,
178            outputs,
179        })
180    }
181
182    /// Returns units in architecture declaration order.
183    pub fn units(&self) -> &[ComponentSpec] {
184        &self.units
185    }
186
187    /// Returns units in dependency-safe execution order.
188    pub fn execution_order(&self) -> impl Iterator<Item = &ComponentSpec> {
189        self.execution_order.iter().map(|index| &self.units[*index])
190    }
191
192    /// Returns graph outputs in declared order.
193    pub fn outputs(&self) -> impl Iterator<Item = &ComponentSpec> {
194        self.outputs.iter().map(|index| &self.units[*index])
195    }
196
197    /// Counts components in one residency class.
198    pub fn residency_count(&self, class: ComponentResidencyClass) -> usize {
199        self.units
200            .iter()
201            .filter(|unit| unit.residency == class)
202            .count()
203    }
204}
205
206/// Invalid component graph declaration.
207#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
208pub enum ComponentGraphError {
209    /// No components were declared.
210    #[error("component graph cannot be empty")]
211    Empty,
212    /// One component identity was empty.
213    #[error("component identity cannot be empty")]
214    EmptyId,
215    /// A component identity occurred more than once.
216    #[error("component {0:?} was declared more than once")]
217    DuplicateId(String),
218    /// Dependency and expected-domain counts differ.
219    #[error("component {unit:?} has {dependencies} dependencies but {domains} dependency domains")]
220    DependencyArity {
221        /// Component identity.
222        unit: String,
223        /// Dependency count.
224        dependencies: usize,
225        /// Expected-domain count.
226        domains: usize,
227    },
228    /// A dependency identity was not declared.
229    #[error("component {unit:?} depends on unknown component {dependency:?}")]
230    UnknownDependency {
231        /// Consumer component.
232        unit: String,
233        /// Missing dependency.
234        dependency: String,
235    },
236    /// The same dependency was listed twice.
237    #[error("component {unit:?} repeats dependency {dependency:?}")]
238    DuplicateDependency {
239        /// Consumer component.
240        unit: String,
241        /// Repeated dependency.
242        dependency: String,
243    },
244    /// Producer and consumer domains disagree.
245    #[error("component {unit:?} expects {expected:?} from {dependency:?}, got {actual:?}")]
246    DomainMismatch {
247        /// Consumer component.
248        unit: String,
249        /// Producer component.
250        dependency: String,
251        /// Required domain.
252        expected: ComponentDomain,
253        /// Produced domain.
254        actual: ComponentDomain,
255    },
256    /// Dependency graph contains a cycle.
257    #[error("component graph contains a dependency cycle")]
258    Cycle,
259    /// No observable output was declared.
260    #[error("component graph requires at least one output")]
261    NoOutputs,
262    /// One output identity was repeated.
263    #[error("component graph repeats an output")]
264    DuplicateOutput,
265    /// An output identity was not declared.
266    #[error("unknown component output {0:?}")]
267    UnknownOutput(String),
268}
269
270#[cfg(test)]
271mod tests {
272    use super::*;
273
274    fn unit(id: &str, output: ComponentDomain) -> ComponentSpec {
275        ComponentSpec {
276            id: id.into(),
277            kind: ComponentKind::StaticText,
278            external_inputs: vec![],
279            dependencies: vec![],
280            dependency_inputs: vec![],
281            output,
282            residency: ComponentResidencyClass::Static,
283        }
284    }
285
286    #[test]
287    fn component_graph_validates_domains_and_preserves_outputs() {
288        let embedding = ComponentSpec {
289            external_inputs: vec![ComponentDomain::TokenIds],
290            ..unit("embedding", ComponentDomain::HiddenStates)
291        };
292        let decoder = ComponentSpec {
293            id: "decoder.0".into(),
294            kind: ComponentKind::Decoder,
295            dependencies: vec!["embedding".into()],
296            dependency_inputs: vec![ComponentDomain::HiddenStates],
297            output: ComponentDomain::HiddenStates,
298            residency: ComponentResidencyClass::Decoder,
299            external_inputs: vec![],
300        };
301        let graph = ComponentGraph::new(vec![decoder, embedding], ["decoder.0"]).unwrap();
302        assert_eq!(
303            graph
304                .execution_order()
305                .map(|unit| unit.id.as_str())
306                .collect::<Vec<_>>(),
307            ["embedding", "decoder.0"]
308        );
309        assert_eq!(graph.residency_count(ComponentResidencyClass::Decoder), 1);
310    }
311
312    #[test]
313    fn component_graph_rejects_domain_mismatch_and_cycles() {
314        let mut decoder = unit("decoder", ComponentDomain::HiddenStates);
315        decoder.dependencies = vec!["tokens".into()];
316        decoder.dependency_inputs = vec![ComponentDomain::HiddenStates];
317        let tokens = unit("tokens", ComponentDomain::TokenIds);
318        assert!(matches!(
319            ComponentGraph::new(vec![tokens, decoder], ["decoder"]),
320            Err(ComponentGraphError::DomainMismatch { .. })
321        ));
322        let mut left = unit("left", ComponentDomain::HiddenStates);
323        left.dependencies = vec!["right".into()];
324        left.dependency_inputs = vec![ComponentDomain::HiddenStates];
325        let mut right = unit("right", ComponentDomain::HiddenStates);
326        right.dependencies = vec!["left".into()];
327        right.dependency_inputs = vec![ComponentDomain::HiddenStates];
328        assert_eq!(
329            ComponentGraph::new(vec![left, right], ["left"]).unwrap_err(),
330            ComponentGraphError::Cycle
331        );
332    }
333}