1use 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#[derive(Debug, Clone, Default, Eq, PartialEq)]
20pub struct WeightMaterializationReport {
21 pub admitted_working_set_bytes: u64,
23 pub transformed_weights: usize,
25 pub source_tiles: usize,
27 pub peak_in_flight_tiles: usize,
29 pub source_bytes_read: u64,
31 pub output_bytes: u64,
33 pub peak_planned_working_set_bytes: u64,
35 pub largest_source_tile_bytes: u64,
37 pub largest_output_tile_bytes: u64,
39}
40
41#[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 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 pub const fn initialized(&self) -> bool {
73 self.initialized
74 }
75
76 pub const fn offload(&self) -> &OffloadReport {
78 &self.offload
79 }
80
81 pub fn units(&self) -> &[UnitResidencyReport] {
83 &self.units
84 }
85
86 pub fn active_window(&self) -> &[OffloadUnitId] {
88 &self.active_window
89 }
90
91 pub const fn weight_store(&self) -> &WeightStoreDiagnostics {
93 &self.weight_store
94 }
95
96 pub const fn materialization(&self) -> Option<&WeightMaterializationReport> {
98 self.materialization.as_ref()
99 }
100
101 pub fn with_materialization(
103 mut self,
104 materialization: Option<WeightMaterializationReport>,
105 ) -> Self {
106 self.materialization = materialization;
107 self
108 }
109}
110
111#[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 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 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 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 pub fn name(&self) -> &str {
195 &self.name
196 }
197
198 pub fn alias_of(&self) -> Option<&str> {
200 self.alias_of.as_deref()
201 }
202
203 pub const fn is_alias(&self) -> bool {
206 self.alias_of.is_some()
207 }
208
209 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 pub fn logical_target(&self) -> Option<&str> {
217 self.logical_target.as_deref()
218 }
219
220 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 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 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 pub fn checkpoint_key(&self) -> &str {
260 &self.checkpoint_key
261 }
262
263 pub fn selection(&self) -> &TensorSelection {
265 &self.selection
266 }
267
268 pub const fn recipe(&self) -> Option<&DerivedWeightRecipe> {
270 self.recipe.as_ref()
271 }
272
273 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 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 pub const fn expected_bytes(&self) -> u64 {
297 self.expected_bytes
298 }
299
300 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 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#[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 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 pub fn owners(&self) -> impl Iterator<Item = &'a WeightBinding> + '_ {
400 self.owners.iter().copied()
401 }
402
403 pub fn aliases(&self) -> impl Iterator<Item = (&'a WeightBinding, &'a WeightBinding)> + '_ {
405 self.aliases.iter().copied()
406 }
407}
408
409#[derive(Debug, thiserror::Error)]
411pub enum WeightBindingSelectionError {
412 #[error(transparent)]
414 Recipe(#[from] RecipeError),
415 #[error(transparent)]
417 Declaration(#[from] ResidencyDeclarationError),
418}
419
420#[derive(Debug, Clone, Eq, PartialEq)]
422pub struct OffloadUnit {
423 id: OffloadUnitId,
424 bindings: Vec<WeightBinding>,
425}
426
427#[derive(Debug)]
434pub struct ResidencyController {
435 ledger: ResidencyLedger,
436 units: BTreeMap<OffloadUnitId, OffloadUnit>,
437 alias_owners: BTreeMap<(OffloadUnitId, String), (OffloadUnitId, String)>,
438}
439
440#[derive(Debug, Clone, Eq, PartialEq)]
442pub struct ResidencyAcquisition {
443 ids: Vec<OffloadUnitId>,
444 missing: Vec<bool>,
445}
446
447pub trait ResidencyLeaseStorage {
449 type DeviceValue;
451 type HostValue;
453 type Error;
455 type BindingNames<'a>: Iterator<Item = &'a str>
457 where
458 Self: 'a;
459
460 fn device_value<'a>(
462 &'a self,
463 id: &OffloadUnitId,
464 name: &str,
465 ) -> Result<&'a Self::DeviceValue, Self::Error>;
466
467 fn host_value<'a>(
469 &'a self,
470 id: &OffloadUnitId,
471 name: &str,
472 ) -> Result<&'a Self::HostValue, Self::Error>;
473
474 fn binding_names(&self) -> Self::BindingNames<'_>;
476}
477
478pub trait ResidencyLeaseOwner: Sized {
480 fn release_residency_pin(&self, id: &OffloadUnitId, tier: MemoryTier);
482}
483
484pub trait ResidencyTransferOwner<C, R>: Sized {
486 type Executor: ?Sized;
488 type Error;
490
491 fn order_after(
493 completion: &C,
494 executor: &Self::Executor,
495 id: &OffloadUnitId,
496 ) -> Result<(), Self::Error>;
497
498 fn is_complete(completion: &C, id: &OffloadUnitId) -> Result<bool, Self::Error>;
500
501 fn wait(completion: &C, id: &OffloadUnitId) -> Result<(), Self::Error>;
503
504 fn finish_resources(resources: R, succeeded: bool);
506
507 fn resolve_transfer(
509 &self,
510 ids: &[OffloadUnitId],
511 tier: MemoryTier,
512 generation: u64,
513 succeeded: bool,
514 ) -> Result<(), Self::Error>;
515}
516
517pub 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 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 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 pub fn leases(&self) -> &[L] {
575 &self.leases
576 }
577
578 pub fn is_empty(&self) -> bool {
580 self.leases.is_empty()
581 }
582
583 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 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 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
651pub 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 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 pub fn id(&self) -> &OffloadUnitId {
680 &self.id
681 }
682
683 pub const fn tier(&self) -> MemoryTier {
685 self.tier
686 }
687
688 pub fn device_value(&self, name: &str) -> Result<&S::DeviceValue, S::Error> {
690 self.storage.device_value(&self.id, name)
691 }
692
693 pub fn host_value(&self, name: &str) -> Result<&S::HostValue, S::Error> {
695 self.storage.host_value(&self.id, name)
696 }
697
698 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 pub fn ids(&self) -> &[OffloadUnitId] {
719 &self.ids
720 }
721
722 pub fn missing(&self) -> &[bool] {
724 &self.missing
725 }
726
727 pub fn is_hit(&self) -> bool {
729 self.missing.iter().all(|missing| !missing)
730 }
731
732 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 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 pub fn unit(&self, id: &OffloadUnitId) -> Option<&OffloadUnit> {
824 self.units.get(id)
825 }
826
827 pub fn units(&self) -> impl ExactSizeIterator<Item = &OffloadUnit> {
829 self.units.values()
830 }
831
832 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 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 pub const fn ledger(&self) -> &ResidencyLedger {
873 &self.ledger
874 }
875
876 pub fn ledger_mut(&mut self) -> &mut ResidencyLedger {
878 &mut self.ledger
879 }
880
881 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 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 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 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 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 #[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 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 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 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 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#[derive(Debug, thiserror::Error)]
1124pub enum ResidencyControllerError {
1125 #[error(transparent)]
1127 Declaration(#[from] ResidencyDeclarationError),
1128 #[error("duplicate residency unit definition: {id}")]
1130 DuplicateUnitDefinition {
1131 id: OffloadUnitId,
1133 },
1134 #[error("offload plan unit {id} has no residency unit definition")]
1136 MissingUnitDefinition {
1137 id: OffloadUnitId,
1139 },
1140 #[error("residency unit {id} is absent from the offload plan")]
1142 UnexpectedUnitDefinition {
1143 id: OffloadUnitId,
1145 },
1146 #[error(
1148 "residency unit {id} defines {actual_bytes} bytes but its plan reserves {planned_bytes}"
1149 )]
1150 UnitByteMismatch {
1151 id: OffloadUnitId,
1153 planned_bytes: u64,
1155 actual_bytes: u64,
1157 },
1158 #[error(
1160 "binding {binding:?} in unit {id} selects {actual_bytes} bytes but declares {expected_bytes}"
1161 )]
1162 BindingByteMismatch {
1163 id: OffloadUnitId,
1165 binding: String,
1167 expected_bytes: u64,
1169 actual_bytes: u64,
1171 },
1172 #[error("derived-weight recipe for binding {binding:?} failed: {source}")]
1174 Recipe {
1175 binding: String,
1177 #[source]
1179 source: RecipeError,
1180 },
1181 #[error("residency arithmetic overflow: {context}")]
1183 ArithmeticOverflow {
1184 context: &'static str,
1186 },
1187}
1188
1189pub trait ResidencyWindowManager {
1191 type Error: std::error::Error + From<ResidencyWindowError>;
1193
1194 fn prepare_window(
1196 &self,
1197 active: &[OffloadUnitId],
1198 upcoming: &[OffloadUnitId],
1199 tier: MemoryTier,
1200 ) -> Result<Vec<(OffloadUnitId, PrefetchOutcome)>, Self::Error>;
1201
1202 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 fn evict(&self, id: &OffloadUnitId, tier: MemoryTier) -> Result<bool, Self::Error>;
1213
1214 fn unit_reports(&self) -> Result<Vec<UnitResidencyReport>, Self::Error>;
1216}
1217
1218#[derive(Debug, Clone)]
1220pub struct DeviceLayerWindow {
1221 units: Vec<OffloadUnitId>,
1222 depth: usize,
1223}
1224
1225impl DeviceLayerWindow {
1226 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 pub const fn depth(&self) -> usize {
1256 self.depth
1257 }
1258
1259 pub fn units(&self) -> &[OffloadUnitId] {
1261 &self.units
1262 }
1263
1264 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 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 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 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#[derive(Debug, Clone)]
1312pub struct ResidentLayerGroup {
1313 id: String,
1314 window: DeviceLayerWindow,
1315}
1316
1317impl ResidentLayerGroup {
1318 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 pub fn id(&self) -> &str {
1336 &self.id
1337 }
1338
1339 pub fn units(&self) -> &[OffloadUnitId] {
1341 self.window.units()
1342 }
1343
1344 pub const fn depth(&self) -> usize {
1346 self.window.depth()
1347 }
1348
1349 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 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 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 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#[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 pub fn id(&self) -> &str {
1434 &self.id
1435 }
1436 pub const fn ordered_units(&self) -> usize {
1438 self.ordered_units
1439 }
1440 pub const fn window_depth(&self) -> usize {
1442 self.window_depth
1443 }
1444 pub const fn host_bytes(&self) -> u64 {
1446 self.host_bytes
1447 }
1448 pub const fn device_bytes(&self) -> u64 {
1450 self.device_bytes
1451 }
1452 pub const fn device_units(&self) -> usize {
1454 self.device_units
1455 }
1456}
1457
1458#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1460pub enum ResidencyWindowError {
1461 #[error("device layer window requires at least one ordered unit")]
1463 EmptyLayerWindow,
1464 #[error("device layer window depth {depth} exceeds {layer_count} ordered units")]
1466 OversizedLayerWindow {
1467 depth: usize,
1469 layer_count: usize,
1471 },
1472 #[error("device layer index {index} is outside {layer_count} ordered units")]
1474 InvalidLayerIndex {
1475 index: usize,
1477 layer_count: usize,
1479 },
1480 #[error("device layer window contains duplicate unit {id}")]
1482 DuplicateLayerWindowUnit {
1483 id: OffloadUnitId,
1485 },
1486 #[error("residency window group identifiers must not be empty")]
1488 InvalidGroupId,
1489 #[error("residency window arithmetic overflow: {context}")]
1491 ArithmeticOverflow {
1492 context: &'static str,
1494 },
1495}
1496
1497impl OffloadUnit {
1498 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 pub fn id(&self) -> &OffloadUnitId {
1522 &self.id
1523 }
1524
1525 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#[derive(Debug, Clone, Eq, PartialEq, thiserror::Error)]
1564pub enum ResidencyDeclarationError {
1565 #[error("weight binding names must not be empty")]
1567 InvalidBindingName,
1568 #[error("weight binding {name:?} has an empty checkpoint key")]
1570 InvalidCheckpointKey {
1571 name: String,
1573 },
1574 #[error("weight binding {name:?} has no checkpoint recipe source")]
1576 EmptyRecipeSources {
1577 name: String,
1579 },
1580 #[error(
1582 "weight binding {name:?} has invalid quantization companions {scales:?} and {biases:?}"
1583 )]
1584 InvalidQuantizationCompanions {
1585 name: String,
1587 scales: String,
1589 biases: String,
1591 },
1592 #[error("weight binding {name:?} must contain at least one byte")]
1594 ZeroSizedBinding {
1595 name: String,
1597 },
1598 #[error("residency unit {id} must contain at least one binding")]
1600 EmptyUnit {
1601 id: OffloadUnitId,
1603 },
1604 #[error("residency unit {id} has duplicate binding name {name:?}")]
1606 DuplicateBindingName {
1607 id: OffloadUnitId,
1609 name: String,
1611 },
1612 #[error("duplicate logical weight binding {name:?}")]
1614 DuplicateLogicalBinding {
1615 name: String,
1617 },
1618 #[error("weight binding alias {alias:?} has unknown owner {owner:?}")]
1620 UnknownBindingAliasOwner {
1621 alias: String,
1623 owner: String,
1625 },
1626 #[error("weight binding alias {alias:?} has ambiguous owner {owner:?}")]
1628 AmbiguousBindingAliasOwner {
1629 alias: String,
1631 owner: String,
1633 },
1634 #[error("weight binding alias cycle contains {name:?}")]
1636 BindingAliasCycle {
1637 name: String,
1639 },
1640 #[error("weight binding alias {alias:?} declares {alias_bytes} bytes but owner {owner:?} declares {owner_bytes}")]
1642 BindingAliasByteMismatch {
1643 alias: String,
1645 owner: String,
1647 alias_bytes: u64,
1649 owner_bytes: u64,
1651 },
1652 #[error("weight binding alias {name:?} cannot own a physical checkpoint source")]
1654 AliasHasPhysicalSource {
1655 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}