Skip to main content

pot_o_extensions/
device_protocol.rs

1//! Device protocol abstraction: native, ESP32-S, ESP8266, and WASM device types and status.
2
3use pot_o_core::TribeResult;
4use serde::{Deserialize, Serialize};
5use std::sync::atomic::{AtomicU64, Ordering};
6use std::time::Instant;
7
8// ---------------------------------------------------------------------------
9// Trait
10// ---------------------------------------------------------------------------
11
12#[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
29/// How a mining device communicates with the validator.
30pub 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
38// ---------------------------------------------------------------------------
39// NativeDevice
40// ---------------------------------------------------------------------------
41
42pub 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 // 512 MB
69    }
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
91// ---------------------------------------------------------------------------
92// ESP32SDevice
93// ---------------------------------------------------------------------------
94
95/// ESP32-S device protocol handler.
96/// Tracks registered devices by ID and provides heartbeat status.
97/// The actual mining runs on the ESP firmware; this represents
98/// the validator-side view of a connected ESP32-S.
99pub 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 // 320 KB
145    }
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
167// ---------------------------------------------------------------------------
168// ESP8266Device
169// ---------------------------------------------------------------------------
170
171/// ESP8266 device protocol handler.
172/// Reduced tensor dimensions (32x32) and limited operation set.
173pub 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 // 80 KB
219    }
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
234// ---------------------------------------------------------------------------
235// WasmDevice (stubbed – pending wasm-bindgen integration)
236// ---------------------------------------------------------------------------
237
238pub 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 // 64 MB WASM linear memory
249    }
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}