laddu_memory/
resource.rs

1use serde::{Deserialize, Serialize};
2
3use crate::discovery::{CapacitySnapshot, MemoryProbe, SystemMemoryProbe};
4
5/// The basis used to obtain a resource's capacity information.
6#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
7#[serde(rename_all = "snake_case")]
8pub enum CapacitySource {
9    /// Host operating-system telemetry.
10    OperatingSystem,
11    /// A process or container limit.
12    Cgroup,
13    /// NVIDIA Management Library telemetry.
14    Nvml,
15    /// Linux DRM/sysfs telemetry.
16    Drm,
17    /// Windows DXGI telemetry.
18    Dxgi,
19    /// Apple Metal working-set telemetry.
20    Metal,
21    /// Capacity supplied by the user.
22    User,
23    /// Capacity is not observable and planning is adaptive.
24    Adaptive,
25}
26
27/// Kind of physical memory resource.
28#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
29#[serde(rename_all = "snake_case")]
30pub enum MemoryResourceKind {
31    /// Host RAM.
32    Host,
33    /// Accelerator-local or unified memory.
34    Device,
35}
36
37/// Stable information used to match a runtime accelerator to platform telemetry.
38#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
39pub struct DeviceIdentity {
40    /// Backend adapter index.
41    pub adapter_index: usize,
42    /// PCI vendor identifier, or zero when unavailable.
43    pub vendor_id: u32,
44    /// PCI device identifier, or zero when unavailable.
45    pub device_id: u32,
46    /// PCI bus identifier, or an empty string when unavailable.
47    pub pci_bus_id: String,
48}
49
50/// Snapshot of one physical memory resource.
51#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
52pub struct MemoryResource {
53    /// Stable resource identifier.
54    pub id: String,
55    /// Human-readable resource name.
56    pub name: String,
57    /// Host or accelerator memory.
58    pub kind: MemoryResourceKind,
59    /// Total capacity when observable.
60    pub total_bytes: Option<u64>,
61    /// Currently available capacity when observable.
62    pub available_bytes: Option<u64>,
63    /// Source of capacity information.
64    pub capacity_source: CapacitySource,
65    /// Accelerator identity used to refresh platform telemetry.
66    pub device_identity: Option<DeviceIdentity>,
67}
68
69#[derive(Clone, Copy, Debug, PartialEq, Eq)]
70pub(crate) enum ResourceValidationError {
71    EmptyId,
72    AvailableExceedsTotal,
73    HostHasDeviceIdentity,
74    HostHasDeviceTelemetry,
75    DeviceHasHostTelemetry,
76    TelemetryHasNoCapacity,
77}
78
79impl MemoryResource {
80    fn adaptive_device(id: impl Into<String>, name: impl Into<String>) -> Self {
81        Self {
82            id: id.into(),
83            name: name.into(),
84            kind: MemoryResourceKind::Device,
85            total_bytes: None,
86            available_bytes: None,
87            capacity_source: CapacitySource::Adaptive,
88            device_identity: None,
89        }
90    }
91
92    pub(crate) fn discover_device(
93        id: impl Into<String>,
94        name: impl Into<String>,
95        identity: DeviceIdentity,
96        fallback_bytes: u64,
97    ) -> Self {
98        let mut resource = Self::adaptive_device(id, name);
99        resource.device_identity = Some(identity);
100        if let Ok(snapshot) = SystemMemoryProbe.probe_device(&resource) {
101            resource.apply_capacity_snapshot(snapshot);
102            return resource;
103        }
104        resource.total_bytes = Some(fallback_bytes);
105        resource.available_bytes = Some(fallback_bytes);
106        resource
107    }
108
109    pub(crate) fn user_device(
110        id: impl Into<String>,
111        name: impl Into<String>,
112        total_bytes: u64,
113        available_bytes: Option<u64>,
114    ) -> Self {
115        Self {
116            id: id.into(),
117            name: name.into(),
118            kind: MemoryResourceKind::Device,
119            total_bytes: Some(total_bytes),
120            available_bytes: Some(available_bytes.unwrap_or(total_bytes).min(total_bytes)),
121            capacity_source: CapacitySource::User,
122            device_identity: None,
123        }
124    }
125
126    pub(crate) fn apply_capacity_snapshot(&mut self, snapshot: CapacitySnapshot) {
127        self.total_bytes = Some(snapshot.total_bytes);
128        self.available_bytes = Some(snapshot.available_bytes);
129        self.capacity_source = snapshot.source;
130    }
131
132    pub(crate) fn validate(&self) -> Result<(), ResourceValidationError> {
133        if self.id.is_empty() {
134            return Err(ResourceValidationError::EmptyId);
135        }
136        if let (Some(total), Some(available)) = (self.total_bytes, self.available_bytes)
137            && available > total
138        {
139            return Err(ResourceValidationError::AvailableExceedsTotal);
140        }
141        match self.kind {
142            MemoryResourceKind::Host => {
143                if self.device_identity.is_some() {
144                    return Err(ResourceValidationError::HostHasDeviceIdentity);
145                }
146                if matches!(
147                    self.capacity_source,
148                    CapacitySource::Nvml
149                        | CapacitySource::Drm
150                        | CapacitySource::Dxgi
151                        | CapacitySource::Metal
152                        | CapacitySource::Adaptive
153                ) {
154                    return Err(ResourceValidationError::HostHasDeviceTelemetry);
155                }
156            }
157            MemoryResourceKind::Device => {
158                if matches!(
159                    self.capacity_source,
160                    CapacitySource::OperatingSystem | CapacitySource::Cgroup
161                ) {
162                    return Err(ResourceValidationError::DeviceHasHostTelemetry);
163                }
164            }
165        }
166        if self.capacity_source != CapacitySource::Adaptive
167            && (self.total_bytes.is_none() || self.available_bytes.is_none())
168        {
169            return Err(ResourceValidationError::TelemetryHasNoCapacity);
170        }
171        Ok(())
172    }
173
174    pub(crate) fn normalize_capacity(&mut self) {
175        if let Some(total) = self.total_bytes
176            && let Some(available) = self.available_bytes.as_mut()
177        {
178            *available = (*available).min(total);
179        }
180    }
181
182    pub(crate) fn effective_available(&self) -> u64 {
183        self.available_bytes
184            .or(self.total_bytes)
185            .unwrap_or(u64::MAX)
186    }
187
188    pub(crate) fn is_refreshable(&self) -> bool {
189        self.kind == MemoryResourceKind::Device
190            && self.capacity_source != CapacitySource::User
191            && self.device_identity.is_some()
192    }
193}
194
195#[cfg(test)]
196mod tests {
197    use super::*;
198    use crate::discovery::discover_host;
199
200    fn device() -> MemoryResource {
201        MemoryResource {
202            id: "test".into(),
203            name: "Test".into(),
204            kind: MemoryResourceKind::Device,
205            total_bytes: Some(1_000),
206            available_bytes: Some(500),
207            capacity_source: CapacitySource::User,
208            device_identity: None,
209        }
210    }
211
212    #[test]
213    fn validation_names_invalid_field_combinations() {
214        let mut invalid = device();
215        invalid.id.clear();
216        assert_eq!(invalid.validate(), Err(ResourceValidationError::EmptyId));
217        let mut invalid = device();
218        invalid.available_bytes = Some(1_001);
219        assert_eq!(
220            invalid.validate(),
221            Err(ResourceValidationError::AvailableExceedsTotal)
222        );
223        let mut invalid = discover_host();
224        invalid.device_identity = Some(DeviceIdentity {
225            adapter_index: 0,
226            vendor_id: 1,
227            device_id: 2,
228            pci_bus_id: "bus-0".into(),
229        });
230        assert_eq!(
231            invalid.validate(),
232            Err(ResourceValidationError::HostHasDeviceIdentity)
233        );
234        let mut invalid = device();
235        invalid.capacity_source = CapacitySource::OperatingSystem;
236        assert_eq!(
237            invalid.validate(),
238            Err(ResourceValidationError::DeviceHasHostTelemetry)
239        );
240
241        let mut invalid = device();
242        invalid.capacity_source = CapacitySource::Nvml;
243        invalid.total_bytes = None;
244        assert_eq!(
245            invalid.validate(),
246            Err(ResourceValidationError::TelemetryHasNoCapacity)
247        );
248
249        let mut invalid = discover_host();
250        invalid.capacity_source = CapacitySource::Metal;
251        assert_eq!(
252            invalid.validate(),
253            Err(ResourceValidationError::HostHasDeviceTelemetry)
254        );
255    }
256}
laddu_memory/resource.rs

laddu_memory/
resource.rs
