1use pot_o_core::TribeResult;
4use serde::{Deserialize, Serialize};
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Instant;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
13pub enum DeviceType {
14 Native,
15 ESP32S,
16 ESP8266,
17 WASM,
18 Custom,
19}
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct DeviceStatus {
23 pub device_type: DeviceType,
24 pub online: bool,
25 pub uptime_secs: u64,
26 pub last_heartbeat: chrono::DateTime<chrono::Utc>,
27}
28
29pub trait DeviceProtocol: Send + Sync {
31 fn device_type(&self) -> DeviceType;
32 fn max_tensor_dims(&self) -> (usize, usize);
33 fn max_working_memory(&self) -> usize;
34 fn heartbeat(&self) -> TribeResult<DeviceStatus>;
35 fn supported_operations(&self) -> Vec<&'static str>;
36}
37
38pub struct NativeDevice {
43 started_at: Instant,
44}
45
46impl NativeDevice {
47 pub fn new() -> Self {
48 Self {
49 started_at: Instant::now(),
50 }
51 }
52}
53
54impl Default for NativeDevice {
55 fn default() -> Self {
56 Self::new()
57 }
58}
59
60impl DeviceProtocol for NativeDevice {
61 fn device_type(&self) -> DeviceType {
62 DeviceType::Native
63 }
64 fn max_tensor_dims(&self) -> (usize, usize) {
65 (1024, 1024)
66 }
67 fn max_working_memory(&self) -> usize {
68 1024 * 1024 * 512 }
70 fn heartbeat(&self) -> TribeResult<DeviceStatus> {
71 Ok(DeviceStatus {
72 device_type: DeviceType::Native,
73 online: true,
74 uptime_secs: self.started_at.elapsed().as_secs(),
75 last_heartbeat: chrono::Utc::now(),
76 })
77 }
78 fn supported_operations(&self) -> Vec<&'static str> {
79 vec![
80 "matrix_multiply",
81 "convolution",
82 "relu",
83 "sigmoid",
84 "tanh",
85 "dot_product",
86 "normalize",
87 ]
88 }
89}
90
91pub struct ESP32SDevice {
100 pub device_id: String,
101 started_at: Instant,
102 last_seen: AtomicU64,
103}
104
105impl ESP32SDevice {
106 pub fn new(device_id: String) -> Self {
107 let now = std::time::SystemTime::now()
108 .duration_since(std::time::UNIX_EPOCH)
109 .unwrap_or_default()
110 .as_secs();
111 Self {
112 device_id,
113 started_at: Instant::now(),
114 last_seen: AtomicU64::new(now),
115 }
116 }
117
118 pub fn record_heartbeat(&self) {
119 let now = std::time::SystemTime::now()
120 .duration_since(std::time::UNIX_EPOCH)
121 .unwrap_or_default()
122 .as_secs();
123 self.last_seen.store(now, Ordering::Relaxed);
124 }
125
126 pub fn is_stale(&self, timeout_secs: u64) -> bool {
127 let now = std::time::SystemTime::now()
128 .duration_since(std::time::UNIX_EPOCH)
129 .unwrap_or_default()
130 .as_secs();
131 let last = self.last_seen.load(Ordering::Relaxed);
132 now.saturating_sub(last) > timeout_secs
133 }
134}
135
136impl DeviceProtocol for ESP32SDevice {
137 fn device_type(&self) -> DeviceType {
138 DeviceType::ESP32S
139 }
140 fn max_tensor_dims(&self) -> (usize, usize) {
141 (64, 64)
142 }
143 fn max_working_memory(&self) -> usize {
144 320 * 1024 }
146 fn heartbeat(&self) -> TribeResult<DeviceStatus> {
147 self.record_heartbeat();
148 Ok(DeviceStatus {
149 device_type: DeviceType::ESP32S,
150 online: !self.is_stale(90),
151 uptime_secs: self.started_at.elapsed().as_secs(),
152 last_heartbeat: chrono::Utc::now(),
153 })
154 }
155 fn supported_operations(&self) -> Vec<&'static str> {
156 vec![
157 "matrix_multiply",
158 "convolution",
159 "relu",
160 "sigmoid",
161 "dot_product",
162 "normalize",
163 ]
164 }
165}
166
167pub struct ESP8266Device {
174 pub device_id: String,
175 started_at: Instant,
176 last_seen: AtomicU64,
177}
178
179impl ESP8266Device {
180 pub fn new(device_id: String) -> Self {
181 let now = std::time::SystemTime::now()
182 .duration_since(std::time::UNIX_EPOCH)
183 .unwrap_or_default()
184 .as_secs();
185 Self {
186 device_id,
187 started_at: Instant::now(),
188 last_seen: AtomicU64::new(now),
189 }
190 }
191
192 pub fn record_heartbeat(&self) {
193 let now = std::time::SystemTime::now()
194 .duration_since(std::time::UNIX_EPOCH)
195 .unwrap_or_default()
196 .as_secs();
197 self.last_seen.store(now, Ordering::Relaxed);
198 }
199
200 pub fn is_stale(&self, timeout_secs: u64) -> bool {
201 let now = std::time::SystemTime::now()
202 .duration_since(std::time::UNIX_EPOCH)
203 .unwrap_or_default()
204 .as_secs();
205 let last = self.last_seen.load(Ordering::Relaxed);
206 now.saturating_sub(last) > timeout_secs
207 }
208}
209
210impl DeviceProtocol for ESP8266Device {
211 fn device_type(&self) -> DeviceType {
212 DeviceType::ESP8266
213 }
214 fn max_tensor_dims(&self) -> (usize, usize) {
215 (32, 32)
216 }
217 fn max_working_memory(&self) -> usize {
218 80 * 1024 }
220 fn heartbeat(&self) -> TribeResult<DeviceStatus> {
221 self.record_heartbeat();
222 Ok(DeviceStatus {
223 device_type: DeviceType::ESP8266,
224 online: !self.is_stale(90),
225 uptime_secs: self.started_at.elapsed().as_secs(),
226 last_heartbeat: chrono::Utc::now(),
227 })
228 }
229 fn supported_operations(&self) -> Vec<&'static str> {
230 vec!["relu", "sigmoid", "dot_product", "normalize"]
231 }
232}
233
234pub struct WasmDevice;
239
240impl DeviceProtocol for WasmDevice {
241 fn device_type(&self) -> DeviceType {
242 DeviceType::WASM
243 }
244 fn max_tensor_dims(&self) -> (usize, usize) {
245 (256, 256)
246 }
247 fn max_working_memory(&self) -> usize {
248 64 * 1024 * 1024 }
250 fn heartbeat(&self) -> TribeResult<DeviceStatus> {
251 Ok(DeviceStatus {
252 device_type: DeviceType::WASM,
253 online: false,
254 uptime_secs: 0,
255 last_heartbeat: chrono::Utc::now(),
256 })
257 }
258 fn supported_operations(&self) -> Vec<&'static str> {
259 vec![
260 "matrix_multiply",
261 "convolution",
262 "relu",
263 "sigmoid",
264 "tanh",
265 "dot_product",
266 "normalize",
267 ]
268 }
269}