tatara_core/domain/
node.rs

1use chrono::{DateTime, Utc};
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use uuid::Uuid;
5
6use super::job::{DriverType, Resources};
7
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
9#[serde(rename_all = "snake_case")]
10pub enum NodeStatus {
11    Ready,
12    Down,
13    Draining,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize)]
17pub struct Node {
18    pub id: String,
19    pub address: String,
20    pub status: NodeStatus,
21    #[serde(default = "default_eligible")]
22    pub eligible: bool,
23    pub total_resources: Resources,
24    pub available_resources: Resources,
25    pub attributes: HashMap<String, String>,
26    pub drivers: Vec<DriverType>,
27    pub last_heartbeat: DateTime<Utc>,
28    pub allocations: Vec<Uuid>,
29}
30
31fn default_eligible() -> bool {
32    true
33}
34
35impl Node {
36    pub fn local() -> Self {
37        let os = std::env::consts::OS.to_string();
38        let arch = std::env::consts::ARCH.to_string();
39        let hostname = hostname::get()
40            .map(|h| h.to_string_lossy().to_string())
41            .unwrap_or_else(|_| "unknown".to_string());
42
43        let mut attributes = HashMap::new();
44        attributes.insert("os".to_string(), os.clone());
45        attributes.insert("arch".to_string(), arch);
46        attributes.insert("hostname".to_string(), hostname.clone());
47
48        let (cpu_mhz, memory_mb) = detect_resources();
49
50        Self {
51            id: hostname,
52            address: "127.0.0.1:4647".to_string(),
53            status: NodeStatus::Ready,
54            eligible: true,
55            total_resources: Resources { cpu_mhz, memory_mb },
56            available_resources: Resources { cpu_mhz, memory_mb },
57            attributes,
58            drivers: vec![DriverType::Exec],
59            last_heartbeat: Utc::now(),
60            allocations: Vec::new(),
61        }
62    }
63}
64
65fn detect_resources() -> (u64, u64) {
66    let cpu_mhz = (num_cpus() as u64) * 1000;
67
68    #[cfg(target_os = "macos")]
69    let memory_mb = {
70        use std::process::Command;
71        Command::new("sysctl")
72            .args(["-n", "hw.memsize"])
73            .output()
74            .ok()
75            .and_then(|o| {
76                String::from_utf8(o.stdout)
77                    .ok()
78                    .and_then(|s| s.trim().parse::<u64>().ok())
79            })
80            .unwrap_or(0)
81            / (1024 * 1024)
82    };
83
84    #[cfg(target_os = "linux")]
85    let memory_mb = {
86        std::fs::read_to_string("/proc/meminfo")
87            .ok()
88            .and_then(|s| {
89                s.lines()
90                    .find(|l| l.starts_with("MemTotal:"))
91                    .and_then(|l| {
92                        l.split_whitespace()
93                            .nth(1)
94                            .and_then(|v| v.parse::<u64>().ok())
95                    })
96            })
97            .unwrap_or(0)
98            / 1024
99    };
100
101    #[cfg(not(any(target_os = "macos", target_os = "linux")))]
102    let memory_mb = 0;
103
104    (cpu_mhz, memory_mb)
105}
106
107fn num_cpus() -> usize {
108    std::thread::available_parallelism()
109        .map(|p| p.get())
110        .unwrap_or(1)
111}
tatara_core/domain/node.rs

tatara_core/domain/
node.rs
