eredu_runtime/
parameter.rs

1//! Backend-neutral checkpoint materialization and parameter binding.
2
3use std::collections::BTreeMap;
4
5use eredu_checkpoint::{
6    recipe::{AtomicRecipeSet, RecipeCatalog, RecipeError},
7    store::{CheckpointSource, ReadPolicy, StoreError, TensorReadRequest},
8};
9use eredu_nn::{
10    ParameterId, ParameterMetadata, ParameterVisitor, ParameterVisitorMut, Parameterized,
11};
12
13use crate::{ParameterBackend, ResidencyDeclarationError, WeightBinding, WeightBindingPlan};
14
15/// Consumes one atomic recipe publication into owner bindings and lightweight
16/// logical aliases without cloning any derived recipe.
17pub fn bindings_from_recipe_set<C: RecipeCatalog + ?Sized>(
18    catalog: &C,
19    set: AtomicRecipeSet,
20) -> Result<Vec<WeightBinding>, RecipeBindingError> {
21    let (outputs, aliases) = set.into_parts();
22    let mut bytes = BTreeMap::new();
23    for (name, recipe) in &outputs {
24        bytes.insert(name.clone(), recipe.infer(catalog)?.byte_len());
25    }
26    let mut bindings = outputs
27        .into_iter()
28        .map(|(name, recipe)| {
29            let expected = bytes[&name];
30            WeightBinding::from_recipe(name, recipe, expected)
31        })
32        .collect::<Result<Vec<_>, _>>()?;
33    for (alias, owner) in aliases {
34        bindings.push(WeightBinding::alias(alias, owner.clone(), bytes[&owner])?);
35    }
36    WeightBindingPlan::new(&bindings)?;
37    Ok(bindings)
38}
39
40/// Failure while lowering an atomic neutral recipe set into runtime bindings.
41#[derive(Debug, thiserror::Error)]
42pub enum RecipeBindingError {
43    /// Recipe metadata inference failed.
44    #[error(transparent)]
45    Recipe(#[from] RecipeError),
46    /// A runtime binding declaration was invalid.
47    #[error(transparent)]
48    Declaration(#[from] ResidencyDeclarationError),
49}
50
51/// One fully realized atomic unit keyed by stable parameter identity.
52pub struct MaterializedUnit<B: ParameterBackend> {
53    weights: BTreeMap<ParameterId, B::MaterializedWeight>,
54}
55
56impl<B: ParameterBackend> MaterializedUnit<B> {
57    /// Returns the number of realized parameter values.
58    pub fn len(&self) -> usize {
59        self.weights.len()
60    }
61
62    /// Returns whether this unit contains no values.
63    pub fn is_empty(&self) -> bool {
64        self.weights.is_empty()
65    }
66
67    /// Returns whether a stable parameter identity is present.
68    pub fn contains(&self, id: &ParameterId) -> bool {
69        self.weights.contains_key(id)
70    }
71}
72
73/// Materializes one atomic unit without inspecting backend-native values.
74///
75/// Each encoded source remains retained by its backend materialization guard
76/// until exact completion. The returned weights are independently owned native
77/// handles ready for binding.
78pub fn materialize_bindings<B: ParameterBackend>(
79    source: &dyn CheckpointSource,
80    bindings: &[WeightBinding],
81    context: &B::MaterializationContext,
82) -> Result<MaterializedUnit<B>, ParameterOrchestrationError<B::ParameterError>> {
83    let plan = WeightBindingPlan::new(bindings)?;
84    for binding in plan.owners() {
85        let inferred = binding.source_recipe().infer(source)?;
86        if inferred.byte_len() != binding.expected_bytes() {
87            return Err(ParameterOrchestrationError::ByteMismatch {
88                parameter: binding.name().to_owned(),
89                expected: binding.expected_bytes(),
90                actual: inferred.byte_len(),
91            });
92        }
93    }
94    let mut weights = BTreeMap::new();
95    for binding in plan.owners() {
96        let materialization = match binding.recipe() {
97            Some(recipe) => B::materialize_recipe(recipe, source, context),
98            None => {
99                let lease = source.acquire_lease(TensorReadRequest {
100                    key: binding.checkpoint_key().to_owned(),
101                    selection: binding.selection().clone(),
102                    policy: ReadPolicy::RequireBounded,
103                })?;
104                B::materialize(lease, context)
105            }
106        }
107        .map_err(ParameterOrchestrationError::Backend)?;
108        let weight = B::finish_materialization(materialization)
109            .map_err(ParameterOrchestrationError::Backend)?;
110        let id = ParameterId::new(binding.name()).map_err(|error| {
111            ParameterOrchestrationError::InvalidParameterIdentity(error.to_string())
112        })?;
113        if weights.insert(id.clone(), weight).is_some() {
114            return Err(ParameterOrchestrationError::DuplicateBinding { parameter: id });
115        }
116    }
117    for (alias, owner) in plan.aliases() {
118        let owner_id = ParameterId::new(owner.name()).map_err(|error| {
119            ParameterOrchestrationError::InvalidParameterIdentity(error.to_string())
120        })?;
121        let weight = weights
122            .get(&owner_id)
123            .expect("validated owner was materialized before aliases");
124        let weight =
125            B::share_materialized_weight(weight).map_err(ParameterOrchestrationError::Backend)?;
126        let alias_id = ParameterId::new(alias.name()).map_err(|error| {
127            ParameterOrchestrationError::InvalidParameterIdentity(error.to_string())
128        })?;
129        weights.insert(alias_id, weight);
130    }
131    Ok(MaterializedUnit { weights })
132}
133
134/// Binds a realized unit into a statically traversed native module.
135///
136/// Binding is keyed exclusively by stable parameter identity. Missing and
137/// unexpected values fail closed; native tensor shape or storage remains
138/// backend-owned.
139pub fn bind_materialized_unit<B, M>(
140    module: &mut M,
141    mut unit: MaterializedUnit<B>,
142) -> Result<(), ParameterOrchestrationError<B::ParameterError>>
143where
144    B: ParameterBackend,
145    M: Parameterized<B::Parameter>,
146{
147    struct Validator<'a, B: ParameterBackend> {
148        weights: &'a BTreeMap<ParameterId, B::MaterializedWeight>,
149        visited: BTreeMap<ParameterId, ()>,
150        error: Option<ParameterOrchestrationError<B::ParameterError>>,
151    }
152
153    impl<'a, 'value, B: ParameterBackend> ParameterVisitor<'value, B::Parameter> for Validator<'a, B> {
154        fn visit(&mut self, metadata: ParameterMetadata, parameter: &'value B::Parameter) {
155            if self.error.is_some() {
156                return;
157            }
158            let Some(weight) = self.weights.get(&metadata.id) else {
159                self.error = Some(ParameterOrchestrationError::MissingBinding {
160                    parameter: metadata.id,
161                });
162                return;
163            };
164            if let Err(error) = B::validate_bind(parameter, weight) {
165                self.error = Some(ParameterOrchestrationError::Backend(error));
166                return;
167            }
168            self.visited.insert(metadata.id, ());
169        }
170    }
171
172    let mut validator = Validator::<B> {
173        weights: &unit.weights,
174        visited: BTreeMap::new(),
175        error: None,
176    };
177    module.visit_parameters(&mut validator);
178    if let Some(error) = validator.error {
179        return Err(error);
180    }
181    let unexpected = unit
182        .weights
183        .keys()
184        .filter(|id| !validator.visited.contains_key(*id))
185        .cloned()
186        .collect::<Vec<_>>();
187    if !unexpected.is_empty() {
188        return Err(ParameterOrchestrationError::UnexpectedBindings {
189            parameters: unexpected,
190        });
191    }
192
193    struct Binder<'a, B: ParameterBackend> {
194        weights: &'a mut BTreeMap<ParameterId, B::MaterializedWeight>,
195        error: Option<ParameterOrchestrationError<B::ParameterError>>,
196    }
197
198    impl<'a, 'value, B: ParameterBackend> ParameterVisitorMut<'value, B::Parameter> for Binder<'a, B> {
199        fn visit_mut(&mut self, metadata: ParameterMetadata, parameter: &'value mut B::Parameter) {
200            if self.error.is_some() {
201                return;
202            }
203            let Some(weight) = self.weights.remove(&metadata.id) else {
204                self.error = Some(ParameterOrchestrationError::MissingBinding {
205                    parameter: metadata.id,
206                });
207                return;
208            };
209            if let Err(error) = B::bind(parameter, weight) {
210                self.error = Some(ParameterOrchestrationError::Backend(error));
211            }
212        }
213    }
214
215    let mut binder = Binder::<B> {
216        weights: &mut unit.weights,
217        error: None,
218    };
219    module.visit_parameters_mut(&mut binder);
220    if let Some(error) = binder.error {
221        return Err(error);
222    }
223    if !unit.weights.is_empty() {
224        return Err(ParameterOrchestrationError::UnexpectedBindings {
225            parameters: unit.weights.into_keys().collect(),
226        });
227    }
228    Ok(())
229}
230
231/// Failure in backend-neutral parameter materialization or binding.
232#[derive(Debug, thiserror::Error)]
233pub enum ParameterOrchestrationError<E>
234where
235    E: std::error::Error + Send + Sync + 'static,
236{
237    /// Binding aliases were invalid before materialization began.
238    #[error(transparent)]
239    Declaration(#[from] ResidencyDeclarationError),
240    /// Neutral checkpoint access failed.
241    #[error(transparent)]
242    Store(#[from] StoreError),
243    /// Neutral recipe validation failed.
244    #[error(transparent)]
245    Recipe(#[from] RecipeError),
246    /// A stable runtime parameter identity was invalid.
247    #[error("invalid runtime parameter identity: {0}")]
248    InvalidParameterIdentity(String),
249    /// Two bindings targeted the same stable parameter.
250    #[error("duplicate materialized binding for parameter {parameter}")]
251    DuplicateBinding {
252        /// Duplicated identity.
253        parameter: ParameterId,
254    },
255    /// Inferred and declared materialized byte sizes disagreed.
256    #[error("parameter {parameter:?} declares {expected} bytes but its recipe produces {actual}")]
257    ByteMismatch {
258        /// Stable binding name.
259        parameter: String,
260        /// Declared bytes.
261        expected: u64,
262        /// Inferred bytes.
263        actual: u64,
264    },
265    /// A traversed parameter had no realized value.
266    #[error("materialized unit has no value for parameter {parameter}")]
267    MissingBinding {
268        /// Missing parameter identity.
269        parameter: ParameterId,
270    },
271    /// Realized values remained after complete parameter traversal.
272    #[error("materialized unit contains values for unknown parameters: {parameters:?}")]
273    UnexpectedBindings {
274        /// Unknown parameter identities.
275        parameters: Vec<ParameterId>,
276    },
277    /// Backend-native realization or binding failed.
278    #[error("backend parameter operation failed: {0}")]
279    Backend(E),
280}
eredu_runtime/parameter.rs

eredu_runtime/
parameter.rs
