Skip to main content

eredu_runtime/
residency.rs

1//! Backend-neutral immutable-weight residency declarations and control state.
2
3use std::{
4    collections::{BTreeMap, BTreeSet},
5    sync::Weak,
6    time::Duration,
7};
8
9use eredu_checkpoint::{
10    recipe::{DerivedWeightRecipe, RecipeCatalog, RecipeError},
11    store::{TensorSelection, WeightStoreDiagnostics},
12};
13use eredu_core::residency::{
14    EvictedResidencyCopy, MemoryTier, OffloadPlan, OffloadReport, OffloadUnitId, PrefetchOutcome,
15    ResidencyLedger, ResidencyLedgerError, UnitResidencyReport,
16};
17
18/// Deterministic telemetry from one bounded weight-materialization pass.
19#[derive(Debug, Clone, Default, Eq, PartialEq)]
20pub struct WeightMaterializationReport {
21    /// Planner ceiling for simultaneously live conversion data.
22    pub admitted_working_set_bytes: u64,
23    /// Number of dense semantic matrices transformed.
24    pub transformed_weights: usize,
25    /// Number of independently evaluated source row tiles.
26    pub source_tiles: usize,
27    /// Largest number of submitted tile completions retained simultaneously.
28    pub peak_in_flight_tiles: usize,
29    /// Total logical dense bytes selected from the source store.
30    pub source_bytes_read: u64,
31    /// Total encoded output bytes written to persistent storage.
32    pub output_bytes: u64,
33    /// Largest conservative conversion working set admitted for one tile.
34    pub peak_planned_working_set_bytes: u64,
35    /// Largest source recipe output tile.
36    pub largest_source_tile_bytes: u64,
37    /// Largest encoded output tile written together.
38    pub largest_output_tile_bytes: u64,
39}
40
41/// Immutable residency-control and checkpoint-storage telemetry snapshot.
42#[derive(Debug, Clone, Eq, PartialEq)]
43pub struct ResidencyReport {
44    initialized: bool,
45    offload: OffloadReport,
46    units: Vec<UnitResidencyReport>,
47    active_window: Vec<OffloadUnitId>,
48    weight_store: WeightStoreDiagnostics,
49    materialization: Option<WeightMaterializationReport>,
50}
51
52impl ResidencyReport {
53    /// Creates a report from one coherent controller and storage snapshot.
54    pub fn new(
55        initialized: bool,
56        offload: OffloadReport,
57        units: Vec<UnitResidencyReport>,
58        active_window: Vec<OffloadUnitId>,
59        weight_store: WeightStoreDiagnostics,
60    ) -> Self {
61        Self {
62            initialized,
63            offload,
64            units,
65            active_window,
66            weight_store,
67            materialization: None,
68        }
69    }
70
71    /// Returns whether explicit initialization completed successfully.
72    pub const fn initialized(&self) -> bool {
73        self.initialized
74    }
75
76    /// Returns the immutable offload telemetry snapshot.
77    pub const fn offload(&self) -> &OffloadReport {
78        &self.offload
79    }
80
81    /// Returns unit states in identifier order.
82    pub fn units(&self) -> &[UnitResidencyReport] {
83        &self.units
84    }
85
86    /// Returns the protected execution window in identifier order.
87    pub fn active_window(&self) -> &[OffloadUnitId] {
88        &self.active_window
89    }
90
91    /// Returns storage diagnostics, distinct from logical residency telemetry.
92    pub const fn weight_store(&self) -> &WeightStoreDiagnostics {
93        &self.weight_store
94    }
95
96    /// Returns bounded load-time materialization telemetry for these units.
97    pub const fn materialization(&self) -> Option<&WeightMaterializationReport> {
98        self.materialization.as_ref()
99    }
100
101    /// Attaches bounded load-time materialization telemetry.
102    pub fn with_materialization(
103        mut self,
104        materialization: Option<WeightMaterializationReport>,
105    ) -> Self {
106        self.materialization = materialization;
107        self
108    }
109}
110
111/// One named checkpoint selection within an atomic resident unit.
112#[derive(Debug, Clone, Eq, PartialEq)]
113pub struct WeightBinding {
114    name: String,
115    alias_of: Option<String>,
116    logical_target: Option<String>,
117    checkpoint_key: String,
118    selection: TensorSelection,
119    recipe: Option<DerivedWeightRecipe>,
120    quantization_companions: Option<(String, String)>,
121    expected_bytes: u64,
122}
123
124impl WeightBinding {
125    /// Creates a direct binding with a stable local name and selected size.
126    pub fn new(
127        name: impl Into<String>,
128        checkpoint_key: impl Into<String>,
129        selection: TensorSelection,
130        expected_bytes: u64,
131    ) -> Result<Self, ResidencyDeclarationError> {
132        let name = validate_name(name.into())?;
133        let checkpoint_key = checkpoint_key.into();
134        if checkpoint_key.trim().is_empty() {
135            return Err(ResidencyDeclarationError::InvalidCheckpointKey { name });
136        }
137        validate_size(&name, expected_bytes)?;
138        Ok(Self {
139            name,
140            alias_of: None,
141            logical_target: None,
142            checkpoint_key,
143            selection,
144            recipe: None,
145            quantization_companions: None,
146            expected_bytes,
147        })
148    }
149
150    /// Creates a binding backed by a composable derived-weight recipe.
151    pub fn from_recipe(
152        name: impl Into<String>,
153        recipe: DerivedWeightRecipe,
154        expected_bytes: u64,
155    ) -> Result<Self, ResidencyDeclarationError> {
156        let name = validate_name(name.into())?;
157        validate_size(&name, expected_bytes)?;
158        let checkpoint_key = first_source(&name, &recipe)?;
159        Ok(Self {
160            name,
161            alias_of: None,
162            logical_target: None,
163            checkpoint_key,
164            selection: TensorSelection::Full,
165            recipe: Some(recipe),
166            quantization_companions: None,
167            expected_bytes,
168        })
169    }
170
171    /// Creates one logical destination that shares an already materialized
172    /// owner binding in the same atomic unit.
173    pub fn alias(
174        name: impl Into<String>,
175        owner: impl Into<String>,
176        expected_bytes: u64,
177    ) -> Result<Self, ResidencyDeclarationError> {
178        let name = validate_name(name.into())?;
179        let owner = validate_name(owner.into())?;
180        validate_size(&name, expected_bytes)?;
181        Ok(Self {
182            name,
183            alias_of: Some(owner),
184            logical_target: None,
185            checkpoint_key: String::new(),
186            selection: TensorSelection::Full,
187            recipe: None,
188            quantization_companions: None,
189            expected_bytes,
190        })
191    }
192
193    /// Returns the stable name used to look up a resident value.
194    pub fn name(&self) -> &str {
195        &self.name
196    }
197
198    /// Returns the logical owner when this binding is an alias.
199    pub fn alias_of(&self) -> Option<&str> {
200        self.alias_of.as_deref()
201    }
202
203    /// Returns whether this binding reuses another logical binding's native
204    /// materialization.
205    pub const fn is_alias(&self) -> bool {
206        self.alias_of.is_some()
207    }
208
209    /// Replaces the stable name used to address this value inside its resident unit.
210    pub fn with_name(mut self, name: impl Into<String>) -> Result<Self, ResidencyDeclarationError> {
211        self.name = validate_name(name.into())?;
212        Ok(self)
213    }
214
215    /// Returns the architecture-logical parameter destination.
216    pub fn logical_target(&self) -> Option<&str> {
217        self.logical_target.as_deref()
218    }
219
220    /// Attaches the architecture-logical parameter destination.
221    pub fn with_logical_target(
222        mut self,
223        target: impl Into<String>,
224    ) -> Result<Self, ResidencyDeclarationError> {
225        self.logical_target = Some(validate_name(target.into())?);
226        Ok(self)
227    }
228
229    /// Declares exact local companion names produced by load-time quantization.
230    pub fn with_quantization_companions(
231        mut self,
232        scales_binding: impl Into<String>,
233        biases_binding: impl Into<String>,
234    ) -> Result<Self, ResidencyDeclarationError> {
235        let scales_binding = validate_name(scales_binding.into())?;
236        let biases_binding = validate_name(biases_binding.into())?;
237        if scales_binding == self.name
238            || biases_binding == self.name
239            || scales_binding == biases_binding
240        {
241            return Err(ResidencyDeclarationError::InvalidQuantizationCompanions {
242                name: self.name,
243                scales: scales_binding,
244                biases: biases_binding,
245            });
246        }
247        self.quantization_companions = Some((scales_binding, biases_binding));
248        Ok(self)
249    }
250
251    /// Returns exact local scale and affine-bias bindings for load-time quantization.
252    pub fn quantization_companions(&self) -> Option<(&str, &str)> {
253        self.quantization_companions
254            .as_ref()
255            .map(|(scales, biases)| (scales.as_str(), biases.as_str()))
256    }
257
258    /// Returns the first physical checkpoint source.
259    pub fn checkpoint_key(&self) -> &str {
260        &self.checkpoint_key
261    }
262
263    /// Returns the direct checkpoint selection.
264    pub fn selection(&self) -> &TensorSelection {
265        &self.selection
266    }
267
268    /// Returns the derived recipe when this is not a direct binding.
269    pub const fn recipe(&self) -> Option<&DerivedWeightRecipe> {
270        self.recipe.as_ref()
271    }
272
273    /// Returns the complete source recipe represented by this binding.
274    pub fn source_recipe(&self) -> DerivedWeightRecipe {
275        assert!(
276            self.alias_of.is_none(),
277            "logical aliases have no physical recipe"
278        );
279        self.recipe.clone().unwrap_or_else(|| {
280            DerivedWeightRecipe::source(self.checkpoint_key.clone(), self.selection.clone())
281        })
282    }
283
284    /// Returns every checkpoint key consumed by this binding.
285    pub fn checkpoint_keys(&self) -> Vec<&str> {
286        if self.alias_of.is_some() {
287            return Vec::new();
288        }
289        self.recipe.as_ref().map_or_else(
290            || vec![self.checkpoint_key.as_str()],
291            DerivedWeightRecipe::source_keys,
292        )
293    }
294
295    /// Returns the exact logical materialized byte length.
296    pub const fn expected_bytes(&self) -> u64 {
297        self.expected_bytes
298    }
299
300    /// Replaces the physical source with an equivalent validated recipe.
301    pub fn with_source_recipe(
302        mut self,
303        recipe: DerivedWeightRecipe,
304        expected_bytes: u64,
305    ) -> Result<Self, ResidencyDeclarationError> {
306        if self.alias_of.is_some() {
307            return Err(ResidencyDeclarationError::AliasHasPhysicalSource { name: self.name });
308        }
309        validate_size(&self.name, expected_bytes)?;
310        self.checkpoint_key = first_source(&self.name, &recipe)?;
311        self.selection = TensorSelection::Full;
312        self.recipe = Some(recipe);
313        self.expected_bytes = expected_bytes;
314        Ok(self)
315    }
316
317    /// Rewrites one logical output selection into bounded physical sources.
318    pub fn select_bounded_output<C: RecipeCatalog + ?Sized>(
319        self,
320        catalog: &C,
321        selection: TensorSelection,
322    ) -> Result<Self, WeightBindingSelectionError> {
323        let recipe = self.source_recipe().select_bounded(catalog, selection)?;
324        let bytes = recipe.infer(catalog)?.byte_len();
325        Ok(self.with_source_recipe(recipe, bytes)?)
326    }
327}
328
329/// Validated owner/alias partition for one atomic binding unit.
330#[derive(Debug)]
331pub struct WeightBindingPlan<'a> {
332    owners: Vec<&'a WeightBinding>,
333    aliases: Vec<(&'a WeightBinding, &'a WeightBinding)>,
334}
335
336impl<'a> WeightBindingPlan<'a> {
337    /// Validates unique identities, owner existence, cycles, and byte geometry.
338    pub fn new(bindings: &'a [WeightBinding]) -> Result<Self, ResidencyDeclarationError> {
339        let by_name = bindings
340            .iter()
341            .map(|binding| (binding.name(), binding))
342            .collect::<BTreeMap<_, _>>();
343        if by_name.len() != bindings.len() {
344            let duplicate = bindings
345                .iter()
346                .map(WeightBinding::name)
347                .find(|name| bindings.iter().filter(|item| item.name() == *name).count() > 1)
348                .unwrap_or("<unknown>");
349            return Err(ResidencyDeclarationError::DuplicateLogicalBinding {
350                name: duplicate.to_owned(),
351            });
352        }
353
354        fn resolve<'a>(
355            binding: &'a WeightBinding,
356            by_name: &BTreeMap<&str, &'a WeightBinding>,
357            visiting: &mut BTreeSet<String>,
358        ) -> Result<&'a WeightBinding, ResidencyDeclarationError> {
359            let Some(owner_name) = binding.alias_of() else {
360                return Ok(binding);
361            };
362            if !visiting.insert(binding.name().to_owned()) {
363                return Err(ResidencyDeclarationError::BindingAliasCycle {
364                    name: binding.name().to_owned(),
365                });
366            }
367            let owner = by_name.get(owner_name).copied().ok_or_else(|| {
368                ResidencyDeclarationError::UnknownBindingAliasOwner {
369                    alias: binding.name().to_owned(),
370                    owner: owner_name.to_owned(),
371                }
372            })?;
373            let resolved = resolve(owner, by_name, visiting)?;
374            visiting.remove(binding.name());
375            Ok(resolved)
376        }
377
378        let owners = bindings
379            .iter()
380            .filter(|binding| !binding.is_alias())
381            .collect::<Vec<_>>();
382        let mut aliases = Vec::new();
383        for alias in bindings.iter().filter(|binding| binding.is_alias()) {
384            let owner = resolve(alias, &by_name, &mut BTreeSet::new())?;
385            if alias.expected_bytes() != owner.expected_bytes() {
386                return Err(ResidencyDeclarationError::BindingAliasByteMismatch {
387                    alias: alias.name().to_owned(),
388                    owner: owner.name().to_owned(),
389                    alias_bytes: alias.expected_bytes(),
390                    owner_bytes: owner.expected_bytes(),
391                });
392            }
393            aliases.push((alias, owner));
394        }
395        Ok(Self { owners, aliases })
396    }
397
398    /// Canonical bindings which require physical materialization.
399    pub fn owners(&self) -> impl Iterator<Item = &'a WeightBinding> + '_ {
400        self.owners.iter().copied()
401    }
402
403    /// Logical aliases paired with their resolved canonical owners.
404    pub fn aliases(&self) -> impl Iterator<Item = (&'a WeightBinding, &'a WeightBinding)> + '_ {
405        self.aliases.iter().copied()
406    }
407}
408
409/// Failure while rewriting a binding into bounded physical selections.
410#[derive(Debug, thiserror::Error)]
411pub enum WeightBindingSelectionError {
412    /// Neutral recipe inference or selection pushdown failed.
413    #[error(transparent)]
414    Recipe(#[from] RecipeError),
415    /// The rewritten binding declaration was invalid.
416    #[error(transparent)]
417    Declaration(#[from] ResidencyDeclarationError),
418}
419
420/// A deterministic group of weight bindings managed as one atomic unit.
421#[derive(Debug, Clone, Eq, PartialEq)]
422pub struct OffloadUnit {
423    id: OffloadUnitId,
424    bindings: Vec<WeightBinding>,
425}
426
427/// Validated backend-neutral declarations paired with residency control state.
428///
429/// Concrete backends own their native values separately and mirror every
430/// publication or eviction through this controller's ledger. Keeping the
431/// declarations here ensures checkpoint shape policy and plan identity are
432/// validated before any backend allocation begins.
433#[derive(Debug)]
434pub struct ResidencyController {
435    ledger: ResidencyLedger,
436    units: BTreeMap<OffloadUnitId, OffloadUnit>,
437    alias_owners: BTreeMap<(OffloadUnitId, String), (OffloadUnitId, String)>,
438}
439
440/// One validated immutable-weight acquisition batch and its initial hit/miss state.
441#[derive(Debug, Clone, Eq, PartialEq)]
442pub struct ResidencyAcquisition {
443    ids: Vec<OffloadUnitId>,
444    missing: Vec<bool>,
445}
446
447/// Backend-native host/device storage exposed through a neutral residency lease.
448pub trait ResidencyLeaseStorage {
449    /// Backend-native executable value.
450    type DeviceValue;
451    /// Backend-native host-resident value.
452    type HostValue;
453    /// Concrete lookup failure.
454    type Error;
455    /// Allocation-free or cold-path iterator over stable binding names.
456    type BindingNames<'a>: Iterator<Item = &'a str>
457    where
458        Self: 'a;
459
460    /// Looks up one executable binding.
461    fn device_value<'a>(
462        &'a self,
463        id: &OffloadUnitId,
464        name: &str,
465    ) -> Result<&'a Self::DeviceValue, Self::Error>;
466
467    /// Looks up one host-resident binding.
468    fn host_value<'a>(
469        &'a self,
470        id: &OffloadUnitId,
471        name: &str,
472    ) -> Result<&'a Self::HostValue, Self::Error>;
473
474    /// Returns binding names in stable order.
475    fn binding_names(&self) -> Self::BindingNames<'_>;
476}
477
478/// Concrete manager hook used when a residency lease releases its exact pin.
479pub trait ResidencyLeaseOwner: Sized {
480    /// Releases one tier pin. Drop paths must tolerate an already-destroyed manager.
481    fn release_residency_pin(&self, id: &OffloadUnitId, tier: MemoryTier);
482}
483
484/// Backend hook used by the neutral exact-transfer ownership lifecycle.
485pub trait ResidencyTransferOwner<C, R>: Sized {
486    /// Backend executor on which a dependent submission can be ordered.
487    type Executor: ?Sized;
488    /// Concrete completion or publication failure.
489    type Error;
490
491    /// Orders an executor after one exact transfer without blocking the host.
492    fn order_after(
493        completion: &C,
494        executor: &Self::Executor,
495        id: &OffloadUnitId,
496    ) -> Result<(), Self::Error>;
497
498    /// Observes exact completion without blocking.
499    fn is_complete(completion: &C, id: &OffloadUnitId) -> Result<bool, Self::Error>;
500
501    /// Waits for this exact transfer only.
502    fn wait(completion: &C, id: &OffloadUnitId) -> Result<(), Self::Error>;
503
504    /// Releases retained backend resources according to transfer success.
505    fn finish_resources(resources: R, succeeded: bool);
506
507    /// Publishes or rolls back the exact transfer generation.
508    fn resolve_transfer(
509        &self,
510        ids: &[OffloadUnitId],
511        tier: MemoryTier,
512        generation: u64,
513        succeeded: bool,
514    ) -> Result<(), Self::Error>;
515}
516
517/// Caller-owned exact transfer and source-lifetime guard.
518pub struct ResidencyTransfer<L, C, R, O>
519where
520    O: ResidencyTransferOwner<C, R>,
521{
522    leases: Vec<L>,
523    completion: Option<C>,
524    resources: Option<R>,
525    owner: Weak<O>,
526    ids: Vec<OffloadUnitId>,
527    tier: MemoryTier,
528    generation: u64,
529}
530
531impl<L, C, R, O> ResidencyTransfer<L, C, R, O>
532where
533    O: ResidencyTransferOwner<C, R>,
534{
535    /// Creates an already-complete transfer containing only resident leases.
536    pub fn immediate(leases: Vec<L>, tier: MemoryTier) -> Self {
537        Self {
538            leases,
539            completion: None,
540            resources: None,
541            owner: Weak::new(),
542            ids: Vec::new(),
543            tier,
544            generation: 0,
545        }
546    }
547
548    /// Creates an in-flight transfer owning its exact completion and retained resources.
549    pub fn submitted(
550        leases: Vec<L>,
551        completion: C,
552        resources: R,
553        owner: Weak<O>,
554        ids: Vec<OffloadUnitId>,
555        tier: MemoryTier,
556        generation: u64,
557    ) -> Self {
558        assert!(
559            !ids.is_empty(),
560            "an in-flight residency transfer must contain at least one unit"
561        );
562        Self {
563            leases,
564            completion: Some(completion),
565            resources: Some(resources),
566            owner,
567            ids,
568            tier,
569            generation,
570        }
571    }
572
573    /// Returns resident unit leases protected by this transfer.
574    pub fn leases(&self) -> &[L] {
575        &self.leases
576    }
577
578    /// Returns whether the transfer contains no resident units.
579    pub fn is_empty(&self) -> bool {
580        self.leases.is_empty()
581    }
582
583    /// Orders dependent backend work after this exact transfer.
584    pub fn order_after(&self, executor: &O::Executor) -> Result<(), O::Error> {
585        match &self.completion {
586            Some(completion) => O::order_after(completion, executor, self.primary_id()),
587            None => Ok(()),
588        }
589    }
590
591    /// Returns whether this exact transfer completed without blocking.
592    pub fn is_complete(&self) -> Result<bool, O::Error> {
593        match &self.completion {
594            Some(completion) => O::is_complete(completion, self.primary_id()),
595            None => Ok(true),
596        }
597    }
598
599    /// Waits for exact completion and publishes or rolls back the transfer.
600    pub fn synchronize(&mut self) -> Result<(), O::Error> {
601        self.finish(true)
602    }
603
604    fn primary_id(&self) -> &OffloadUnitId {
605        self.ids
606            .first()
607            .expect("an in-flight residency transfer has at least one unit")
608    }
609
610    fn finish(&mut self, report_error: bool) -> Result<(), O::Error> {
611        let result = match &self.completion {
612            Some(completion) => O::wait(completion, self.primary_id()),
613            None => Ok(()),
614        };
615        let succeeded = result.is_ok();
616        if let Some(resources) = self.resources.take() {
617            O::finish_resources(resources, succeeded);
618        }
619        if succeeded {
620            self.completion = None;
621        }
622        self.publish(succeeded)?;
623        match result {
624            Ok(()) => Ok(()),
625            Err(error) if report_error => Err(error),
626            Err(_) => Ok(()),
627        }
628    }
629
630    fn publish(&mut self, succeeded: bool) -> Result<(), O::Error> {
631        if self.generation == 0 {
632            return Ok(());
633        }
634        let generation = std::mem::take(&mut self.generation);
635        if let Some(owner) = self.owner.upgrade() {
636            owner.resolve_transfer(&self.ids, self.tier, generation, succeeded)?;
637        }
638        Ok(())
639    }
640}
641
642impl<L, C, R, O> Drop for ResidencyTransfer<L, C, R, O>
643where
644    O: ResidencyTransferOwner<C, R>,
645{
646    fn drop(&mut self) {
647        let _ = self.finish(false);
648    }
649}
650
651/// Statically dispatched lease retaining one backend-native resident unit.
652pub struct ResidencyLease<S, O>
653where
654    S: ResidencyLeaseStorage,
655    O: ResidencyLeaseOwner,
656{
657    id: OffloadUnitId,
658    tier: MemoryTier,
659    storage: S,
660    owner: Weak<O>,
661}
662
663impl<S, O> ResidencyLease<S, O>
664where
665    S: ResidencyLeaseStorage,
666    O: ResidencyLeaseOwner,
667{
668    /// Creates a lease after the neutral controller has pinned the resident copy.
669    pub fn new(id: OffloadUnitId, tier: MemoryTier, storage: S, owner: Weak<O>) -> Self {
670        Self {
671            id,
672            tier,
673            storage,
674            owner,
675        }
676    }
677
678    /// Returns the acquired unit identifier.
679    pub fn id(&self) -> &OffloadUnitId {
680        &self.id
681    }
682
683    /// Returns the protected resident tier.
684    pub const fn tier(&self) -> MemoryTier {
685        self.tier
686    }
687
688    /// Looks up one backend-native executable binding without cloning it.
689    pub fn device_value(&self, name: &str) -> Result<&S::DeviceValue, S::Error> {
690        self.storage.device_value(&self.id, name)
691    }
692
693    /// Looks up one backend-native host binding without cloning it.
694    pub fn host_value(&self, name: &str) -> Result<&S::HostValue, S::Error> {
695        self.storage.host_value(&self.id, name)
696    }
697
698    /// Returns binding names in stable order.
699    pub fn binding_names(&self) -> S::BindingNames<'_> {
700        self.storage.binding_names()
701    }
702}
703
704impl<S, O> Drop for ResidencyLease<S, O>
705where
706    S: ResidencyLeaseStorage,
707    O: ResidencyLeaseOwner,
708{
709    fn drop(&mut self) {
710        if let Some(owner) = self.owner.upgrade() {
711            owner.release_residency_pin(&self.id, self.tier);
712        }
713    }
714}
715
716impl ResidencyAcquisition {
717    /// Returns requested units in caller order.
718    pub fn ids(&self) -> &[OffloadUnitId] {
719        &self.ids
720    }
721
722    /// Returns one flag per requested unit; `true` requires backend realization.
723    pub fn missing(&self) -> &[bool] {
724        &self.missing
725    }
726
727    /// Returns whether every requested copy was already resident.
728    pub fn is_hit(&self) -> bool {
729        self.missing.iter().all(|missing| !missing)
730    }
731
732    /// Returns missing units in caller order.
733    pub fn missing_ids(&self) -> impl Iterator<Item = &OffloadUnitId> {
734        self.ids
735            .iter()
736            .zip(&self.missing)
737            .filter_map(|(id, &missing)| missing.then_some(id))
738    }
739
740    fn temporary_protection(&self) -> BTreeSet<OffloadUnitId> {
741        self.ids.iter().cloned().collect()
742    }
743}
744
745impl ResidencyController {
746    /// Validates declarations against checkpoint metadata and an explicit plan.
747    pub fn new<C: RecipeCatalog + ?Sized>(
748        catalog: &C,
749        plan: OffloadPlan,
750        units: impl IntoIterator<Item = OffloadUnit>,
751    ) -> Result<Self, ResidencyControllerError> {
752        let mut definitions = BTreeMap::new();
753        for unit in units {
754            let id = unit.id().clone();
755            if definitions.insert(id.clone(), unit).is_some() {
756                return Err(ResidencyControllerError::DuplicateUnitDefinition { id });
757            }
758        }
759        for spec in plan.units() {
760            if !definitions.contains_key(spec.id()) {
761                return Err(ResidencyControllerError::MissingUnitDefinition {
762                    id: spec.id().clone(),
763                });
764            }
765        }
766        if let Some(id) = definitions
767            .keys()
768            .find(|id| plan.unit(id).is_none())
769            .cloned()
770        {
771            return Err(ResidencyControllerError::UnexpectedUnitDefinition { id });
772        }
773
774        let alias_owners = validate_global_binding_aliases(&definitions)?;
775
776        for spec in plan.units() {
777            let unit = definitions
778                .get(spec.id())
779                .expect("definition identity validated above");
780            let mut total = 0u64;
781            for binding in unit.bindings().iter().filter(|binding| !binding.is_alias()) {
782                total = total.checked_add(binding.expected_bytes()).ok_or(
783                    ResidencyControllerError::ArithmeticOverflow {
784                        context: "unit binding byte total",
785                    },
786                )?;
787                if !binding.is_alias() {
788                    let actual = binding
789                        .source_recipe()
790                        .infer(catalog)
791                        .map_err(|source| ResidencyControllerError::Recipe {
792                            binding: binding.name().to_owned(),
793                            source,
794                        })?
795                        .byte_len();
796                    if actual != binding.expected_bytes() {
797                        return Err(ResidencyControllerError::BindingByteMismatch {
798                            id: unit.id().clone(),
799                            binding: binding.name().to_owned(),
800                            expected_bytes: binding.expected_bytes(),
801                            actual_bytes: actual,
802                        });
803                    }
804                }
805            }
806            if total != spec.bytes() {
807                return Err(ResidencyControllerError::UnitByteMismatch {
808                    id: unit.id().clone(),
809                    planned_bytes: spec.bytes(),
810                    actual_bytes: total,
811                });
812            }
813        }
814
815        Ok(Self {
816            ledger: ResidencyLedger::new(plan),
817            units: definitions,
818            alias_owners,
819        })
820    }
821
822    /// Returns the validated declaration for one planned unit.
823    pub fn unit(&self, id: &OffloadUnitId) -> Option<&OffloadUnit> {
824        self.units.get(id)
825    }
826
827    /// Returns declarations in stable unit-identifier order.
828    pub fn units(&self) -> impl ExactSizeIterator<Item = &OffloadUnit> {
829        self.units.values()
830    }
831
832    /// Resolves a logical alias to its canonical owner unit and binding.
833    pub fn binding_owner(
834        &self,
835        unit: &OffloadUnitId,
836        binding: &WeightBinding,
837    ) -> Option<(&OffloadUnitId, &WeightBinding)> {
838        let (owner_unit, owner_name) = self
839            .alias_owners
840            .get(&(unit.clone(), binding.name().to_owned()))?;
841        let owner = self.units.get(owner_unit)?;
842        let binding = owner
843            .bindings()
844            .iter()
845            .find(|binding| binding.name() == owner_name)?;
846        Some((owner_unit, binding))
847    }
848
849    /// Returns the canonical owner location when a binding participates in a
850    /// shared alias family, including the canonical owner itself.
851    pub fn shared_binding_owner(
852        &self,
853        unit: &OffloadUnitId,
854        binding: &WeightBinding,
855    ) -> Option<(&OffloadUnitId, &WeightBinding)> {
856        if binding.is_alias() {
857            return self.binding_owner(unit, binding);
858        }
859        let location = (unit.clone(), binding.name().to_owned());
860        if !self.alias_owners.values().any(|owner| owner == &location) {
861            return None;
862        }
863        let (owner_unit, owner) = self.units.get_key_value(unit)?;
864        let owner = owner
865            .bindings()
866            .iter()
867            .find(|candidate| candidate.name() == binding.name())?;
868        Some((owner_unit, owner))
869    }
870
871    /// Returns immutable ownership, capacity, and telemetry state.
872    pub const fn ledger(&self) -> &ResidencyLedger {
873        &self.ledger
874    }
875
876    /// Returns mutable ownership, capacity, and telemetry state.
877    pub fn ledger_mut(&mut self) -> &mut ResidencyLedger {
878        &mut self.ledger
879    }
880
881    /// Validates one batch and snapshots which requested copies need realization.
882    pub fn plan_acquisition(
883        &mut self,
884        ids: &[OffloadUnitId],
885        tier: MemoryTier,
886    ) -> Result<ResidencyAcquisition, ResidencyLedgerError> {
887        self.ledger.require_initialized()?;
888        self.plan_acquisition_inner(ids, tier)
889    }
890
891    /// Validates a batch while the manager is realizing its initial planned tiers.
892    pub fn plan_initialization_acquisition(
893        &mut self,
894        ids: &[OffloadUnitId],
895        tier: MemoryTier,
896    ) -> Result<ResidencyAcquisition, ResidencyLedgerError> {
897        self.plan_acquisition_inner(ids, tier)
898    }
899
900    fn plan_acquisition_inner(
901        &mut self,
902        ids: &[OffloadUnitId],
903        tier: MemoryTier,
904    ) -> Result<ResidencyAcquisition, ResidencyLedgerError> {
905        self.ledger.validate_batch(ids, tier)?;
906        let missing = ids
907            .iter()
908            .map(|id| self.ledger.is_resident(id, tier).map(|resident| !resident))
909            .collect::<Result<Vec<_>, _>>()?;
910        Ok(ResidencyAcquisition {
911            ids: ids.to_vec(),
912            missing,
913        })
914    }
915
916    /// Reserves backend-supplied physical capacities while protecting the complete batch.
917    pub fn reserve_acquisition(
918        &mut self,
919        acquisition: &ResidencyAcquisition,
920        reservations: &[(OffloadUnitId, u64)],
921        tier: MemoryTier,
922    ) -> Result<Vec<EvictedResidencyCopy>, ResidencyLedgerError> {
923        self.ledger
924            .reserve_copies(reservations, tier, &acquisition.temporary_protection())
925    }
926
927    /// Updates recency for copies which were hits when the batch began.
928    pub fn touch_acquisition_hits(
929        &mut self,
930        acquisition: &ResidencyAcquisition,
931        tier: MemoryTier,
932    ) -> Result<(), ResidencyLedgerError> {
933        for (id, &missing) in acquisition.ids.iter().zip(&acquisition.missing) {
934            if !missing {
935                self.ledger.touch(id, tier)?;
936            }
937        }
938        Ok(())
939    }
940
941    /// Rolls back every missing copy which remains an unpublished reservation.
942    pub fn rollback_acquisition(
943        &mut self,
944        acquisition: &ResidencyAcquisition,
945        tier: MemoryTier,
946    ) -> Result<(), ResidencyLedgerError> {
947        for id in acquisition.missing_ids() {
948            self.ledger.rollback_reserved(id, tier)?;
949        }
950        Ok(())
951    }
952
953    /// Publishes one realized copy and records its backend transfer observation.
954    #[allow(clippy::too_many_arguments)]
955    pub fn publish_acquisition_copy(
956        &mut self,
957        id: &OffloadUnitId,
958        tier: MemoryTier,
959        actual_bytes: u64,
960        transferred_bytes: u64,
961        transfer_generation: Option<u64>,
962        direction: eredu_core::residency::TransferDirection,
963        duration: Duration,
964    ) -> Result<(), ResidencyLedgerError> {
965        self.ledger
966            .publish_reserved(id, tier, actual_bytes, transfer_generation)?;
967        self.ledger
968            .record_transfer(direction, transferred_bytes, duration);
969        Ok(())
970    }
971
972    /// Records a prefetch hit or miss before backend realization begins.
973    pub fn begin_prefetch(
974        &mut self,
975        id: &OffloadUnitId,
976        tier: MemoryTier,
977    ) -> Result<PrefetchOutcome, ResidencyLedgerError> {
978        self.ledger.require_initialized()?;
979        let outcome = if self.ledger.is_resident(id, tier)? {
980            PrefetchOutcome::Hit
981        } else {
982            PrefetchOutcome::Miss
983        };
984        self.ledger.record_prefetch(tier, outcome);
985        Ok(outcome)
986    }
987
988    /// Resolves one exact transfer and returns backend copies invalidated by failure.
989    pub fn resolve_transfer(
990        &mut self,
991        ids: &[OffloadUnitId],
992        tier: MemoryTier,
993        generation: u64,
994        succeeded: bool,
995    ) -> Result<Vec<EvictedResidencyCopy>, ResidencyLedgerError> {
996        self.ledger
997            .resolve_transfer(ids, tier, generation, succeeded)
998    }
999
1000    /// Replaces one protected window and selects unique bounded lookahead in caller order.
1001    ///
1002    /// A concrete backend calls this after any in-flight copies touching the requested
1003    /// units have reached a stable state.
1004    pub fn commit_group_window(
1005        &mut self,
1006        group: &str,
1007        active: &[OffloadUnitId],
1008        upcoming: &[OffloadUnitId],
1009        tier: MemoryTier,
1010    ) -> Result<Vec<OffloadUnitId>, eredu_core::residency::ResidencyLedgerError> {
1011        self.ledger.require_initialized()?;
1012        for id in active.iter().chain(upcoming) {
1013            self.ledger.spec(id)?;
1014        }
1015        self.ledger.set_group_window(group, active, tier)?;
1016        let depth = self.ledger.plan().config().prefetch_depth();
1017        let mut seen = BTreeSet::new();
1018        Ok(upcoming
1019            .iter()
1020            .filter(|id| seen.insert((*id).clone()))
1021            .take(depth)
1022            .cloned()
1023            .collect())
1024    }
1025
1026    /// Replaces one protected window without selecting or materializing lookahead.
1027    pub fn protect_group_window(
1028        &mut self,
1029        group: &str,
1030        active: &[OffloadUnitId],
1031        tier: MemoryTier,
1032    ) -> Result<(), eredu_core::residency::ResidencyLedgerError> {
1033        self.commit_group_window(group, active, &[], tier)
1034            .map(|_| ())
1035    }
1036}
1037
1038type BindingLocation = (OffloadUnitId, String);
1039type BindingAliasMap = BTreeMap<BindingLocation, BindingLocation>;
1040
1041fn validate_global_binding_aliases(
1042    units: &BTreeMap<OffloadUnitId, OffloadUnit>,
1043) -> Result<BindingAliasMap, ResidencyControllerError> {
1044    let mut identities = BTreeMap::<String, Vec<BindingLocation>>::new();
1045    let mut aliases = BTreeMap::<BindingLocation, String>::new();
1046    let mut bytes = BTreeMap::<BindingLocation, u64>::new();
1047    for (unit_id, unit) in units {
1048        for binding in unit.bindings() {
1049            let location = (unit_id.clone(), binding.name().to_owned());
1050            let identity = binding
1051                .logical_target()
1052                .unwrap_or(binding.name())
1053                .to_owned();
1054            identities
1055                .entry(identity)
1056                .or_default()
1057                .push(location.clone());
1058            bytes.insert(location.clone(), binding.expected_bytes());
1059            if let Some(owner) = binding.alias_of() {
1060                aliases.insert(location, owner.to_owned());
1061            }
1062        }
1063    }
1064
1065    fn resolve(
1066        location: &BindingLocation,
1067        identities: &BTreeMap<String, Vec<BindingLocation>>,
1068        aliases: &BTreeMap<BindingLocation, String>,
1069        visiting: &mut BTreeSet<BindingLocation>,
1070    ) -> Result<BindingLocation, ResidencyControllerError> {
1071        let Some(destination) = aliases.get(location) else {
1072            return Ok(location.clone());
1073        };
1074        if !visiting.insert(location.clone()) {
1075            return Err(ResidencyControllerError::Declaration(
1076                ResidencyDeclarationError::BindingAliasCycle {
1077                    name: location.1.clone(),
1078                },
1079            ));
1080        }
1081        let candidates = identities.get(destination).ok_or_else(|| {
1082            ResidencyControllerError::Declaration(
1083                ResidencyDeclarationError::UnknownBindingAliasOwner {
1084                    alias: location.1.clone(),
1085                    owner: destination.clone(),
1086                },
1087            )
1088        })?;
1089        if candidates.len() != 1 {
1090            return Err(ResidencyControllerError::Declaration(
1091                ResidencyDeclarationError::AmbiguousBindingAliasOwner {
1092                    alias: location.1.clone(),
1093                    owner: destination.clone(),
1094                },
1095            ));
1096        }
1097        let owner = resolve(&candidates[0], identities, aliases, visiting)?;
1098        visiting.remove(location);
1099        Ok(owner)
1100    }
1101
1102    let mut resolved = BTreeMap::new();
1103    for alias in aliases.keys() {
1104        let owner = resolve(alias, &identities, &aliases, &mut BTreeSet::new())?;
1105        let alias_bytes = bytes[alias];
1106        let owner_bytes = bytes[&owner];
1107        if alias_bytes != owner_bytes {
1108            return Err(ResidencyControllerError::Declaration(
1109                ResidencyDeclarationError::BindingAliasByteMismatch {
1110                    alias: alias.1.clone(),
1111                    owner: owner.1.clone(),
1112                    alias_bytes,
1113                    owner_bytes,
1114                },
1115            ));
1116        }
1117        resolved.insert(alias.clone(), owner);
1118    }
1119    Ok(resolved)
1120}
1121
1122/// Failure while validating a residency control plane.
1123#[derive(Debug, thiserror::Error)]
1124pub enum ResidencyControllerError {
1125    /// A binding alias graph was invalid.
1126    #[error(transparent)]
1127    Declaration(#[from] ResidencyDeclarationError),
1128    /// More than one definition used the same plan identifier.
1129    #[error("duplicate residency unit definition: {id}")]
1130    DuplicateUnitDefinition {
1131        /// Duplicated identifier.
1132        id: OffloadUnitId,
1133    },
1134    /// The plan had no matching unit definition.
1135    #[error("offload plan unit {id} has no residency unit definition")]
1136    MissingUnitDefinition {
1137        /// Missing identifier.
1138        id: OffloadUnitId,
1139    },
1140    /// A definition had no matching plan entry.
1141    #[error("residency unit {id} is absent from the offload plan")]
1142    UnexpectedUnitDefinition {
1143        /// Unexpected identifier.
1144        id: OffloadUnitId,
1145    },
1146    /// Binding sizes did not sum to the plan's unit size.
1147    #[error(
1148        "residency unit {id} defines {actual_bytes} bytes but its plan reserves {planned_bytes}"
1149    )]
1150    UnitByteMismatch {
1151        /// Unit identifier.
1152        id: OffloadUnitId,
1153        /// Bytes reserved by the plan.
1154        planned_bytes: u64,
1155        /// Sum of binding sizes.
1156        actual_bytes: u64,
1157    },
1158    /// A binding's selected checkpoint size contradicted its declaration.
1159    #[error(
1160        "binding {binding:?} in unit {id} selects {actual_bytes} bytes but declares {expected_bytes}"
1161    )]
1162    BindingByteMismatch {
1163        /// Unit identifier.
1164        id: OffloadUnitId,
1165        /// Binding name.
1166        binding: String,
1167        /// Declared size.
1168        expected_bytes: u64,
1169        /// Catalog-validated size.
1170        actual_bytes: u64,
1171    },
1172    /// A derived-weight recipe was invalid.
1173    #[error("derived-weight recipe for binding {binding:?} failed: {source}")]
1174    Recipe {
1175        /// Local binding name.
1176        binding: String,
1177        /// Invalid recipe.
1178        #[source]
1179        source: RecipeError,
1180    },
1181    /// Checked byte arithmetic overflowed.
1182    #[error("residency arithmetic overflow: {context}")]
1183    ArithmeticOverflow {
1184        /// Calculation that overflowed.
1185        context: &'static str,
1186    },
1187}
1188
1189/// Backend-independent operations required by ordered residency windows.
1190pub trait ResidencyWindowManager {
1191    /// Manager-specific failure including neutral window validation failures.
1192    type Error: std::error::Error + From<ResidencyWindowError>;
1193
1194    /// Replaces the default protected window and prepares bounded lookahead.
1195    fn prepare_window(
1196        &self,
1197        active: &[OffloadUnitId],
1198        upcoming: &[OffloadUnitId],
1199        tier: MemoryTier,
1200    ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, Self::Error>;
1201
1202    /// Replaces one named protected window and prepares bounded lookahead.
1203    fn prepare_group_window(
1204        &self,
1205        group: &str,
1206        active: &[OffloadUnitId],
1207        upcoming: &[OffloadUnitId],
1208        tier: MemoryTier,
1209    ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, Self::Error>;
1210
1211    /// Removes one concrete resident copy if present.
1212    fn evict(&self, id: &OffloadUnitId, tier: MemoryTier) -> Result<bool, Self::Error>;
1213
1214    /// Returns logical unit state in stable identifier order.
1215    fn unit_reports(&self) -> Result<Vec<UnitResidencyReport>, Self::Error>;
1216}
1217
1218/// Deterministic controller for a bounded ordered device-layer window.
1219#[derive(Debug, Clone)]
1220pub struct DeviceLayerWindow {
1221    units: Vec<OffloadUnitId>,
1222    depth: usize,
1223}
1224
1225impl DeviceLayerWindow {
1226    /// Creates a controller for a non-empty, duplicate-free unit sequence.
1227    pub fn new(
1228        units: impl IntoIterator<Item = OffloadUnitId>,
1229        depth: usize,
1230    ) -> Result<Self, ResidencyWindowError> {
1231        let units = units.into_iter().collect::<Vec<_>>();
1232        if units.is_empty() {
1233            return Err(ResidencyWindowError::EmptyLayerWindow);
1234        }
1235        if depth == 0 || depth > units.len() {
1236            return Err(ResidencyWindowError::OversizedLayerWindow {
1237                depth,
1238                layer_count: units.len(),
1239            });
1240        }
1241        let unique = units.iter().collect::<BTreeSet<_>>();
1242        if unique.len() != units.len() {
1243            return Err(ResidencyWindowError::DuplicateLayerWindowUnit {
1244                id: units
1245                    .iter()
1246                    .find(|id| units.iter().filter(|candidate| *candidate == *id).count() > 1)
1247                    .expect("duplicate exists")
1248                    .clone(),
1249            });
1250        }
1251        Ok(Self { units, depth })
1252    }
1253
1254    /// Returns the maximum number of ordered units kept on the device.
1255    pub const fn depth(&self) -> usize {
1256        self.depth
1257    }
1258
1259    /// Returns units in execution order.
1260    pub fn units(&self) -> &[OffloadUnitId] {
1261        &self.units
1262    }
1263
1264    /// Returns the desired window beginning at `current`.
1265    pub fn desired(&self, current: usize) -> Result<&[OffloadUnitId], ResidencyWindowError> {
1266        if current >= self.units.len() {
1267            return Err(ResidencyWindowError::InvalidLayerIndex {
1268                index: current,
1269                layer_count: self.units.len(),
1270            });
1271        }
1272        let end = current.saturating_add(self.depth).min(self.units.len());
1273        Ok(&self.units[current..end])
1274    }
1275
1276    /// Prepares and trims the default window beginning at `current`.
1277    pub fn prepare<M: ResidencyWindowManager>(
1278        &self,
1279        manager: &M,
1280        current: usize,
1281    ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, M::Error> {
1282        let desired = self.desired(current).map_err(M::Error::from)?;
1283        let outcomes = manager.prepare_window(desired, desired, MemoryTier::Device)?;
1284        self.trim_to(manager, desired)?;
1285        Ok(outcomes)
1286    }
1287
1288    /// Explicitly evicts every managed device copy outside `desired`.
1289    pub fn trim_to<M: ResidencyWindowManager>(
1290        &self,
1291        manager: &M,
1292        desired: &[OffloadUnitId],
1293    ) -> Result<(), M::Error> {
1294        let desired = desired.iter().collect::<BTreeSet<_>>();
1295        for id in &self.units {
1296            if !desired.contains(id) {
1297                manager.evict(id, MemoryTier::Device)?;
1298            }
1299        }
1300        Ok(())
1301    }
1302
1303    /// Clears protection and removes every managed device copy.
1304    pub fn clear<M: ResidencyWindowManager>(&self, manager: &M) -> Result<(), M::Error> {
1305        manager.prepare_window(&[], &[], MemoryTier::Device)?;
1306        self.trim_to(manager, &[])
1307    }
1308}
1309
1310/// A named sequential execution stack with an independent device window.
1311#[derive(Debug, Clone)]
1312pub struct ResidentLayerGroup {
1313    id: String,
1314    window: DeviceLayerWindow,
1315}
1316
1317impl ResidentLayerGroup {
1318    /// Creates a named group over ordered residency units.
1319    pub fn new(
1320        id: impl Into<String>,
1321        units: impl IntoIterator<Item = OffloadUnitId>,
1322        depth: usize,
1323    ) -> Result<Self, ResidencyWindowError> {
1324        let id = id.into();
1325        if id.trim().is_empty() {
1326            return Err(ResidencyWindowError::InvalidGroupId);
1327        }
1328        Ok(Self {
1329            id,
1330            window: DeviceLayerWindow::new(units, depth)?,
1331        })
1332    }
1333
1334    /// Returns the stable group identifier.
1335    pub fn id(&self) -> &str {
1336        &self.id
1337    }
1338
1339    /// Returns ordered units in this group.
1340    pub fn units(&self) -> &[OffloadUnitId] {
1341        self.window.units()
1342    }
1343
1344    /// Returns the configured device-unit bound.
1345    pub const fn depth(&self) -> usize {
1346        self.window.depth()
1347    }
1348
1349    /// Prepares this group's window without replacing another group's window.
1350    pub fn prepare<M: ResidencyWindowManager>(
1351        &self,
1352        manager: &M,
1353        current: usize,
1354    ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, M::Error> {
1355        let desired = self.window.desired(current).map_err(M::Error::from)?;
1356        let outcomes =
1357            manager.prepare_group_window(&self.id, desired, desired, MemoryTier::Device)?;
1358        self.window.trim_to(manager, desired)?;
1359        Ok(outcomes)
1360    }
1361
1362    /// Trims this group to the desired window.
1363    pub fn trim_to<M: ResidencyWindowManager>(
1364        &self,
1365        manager: &M,
1366        desired: &[OffloadUnitId],
1367    ) -> Result<(), M::Error> {
1368        self.window.trim_to(manager, desired)
1369    }
1370
1371    /// Clears only this group's protection and device copies.
1372    pub fn clear<M: ResidencyWindowManager>(&self, manager: &M) -> Result<(), M::Error> {
1373        manager.prepare_group_window(&self.id, &[], &[], MemoryTier::Device)?;
1374        self.window.trim_to(manager, &[])
1375    }
1376
1377    /// Returns current logical residency attributed to this group's units.
1378    pub fn report<M: ResidencyWindowManager>(
1379        &self,
1380        manager: &M,
1381    ) -> Result<ResidentLayerGroupReport, M::Error> {
1382        let ids = self.units().iter().collect::<BTreeSet<_>>();
1383        let mut host_bytes = 0u64;
1384        let mut device_bytes = 0u64;
1385        let mut device_units = 0usize;
1386        for unit in manager
1387            .unit_reports()?
1388            .iter()
1389            .filter(|unit| ids.contains(unit.id()))
1390        {
1391            if unit.host_resident() {
1392                host_bytes = host_bytes
1393                    .checked_add(unit.host_allocated_bytes())
1394                    .ok_or(ResidencyWindowError::ArithmeticOverflow {
1395                        context: "execution group host bytes",
1396                    })
1397                    .map_err(M::Error::from)?;
1398            }
1399            if unit.device_resident() {
1400                device_bytes = device_bytes
1401                    .checked_add(unit.device_allocated_bytes())
1402                    .ok_or(ResidencyWindowError::ArithmeticOverflow {
1403                        context: "execution group device bytes",
1404                    })
1405                    .map_err(M::Error::from)?;
1406                device_units += 1;
1407            }
1408        }
1409        Ok(ResidentLayerGroupReport {
1410            id: self.id.clone(),
1411            ordered_units: self.units().len(),
1412            window_depth: self.depth(),
1413            host_bytes,
1414            device_bytes,
1415            device_units,
1416        })
1417    }
1418}
1419
1420/// Logical residency attributed to one named execution group.
1421#[derive(Debug, Clone, Eq, PartialEq)]
1422pub struct ResidentLayerGroupReport {
1423    id: String,
1424    ordered_units: usize,
1425    window_depth: usize,
1426    host_bytes: u64,
1427    device_bytes: u64,
1428    device_units: usize,
1429}
1430
1431impl ResidentLayerGroupReport {
1432    /// Returns the group identifier.
1433    pub fn id(&self) -> &str {
1434        &self.id
1435    }
1436    /// Returns the number of ordered units.
1437    pub const fn ordered_units(&self) -> usize {
1438        self.ordered_units
1439    }
1440    /// Returns the configured maximum device-unit count.
1441    pub const fn window_depth(&self) -> usize {
1442        self.window_depth
1443    }
1444    /// Returns current physical host allocation capacity for group units.
1445    pub const fn host_bytes(&self) -> u64 {
1446        self.host_bytes
1447    }
1448    /// Returns current device-resident bytes for group units.
1449    pub const fn device_bytes(&self) -> u64 {
1450        self.device_bytes
1451    }
1452    /// Returns current device-resident group units.
1453    pub const fn device_units(&self) -> usize {
1454        self.device_units
1455    }
1456}
1457
1458/// Invalid ordered residency-window configuration or accounting.
1459#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1460pub enum ResidencyWindowError {
1461    /// A layer window had no units.
1462    #[error("device layer window requires at least one ordered unit")]
1463    EmptyLayerWindow,
1464    /// A device window depth was zero or exceeded its unit count.
1465    #[error("device layer window depth {depth} exceeds {layer_count} ordered units")]
1466    OversizedLayerWindow {
1467        /// Requested resident-unit bound.
1468        depth: usize,
1469        /// Available ordered units.
1470        layer_count: usize,
1471    },
1472    /// A requested current unit was outside the sequence.
1473    #[error("device layer index {index} is outside {layer_count} ordered units")]
1474    InvalidLayerIndex {
1475        /// Requested index.
1476        index: usize,
1477        /// Available ordered units.
1478        layer_count: usize,
1479    },
1480    /// An ordered layer window repeated a unit identifier.
1481    #[error("device layer window contains duplicate unit {id}")]
1482    DuplicateLayerWindowUnit {
1483        /// Duplicated identifier.
1484        id: OffloadUnitId,
1485    },
1486    /// A named execution group had an empty identifier.
1487    #[error("residency window group identifiers must not be empty")]
1488    InvalidGroupId,
1489    /// Checked group accounting overflowed.
1490    #[error("residency window arithmetic overflow: {context}")]
1491    ArithmeticOverflow {
1492        /// Calculation that overflowed.
1493        context: &'static str,
1494    },
1495}
1496
1497impl OffloadUnit {
1498    /// Creates a non-empty unit and sorts bindings by local name.
1499    pub fn new(
1500        id: OffloadUnitId,
1501        bindings: impl IntoIterator<Item = WeightBinding>,
1502    ) -> Result<Self, ResidencyDeclarationError> {
1503        let mut bindings = bindings.into_iter().collect::<Vec<_>>();
1504        if bindings.is_empty() {
1505            return Err(ResidencyDeclarationError::EmptyUnit { id });
1506        }
1507        bindings.sort_by(|left, right| left.name.cmp(&right.name));
1508        if let Some(pair) = bindings
1509            .windows(2)
1510            .find(|pair| pair[0].name == pair[1].name)
1511        {
1512            return Err(ResidencyDeclarationError::DuplicateBindingName {
1513                id,
1514                name: pair[0].name.clone(),
1515            });
1516        }
1517        Ok(Self { id, bindings })
1518    }
1519
1520    /// Returns the plan identifier for this unit.
1521    pub fn id(&self) -> &OffloadUnitId {
1522        &self.id
1523    }
1524
1525    /// Returns bindings in stable local-name order.
1526    pub fn bindings(&self) -> &[WeightBinding] {
1527        &self.bindings
1528    }
1529}
1530
1531fn validate_name(name: String) -> Result<String, ResidencyDeclarationError> {
1532    if name.trim().is_empty() {
1533        Err(ResidencyDeclarationError::InvalidBindingName)
1534    } else {
1535        Ok(name)
1536    }
1537}
1538
1539fn validate_size(name: &str, expected_bytes: u64) -> Result<(), ResidencyDeclarationError> {
1540    if expected_bytes == 0 {
1541        Err(ResidencyDeclarationError::ZeroSizedBinding {
1542            name: name.to_owned(),
1543        })
1544    } else {
1545        Ok(())
1546    }
1547}
1548
1549fn first_source(
1550    name: &str,
1551    recipe: &DerivedWeightRecipe,
1552) -> Result<String, ResidencyDeclarationError> {
1553    recipe
1554        .source_keys()
1555        .first()
1556        .map(|key| (*key).to_owned())
1557        .ok_or_else(|| ResidencyDeclarationError::EmptyRecipeSources {
1558            name: name.to_owned(),
1559        })
1560}
1561
1562/// Invalid backend-neutral residency declaration.
1563#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1564pub enum ResidencyDeclarationError {
1565    /// A binding name was empty.
1566    #[error("weight binding names must not be empty")]
1567    InvalidBindingName,
1568    /// A binding checkpoint key was empty.
1569    #[error("weight binding {name:?} has an empty checkpoint key")]
1570    InvalidCheckpointKey {
1571        /// Invalid local name.
1572        name: String,
1573    },
1574    /// A recipe had no physical source.
1575    #[error("weight binding {name:?} has no checkpoint recipe source")]
1576    EmptyRecipeSources {
1577        /// Invalid local name.
1578        name: String,
1579    },
1580    /// Load-time quantization companion names collide with one another.
1581    #[error(
1582        "weight binding {name:?} has invalid quantization companions {scales:?} and {biases:?}"
1583    )]
1584    InvalidQuantizationCompanions {
1585        /// Quantizable weight binding.
1586        name: String,
1587        /// Declared scale binding.
1588        scales: String,
1589        /// Declared affine-bias binding.
1590        biases: String,
1591    },
1592    /// A binding declared no bytes.
1593    #[error("weight binding {name:?} must contain at least one byte")]
1594    ZeroSizedBinding {
1595        /// Invalid local name.
1596        name: String,
1597    },
1598    /// A unit had no bindings.
1599    #[error("residency unit {id} must contain at least one binding")]
1600    EmptyUnit {
1601        /// Unit identifier.
1602        id: OffloadUnitId,
1603    },
1604    /// Two bindings in one unit had the same local name.
1605    #[error("residency unit {id} has duplicate binding name {name:?}")]
1606    DuplicateBindingName {
1607        /// Unit identifier.
1608        id: OffloadUnitId,
1609        /// Duplicated local name.
1610        name: String,
1611    },
1612    /// A general binding list repeated one logical identity.
1613    #[error("duplicate logical weight binding {name:?}")]
1614    DuplicateLogicalBinding {
1615        /// Repeated logical identity.
1616        name: String,
1617    },
1618    /// An alias named no binding in its atomic unit.
1619    #[error("weight binding alias {alias:?} has unknown owner {owner:?}")]
1620    UnknownBindingAliasOwner {
1621        /// Alias identity.
1622        alias: String,
1623        /// Missing owner identity.
1624        owner: String,
1625    },
1626    /// An alias destination did not identify one unique global owner.
1627    #[error("weight binding alias {alias:?} has ambiguous owner {owner:?}")]
1628    AmbiguousBindingAliasOwner {
1629        /// Alias identity.
1630        alias: String,
1631        /// Ambiguous owner identity.
1632        owner: String,
1633    },
1634    /// Alias declarations formed a cycle.
1635    #[error("weight binding alias cycle contains {name:?}")]
1636    BindingAliasCycle {
1637        /// One member of the cycle.
1638        name: String,
1639    },
1640    /// Alias and resolved owner disagreed on materialized byte geometry.
1641    #[error("weight binding alias {alias:?} declares {alias_bytes} bytes but owner {owner:?} declares {owner_bytes}")]
1642    BindingAliasByteMismatch {
1643        /// Alias identity.
1644        alias: String,
1645        /// Resolved canonical owner.
1646        owner: String,
1647        /// Alias byte declaration.
1648        alias_bytes: u64,
1649        /// Owner byte declaration.
1650        owner_bytes: u64,
1651    },
1652    /// An alias was incorrectly rewritten with a physical source.
1653    #[error("weight binding alias {name:?} cannot own a physical checkpoint source")]
1654    AliasHasPhysicalSource {
1655        /// Alias identity.
1656        name: String,
1657    },
1658}
1659
1660#[cfg(test)]
1661mod tests {
1662    use std::{
1663        cell::RefCell,
1664        sync::{Arc, Mutex},
1665    };
1666
1667    use eredu_checkpoint::{store::TensorMetadata, StoredDtype};
1668    use eredu_core::residency::{MemoryTier, OffloadConfig, OffloadUnitSpec, ResidencyPolicy};
1669
1670    use super::*;
1671
1672    struct Catalog(BTreeMap<String, TensorMetadata>);
1673
1674    struct TestLeaseStorage(BTreeMap<String, u32>);
1675
1676    impl ResidencyLeaseStorage for TestLeaseStorage {
1677        type DeviceValue = u32;
1678        type HostValue = u32;
1679        type Error = &'static str;
1680        type BindingNames<'a> = std::iter::Map<
1681            std::collections::btree_map::Keys<'a, String, u32>,
1682            fn(&'a String) -> &'a str,
1683        >;
1684
1685        fn device_value<'a>(
1686            &'a self,
1687            _: &OffloadUnitId,
1688            name: &str,
1689        ) -> Result<&'a Self::DeviceValue, Self::Error> {
1690            self.0.get(name).ok_or("unknown binding")
1691        }
1692
1693        fn host_value<'a>(
1694            &'a self,
1695            _: &OffloadUnitId,
1696            name: &str,
1697        ) -> Result<&'a Self::HostValue, Self::Error> {
1698            self.0.get(name).ok_or("unknown binding")
1699        }
1700
1701        fn binding_names(&self) -> Self::BindingNames<'_> {
1702            self.0.keys().map(String::as_str)
1703        }
1704    }
1705
1706    #[derive(Default)]
1707    struct TestLeaseOwner(Mutex<Vec<(OffloadUnitId, MemoryTier)>>);
1708
1709    impl ResidencyLeaseOwner for TestLeaseOwner {
1710        fn release_residency_pin(&self, id: &OffloadUnitId, tier: MemoryTier) {
1711            self.0.lock().unwrap().push((id.clone(), tier));
1712        }
1713    }
1714
1715    struct TestTransferCompletion {
1716        succeeds: bool,
1717        waits: Arc<Mutex<usize>>,
1718    }
1719
1720    struct TestTransferResources(Arc<Mutex<Vec<bool>>>);
1721
1722    type TransferCompletionRecord = (Vec<OffloadUnitId>, MemoryTier, u64, bool);
1723
1724    #[derive(Default)]
1725    struct TestTransferOwner(Mutex<Vec<TransferCompletionRecord>>);
1726
1727    impl ResidencyTransferOwner<TestTransferCompletion, TestTransferResources> for TestTransferOwner {
1728        type Executor = Mutex<usize>;
1729        type Error = &'static str;
1730
1731        fn order_after(
1732            _: &TestTransferCompletion,
1733            executor: &Self::Executor,
1734            _: &OffloadUnitId,
1735        ) -> Result<(), Self::Error> {
1736            *executor.lock().unwrap() += 1;
1737            Ok(())
1738        }
1739
1740        fn is_complete(
1741            completion: &TestTransferCompletion,
1742            _: &OffloadUnitId,
1743        ) -> Result<bool, Self::Error> {
1744            Ok(completion.succeeds)
1745        }
1746
1747        fn wait(completion: &TestTransferCompletion, _: &OffloadUnitId) -> Result<(), Self::Error> {
1748            *completion.waits.lock().unwrap() += 1;
1749            completion.succeeds.then_some(()).ok_or("transfer failed")
1750        }
1751
1752        fn finish_resources(resources: TestTransferResources, succeeded: bool) {
1753            resources.0.lock().unwrap().push(succeeded);
1754        }
1755
1756        fn resolve_transfer(
1757            &self,
1758            ids: &[OffloadUnitId],
1759            tier: MemoryTier,
1760            generation: u64,
1761            succeeded: bool,
1762        ) -> Result<(), Self::Error> {
1763            self.0
1764                .lock()
1765                .unwrap()
1766                .push((ids.to_vec(), tier, generation, succeeded));
1767            Ok(())
1768        }
1769    }
1770
1771    #[derive(Default)]
1772    struct WindowManager {
1773        prepared: RefCell<Vec<(String, Vec<OffloadUnitId>)>>,
1774        evicted: RefCell<Vec<OffloadUnitId>>,
1775    }
1776
1777    impl ResidencyWindowManager for WindowManager {
1778        type Error = ResidencyWindowError;
1779
1780        fn prepare_window(
1781            &self,
1782            active: &[OffloadUnitId],
1783            _: &[OffloadUnitId],
1784            _: MemoryTier,
1785        ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, Self::Error> {
1786            self.prepared
1787                .borrow_mut()
1788                .push(("default".into(), active.to_vec()));
1789            Ok(active
1790                .iter()
1791                .cloned()
1792                .map(|id| (id, PrefetchOutcome::Hit))
1793                .collect())
1794        }
1795
1796        fn prepare_group_window(
1797            &self,
1798            group: &str,
1799            active: &[OffloadUnitId],
1800            _: &[OffloadUnitId],
1801            _: MemoryTier,
1802        ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, Self::Error> {
1803            self.prepared
1804                .borrow_mut()
1805                .push((group.to_owned(), active.to_vec()));
1806            Ok(active
1807                .iter()
1808                .cloned()
1809                .map(|id| (id, PrefetchOutcome::Miss))
1810                .collect())
1811        }
1812
1813        fn evict(&self, id: &OffloadUnitId, _: MemoryTier) -> Result<bool, Self::Error> {
1814            self.evicted.borrow_mut().push(id.clone());
1815            Ok(true)
1816        }
1817
1818        fn unit_reports(&self) -> Result<Vec<UnitResidencyReport>, Self::Error> {
1819            Ok(Vec::new())
1820        }
1821    }
1822
1823    impl RecipeCatalog for Catalog {
1824        fn tensor_metadata(
1825            &self,
1826            key: &str,
1827        ) -> Result<TensorMetadata, eredu_checkpoint::store::StoreError> {
1828            self.0.get(key).cloned().ok_or_else(|| {
1829                eredu_checkpoint::store::StoreError::UnknownTensor {
1830                    key: key.to_owned(),
1831                }
1832            })
1833        }
1834    }
1835
1836    fn metadata(name: &str, shape: Vec<usize>) -> TensorMetadata {
1837        TensorMetadata {
1838            name: name.to_owned(),
1839            logical_shape: shape.clone(),
1840            physical_shape: shape,
1841            stored_dtype: StoredDtype::F32,
1842            encoded_byte_len: 0,
1843            backing_shard: None,
1844        }
1845    }
1846
1847    #[test]
1848    fn declarations_are_validated_and_deterministic() {
1849        let b = WeightBinding::new("b", "b.weight", TensorSelection::Full, 4).unwrap();
1850        let a = WeightBinding::new("a", "a.weight", TensorSelection::Full, 8).unwrap();
1851        let id = OffloadUnitId::new("layer.0").unwrap();
1852        let unit = OffloadUnit::new(id, [b, a]).unwrap();
1853        assert_eq!(unit.bindings()[0].name(), "a");
1854        assert_eq!(unit.bindings()[1].name(), "b");
1855    }
1856
1857    #[test]
1858    fn neutral_lease_exposes_native_storage_and_releases_exact_pin() {
1859        let owner = Arc::new(TestLeaseOwner::default());
1860        let id = OffloadUnitId::new("layer.0").unwrap();
1861        let lease = ResidencyLease::new(
1862            id.clone(),
1863            MemoryTier::Device,
1864            TestLeaseStorage(BTreeMap::from([("weight".into(), 7)])),
1865            Arc::downgrade(&owner),
1866        );
1867        assert_eq!(lease.device_value("weight"), Ok(&7));
1868        assert_eq!(lease.binding_names().collect::<Vec<_>>(), vec!["weight"]);
1869        drop(lease);
1870        assert_eq!(*owner.0.lock().unwrap(), vec![(id, MemoryTier::Device)]);
1871    }
1872
1873    #[test]
1874    fn neutral_transfer_orders_publishes_and_releases_resources() {
1875        let owner = Arc::new(TestTransferOwner::default());
1876        let waits = Arc::new(Mutex::new(0));
1877        let resources = Arc::new(Mutex::new(Vec::new()));
1878        let executor = Mutex::new(0);
1879        let id = OffloadUnitId::new("layer.0").unwrap();
1880        let mut transfer = ResidencyTransfer::submitted(
1881            vec![7],
1882            TestTransferCompletion {
1883                succeeds: true,
1884                waits: Arc::clone(&waits),
1885            },
1886            TestTransferResources(Arc::clone(&resources)),
1887            Arc::downgrade(&owner),
1888            vec![id.clone()],
1889            MemoryTier::Device,
1890            11,
1891        );
1892
1893        assert_eq!(transfer.leases(), &[7]);
1894        transfer.order_after(&executor).unwrap();
1895        assert_eq!(*executor.lock().unwrap(), 1);
1896        transfer.synchronize().unwrap();
1897        assert!(transfer.is_complete().unwrap());
1898        assert_eq!(*waits.lock().unwrap(), 1);
1899        assert_eq!(*resources.lock().unwrap(), vec![true]);
1900        assert_eq!(
1901            *owner.0.lock().unwrap(),
1902            vec![(vec![id], MemoryTier::Device, 11, true)]
1903        );
1904    }
1905
1906    #[test]
1907    fn neutral_transfer_failure_remains_observable_and_resolves_once() {
1908        let owner = Arc::new(TestTransferOwner::default());
1909        let waits = Arc::new(Mutex::new(0));
1910        let resources = Arc::new(Mutex::new(Vec::new()));
1911        let id = OffloadUnitId::new("layer.0").unwrap();
1912        let mut transfer = ResidencyTransfer::submitted(
1913            Vec::<u8>::new(),
1914            TestTransferCompletion {
1915                succeeds: false,
1916                waits: Arc::clone(&waits),
1917            },
1918            TestTransferResources(Arc::clone(&resources)),
1919            Arc::downgrade(&owner),
1920            vec![id.clone()],
1921            MemoryTier::Host,
1922            4,
1923        );
1924
1925        assert_eq!(transfer.synchronize(), Err("transfer failed"));
1926        assert_eq!(transfer.synchronize(), Err("transfer failed"));
1927        assert_eq!(*resources.lock().unwrap(), vec![false]);
1928        assert_eq!(
1929            *owner.0.lock().unwrap(),
1930            vec![(vec![id], MemoryTier::Host, 4, false)]
1931        );
1932        drop(transfer);
1933        assert_eq!(*waits.lock().unwrap(), 3);
1934    }
1935
1936    #[test]
1937    fn binding_selection_rewrites_sources_and_exact_bytes_neutrally() {
1938        let catalog = Catalog(BTreeMap::from([(
1939            "weight".into(),
1940            metadata("weight", vec![2, 2]),
1941        )]));
1942        let binding = WeightBinding::new("weight", "weight", TensorSelection::Full, 16)
1943            .unwrap()
1944            .select_bounded_output(
1945                &catalog,
1946                TensorSelection::Range {
1947                    axis: 0,
1948                    start: 1,
1949                    end: 2,
1950                },
1951            )
1952            .unwrap();
1953
1954        assert_eq!(binding.expected_bytes(), 8);
1955        assert!(matches!(
1956            binding.source_recipe(),
1957            DerivedWeightRecipe::Source {
1958                selection: TensorSelection::Range {
1959                    axis: 0,
1960                    start: 1,
1961                    end: 2,
1962                },
1963                ..
1964            }
1965        ));
1966    }
1967
1968    #[test]
1969    fn controller_validates_catalog_bytes_before_allocating_backend_storage() {
1970        let catalog = Catalog(BTreeMap::from([
1971            ("a.weight".into(), metadata("a.weight", vec![2])),
1972            ("b.weight".into(), metadata("b.weight", vec![1])),
1973        ]));
1974        let id = OffloadUnitId::new("layer.0").unwrap();
1975        let unit = OffloadUnit::new(
1976            id.clone(),
1977            [
1978                WeightBinding::new("a", "a.weight", TensorSelection::Full, 8).unwrap(),
1979                WeightBinding::new("b", "b.weight", TensorSelection::Full, 4).unwrap(),
1980            ],
1981        )
1982        .unwrap();
1983        let plan = OffloadPlan::new(
1984            OffloadConfig::default(),
1985            [
1986                OffloadUnitSpec::new(id.clone(), 12, ResidencyPolicy::Windowed, MemoryTier::Disk)
1987                    .unwrap(),
1988            ],
1989        )
1990        .unwrap();
1991
1992        let controller = ResidencyController::new(&catalog, plan, [unit]).unwrap();
1993        assert_eq!(controller.units().len(), 1);
1994        assert_eq!(controller.unit(&id).unwrap().bindings().len(), 2);
1995        assert!(!controller.ledger().initialized());
1996    }
1997
1998    #[test]
1999    fn controller_resolves_aliases_across_independent_units() {
2000        let catalog = Catalog(BTreeMap::from([
2001            ("physical.owner".into(), metadata("physical.owner", vec![1])),
2002            ("slice.local".into(), metadata("slice.local", vec![1])),
2003        ]));
2004        let owner_id = OffloadUnitId::new("slice.0").unwrap();
2005        let alias_id = OffloadUnitId::new("slice.1").unwrap();
2006        let owner_binding =
2007            WeightBinding::new("weight", "physical.owner", TensorSelection::Full, 4)
2008                .unwrap()
2009                .with_logical_target("shared.owner")
2010                .unwrap();
2011        let alias_binding = WeightBinding::alias("weight", "shared.owner", 4)
2012            .unwrap()
2013            .with_logical_target("slice.1.weight")
2014            .unwrap();
2015        let local_binding =
2016            WeightBinding::new("local", "slice.local", TensorSelection::Full, 4).unwrap();
2017        let units = [
2018            OffloadUnit::new(owner_id.clone(), [owner_binding]).unwrap(),
2019            OffloadUnit::new(alias_id.clone(), [alias_binding, local_binding]).unwrap(),
2020        ];
2021        let plan = OffloadPlan::new(
2022            OffloadConfig::default(),
2023            [
2024                OffloadUnitSpec::new(
2025                    owner_id.clone(),
2026                    4,
2027                    ResidencyPolicy::Windowed,
2028                    MemoryTier::Disk,
2029                )
2030                .unwrap(),
2031                OffloadUnitSpec::new(
2032                    alias_id.clone(),
2033                    4,
2034                    ResidencyPolicy::Windowed,
2035                    MemoryTier::Disk,
2036                )
2037                .unwrap(),
2038            ],
2039        )
2040        .unwrap();
2041        let controller = ResidencyController::new(&catalog, plan, units).unwrap();
2042        let alias = controller
2043            .unit(&alias_id)
2044            .unwrap()
2045            .bindings()
2046            .iter()
2047            .find(|binding| binding.is_alias())
2048            .unwrap();
2049        let (resolved_unit, resolved) = controller.binding_owner(&alias_id, alias).unwrap();
2050        assert_eq!(resolved_unit, &owner_id);
2051        assert_eq!(resolved.logical_target(), Some("shared.owner"));
2052    }
2053
2054    #[test]
2055    fn controller_owns_named_window_and_unique_lookahead_selection() {
2056        let ids = ["a", "b", "c"].map(|name| OffloadUnitId::new(format!("layer.{name}")).unwrap());
2057        let catalog = Catalog(BTreeMap::from([
2058            ("a".into(), metadata("a", vec![1])),
2059            ("b".into(), metadata("b", vec![1])),
2060            ("c".into(), metadata("c", vec![1])),
2061        ]));
2062        let units = ids.iter().zip(["a", "b", "c"]).map(|(id, key)| {
2063            OffloadUnit::new(
2064                id.clone(),
2065                [WeightBinding::new("weight", key, TensorSelection::Full, 4).unwrap()],
2066            )
2067            .unwrap()
2068        });
2069        let plan = OffloadPlan::new(
2070            OffloadConfig::new(None, None, 2).unwrap(),
2071            ids.iter().map(|id| {
2072                OffloadUnitSpec::new(id.clone(), 4, ResidencyPolicy::Windowed, MemoryTier::Disk)
2073                    .unwrap()
2074            }),
2075        )
2076        .unwrap();
2077        let mut controller = ResidencyController::new(&catalog, plan, units).unwrap();
2078        controller.ledger_mut().mark_initialized();
2079        let selected = controller
2080            .commit_group_window(
2081                "decoder",
2082                &[ids[0].clone()],
2083                &[ids[1].clone(), ids[1].clone(), ids[2].clone()],
2084                MemoryTier::Device,
2085            )
2086            .unwrap();
2087        assert_eq!(selected, vec![ids[1].clone(), ids[2].clone()]);
2088        assert_eq!(
2089            controller.ledger().active_window(),
2090            BTreeSet::from([ids[0].clone()])
2091        );
2092    }
2093
2094    #[test]
2095    fn controller_owns_acquisition_reservation_and_rollback() {
2096        let ids = ["a", "b"].map(|name| OffloadUnitId::new(format!("layer.{name}")).unwrap());
2097        let catalog = Catalog(BTreeMap::from([
2098            ("a".into(), metadata("a", vec![1])),
2099            ("b".into(), metadata("b", vec![1])),
2100        ]));
2101        let units = ids.iter().zip(["a", "b"]).map(|(id, key)| {
2102            OffloadUnit::new(
2103                id.clone(),
2104                [WeightBinding::new("weight", key, TensorSelection::Full, 4).unwrap()],
2105            )
2106            .unwrap()
2107        });
2108        let plan = OffloadPlan::new(
2109            OffloadConfig::new(Some(8), Some(8), 1).unwrap(),
2110            ids.iter().map(|id| {
2111                OffloadUnitSpec::new(id.clone(), 4, ResidencyPolicy::Cacheable, MemoryTier::Disk)
2112                    .unwrap()
2113            }),
2114        )
2115        .unwrap();
2116        let mut controller = ResidencyController::new(&catalog, plan, units).unwrap();
2117        controller.ledger_mut().mark_initialized();
2118
2119        let acquisition = controller
2120            .plan_acquisition(&ids, MemoryTier::Device)
2121            .unwrap();
2122        assert_eq!(acquisition.missing(), &[true, true]);
2123        assert!(controller
2124            .reserve_acquisition(
2125                &acquisition,
2126                &[(ids[0].clone(), 4), (ids[1].clone(), 4)],
2127                MemoryTier::Device,
2128            )
2129            .unwrap()
2130            .is_empty());
2131        controller
2132            .rollback_acquisition(&acquisition, MemoryTier::Device)
2133            .unwrap();
2134        assert!(!controller
2135            .ledger()
2136            .is_resident(&ids[0], MemoryTier::Device)
2137            .unwrap());
2138        assert!(!controller
2139            .ledger()
2140            .is_resident(&ids[1], MemoryTier::Device)
2141            .unwrap());
2142
2143        let acquisition = controller
2144            .plan_acquisition(&[ids[0].clone()], MemoryTier::Device)
2145            .unwrap();
2146        controller
2147            .reserve_acquisition(&acquisition, &[(ids[0].clone(), 4)], MemoryTier::Device)
2148            .unwrap();
2149        controller
2150            .publish_acquisition_copy(
2151                &ids[0],
2152                MemoryTier::Device,
2153                4,
2154                4,
2155                None,
2156                eredu_core::residency::TransferDirection::DiskToDevice,
2157                Duration::from_millis(2),
2158            )
2159            .unwrap();
2160        assert!(controller
2161            .ledger()
2162            .is_resident(&ids[0], MemoryTier::Device)
2163            .unwrap());
2164        assert_eq!(
2165            controller
2166                .ledger()
2167                .telemetry()
2168                .transfer(eredu_core::residency::TransferDirection::DiskToDevice)
2169                .bytes(),
2170            4
2171        );
2172    }
2173
2174    #[test]
2175    fn controller_rejects_binding_and_plan_byte_mismatches() {
2176        let catalog = Catalog(BTreeMap::from([(
2177            "weight".into(),
2178            metadata("weight", vec![2]),
2179        )]));
2180        let id = OffloadUnitId::new("layer.0").unwrap();
2181        let plan = |bytes| {
2182            OffloadPlan::new(
2183                OffloadConfig::default(),
2184                [OffloadUnitSpec::new(
2185                    id.clone(),
2186                    bytes,
2187                    ResidencyPolicy::Windowed,
2188                    MemoryTier::Disk,
2189                )
2190                .unwrap()],
2191            )
2192            .unwrap()
2193        };
2194
2195        let wrong_binding = OffloadUnit::new(
2196            id.clone(),
2197            [WeightBinding::new("weight", "weight", TensorSelection::Full, 4).unwrap()],
2198        )
2199        .unwrap();
2200        assert!(matches!(
2201            ResidencyController::new(&catalog, plan(4), [wrong_binding]),
2202            Err(ResidencyControllerError::BindingByteMismatch { .. })
2203        ));
2204
2205        let wrong_plan = OffloadUnit::new(
2206            id.clone(),
2207            [WeightBinding::new("weight", "weight", TensorSelection::Full, 8).unwrap()],
2208        )
2209        .unwrap();
2210        assert!(matches!(
2211            ResidencyController::new(&catalog, plan(16), [wrong_plan]),
2212            Err(ResidencyControllerError::UnitByteMismatch { .. })
2213        ));
2214    }
2215
2216    #[test]
2217    fn named_windows_prepare_and_trim_without_backend_types() {
2218        let ids = ["layer.0", "layer.1", "layer.2"].map(|id| OffloadUnitId::new(id).unwrap());
2219        let group = ResidentLayerGroup::new("decoder", ids.clone(), 2).unwrap();
2220        let manager = WindowManager::default();
2221
2222        assert_eq!(group.prepare(&manager, 1).unwrap().len(), 2);
2223        assert_eq!(
2224            manager.prepared.borrow()[0],
2225            ("decoder".into(), vec![ids[1].clone(), ids[2].clone()])
2226        );
2227        assert_eq!(manager.evicted.borrow().as_slice(), &[ids[0].clone()]);
2228        assert_eq!(group.report(&manager).unwrap().device_units(), 0);
2229    }
2230}