Skip to main content

tenzro_types/
hardware.rs

1//! Hardware capability descriptors used across the Tenzro stack.
2//!
3//! Both `tenzro-network` (for `ProviderAnnouncementMessage`) and `tenzro-model`
4//! (for `ModelProvisioner`) need to describe the hardware envelope of a node.
5//! Lifting the type into `tenzro-types` keeps the dependency graph acyclic
6//! (`tenzro-network -> tenzro-types`, `tenzro-model -> tenzro-network -> tenzro-types`).
7
8use serde::{Deserialize, Serialize};
9
10/// Hardware capabilities advertised by a provider node.
11///
12/// Used both as a local provisioning input (in `tenzro-model::ModelProvisioner`)
13/// and as part of the gossiped `ProviderAnnouncementMessage` payload so that
14/// consumers can route inference / TEE work by available memory, GPU presence,
15/// and CPU shape without an extra round-trip RPC.
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
17pub struct HardwareCapabilities {
18    /// Total RAM in GiB.
19    pub ram_gb: u32,
20    /// Total VRAM in GiB across all GPUs visible to the node (0 if no GPU).
21    pub vram_gb: u32,
22    /// Free disk space in GiB at the configured data directory.
23    pub disk_gb: u32,
24    /// Whether a usable TEE (TDX / SEV-SNP / Nitro / NVIDIA CC) was detected.
25    pub tee_available: bool,
26    /// CPU architecture string (`x86_64`, `aarch64`, …) per `std::env::consts::ARCH`.
27    pub cpu_arch: String,
28    /// Number of CPU cores reported by `std::thread::available_parallelism()`.
29    pub cpu_cores: u32,
30}
31
32impl Default for HardwareCapabilities {
33    fn default() -> Self {
34        Self {
35            ram_gb: 8,
36            vram_gb: 0,
37            disk_gb: 50,
38            tee_available: false,
39            cpu_arch: std::env::consts::ARCH.to_string(),
40            cpu_cores: std::thread::available_parallelism()
41                .map(|n| n.get() as u32)
42                .unwrap_or(4),
43        }
44    }
45}
46
47impl HardwareCapabilities {
48    /// Detect hardware capabilities from the current system.
49    ///
50    /// Linux: parses `/proc/meminfo` for `MemTotal:`.
51    /// macOS: shells out to `sysctl -n hw.memsize`.
52    /// All other fields fall back to the `Default` values; richer detection
53    /// (GPU / disk free / TEE probe) is layered in by `tenzro-tee` and the
54    /// node's startup sequence before the value is broadcast.
55    pub fn detect() -> Self {
56        let mut caps = Self::default();
57
58        #[cfg(target_os = "linux")]
59        {
60            if let Ok(meminfo) = std::fs::read_to_string("/proc/meminfo") {
61                for line in meminfo.lines() {
62                    if line.starts_with("MemTotal:")
63                        && let Some(kb_str) = line.split_whitespace().nth(1)
64                        && let Ok(kb) = kb_str.parse::<u64>()
65                    {
66                        caps.ram_gb = (kb / 1_048_576) as u32;
67                    }
68                }
69            }
70        }
71
72        #[cfg(target_os = "macos")]
73        {
74            use std::process::Command;
75            if let Ok(output) = Command::new("sysctl").arg("-n").arg("hw.memsize").output()
76                && let Ok(bytes_str) = String::from_utf8(output.stdout)
77                && let Ok(bytes) = bytes_str.trim().parse::<u64>()
78            {
79                caps.ram_gb = (bytes / 1_073_741_824) as u32;
80            }
81        }
82
83        caps
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn default_has_arch_and_cores() {
93        let caps = HardwareCapabilities::default();
94        assert!(!caps.cpu_arch.is_empty());
95        assert!(caps.cpu_cores >= 1);
96    }
97
98    #[test]
99    fn detect_returns_nonzero_ram() {
100        let caps = HardwareCapabilities::detect();
101        // Detection may fall back to default on unusual platforms; the
102        // default still reports a positive value (8 GiB), so this asserts
103        // the contract in either branch.
104        assert!(caps.ram_gb > 0);
105    }
106}