1use crate::cgroup_memory::detect_cgroup_memory;
2pub use crate::cgroup_memory::{
3 CgroupMemoryAvailability, CgroupMemoryLimit, CgroupMemoryObservation, CgroupMemoryProbeFailure,
4 CgroupMemoryProbeFailureKind,
5};
6
7pub const LIBRARY_ROW_CHUNK_TARGET_BYTES: usize = 8 * 1024 * 1024;
30
31#[derive(Clone, Debug)]
32pub struct ResourcePolicy {
33 pub max_single_materialization_bytes: usize,
34 pub max_operator_cache_bytes: usize,
35 pub max_spatial_distance_cache_bytes: usize,
36 pub max_owned_data_cache_bytes: usize,
37 pub row_chunk_target_bytes: usize,
38 pub derivative_storage_mode: DerivativeStorageMode,
39}
40
41pub const OWNED_DATA_CACHE_MAX_ENTRIES: usize = 2;
42
43const GOVERNOR_BUDGET_NUMERATOR: u128 = 3;
62const GOVERNOR_BUDGET_DENOMINATOR: u128 = 4;
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq)]
67pub enum MemoryAvailabilitySource {
68 Host,
69 Cgroup,
70 HostAndCgroup,
71 CgroupProbeFailure,
72}
73
74impl std::fmt::Display for MemoryAvailabilitySource {
75 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 match self {
77 Self::Host => formatter.write_str("host"),
78 Self::Cgroup => formatter.write_str("cgroup"),
79 Self::HostAndCgroup => formatter.write_str("host and cgroup equally"),
80 Self::CgroupProbeFailure => formatter.write_str("cgroup probe failure"),
81 }
82 }
83}
84
85#[derive(Clone, Debug, PartialEq, Eq)]
113pub struct MemoryAvailability {
114 host_available_bytes: u64,
115 host_total_bytes: u64,
116 cgroup: CgroupMemoryObservation,
117 available_bytes: u64,
118 capacity_bytes: u64,
119 limiting_source: MemoryAvailabilitySource,
120}
121
122impl MemoryAvailability {
123 pub(crate) fn from_observation(
124 host_available_bytes: u64,
125 host_total_bytes: u64,
126 cgroup: CgroupMemoryObservation,
127 ) -> Self {
128 use std::cmp::Ordering;
129
130 let capacity_bytes = match &cgroup {
137 CgroupMemoryObservation::NotPresent | CgroupMemoryObservation::V2Unbounded { .. } => {
138 host_total_bytes.max(host_available_bytes)
139 }
140 CgroupMemoryObservation::V2Limited(observation)
141 | CgroupMemoryObservation::V1Limited(observation) => host_total_bytes
142 .max(host_available_bytes)
143 .min(observation.limit_bytes()),
144 CgroupMemoryObservation::ProbeFailed(_) => 0,
145 };
146
147 let (available_bytes, limiting_source) = match &cgroup {
148 CgroupMemoryObservation::NotPresent | CgroupMemoryObservation::V2Unbounded { .. } => {
149 (host_available_bytes, MemoryAvailabilitySource::Host)
150 }
151 CgroupMemoryObservation::V2Limited(observation)
152 | CgroupMemoryObservation::V1Limited(observation) => {
153 match observation.available_bytes().cmp(&host_available_bytes) {
154 Ordering::Less => (
155 observation.available_bytes(),
156 MemoryAvailabilitySource::Cgroup,
157 ),
158 Ordering::Equal => (
159 host_available_bytes,
160 MemoryAvailabilitySource::HostAndCgroup,
161 ),
162 Ordering::Greater => (host_available_bytes, MemoryAvailabilitySource::Host),
163 }
164 }
165 CgroupMemoryObservation::ProbeFailed(_) => {
166 (0, MemoryAvailabilitySource::CgroupProbeFailure)
167 }
168 };
169 Self {
170 host_available_bytes,
171 host_total_bytes,
172 cgroup,
173 available_bytes,
174 capacity_bytes,
175 limiting_source,
176 }
177 }
178
179 pub const fn host_available_bytes(&self) -> u64 {
180 self.host_available_bytes
181 }
182
183 pub const fn host_total_bytes(&self) -> u64 {
184 self.host_total_bytes
185 }
186
187 pub const fn capacity_bytes(&self) -> u64 {
193 self.capacity_bytes
194 }
195
196 pub fn capacity_bytes_usize(&self) -> usize {
197 usize::try_from(self.capacity_bytes).unwrap_or(usize::MAX)
198 }
199
200 pub const fn cgroup(&self) -> &CgroupMemoryObservation {
201 &self.cgroup
202 }
203
204 pub const fn available_bytes(&self) -> u64 {
205 self.available_bytes
206 }
207
208 pub fn available_bytes_usize(&self) -> usize {
209 usize::try_from(self.available_bytes).unwrap_or(usize::MAX)
210 }
211
212 pub const fn limiting_source(&self) -> MemoryAvailabilitySource {
213 self.limiting_source
214 }
215}
216
217impl std::fmt::Display for MemoryAvailability {
218 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
219 match &self.cgroup {
220 CgroupMemoryObservation::ProbeFailed(failure) => write!(
221 formatter,
222 "0 bytes admitted because the active cgroup probe failed closed (host_available={}, failure={})",
223 self.host_available_bytes, failure,
224 ),
225 observation => write!(
226 formatter,
227 "{} bytes limited by {} (capacity={}, host_available={}, host_total={}, {})",
228 self.available_bytes,
229 self.limiting_source,
230 self.capacity_bytes,
231 self.host_available_bytes,
232 self.host_total_bytes,
233 observation,
234 ),
235 }
236 }
237}
238
239pub fn memory_availability_probe_count() -> u64 {
247 MEMORY_AVAILABILITY_PROBES.load(std::sync::atomic::Ordering::Relaxed)
248}
249
250static MEMORY_AVAILABILITY_PROBES: std::sync::atomic::AtomicU64 =
251 std::sync::atomic::AtomicU64::new(0);
252
253pub fn resample_memory_availability() -> MemoryAvailability {
268 static SYSTEM: OnceLock<Mutex<sysinfo::System>> = OnceLock::new();
269 let system = SYSTEM.get_or_init(|| Mutex::new(sysinfo::System::new()));
270 let mut system = system.lock().expect("sysinfo system mutex poisoned");
271 system.refresh_memory();
272 let cgroup = detect_cgroup_memory();
273 MEMORY_AVAILABILITY_PROBES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
274 MemoryAvailability::from_observation(
275 system.available_memory(),
276 system.total_memory(),
277 cgroup,
278 )
279}
280
281pub fn process_memory_availability() -> &'static MemoryAvailability {
298 &MemoryGovernor::global().ledger.availability
299}
300
301pub fn process_available_memory_bytes() -> usize {
305 process_memory_availability().available_bytes_usize()
306}
307
308fn governor_budget_from_availability(availability: &MemoryAvailability) -> usize {
339 stationary_headroom_of_capacity(availability)
340}
341
342fn governor_materialization_cap_from_availability(availability: &MemoryAvailability) -> usize {
347 stationary_headroom_of_capacity(availability)
348}
349
350fn stationary_headroom_of_capacity(availability: &MemoryAvailability) -> usize {
361 let scaled = u128::from(availability.capacity_bytes()) * GOVERNOR_BUDGET_NUMERATOR
362 / GOVERNOR_BUDGET_DENOMINATOR;
363 usize::try_from(scaled).unwrap_or(usize::MAX)
364}
365
366#[derive(Debug, Clone, thiserror::Error, PartialEq, Eq)]
374pub enum MemoryReservationError {
375 #[error(
376 "{context}: cannot reserve {requested_bytes} bytes; {reserved_bytes} of {budget_bytes} bytes already reserved process-wide; detected availability: {availability}"
377 )]
378 BudgetExceeded {
379 context: Box<str>,
380 requested_bytes: usize,
381 reserved_bytes: usize,
382 budget_bytes: usize,
383 availability: MemoryAvailability,
384 },
385
386 #[error(
387 "{context}: dense allocation size overflow for {copies} copies of a {nrows}x{ncols} f64 matrix"
388 )]
389 SizeOverflow {
390 context: Box<str>,
391 nrows: usize,
392 ncols: usize,
393 copies: usize,
394 },
395}
396
397#[derive(Debug)]
398struct GovernorLedger {
399 budget_bytes: usize,
400 materialization_cap_bytes: usize,
401 availability: MemoryAvailability,
402 reserved_bytes: std::sync::atomic::AtomicUsize,
403}
404
405#[derive(Debug, Clone)]
424pub struct MemoryGovernor {
425 ledger: Arc<GovernorLedger>,
426}
427
428impl MemoryGovernor {
429 pub fn global() -> &'static MemoryGovernor {
435 static GLOBAL: OnceLock<MemoryGovernor> = OnceLock::new();
436 GLOBAL.get_or_init(|| {
437 let availability = resample_memory_availability();
438 MemoryGovernor::with_detected_availability(availability)
439 })
440 }
441
442 fn with_detected_availability(availability: MemoryAvailability) -> Self {
443 let budget_bytes = governor_budget_from_availability(&availability);
444 let materialization_cap_bytes = governor_materialization_cap_from_availability(&availability);
445 Self {
446 ledger: Arc::new(GovernorLedger {
447 budget_bytes,
448 materialization_cap_bytes,
449 availability,
450 reserved_bytes: std::sync::atomic::AtomicUsize::new(0),
451 }),
452 }
453 }
454
455 pub fn budget_bytes(&self) -> usize {
459 self.ledger.budget_bytes
460 }
461
462 pub fn availability(&self) -> MemoryAvailability {
463 self.ledger.availability.clone()
464 }
465
466 pub fn reserved_bytes(&self) -> usize {
467 self.ledger
468 .reserved_bytes
469 .load(std::sync::atomic::Ordering::Acquire)
470 }
471
472 pub fn remaining_bytes(&self) -> usize {
477 self.ledger
478 .budget_bytes
479 .saturating_sub(self.reserved_bytes())
480 }
481
482 pub fn single_materialization_cap_bytes(&self) -> usize {
504 self.ledger.materialization_cap_bytes
505 }
506
507 pub fn try_reserve(
514 &self,
515 bytes: usize,
516 context: &str,
517 ) -> Result<MemoryReservation, MemoryReservationError> {
518 use std::sync::atomic::Ordering;
519 let mut current = self.ledger.reserved_bytes.load(Ordering::Relaxed);
520 loop {
521 let next = match current.checked_add(bytes) {
522 Some(next) if next <= self.ledger.budget_bytes => next,
523 _ => {
524 return Err(MemoryReservationError::BudgetExceeded {
525 context: context.into(),
526 requested_bytes: bytes,
527 reserved_bytes: current,
528 budget_bytes: self.ledger.budget_bytes,
529 availability: self.ledger.availability.clone(),
530 });
531 }
532 };
533 match self.ledger.reserved_bytes.compare_exchange_weak(
534 current,
535 next,
536 Ordering::AcqRel,
537 Ordering::Relaxed,
538 ) {
539 Ok(_) => {
540 return Ok(MemoryReservation {
541 ledger: Arc::clone(&self.ledger),
542 bytes,
543 });
544 }
545 Err(observed) => current = observed,
546 }
547 }
548 }
549
550 pub fn try_reserve_dense_f64(
554 &self,
555 nrows: usize,
556 ncols: usize,
557 context: &str,
558 ) -> Result<MemoryReservation, MemoryReservationError> {
559 self.try_reserve_dense_f64_copies(nrows, ncols, 1, context)
560 }
561
562 pub fn try_reserve_dense_f64_copies(
565 &self,
566 nrows: usize,
567 ncols: usize,
568 copies: usize,
569 context: &str,
570 ) -> Result<MemoryReservation, MemoryReservationError> {
571 let bytes = dense_f64_bytes(nrows, ncols)
572 .and_then(|one| one.checked_mul(copies))
573 .ok_or_else(|| MemoryReservationError::SizeOverflow {
574 context: context.into(),
575 nrows,
576 ncols,
577 copies,
578 })?;
579 self.try_reserve(bytes, context)
580 }
581}
582
583pub const fn dense_f64_bytes(nrows: usize, ncols: usize) -> Option<usize> {
585 match nrows.checked_mul(ncols) {
586 Some(cells) => cells.checked_mul(std::mem::size_of::<f64>()),
587 None => None,
588 }
589}
590
591#[derive(Debug)]
595#[must_use = "dropping a memory reservation immediately releases its ledger charge"]
596pub struct MemoryReservation {
597 ledger: Arc<GovernorLedger>,
598 bytes: usize,
599}
600
601impl MemoryReservation {
602 pub fn bytes(&self) -> usize {
603 self.bytes
604 }
605
606 pub fn bind<T>(self, value: T) -> Governed<T> {
608 Governed {
609 value,
610 reservation: self,
611 }
612 }
613}
614
615#[derive(Debug)]
621#[must_use = "the governed value owns a live process-wide memory reservation"]
622pub struct Governed<T> {
623 value: T,
624 reservation: MemoryReservation,
625}
626
627impl<T> Governed<T> {
628 pub fn reserved_bytes(&self) -> usize {
629 self.reservation.bytes()
630 }
631}
632
633impl<T> std::ops::Deref for Governed<T> {
634 type Target = T;
635
636 fn deref(&self) -> &Self::Target {
637 &self.value
638 }
639}
640
641impl<T> std::ops::DerefMut for Governed<T> {
642 fn deref_mut(&mut self) -> &mut Self::Target {
643 &mut self.value
644 }
645}
646
647impl<T> AsRef<T> for Governed<T> {
648 fn as_ref(&self) -> &T {
649 &self.value
650 }
651}
652
653impl<T> AsMut<T> for Governed<T> {
654 fn as_mut(&mut self) -> &mut T {
655 &mut self.value
656 }
657}
658
659impl Drop for MemoryReservation {
660 fn drop(&mut self) {
661 self.ledger
662 .reserved_bytes
663 .fetch_sub(self.bytes, std::sync::atomic::Ordering::AcqRel);
664 }
665}
666
667#[derive(Clone, Copy, Debug, Default)]
670pub struct ProblemHints {
671 pub marginal_slope_large_scale_active: bool,
672}
673
674#[derive(Clone, Copy, Debug, PartialEq, Eq)]
675pub enum DerivativeStorageMode {
676 AnalyticOperatorRequired,
678 MaterializeIfSmall,
680 DiagnosticsOnly,
682}
683
684#[derive(Clone, Debug)]
685pub struct MaterializationPolicy {
686 pub max_single_dense_bytes: usize,
687 pub max_cached_dense_bytes: usize,
688 pub row_chunk_target_bytes: usize,
689 pub allow_operator_materialization: bool,
690 pub allow_diagnostic_materialization: bool,
691}
692
693#[derive(Debug, thiserror::Error)]
694pub enum MatrixMaterializationError {
695 #[error(
696 "{context}: dense materialization of {nrows}x{ncols} requires {bytes} bytes (limit {limit_bytes})"
697 )]
698 TooLarge {
699 context: &'static str,
700 nrows: usize,
701 ncols: usize,
702 bytes: usize,
703 limit_bytes: usize,
704 },
705
706 #[error("{context}: operator does not implement chunked row access")]
707 MissingRowChunk { context: &'static str },
708
709 #[error("{context}: row materialization failed: {reason}")]
710 RowMaterializationFailed {
711 context: &'static str,
712 reason: String,
713 },
714
715 #[error("{context}: materialization forbidden by policy (mode={mode:?})")]
716 Forbidden {
717 context: &'static str,
718 mode: DerivativeStorageMode,
719 },
720
721 #[error(transparent)]
725 Reservation(#[from] MemoryReservationError),
726}
727
728pub trait ResidentBytes {
729 fn resident_bytes(&self) -> usize;
730}
731
732impl ResourcePolicy {
733 pub fn default_library() -> Self {
750 Self::for_observed_memory(process_memory_availability())
751 }
752
753 pub fn for_observed_memory(availability: &MemoryAvailability) -> Self {
771 let single_cap = governor_materialization_cap_from_availability(availability);
772 Self {
773 max_single_materialization_bytes: single_cap,
774 max_operator_cache_bytes: single_cap,
775 max_spatial_distance_cache_bytes: single_cap,
776 max_owned_data_cache_bytes: single_cap,
777 row_chunk_target_bytes: LIBRARY_ROW_CHUNK_TARGET_BYTES,
778 derivative_storage_mode: DerivativeStorageMode::MaterializeIfSmall,
779 }
780 }
781
782 pub fn analytic_operator_required() -> Self {
788 let base = Self::default_library();
789 Self {
790 derivative_storage_mode: DerivativeStorageMode::AnalyticOperatorRequired,
791 ..base
792 }
793 }
794
795 pub fn for_problem(hints: ProblemHints) -> Self {
807 if hints.marginal_slope_large_scale_active {
808 return Self::analytic_operator_required();
809 }
810 Self::default_library()
811 }
812
813 pub fn permissive_small_data() -> Self {
816 let base = Self::default_library();
817 Self {
818 row_chunk_target_bytes: 64 * 1024 * 1024,
819 ..base
820 }
821 }
822
823 pub const fn material_policy(&self) -> MaterializationPolicy {
824 MaterializationPolicy {
825 max_single_dense_bytes: self.max_single_materialization_bytes,
826 max_cached_dense_bytes: self.max_operator_cache_bytes,
827 row_chunk_target_bytes: self.row_chunk_target_bytes,
828 allow_operator_materialization: matches!(
829 self.derivative_storage_mode,
830 DerivativeStorageMode::MaterializeIfSmall
831 ),
832 allow_diagnostic_materialization: !matches!(
833 self.derivative_storage_mode,
834 DerivativeStorageMode::AnalyticOperatorRequired
835 ),
836 }
837 }
838}
839
840pub const fn rows_for_target_bytes(target_bytes: usize, cols: usize) -> usize {
843 let raw_bytes_per_row = cols.saturating_mul(std::mem::size_of::<f64>());
844 let bytes_per_row = if raw_bytes_per_row == 0 {
845 1
846 } else {
847 raw_bytes_per_row
848 };
849 let rows = target_bytes / bytes_per_row;
850 if rows == 0 { 1 } else { rows }
851}
852
853pub fn prediction_chunk_rows(parameter_dim: usize, local_dim: usize, total_rows: usize) -> usize {
860 const MIN_ROWS: usize = 16;
861 const MAX_ROWS: usize = 4096;
862
863 if total_rows == 0 {
864 return 1;
865 }
866 let live_f64_values_per_row = parameter_dim
867 .max(1)
868 .saturating_mul(local_dim.max(1))
869 .saturating_mul(4);
870 rows_for_target_bytes(
871 ResourcePolicy::default_library().row_chunk_target_bytes,
872 live_f64_values_per_row,
873 )
874 .clamp(MIN_ROWS, MAX_ROWS)
875 .min(total_rows)
876}
877
878use std::collections::{HashMap, VecDeque};
879use std::hash::{Hash, Hasher};
880use std::sync::{Arc, Mutex, OnceLock};
881
882pub struct ByteLruCache<K: Eq + Hash + Clone, V> {
893 shards: Box<[Mutex<ByteLruInner<K, V>>]>,
904 shard_bytes: usize,
906 shard_entries: Option<usize>,
908 max_bytes: usize,
909 governor: MemoryGovernor,
910}
911
912struct ByteLruInner<K, V> {
913 map: HashMap<K, (V, usize, MemoryReservation)>,
916 order: VecDeque<K>,
917 resident_bytes: usize,
918}
919
920impl<K: Eq + Hash + Clone, V: Clone + ResidentBytes> ByteLruCache<K, V> {
921 pub fn new(max_bytes: usize) -> Self {
922 Self::build(max_bytes, None, 1)
923 }
924
925 pub fn with_max_entries(max_bytes: usize, max_entries: usize) -> Self {
926 Self::build(max_bytes, Some(max_entries), 1)
927 }
928
929 pub fn new_sharded(max_bytes: usize, shard_count: usize) -> Self {
934 Self::build(max_bytes, None, shard_count)
935 }
936
937 pub fn with_max_entries_sharded(
940 max_bytes: usize,
941 max_entries: usize,
942 shard_count: usize,
943 ) -> Self {
944 Self::build(max_bytes, Some(max_entries), shard_count)
945 }
946
947 fn build(max_bytes: usize, max_entries: Option<usize>, shard_count: usize) -> Self {
948 Self::build_with_governor(
949 max_bytes,
950 max_entries,
951 shard_count,
952 MemoryGovernor::global().clone(),
953 )
954 }
955
956 fn build_with_governor(
957 max_bytes: usize,
958 max_entries: Option<usize>,
959 shard_count: usize,
960 governor: MemoryGovernor,
961 ) -> Self {
962 let shard_count = shard_count.max(1);
963 let shard_bytes = max_bytes.div_ceil(shard_count);
968 let shard_entries = max_entries.map(|m| {
969 if m == 0 {
970 0
971 } else {
972 m.div_ceil(shard_count).max(1)
973 }
974 });
975 let shards = (0..shard_count)
976 .map(|_| {
977 Mutex::new(ByteLruInner {
978 map: HashMap::new(),
979 order: VecDeque::new(),
980 resident_bytes: 0,
981 })
982 })
983 .collect::<Vec<_>>()
984 .into_boxed_slice();
985 Self {
986 shards,
987 shard_bytes,
988 shard_entries,
989 max_bytes,
990 governor,
991 }
992 }
993
994 #[inline]
995 fn shard(&self, key: &K) -> &Mutex<ByteLruInner<K, V>> {
996 if self.shards.len() == 1 {
997 return &self.shards[0];
998 }
999 let mut hasher = std::collections::hash_map::DefaultHasher::new();
1000 key.hash(&mut hasher);
1001 &self.shards[(hasher.finish() as usize) % self.shards.len()]
1002 }
1003
1004 pub fn get(&self, key: &K) -> Option<V> {
1005 let mut g = self.shard(key).lock().unwrap_or_else(|p| p.into_inner());
1007 let v = g.map.get(key)?.0.clone();
1008 if let Some(pos) = g.order.iter().position(|k| k == key) {
1010 let k = g
1011 .order
1012 .remove(pos)
1013 .expect("position() returned an in-bounds index into this same deque");
1014 g.order.push_back(k);
1015 }
1016 Some(v)
1017 }
1018
1019 pub fn insert(&self, key: K, value: V) {
1020 let charge = value.resident_bytes();
1021 let mut g = self.shard(&key).lock().unwrap_or_else(|p| p.into_inner());
1022
1023 if let Some((_old, old_charge, _reservation)) = g.map.remove(&key) {
1026 g.resident_bytes = g.resident_bytes.saturating_sub(old_charge);
1027 if let Some(pos) = g.order.iter().position(|k| k == &key) {
1028 g.order.remove(pos);
1029 }
1030 }
1031
1032 if charge > self.shard_bytes {
1033 return;
1035 }
1036
1037 if let Some(max_entries) = self.shard_entries {
1038 if max_entries == 0 {
1039 return;
1040 }
1041 while g.map.len() >= max_entries {
1042 if let Some(evict_key) = g.order.pop_front() {
1043 if let Some((_v, c, _reservation)) = g.map.remove(&evict_key) {
1044 g.resident_bytes = g.resident_bytes.saturating_sub(c);
1045 }
1046 } else {
1047 break;
1048 }
1049 }
1050 }
1051
1052 while g.resident_bytes + charge > self.shard_bytes {
1053 if let Some(evict_key) = g.order.pop_front() {
1054 if let Some((_v, c, _reservation)) = g.map.remove(&evict_key) {
1055 g.resident_bytes = g.resident_bytes.saturating_sub(c);
1056 }
1057 } else {
1058 break;
1059 }
1060 }
1061
1062 let reservation = match self
1063 .governor
1064 .try_reserve(charge, "ByteLruCache resident entry")
1065 {
1066 Ok(reservation) => reservation,
1067 Err(_) => return,
1068 };
1069 g.map.insert(key.clone(), (value, charge, reservation));
1070 g.order.push_back(key);
1071 g.resident_bytes = g.resident_bytes.saturating_add(charge);
1072 }
1073
1074 pub fn resident_bytes(&self) -> usize {
1075 self.shards
1076 .iter()
1077 .map(|shard| {
1078 shard
1079 .lock()
1080 .unwrap_or_else(|p| p.into_inner())
1081 .resident_bytes
1082 })
1083 .sum()
1084 }
1085
1086 pub const fn max_bytes(&self) -> usize {
1087 self.max_bytes
1088 }
1089
1090 pub fn len(&self) -> usize {
1091 self.shards
1092 .iter()
1093 .map(|shard| shard.lock().unwrap_or_else(|p| p.into_inner()).map.len())
1094 .sum()
1095 }
1096
1097 pub fn is_empty(&self) -> bool {
1098 self.len() == 0
1099 }
1100
1101 pub fn clear(&self) {
1102 for shard in self.shards.iter() {
1103 let mut g = shard.lock().unwrap_or_else(|p| p.into_inner());
1104 g.map.clear();
1105 g.order.clear();
1106 g.resident_bytes = 0;
1107 }
1108 }
1109}
1110
1111impl<K: Eq + Hash + Clone, V: Clone + ResidentBytes> std::fmt::Debug for ByteLruCache<K, V> {
1112 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1113 f.debug_struct("ByteLruCache")
1114 .field("resident_bytes", &self.resident_bytes())
1115 .field("max_bytes", &self.max_bytes)
1116 .field("shard_count", &self.shards.len())
1117 .field("shard_bytes", &self.shard_bytes)
1118 .field("shard_entries", &self.shard_entries)
1119 .finish()
1120 }
1121}
1122
1123impl ResidentBytes for Arc<ndarray::Array2<f64>> {
1130 fn resident_bytes(&self) -> usize {
1131 std::mem::size_of::<f64>()
1132 .saturating_mul(self.nrows())
1133 .saturating_mul(self.ncols())
1134 }
1135}
1136
1137pub struct RayonSafeOnce<T> {
1158 slot: std::sync::OnceLock<T>,
1159}
1160
1161impl<T> RayonSafeOnce<T> {
1162 pub const fn new() -> Self {
1163 Self {
1164 slot: std::sync::OnceLock::new(),
1165 }
1166 }
1167
1168 #[inline]
1170 pub fn get(&self) -> Option<&T> {
1171 self.slot.get()
1172 }
1173
1174 pub fn get_or_compute<F>(&self, init: F) -> &T
1186 where
1187 F: FnOnce() -> T,
1188 {
1189 if let Some(v) = self.slot.get() {
1190 return v;
1191 }
1192 let candidate = init();
1193 if self.slot.set(candidate).is_err() {
1194 log::trace!(
1195 "RayonSafeOnce: a concurrent initializer won the race; \
1196 keeping its value and discarding this candidate"
1197 );
1198 }
1199 self.slot
1200 .get()
1201 .expect("RayonSafeOnce slot populated by set() above")
1202 }
1203}
1204
1205impl<T> Default for RayonSafeOnce<T> {
1206 fn default() -> Self {
1207 Self::new()
1208 }
1209}
1210
1211impl<T: Clone> Clone for RayonSafeOnce<T> {
1212 fn clone(&self) -> Self {
1213 let cloned = Self::new();
1214 if let Some(value) = self.slot.get() {
1215 cloned.slot.get_or_init(|| value.clone());
1221 }
1222 cloned
1223 }
1224}
1225
1226impl<T: std::fmt::Debug> std::fmt::Debug for RayonSafeOnce<T> {
1227 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1228 f.debug_struct("RayonSafeOnce")
1229 .field("slot", &self.slot.get())
1230 .finish()
1231 }
1232}
1233
1234#[cfg(test)]
1235mod byte_lru_tests {
1236 use super::*;
1237
1238 fn cache_test_governor(budget_bytes: usize) -> MemoryGovernor {
1239 let available_bytes = (budget_bytes as u128 * GOVERNOR_BUDGET_DENOMINATOR)
1240 .div_ceil(GOVERNOR_BUDGET_NUMERATOR);
1241 let available_bytes =
1242 u64::try_from(available_bytes).expect("test cache budget must fit in u64");
1243 MemoryGovernor::with_detected_availability(MemoryAvailability::from_observation(
1244 available_bytes,
1245 available_bytes,
1246 CgroupMemoryObservation::NotPresent,
1247 ))
1248 }
1249
1250 #[derive(Clone, PartialEq, Debug)]
1252 struct Payload(u64);
1253 impl ResidentBytes for Payload {
1254 fn resident_bytes(&self) -> usize {
1255 8
1256 }
1257 }
1258
1259 #[test]
1260 fn single_shard_round_trips_and_evicts_by_bytes() {
1261 let cache: ByteLruCache<u64, Payload> =
1263 ByteLruCache::build_with_governor(24, None, 1, cache_test_governor(24));
1264 for k in 0..3 {
1265 cache.insert(k, Payload(k));
1266 }
1267 assert_eq!(cache.len(), 3);
1268 assert_eq!(cache.resident_bytes(), 24);
1269 assert_eq!(cache.get(&0), Some(Payload(0)));
1271 cache.insert(3, Payload(3));
1272 assert_eq!(cache.len(), 3);
1274 assert_eq!(cache.get(&1), None);
1275 assert_eq!(cache.get(&0), Some(Payload(0)));
1276 assert_eq!(cache.get(&3), Some(Payload(3)));
1277 }
1278
1279 #[test]
1280 fn zero_entry_budget_disables_caching_in_every_shard() {
1281 let single: ByteLruCache<u64, Payload> = ByteLruCache::with_max_entries(1 << 20, 0);
1282 single.insert(7, Payload(7));
1283 assert_eq!(single.get(&7), None);
1284 let sharded: ByteLruCache<u64, Payload> =
1285 ByteLruCache::with_max_entries_sharded(1 << 20, 0, 16);
1286 sharded.insert(7, Payload(7));
1287 assert_eq!(sharded.get(&7), None);
1288 }
1289
1290 #[test]
1291 fn sharded_cache_retrieves_all_keys_and_respects_aggregate_budget() {
1292 let shard_count = 8usize;
1296 let max_bytes = 8 * 64; let cache: ByteLruCache<u64, Payload> = ByteLruCache::build_with_governor(
1298 max_bytes,
1299 None,
1300 shard_count,
1301 cache_test_governor(max_bytes),
1302 );
1303 for k in 0..64u64 {
1304 cache.insert(k, Payload(k));
1305 }
1306 assert!(cache.resident_bytes() <= max_bytes.div_ceil(shard_count) * shard_count);
1308 cache.insert(123, Payload(123));
1310 assert_eq!(cache.get(&123), Some(Payload(123)));
1311 assert!(!cache.is_empty());
1312 cache.clear();
1313 assert_eq!(cache.len(), 0);
1314 assert_eq!(cache.resident_bytes(), 0);
1315 }
1316}
1317
1318#[cfg(test)]
1319mod resource_policy_tests {
1320 use super::*;
1321
1322 fn test_governor(budget_bytes: usize) -> MemoryGovernor {
1323 let available_bytes = (budget_bytes as u128 * GOVERNOR_BUDGET_DENOMINATOR)
1324 .div_ceil(GOVERNOR_BUDGET_NUMERATOR);
1325 let available_bytes =
1326 u64::try_from(available_bytes).expect("test budget has a representable observation");
1327 let availability = MemoryAvailability::from_observation(
1328 available_bytes,
1329 available_bytes,
1330 CgroupMemoryObservation::NotPresent,
1331 );
1332 let governor = MemoryGovernor::with_detected_availability(availability);
1333 assert_eq!(governor.budget_bytes(), budget_bytes);
1334 governor
1335 }
1336
1337 #[test]
1340 fn rows_for_target_bytes_exact_fit() {
1341 assert_eq!(rows_for_target_bytes(8, 1), 1);
1343 }
1344
1345 #[test]
1346 fn rows_for_target_bytes_multiple_rows() {
1347 assert_eq!(rows_for_target_bytes(80, 1), 10);
1349 }
1350
1351 #[test]
1352 fn rows_for_target_bytes_multiple_cols() {
1353 assert_eq!(rows_for_target_bytes(128, 4), 4);
1355 }
1356
1357 #[test]
1358 fn rows_for_target_bytes_zero_target_returns_one() {
1359 assert_eq!(rows_for_target_bytes(0, 1), 1);
1361 }
1362
1363 #[test]
1364 fn rows_for_target_bytes_zero_cols_returns_non_zero() {
1365 assert_eq!(rows_for_target_bytes(100, 0), 100);
1367 }
1368
1369 #[test]
1370 fn rows_for_target_bytes_large_target() {
1371 let target = LIBRARY_ROW_CHUNK_TARGET_BYTES;
1375 let cols = 1024_usize;
1376 let expected = target / (cols * std::mem::size_of::<f64>());
1377 assert_eq!(rows_for_target_bytes(target, cols), expected);
1378 }
1379
1380 #[test]
1381 fn prediction_chunks_share_the_runtime_byte_budget() {
1382 assert_eq!(prediction_chunk_rows(1024, 1, 100_000), 256);
1383 assert_eq!(prediction_chunk_rows(32, 2, 100_000), 4096);
1384 }
1385
1386 #[test]
1387 fn prediction_chunks_respect_dataset_bounds() {
1388 assert_eq!(prediction_chunk_rows(1, 1, 7), 7);
1389 assert_eq!(prediction_chunk_rows(1, 1, 0), 1);
1390 }
1391
1392 #[test]
1395 fn for_problem_small_data_uses_materialize_if_small() {
1396 let p = ResourcePolicy::for_problem(ProblemHints::default());
1397 assert_eq!(
1398 p.derivative_storage_mode,
1399 DerivativeStorageMode::MaterializeIfSmall
1400 );
1401 }
1402
1403 #[test]
1404 fn for_problem_has_no_row_or_column_cliff() {
1405 let narrow = ResourcePolicy::for_problem(ProblemHints::default());
1406 let wide = ResourcePolicy::for_problem(ProblemHints::default());
1407 assert_eq!(
1408 narrow.derivative_storage_mode,
1409 DerivativeStorageMode::MaterializeIfSmall
1410 );
1411 assert_eq!(
1412 wide.derivative_storage_mode,
1413 DerivativeStorageMode::MaterializeIfSmall
1414 );
1415 }
1416
1417 #[test]
1418 fn for_problem_dimension_overflow_defers_to_typed_reservation() {
1419 let policy = ResourcePolicy::for_problem(ProblemHints::default());
1420 assert_eq!(
1421 policy.derivative_storage_mode,
1422 DerivativeStorageMode::MaterializeIfSmall
1423 );
1424 }
1425
1426 #[test]
1427 fn for_problem_marginal_slope_hint_is_strict() {
1428 let p = ResourcePolicy::for_problem(ProblemHints {
1429 marginal_slope_large_scale_active: true,
1430 });
1431 assert_eq!(
1432 p.derivative_storage_mode,
1433 DerivativeStorageMode::AnalyticOperatorRequired
1434 );
1435 }
1436
1437 #[test]
1440 fn material_policy_default_library_allows_operator_and_diagnostics() {
1441 let mp = ResourcePolicy::default_library().material_policy();
1442 assert!(mp.allow_operator_materialization);
1443 assert!(mp.allow_diagnostic_materialization);
1444 }
1445
1446 #[test]
1447 fn material_policy_analytic_operator_required_blocks_both() {
1448 let mp = ResourcePolicy::analytic_operator_required().material_policy();
1449 assert!(!mp.allow_operator_materialization);
1450 assert!(!mp.allow_diagnostic_materialization);
1451 }
1452
1453 #[test]
1454 fn material_policy_propagates_byte_limits() {
1455 let policy = ResourcePolicy::default_library();
1456 let mp = policy.material_policy();
1457 assert_eq!(
1458 mp.max_single_dense_bytes,
1459 policy.max_single_materialization_bytes
1460 );
1461 assert_eq!(mp.max_cached_dense_bytes, policy.max_operator_cache_bytes);
1462 assert_eq!(mp.row_chunk_target_bytes, policy.row_chunk_target_bytes);
1463 }
1464
1465 #[test]
1468 fn reservations_account_and_release_on_drop() {
1469 let governor = test_governor(1_000);
1470 assert_eq!(governor.remaining_bytes(), 1_000);
1471 let first = governor.try_reserve(600, "test-first").expect("fits");
1472 assert_eq!(governor.reserved_bytes(), 600);
1473 assert_eq!(governor.remaining_bytes(), 400);
1474 assert_eq!(first.bytes(), 600);
1475 drop(first);
1476 assert_eq!(governor.reserved_bytes(), 0);
1477 assert_eq!(governor.remaining_bytes(), 1_000);
1478 }
1479
1480 #[test]
1481 fn jointly_excessive_reservations_are_refused_with_evidence() {
1482 let governor = test_governor(1_000);
1486 let availability = governor.availability();
1487 let held = governor.try_reserve(600, "test-held").expect("fits alone");
1488 let refusal = governor
1489 .try_reserve(600, "test-joint")
1490 .expect_err("600 + 600 exceeds the 1000-byte budget");
1491 assert_eq!(
1492 refusal,
1493 MemoryReservationError::BudgetExceeded {
1494 context: "test-joint".into(),
1495 requested_bytes: 600,
1496 reserved_bytes: 600,
1497 budget_bytes: 1_000,
1498 availability,
1499 }
1500 );
1501 drop(held);
1504 let refreshed = governor
1505 .try_reserve(600, "test-joint")
1506 .expect("fits after release");
1507 assert_eq!(refreshed.bytes(), 600);
1508 }
1509
1510 #[test]
1511 fn dense_reservation_uses_checked_footprint() {
1512 let governor = test_governor(1 << 20);
1513 let ok = governor
1514 .try_reserve_dense_f64(1024, 64, "test-dense")
1515 .expect("512 KiB fits in 1 MiB");
1516 assert_eq!(ok.bytes(), 1024 * 64 * 8);
1517 drop(ok);
1518 governor
1521 .try_reserve_dense_f64(usize::MAX, 2, "test-overflow")
1522 .expect_err("overflowing footprint cannot be reserved");
1523 assert!(matches!(
1524 governor.try_reserve_dense_f64(usize::MAX, 2, "test-overflow"),
1525 Err(MemoryReservationError::SizeOverflow { .. })
1526 ));
1527 }
1528
1529 #[test]
1530 fn concurrent_reservations_never_oversubscribe() {
1531 let governor = std::sync::Arc::new(test_governor(1_000));
1532 let granted = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
1533 let barrier = std::sync::Arc::new(std::sync::Barrier::new(9));
1534 std::thread::scope(|scope| {
1535 for _ in 0..8 {
1536 let governor = std::sync::Arc::clone(&governor);
1537 let granted = std::sync::Arc::clone(&granted);
1538 let barrier = std::sync::Arc::clone(&barrier);
1539 scope.spawn(move || {
1540 let held = governor.try_reserve(200, "test-race");
1541 if held.is_ok() {
1542 granted.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
1543 }
1544 barrier.wait();
1545 assert!(governor.reserved_bytes() <= governor.budget_bytes());
1546 barrier.wait();
1547 drop(held);
1548 });
1549 }
1550 barrier.wait();
1551 assert_eq!(granted.load(std::sync::atomic::Ordering::SeqCst), 5);
1552 assert_eq!(governor.reserved_bytes(), 1_000);
1553 barrier.wait();
1554 });
1555 assert_eq!(governor.reserved_bytes(), 0);
1556 }
1557
1558 #[test]
1559 fn global_policy_caps_are_one_shared_admission_ceiling() {
1560 let governor = MemoryGovernor::global();
1561 assert_eq!(
1562 governor.single_materialization_cap_bytes(),
1563 governor_materialization_cap_from_availability(&governor.availability())
1564 );
1565 assert_eq!(
1569 governor.single_materialization_cap_bytes(),
1570 governor.budget_bytes(),
1571 "the routing cap and the ledger budget are one ceiling asked two questions"
1572 );
1573 let policy = ResourcePolicy::default_library();
1574 assert_eq!(
1575 policy.max_single_materialization_bytes,
1576 governor.single_materialization_cap_bytes()
1577 );
1578 let strict = ResourcePolicy::analytic_operator_required();
1579 assert_eq!(
1580 strict.max_single_materialization_bytes,
1581 policy.max_single_materialization_bytes
1582 );
1583 }
1584
1585 #[test]
1586 fn governed_value_holds_and_releases_its_charge() {
1587 let governor = test_governor(64);
1588 let governed = governor
1589 .try_reserve(32, "governed-value")
1590 .expect("reservation fits")
1591 .bind(vec![0_u8; 32]);
1592 assert_eq!(governed.len(), 32);
1593 assert_eq!(governed.reserved_bytes(), 32);
1594 assert_eq!(governor.reserved_bytes(), 32);
1595 drop(governed);
1596 assert_eq!(governor.reserved_bytes(), 0);
1597 }
1598
1599 #[test]
1605 fn memory_availability_distinguishes_host_cgroup_and_exhaustion() {
1606 let host_only =
1607 MemoryAvailability::from_observation(1_000, 4_000, CgroupMemoryObservation::NotPresent);
1608 assert_eq!(host_only.available_bytes(), 1_000);
1609 assert_eq!(host_only.limiting_source(), MemoryAvailabilitySource::Host);
1610 assert_eq!(host_only.capacity_bytes(), 4_000);
1611 assert_eq!(governor_budget_from_availability(&host_only), 3_000);
1612
1613 let exhausted_host =
1617 MemoryAvailability::from_observation(0, 4_000, CgroupMemoryObservation::NotPresent);
1618 assert_eq!(exhausted_host.available_bytes(), 0);
1619 assert_eq!(governor_budget_from_availability(&exhausted_host), 3_000);
1620
1621 let finite_cgroup = MemoryAvailability::from_observation(
1622 1_000,
1623 4_000_000_000,
1624 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1625 "/fixture/leaf",
1626 600,
1627 200,
1628 0,
1629 1,
1630 )),
1631 );
1632 assert_eq!(finite_cgroup.available_bytes(), 400);
1633 assert_eq!(
1634 finite_cgroup.limiting_source(),
1635 MemoryAvailabilitySource::Cgroup
1636 );
1637 assert_eq!(finite_cgroup.capacity_bytes(), 600);
1639 assert_eq!(governor_budget_from_availability(&finite_cgroup), 450);
1640
1641 let exhausted_cgroup = MemoryAvailability::from_observation(
1642 1_000,
1643 4_000_000_000,
1644 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1645 "/fixture/leaf",
1646 600,
1647 600,
1648 0,
1649 1,
1650 )),
1651 );
1652 assert_eq!(exhausted_cgroup.available_bytes(), 0);
1653 assert_eq!(
1654 exhausted_cgroup.limiting_source(),
1655 MemoryAvailabilitySource::Cgroup
1656 );
1657 assert_eq!(exhausted_cgroup.capacity_bytes(), 600);
1660 assert_eq!(governor_budget_from_availability(&exhausted_cgroup), 450);
1661
1662 let host_is_tighter = MemoryAvailability::from_observation(
1665 1_000,
1666 4_000_000_000,
1667 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1668 "/fixture/leaf",
1669 8_000,
1670 2_000,
1671 0,
1672 1,
1673 )),
1674 );
1675 assert_eq!(host_is_tighter.available_bytes(), 1_000);
1676 assert_eq!(
1677 host_is_tighter.limiting_source(),
1678 MemoryAvailabilitySource::Host
1679 );
1680
1681 let equal_cgroup_ceiling = MemoryAvailability::from_observation(
1682 1_000,
1683 4_000_000_000,
1684 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1685 "/fixture/leaf",
1686 1_200,
1687 200,
1688 0,
1689 1,
1690 )),
1691 );
1692 assert_eq!(equal_cgroup_ceiling.available_bytes(), 1_000);
1693 assert_eq!(
1694 equal_cgroup_ceiling.limiting_source(),
1695 MemoryAvailabilitySource::HostAndCgroup
1696 );
1697
1698 let both_exhausted = MemoryAvailability::from_observation(
1699 0,
1700 4_000_000_000,
1701 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1702 "/fixture/leaf",
1703 600,
1704 600,
1705 0,
1706 1,
1707 )),
1708 );
1709 assert_eq!(
1710 both_exhausted.limiting_source(),
1711 MemoryAvailabilitySource::HostAndCgroup
1712 );
1713
1714 let zero_ceiling = MemoryAvailability::from_observation(
1717 8_000,
1718 4_000_000_000,
1719 CgroupMemoryObservation::V2Limited(CgroupMemoryAvailability::fixture(
1720 "/fixture/leaf",
1721 0,
1722 0,
1723 0,
1724 1,
1725 )),
1726 );
1727 assert_eq!(zero_ceiling.available_bytes(), 0);
1728 assert_eq!(
1729 zero_ceiling.limiting_source(),
1730 MemoryAvailabilitySource::Cgroup
1731 );
1732 assert_eq!(governor_budget_from_availability(&zero_ceiling), 0);
1733 }
1734
1735 #[test]
1736 fn literal_unlimited_cgroup_defers_to_host_available_memory_2317() {
1737 let unlimited = MemoryAvailability::from_observation(
1738 2_430_926_848,
1739 4_000_000_000,
1740 CgroupMemoryObservation::V2Unbounded {
1741 cgroup_path: "/fixture/leaf".into(),
1742 inspected_levels: 3,
1743 },
1744 );
1745 assert_eq!(unlimited.available_bytes(), 2_430_926_848);
1746 assert_eq!(unlimited.limiting_source(), MemoryAvailabilitySource::Host);
1747 assert_eq!(unlimited.capacity_bytes(), 4_000_000_000);
1750 assert_eq!(governor_budget_from_availability(&unlimited), 3_000_000_000);
1751 assert!(format!("{unlimited}").contains("unbounded cgroup-v2"));
1752 }
1753
1754 #[test]
1755 fn finite_cgroup_v1_headroom_participates_in_the_same_exact_minimum() {
1756 let availability = MemoryAvailability::from_observation(
1757 8_000,
1758 4_000_000_000,
1759 CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1760 "/sys/fs/cgroup/memory/slurm/job",
1761 4_000,
1762 1_500,
1763 500,
1764 4,
1765 )),
1766 );
1767 assert_eq!(availability.available_bytes(), 3_000);
1768 assert_eq!(
1769 availability.limiting_source(),
1770 MemoryAvailabilitySource::Cgroup
1771 );
1772 assert_eq!(availability.capacity_bytes(), 4_000);
1775 assert_eq!(governor_budget_from_availability(&availability), 3_000);
1776 let evidence = format!("{availability}");
1777 assert!(evidence.contains("cgroup-v1"));
1778 assert!(evidence.contains("available=3000"));
1779 }
1780
1781 #[test]
1795 fn a_cgroup_at_its_limit_moves_neither_the_budget_nor_the_materialization_cap_2684_2702() {
1796 const DESIGN_300X12_BYTES: usize = 300 * 12 * 8;
1797 let limit_bytes = 6 * 1024 * 1024 * 1024_u64;
1798 let at_the_limit = MemoryAvailability::from_observation(
1799 448_648_040_448,
1800 527_799_400 * 1024,
1801 CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1802 "/sys/fs/cgroup/memory/slurm/uid_81060/job_14615476",
1803 limit_bytes,
1804 limit_bytes - 53_248,
1805 0,
1806 6,
1807 )),
1808 );
1809 assert_eq!(at_the_limit.available_bytes(), 53_248);
1810 assert_eq!(at_the_limit.capacity_bytes(), limit_bytes);
1811 let governor = MemoryGovernor::with_detected_availability(at_the_limit);
1812 assert_eq!(governor.budget_bytes(), (limit_bytes as usize) / 4 * 3);
1815 assert!(governor.try_reserve(8 << 30, "larger-than-the-job").is_err());
1817 let se_chunk = governor
1821 .try_reserve_dense_f64_copies(64, 1, 2, "coefficient-SE solve chunk")
1822 .expect("a 1,024-byte SE chunk must be admissible in a 6 GiB job");
1823 assert_eq!(se_chunk.bytes(), 1_024);
1824 drop(se_chunk);
1825 assert_eq!(
1827 governor.single_materialization_cap_bytes(),
1828 (limit_bytes as usize) / 4 * 3
1829 );
1830 assert!(governor.single_materialization_cap_bytes() > DESIGN_300X12_BYTES);
1831
1832 let tiny_ceiling = MemoryAvailability::from_observation(
1836 448_648_040_448,
1837 527_799_400 * 1024,
1838 CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1839 "/fixture/tiny",
1840 1_024,
1841 0,
1842 0,
1843 1,
1844 )),
1845 );
1846 assert_eq!(tiny_ceiling.capacity_bytes(), 1_024);
1847 assert_eq!(
1848 MemoryGovernor::with_detected_availability(tiny_ceiling)
1849 .single_materialization_cap_bytes(),
1850 768
1851 );
1852
1853 let idle = MemoryAvailability::from_observation(
1856 448_648_040_448,
1857 527_799_400 * 1024,
1858 CgroupMemoryObservation::V1Limited(CgroupMemoryAvailability::fixture(
1859 "/sys/fs/cgroup/memory/slurm/uid_81060/job_14615476",
1860 limit_bytes,
1861 92_827_648,
1862 28_672,
1863 6,
1864 )),
1865 );
1866 assert!(idle.available_bytes() > 6_000_000_000);
1867 assert_eq!(
1868 MemoryGovernor::with_detected_availability(idle).single_materialization_cap_bytes(),
1869 (limit_bytes as usize) / 4 * 3
1870 );
1871 }
1872
1873 #[test]
1874 fn malformed_active_cgroup_fails_closed_with_typed_evidence() {
1875 let availability = MemoryAvailability::from_observation(
1876 8_000,
1877 4_000_000_000,
1878 CgroupMemoryObservation::ProbeFailed(CgroupMemoryProbeFailure::fixture(
1879 CgroupMemoryProbeFailureKind::InvalidCounter,
1880 "/fixture/leaf/memory.current",
1881 "expected an unsigned byte count",
1882 )),
1883 );
1884 assert_eq!(availability.available_bytes(), 0);
1885 assert_eq!(availability.capacity_bytes(), 0);
1888 assert_eq!(
1889 availability.limiting_source(),
1890 MemoryAvailabilitySource::CgroupProbeFailure
1891 );
1892 assert_eq!(governor_budget_from_availability(&availability), 0);
1893 assert_eq!(
1894 governor_materialization_cap_from_availability(&availability),
1895 0
1896 );
1897 let evidence = format!("{availability}");
1898 assert!(evidence.contains("failed closed"));
1899 assert!(evidence.contains("invalid-counter"));
1900 }
1901
1902 #[test]
1903 fn compressed_macos_observation_keeps_xnu_available_memory_positive() {
1904 let xnu_available = (75_514_u64 + 69_056 + 3_802) * 16_384;
1909 assert_eq!(xnu_available, 2_430_926_848);
1910 let availability = MemoryAvailability::from_observation(
1911 xnu_available,
1912 8 * 1024 * 1024 * 1024,
1913 CgroupMemoryObservation::NotPresent,
1914 );
1915 assert_eq!(availability.available_bytes(), xnu_available);
1916 assert_eq!(availability.capacity_bytes(), 8 * 1024 * 1024 * 1024);
1919 assert_eq!(
1920 governor_budget_from_availability(&availability),
1921 6 * 1024 * 1024 * 1024
1922 );
1923 }
1924}
1925
1926#[cfg(test)]
1946mod governor_budget_is_capacity_determined_2702_tests {
1947 use super::*;
1948
1949 const HOST_AVAILABLE_BYTES: u64 = 448_648_040_448;
1952 const HOST_TOTAL_BYTES: u64 = 527_799_400 * 1024;
1953 const JOB_LIMIT_BYTES: u64 = 8 * 1024 * 1024 * 1024;
1954
1955 const SE_TRANSFORMED_ROWS: usize = 512;
1961 const SE_CHUNK_BYTES: usize = SE_TRANSFORMED_ROWS * 8 * 2;
1963
1964 fn one_job_cgroup_at_load(charged: u64) -> MemoryAvailability {
1966 crate::test_support::simulated_cgroup_memory_environment(
1967 HOST_AVAILABLE_BYTES,
1968 HOST_TOTAL_BYTES,
1969 JOB_LIMIT_BYTES,
1970 charged,
1971 )
1972 }
1973
1974 fn pre_2702_free_denominated_budget(availability: &MemoryAvailability) -> usize {
1978 let scaled = u128::from(availability.available_bytes()) * GOVERNOR_BUDGET_NUMERATOR
1979 / GOVERNOR_BUDGET_DENOMINATOR;
1980 usize::try_from(scaled).unwrap_or(usize::MAX)
1981 }
1982
1983 #[test]
1984 fn one_job_observed_at_two_load_levels_yields_one_budget_and_one_verdict() {
1985 let roomy = one_job_cgroup_at_load(92_827_648);
1989 let pinned = one_job_cgroup_at_load(JOB_LIMIT_BYTES - 4_096);
1990
1991 assert!(roomy.available_bytes() > 7_000_000_000);
1994 assert_eq!(pinned.available_bytes(), 4_096);
1995 assert_eq!(roomy.capacity_bytes(), JOB_LIMIT_BYTES);
1997 assert_eq!(pinned.capacity_bytes(), JOB_LIMIT_BYTES);
1998
1999 assert_eq!(pre_2702_free_denominated_budget(&pinned), 3_072);
2004 assert!(pre_2702_free_denominated_budget(&pinned) < SE_CHUNK_BYTES);
2005 assert!(pre_2702_free_denominated_budget(&roomy) > SE_CHUNK_BYTES);
2006
2007 let roomy_governor = MemoryGovernor::with_detected_availability(roomy);
2008 let pinned_governor = MemoryGovernor::with_detected_availability(pinned);
2009
2010 assert_eq!(roomy_governor.budget_bytes(), pinned_governor.budget_bytes());
2011 assert_eq!(
2012 pinned_governor.budget_bytes(),
2013 (JOB_LIMIT_BYTES as usize) / 4 * 3
2014 );
2015 assert_eq!(
2016 roomy_governor.remaining_bytes(),
2017 pinned_governor.remaining_bytes()
2018 );
2019
2020 for governor in [&roomy_governor, &pinned_governor] {
2023 let reservation = governor
2024 .try_reserve_dense_f64_copies(
2025 SE_TRANSFORMED_ROWS,
2026 1,
2027 2,
2028 "factorized coefficient-SE solve chunk",
2029 )
2030 .expect("an 8 KiB SE chunk is admissible in an 8 GiB job at any load");
2031 assert_eq!(reservation.bytes(), SE_CHUNK_BYTES);
2032 }
2033 }
2034
2035 #[test]
2036 fn a_request_larger_than_the_job_is_still_refused_at_every_load() {
2037 for charged in [0, JOB_LIMIT_BYTES / 2, JOB_LIMIT_BYTES - 4_096] {
2040 let governor = MemoryGovernor::with_detected_availability(one_job_cgroup_at_load(charged));
2041 let refusal = governor
2042 .try_reserve(16 * 1024 * 1024 * 1024, "twice the job's ceiling")
2043 .expect_err("16 GiB cannot be admitted in an 8 GiB job");
2044 match refusal {
2045 MemoryReservationError::BudgetExceeded { budget_bytes, .. } => {
2046 assert_eq!(budget_bytes, (JOB_LIMIT_BYTES as usize) / 4 * 3);
2047 }
2048 other => panic!("expected a budget refusal naming the ceiling, got {other:?}"),
2049 }
2050 }
2051 }
2052
2053 #[test]
2054 fn the_verdict_does_not_depend_on_what_this_process_did_earlier() {
2055 let governor = MemoryGovernor::with_detected_availability(one_job_cgroup_at_load(0));
2060 let request = || {
2061 governor
2062 .try_reserve_dense_f64_copies(
2063 SE_TRANSFORMED_ROWS,
2064 1,
2065 2,
2066 "factorized coefficient-SE solve chunk",
2067 )
2068 .map(|reservation| reservation.bytes())
2069 };
2070
2071 let before = request().expect("admissible on a fresh ledger");
2072 {
2073 let bulk = governor
2076 .try_reserve(
2077 governor.remaining_bytes() - SE_CHUNK_BYTES,
2078 "prior work in this process",
2079 )
2080 .expect("the bulk reservation is exactly the remaining budget");
2081 assert_eq!(governor.remaining_bytes(), SE_CHUNK_BYTES);
2082 assert!(request().is_ok());
2086 drop(bulk);
2087 }
2088 assert_eq!(governor.reserved_bytes(), 0);
2089 assert_eq!(request().expect("admissible again"), before);
2090 }
2091}