somatize_coordinator/
registry.rs

1//! Coordinator — lightweight gateway that manages worker registration,
2//! routing, and health monitoring.
3//!
4//! Can run as:
5//! - **Standalone binary**: `soma-coordinator --token sk-xxx --port 9090`
6//! - **Embedded**: `Coordinator::new().start_local()` for development
7//!
8//! The coordinator does NOT execute plans. It:
9//! 1. Accepts worker registrations (with capabilities + heartbeats)
10//! 2. Authenticates connections via bearer token
11//! 3. Routes client plan submissions to appropriate workers
12//! 4. Forwards worker events back to the client
13
14use chrono::{DateTime, Utc};
15use serde::{Deserialize, Serialize};
16use somatize_worker::protocol::{Capabilities, LoadMetrics, WorkerId};
17use std::collections::HashMap;
18use std::sync::{Arc, RwLock};
19
20/// Status of a registered worker.
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct WorkerStatus {
23    /// The worker's unique id, as it registered itself.
24    pub id: WorkerId,
25    /// Where clients reach the worker (e.g. `ws://host:8080`). The
26    /// coordinator hands this out on `/submit` and steps aside — the
27    /// plan and its tensor payloads travel client→worker direct, never
28    /// through the coordinator.
29    pub address: String,
30    /// What the worker offers: CPUs, RAM, GPUs, Python envs, tags.
31    /// Placement matches required tags against these via
32    /// [`matches_tags`](Self::matches_tags).
33    pub capabilities: Capabilities,
34    /// Load reported with the latest heartbeat; `None` until the first
35    /// one arrives after registration.
36    pub load: Option<LoadMetrics>,
37    /// The plans currently leased to this worker. `/submit` takes a
38    /// lease ([`WorkerRegistry::claim`]), `/complete` releases it
39    /// ([`WorkerRegistry::release`]) — this list is what makes
40    /// [`has_capacity`](Self::has_capacity) and the least-loaded
41    /// tie-break mean anything.
42    pub active_plans: Vec<String>,
43    /// When the worker last beat (workers beat every 10s). What
44    /// [`is_alive`](Self::is_alive) and the reaper compare against.
45    pub last_heartbeat: DateTime<Utc>,
46    /// False after an explicit [`WorkerRegistry::disconnect`];
47    /// re-registering sets it back. A disconnected worker is never
48    /// alive, however fresh its heartbeat.
49    pub connected: bool,
50}
51
52impl WorkerStatus {
53    /// Whether the worker has capacity for more work.
54    pub fn has_capacity(&self, max_concurrent: usize) -> bool {
55        self.connected && self.active_plans.len() < max_concurrent
56    }
57
58    /// Whether the worker matches a set of required tags.
59    pub fn matches_tags(&self, required: &[String]) -> bool {
60        required
61            .iter()
62            .all(|tag| self.capabilities.tags.contains(tag))
63    }
64
65    /// Whether the worker is considered alive (heartbeat within timeout).
66    pub fn is_alive(&self, timeout_secs: i64) -> bool {
67        self.connected && (Utc::now() - self.last_heartbeat).num_seconds() < timeout_secs
68    }
69}
70
71/// The worker registry — tracks all known workers and their status.
72#[derive(Debug, Clone)]
73pub struct WorkerRegistry {
74    workers: Arc<RwLock<HashMap<WorkerId, WorkerStatus>>>,
75    heartbeat_timeout_secs: i64,
76}
77
78impl WorkerRegistry {
79    /// Read the registry, tolerating poisoning.
80    ///
81    /// Every access used `.unwrap()`, so one handler panicking anywhere
82    /// left the whole coordinator unable to answer about any worker for
83    /// the rest of the process. A `HashMap` of statuses has no invariant
84    /// that spans a lock acquisition, so the data behind a poisoned lock
85    /// is still sound. Same policy the worker already uses.
86    fn read(&self) -> std::sync::RwLockReadGuard<'_, HashMap<WorkerId, WorkerStatus>> {
87        self.workers.read().unwrap_or_else(|e| e.into_inner())
88    }
89
90    fn write(&self) -> std::sync::RwLockWriteGuard<'_, HashMap<WorkerId, WorkerStatus>> {
91        self.workers.write().unwrap_or_else(|e| e.into_inner())
92    }
93
94    /// An empty registry with a 30-second heartbeat timeout — three
95    /// missed beats at the workers' 10-second cadence before a worker
96    /// counts as dead.
97    pub fn new() -> Self {
98        Self {
99            workers: Arc::new(RwLock::new(HashMap::new())),
100            heartbeat_timeout_secs: 30,
101        }
102    }
103
104    /// Override the heartbeat timeout (builder-style). Tests use 0 to
105    /// make everything instantly stale and 3600 to make nothing stale.
106    pub fn with_heartbeat_timeout(mut self, secs: i64) -> Self {
107        self.heartbeat_timeout_secs = secs;
108        self
109    }
110
111    /// Register a new worker or update an existing one.
112    pub fn register(
113        &self,
114        id: impl Into<String>,
115        address: impl Into<String>,
116        capabilities: Capabilities,
117    ) {
118        let id = id.into();
119        let mut workers = self.write();
120        workers.insert(
121            id.clone(),
122            WorkerStatus {
123                id,
124                address: address.into(),
125                capabilities,
126                load: None,
127                active_plans: vec![],
128                last_heartbeat: Utc::now(),
129                connected: true,
130            },
131        );
132    }
133
134    /// Update a worker's heartbeat and load metrics.
135    pub fn heartbeat(&self, worker_id: &str, load: LoadMetrics) {
136        let mut workers = self.write();
137        if let Some(w) = workers.get_mut(worker_id) {
138            w.load = Some(load);
139            w.last_heartbeat = Utc::now();
140        }
141    }
142
143    /// Record that `plan_id` has been placed on `worker_id`.
144    ///
145    /// `active_plans` was initialised to `vec![]` and never touched again,
146    /// so `has_capacity` and the "least loaded" tie-break both compared
147    /// zeroes: placement picked an arbitrary worker and called it balanced.
148    /// Returns false if the worker is unknown.
149    pub fn claim(&self, worker_id: &str, plan_id: impl Into<String>) -> bool {
150        let mut workers = self.write();
151        match workers.get_mut(worker_id) {
152            Some(w) => {
153                let plan_id = plan_id.into();
154                if !w.active_plans.contains(&plan_id) {
155                    w.active_plans.push(plan_id);
156                }
157                true
158            }
159            None => false,
160        }
161    }
162
163    /// Release a plan, whether it finished or failed.
164    pub fn release(&self, worker_id: &str, plan_id: &str) -> bool {
165        let mut workers = self.write();
166        match workers.get_mut(worker_id) {
167            Some(w) => {
168                w.active_plans.retain(|p| p != plan_id);
169                true
170            }
171            None => false,
172        }
173    }
174
175    /// Mark a worker as disconnected.
176    pub fn disconnect(&self, worker_id: &str) {
177        let mut workers = self.write();
178        if let Some(w) = workers.get_mut(worker_id) {
179            w.connected = false;
180        }
181    }
182
183    /// Remove a worker entirely.
184    pub fn remove(&self, worker_id: &str) {
185        let mut workers = self.write();
186        workers.remove(worker_id);
187    }
188
189    /// Get all alive, connected workers.
190    pub fn active_workers(&self) -> Vec<WorkerStatus> {
191        let workers = self.read();
192        workers
193            .values()
194            .filter(|w| w.is_alive(self.heartbeat_timeout_secs))
195            .cloned()
196            .collect()
197    }
198
199    /// Get a specific worker by ID.
200    pub fn get(&self, worker_id: &str) -> Option<WorkerStatus> {
201        let workers = self.read();
202        workers.get(worker_id).cloned()
203    }
204
205    /// Find workers matching required tags with available capacity.
206    pub fn find_workers(&self, tags: &[String], max_concurrent: usize) -> Vec<WorkerStatus> {
207        self.active_workers()
208            .into_iter()
209            .filter(|w| w.matches_tags(tags) && w.has_capacity(max_concurrent))
210            .collect()
211    }
212
213    /// Total number of registered workers (including disconnected).
214    pub fn total_count(&self) -> usize {
215        self.read().len()
216    }
217
218    /// Number of alive, connected workers.
219    pub fn active_count(&self) -> usize {
220        self.active_workers().len()
221    }
222
223    /// Human-readable summary.
224    pub fn summary(&self) -> String {
225        let workers = self.active_workers();
226        let total_cpus: usize = workers.iter().map(|w| w.capabilities.cpu_cores).sum();
227        let total_gpus: usize = workers.iter().map(|w| w.capabilities.gpus.len()).sum();
228        let total_ram: u64 = workers.iter().map(|w| w.capabilities.ram_bytes).sum();
229        format!(
230            "{} workers ({} CPUs, {} GPUs, {:.1} GB RAM)",
231            workers.len(),
232            total_cpus,
233            total_gpus,
234            total_ram as f64 / (1024.0 * 1024.0 * 1024.0),
235        )
236    }
237
238    /// Drop workers that have stopped sending heartbeats.
239    ///
240    /// The predicate was `is_alive(timeout) || w.connected`, and
241    /// `is_alive` already requires `connected` — so it reduced to
242    /// `w.connected` and pruned nothing that was still marked connected,
243    /// however long ago it had last been heard from. Which is the only
244    /// case worth pruning. It also had no callers.
245    ///
246    /// Returns the ids that were dropped, so a caller can log them.
247    pub fn prune_stale(&self) -> Vec<WorkerId> {
248        let timeout = self.heartbeat_timeout_secs;
249        let mut workers = self.write();
250        let stale: Vec<WorkerId> = workers
251            .iter()
252            .filter(|(_, w)| !w.is_alive(timeout))
253            .map(|(id, _)| id.clone())
254            .collect();
255        for id in &stale {
256            workers.remove(id);
257        }
258        stale
259    }
260}
261
262impl Default for WorkerRegistry {
263    fn default() -> Self {
264        Self::new()
265    }
266}
267
268#[cfg(test)]
269mod tests {
270    use super::*;
271    use somatize_worker::protocol::GpuInfo;
272
273    fn test_caps(tags: Vec<String>) -> Capabilities {
274        Capabilities {
275            cpu_cores: 4,
276            ram_bytes: 8_000_000_000,
277            gpus: vec![],
278            python_envs: vec![],
279            tags,
280        }
281    }
282
283    fn gpu_caps() -> Capabilities {
284        Capabilities {
285            cpu_cores: 8,
286            ram_bytes: 32_000_000_000,
287            gpus: vec![GpuInfo {
288                name: "A100".into(),
289                memory_bytes: 80_000_000_000,
290            }],
291            python_envs: vec![],
292            tags: vec!["gpu".into(), "training".into()],
293        }
294    }
295
296    #[test]
297    fn register_and_query() {
298        let registry = WorkerRegistry::new();
299        registry.register("w1", "ws://host1:8080", test_caps(vec!["cpu".into()]));
300        registry.register("w2", "ws://host2:8080", gpu_caps());
301
302        assert_eq!(registry.total_count(), 2);
303        assert_eq!(registry.active_count(), 2);
304
305        let w1 = registry.get("w1").unwrap();
306        assert_eq!(w1.address, "ws://host1:8080");
307        assert!(w1.connected);
308    }
309
310    #[test]
311    fn find_by_tags() {
312        let registry = WorkerRegistry::new();
313        registry.register("cpu1", "ws://c1:8080", test_caps(vec!["cpu".into()]));
314        registry.register("gpu1", "ws://g1:8080", gpu_caps());
315
316        let gpu_workers = registry.find_workers(&["gpu".into()], 10);
317        assert_eq!(gpu_workers.len(), 1);
318        assert_eq!(gpu_workers[0].id, "gpu1");
319
320        let cpu_workers = registry.find_workers(&["cpu".into()], 10);
321        assert_eq!(cpu_workers.len(), 1);
322    }
323
324    #[test]
325    fn disconnect_and_reconnect() {
326        let registry = WorkerRegistry::new();
327        registry.register("w1", "ws://host1:8080", test_caps(vec![]));
328        assert_eq!(registry.active_count(), 1);
329
330        registry.disconnect("w1");
331        assert_eq!(registry.active_count(), 0);
332
333        // Re-register = reconnect
334        registry.register("w1", "ws://host1:8080", test_caps(vec![]));
335        assert_eq!(registry.active_count(), 1);
336    }
337
338    #[test]
339    fn summary_format() {
340        let registry = WorkerRegistry::new();
341        registry.register("w1", "ws://h1:8080", test_caps(vec![]));
342        registry.register("w2", "ws://h2:8080", gpu_caps());
343
344        let s = registry.summary();
345        assert!(s.contains("2 workers"));
346        assert!(s.contains("12 CPUs")); // 4 + 8
347        assert!(s.contains("1 GPUs"));
348    }
349
350    /// A worker that stops beating is dropped, and one that keeps beating
351    /// is not.
352    ///
353    /// `prune_stale` had no callers and a predicate that reduced to
354    /// `w.connected`, so it removed nothing that mattered: a worker whose
355    /// process had died stayed in the registry forever.
356    #[test]
357    fn a_silent_worker_is_reaped_and_a_beating_one_is_not() {
358        // Zero-second timeout: everything registered in the past is stale.
359        let registry = WorkerRegistry::new().with_heartbeat_timeout(0);
360        registry.register("gone", "ws://h1:8080", test_caps(vec![]));
361        assert_eq!(registry.total_count(), 1);
362
363        let reaped = registry.prune_stale();
364        assert_eq!(reaped, vec!["gone".to_string()]);
365        assert_eq!(registry.total_count(), 0, "the dead worker is gone");
366
367        // A generous window: a freshly registered worker survives.
368        let registry = WorkerRegistry::new().with_heartbeat_timeout(3600);
369        registry.register("here", "ws://h1:8080", test_caps(vec![]));
370        assert!(registry.prune_stale().is_empty());
371        assert_eq!(registry.total_count(), 1);
372    }
373
374    /// Placement is only balanced if placements are recorded.
375    ///
376    /// `active_plans` was initialised empty and never touched, so
377    /// `has_capacity` and the least-loaded tie-break both compared zeroes.
378    #[test]
379    fn a_placed_plan_counts_against_the_worker_that_took_it() {
380        let registry = WorkerRegistry::new();
381        registry.register("w1", "ws://h1:8080", test_caps(vec![]));
382
383        assert!(registry.claim("w1", "plan-1"));
384        assert_eq!(registry.get("w1").unwrap().active_plans, vec!["plan-1"]);
385
386        // At a cap of one, the worker is now full.
387        assert!(registry.find_workers(&[], 1).is_empty());
388
389        // Claiming the same plan twice is not two plans.
390        registry.claim("w1", "plan-1");
391        assert_eq!(registry.get("w1").unwrap().active_plans.len(), 1);
392
393        assert!(registry.release("w1", "plan-1"));
394        assert_eq!(registry.find_workers(&[], 1).len(), 1, "capacity is back");
395
396        // An unknown worker is reported, not silently accepted.
397        assert!(!registry.claim("nobody", "plan-2"));
398        assert!(!registry.release("nobody", "plan-2"));
399    }
400
401    /// The least-loaded worker is the one with fewest placements.
402    #[test]
403    fn placement_prefers_the_less_loaded_worker() {
404        let registry = WorkerRegistry::new();
405        registry.register("busy", "ws://h1:8080", test_caps(vec![]));
406        registry.register("idle", "ws://h2:8080", test_caps(vec![]));
407        registry.claim("busy", "plan-1");
408        registry.claim("busy", "plan-2");
409
410        let best = registry
411            .find_workers(&[], 4)
412            .into_iter()
413            .min_by_key(|w| w.active_plans.len())
414            .expect("a candidate");
415        assert_eq!(best.id, "idle");
416    }
417
418    #[test]
419    fn capacity_check() {
420        let registry = WorkerRegistry::new();
421        registry.register("w1", "ws://h1:8080", test_caps(vec![]));
422
423        // With max_concurrent=0, no one has capacity
424        let workers = registry.find_workers(&[], 0);
425        assert!(workers.is_empty());
426
427        // With max_concurrent=1, worker with 0 active plans has capacity
428        let workers = registry.find_workers(&[], 1);
429        assert_eq!(workers.len(), 1);
430    }
431}
somatize_coordinator/registry.rs

somatize_coordinator/
registry.rs
