cranpose_services/
device_info.rs

1//! What the device has, and what this process is using of it.
2//!
3//! Applications ask two different questions here. *How much memory does this
4//! device have* sizes a decision made once — whether a large model fits at all.
5//! *How much is this process holding, and how much may it still have* is asked
6//! while work is running, because the answer moves and because the platform
7//! kills a process that gets it wrong.
8//!
9//! Every reading is optional. A platform that will not say reports `None`
10//! rather than a zero an application would divide by, and
11//! [`release_free_memory`] reports whether the platform has such a call at all
12//! instead of quietly doing nothing.
13//!
14//! The default reads what can be read without leaving safe Rust:
15//! `/proc/meminfo` and `/proc/self/statm` on Linux and Android. The platform
16//! backends install richer implementations through
17//! [`set_platform_device_info`].
18
19use std::{cell::RefCell, rc::Rc, time::Duration};
20
21/// Provides device and process information.
22pub trait DeviceInfo {
23    /// Total physical memory in bytes, or `None` if unknown.
24    fn total_memory_bytes(&self) -> Option<u64>;
25
26    /// Memory this process currently holds resident, or `None` where the
27    /// platform will not say.
28    fn resident_memory_bytes(&self) -> Option<u64> {
29        None
30    }
31
32    /// Memory this process may still allocate before the platform stops it.
33    ///
34    /// Not the same as free system memory: a device with gigabytes free may
35    /// still refuse this process another hundred megabytes, and it is the
36    /// second number that decides whether the next allocation is the one that
37    /// gets the application killed.
38    fn available_memory_bytes(&self) -> Option<u64> {
39        None
40    }
41
42    /// Processor time this process has used, user and system together.
43    ///
44    /// Wall-clock time says how long something took; this says how much of a
45    /// core it took, which is what a background lane that must not heat the
46    /// device up is actually rationing.
47    fn process_cpu_time(&self) -> Option<Duration> {
48        None
49    }
50
51    /// Asks the allocator to return free pages to the system.
52    ///
53    /// Returns whether the platform has such a call. Freeing a large buffer
54    /// does not necessarily shrink the process — an allocator keeps the pages
55    /// for the next allocation — and on a platform that kills by resident size
56    /// that is the difference between finishing and being killed.
57    fn release_free_memory(&self) -> bool {
58        false
59    }
60}
61
62pub type DeviceInfoRef = Rc<dyn DeviceInfo>;
63
64struct DefaultDeviceInfo;
65
66impl DeviceInfo for DefaultDeviceInfo {
67    fn total_memory_bytes(&self) -> Option<u64> {
68        #[cfg(any(target_os = "linux", target_os = "android"))]
69        {
70            let text = std::fs::read_to_string("/proc/meminfo").ok()?;
71            for line in text.lines() {
72                if let Some(rest) = line.strip_prefix("MemTotal:") {
73                    let kb: u64 = rest.split_whitespace().next()?.parse().ok()?;
74                    return Some(kb * 1024);
75                }
76            }
77            None
78        }
79        #[cfg(not(any(target_os = "linux", target_os = "android")))]
80        {
81            None
82        }
83    }
84
85    fn resident_memory_bytes(&self) -> Option<u64> {
86        #[cfg(any(target_os = "linux", target_os = "android"))]
87        {
88            let text = std::fs::read_to_string("/proc/self/statm").ok()?;
89            resident_bytes_from_statm(&text, page_size_bytes())
90        }
91        #[cfg(not(any(target_os = "linux", target_os = "android")))]
92        {
93            None
94        }
95    }
96}
97
98#[cfg(any(target_os = "linux", target_os = "android", test))]
99fn resident_bytes_from_statm(text: &str, page_size: u64) -> Option<u64> {
100    text.split_whitespace()
101        .nth(1)?
102        .parse::<u64>()
103        .ok()?
104        .checked_mul(page_size)
105}
106
107#[cfg(any(target_os = "linux", target_os = "android", test))]
108const fn page_size_bytes() -> u64 {
109    4096
110}
111
112thread_local! {
113    static PLATFORM_DEVICE_INFO: RefCell<Option<DeviceInfoRef>> = const { RefCell::new(None) };
114}
115
116/// Installs a platform device-info implementation, replacing any previous one.
117pub fn set_platform_device_info(info: DeviceInfoRef) {
118    PLATFORM_DEVICE_INFO.with(|cell| *cell.borrow_mut() = Some(info));
119}
120
121/// Removes any registered platform device info (tests and teardown).
122pub fn clear_platform_device_info() {
123    PLATFORM_DEVICE_INFO.with(|cell| *cell.borrow_mut() = None);
124}
125
126/// The active device info: the platform implementation if installed, otherwise
127/// the built-in default.
128pub fn device_info() -> DeviceInfoRef {
129    PLATFORM_DEVICE_INFO
130        .with(|cell| cell.borrow().clone())
131        .unwrap_or_else(|| Rc::new(DefaultDeviceInfo))
132}
133
134/// Asks the allocator to return free pages to the system.
135///
136/// Returns whether the platform has such a call — see
137/// [`DeviceInfo::release_free_memory`]. Worth doing after finishing with a
138/// large buffer on a platform that kills by resident size, and worth doing
139/// nowhere else: it walks the allocator's arenas.
140pub fn release_free_memory() -> bool {
141    device_info().release_free_memory()
142}
143
144#[cfg(test)]
145#[path = "tests/device_info_tests.rs"]
146mod tests;
cranpose_services/device_info.rs

cranpose_services/
device_info.rs
