events.rs 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382
  1. //! SSE event bus + state snapshot for the dashboard.
  2. //!
  3. //! One multiplexed `broadcast` channel; `/events` forwards every event and
  4. //! the UI routes by `type`. Emits `state` (supervisor + job finishes),
  5. //! `job` (registry transitions), and `log` (attached window lines).
  6. use crate::config::Config;
  7. use crate::web::probe::ContainerSnap;
  8. use serde::Serialize;
  9. use std::time::{SystemTime, UNIX_EPOCH};
  10. /// One SSE event (serialized as `{"type":"state"|"job", ...}` on the wire).
  11. #[derive(Serialize, Clone, Debug)]
  12. #[serde(tag = "type", rename_all = "lowercase")]
  13. pub enum Event {
  14. State {
  15. ts: u64,
  16. snapshot: Snapshot,
  17. },
  18. Job {
  19. id: String,
  20. /// "serve" | "stop" | "power"
  21. kind: &'static str,
  22. event: String,
  23. cluster: String,
  24. recipe: Option<String>,
  25. message: String,
  26. },
  27. Log {
  28. /// Attach id (matches `POST /api/logs` response).
  29. attach: String,
  30. cluster: String,
  31. node: String,
  32. recipe: String,
  33. rank: u32,
  34. line: String,
  35. },
  36. }
  37. /// Per-node view for the LED grid.
  38. #[derive(Serialize, Clone, Debug)]
  39. pub struct NodeSnap {
  40. pub name: String,
  41. pub ssh_ip: String,
  42. pub reachable: bool,
  43. pub occupied: bool,
  44. pub containers: Vec<ContainerSnap>,
  45. /// Node rank — log windows attach the rank's `docker logs` stream.
  46. pub rank: u32,
  47. }
  48. /// Per-cluster view.
  49. #[derive(Serialize, Clone, Debug)]
  50. pub struct ClusterSnap {
  51. pub name: String,
  52. pub nodes: Vec<NodeSnap>,
  53. /// All nodes reachable AND blocking HA switches on (cluster verdict).
  54. pub up: bool,
  55. /// A blocking `extra_entity` reports anything but "on".
  56. pub switch_off: bool,
  57. /// Non-blocking entities that are off (display only, e.g. a fan).
  58. pub soft_off: Vec<String>,
  59. /// Free nodes; 0 when the cluster is down.
  60. pub free: u32,
  61. }
  62. /// One log window attach (for state snapshots — the UI restores open
  63. /// log windows from `/api/state` after a reload).
  64. #[derive(Serialize, Clone, Debug)]
  65. pub struct AttachSnap {
  66. pub id: String,
  67. pub cluster: String,
  68. pub node: String,
  69. pub recipe: String,
  70. pub rank: u32,
  71. /// Stream ended on its own (container stopped, SSH drop).
  72. pub done: bool,
  73. }
  74. /// The full dashboard snapshot (also the `GET /api/state` body).
  75. #[derive(Serialize, Clone, Debug)]
  76. pub struct Snapshot {
  77. pub ts: u64,
  78. pub clusters: Vec<ClusterSnap>,
  79. pub jobs: Vec<crate::web::jobs::Job>,
  80. pub attaches: Vec<AttachSnap>,
  81. }
  82. fn now_ts() -> u64 {
  83. SystemTime::now()
  84. .duration_since(UNIX_EPOCH)
  85. .map(|d| d.as_secs())
  86. .unwrap_or(0)
  87. }
  88. /// Pure snapshot builder — unit-testable without SSH.
  89. ///
  90. /// Cluster-level semantics via [`crate::cluster_state::verdict`]: one
  91. /// offline node or a switch off takes the whole cluster down, and a down
  92. /// cluster reports no containers (even on nodes that are alive).
  93. pub fn build_snapshot(
  94. config: &Config,
  95. fleet: &crate::web::probe::Fleet,
  96. jobs: &[crate::web::jobs::Job],
  97. attaches: &[AttachSnap],
  98. ) -> Snapshot {
  99. let nodes = &fleet.nodes;
  100. let switch_off = &fleet.switch_off;
  101. let soft_off = &fleet.soft_off;
  102. let mut clusters = Vec::with_capacity(config.clusters.len());
  103. for c in &config.clusters {
  104. let mut node_snaps = Vec::with_capacity(c.nodes.len());
  105. let mut reachables: Vec<(bool, bool)> = Vec::with_capacity(c.nodes.len());
  106. for n in &c.nodes {
  107. let s = nodes.get(&n.name);
  108. let reachable = s.is_some_and(|x| x.reachable);
  109. let occupied = s.is_some_and(|x| x.occupied);
  110. reachables.push((reachable, occupied));
  111. }
  112. let so = switch_off.get(&c.name).copied().unwrap_or(false);
  113. let v = crate::cluster_state::verdict(&reachables, so);
  114. for n in &c.nodes {
  115. let s = nodes.get(&n.name);
  116. let reachable = s.is_some_and(|x| x.reachable);
  117. let occupied = s.is_some_and(|x| x.occupied);
  118. // A down cluster reports no containers — the cluster is off as
  119. // a whole, whatever individual nodes still run.
  120. let containers = if v.up {
  121. s.map(|x| x.containers.clone()).unwrap_or_default()
  122. } else {
  123. Vec::new()
  124. };
  125. node_snaps.push(NodeSnap {
  126. name: n.name.clone(),
  127. ssh_ip: n.ssh_ip.clone(),
  128. reachable,
  129. occupied,
  130. containers,
  131. rank: n.rank,
  132. });
  133. }
  134. clusters.push(ClusterSnap {
  135. name: c.name.clone(),
  136. nodes: node_snaps,
  137. up: v.up,
  138. switch_off: so,
  139. soft_off: soft_off.get(&c.name).cloned().unwrap_or_default(),
  140. free: v.free,
  141. });
  142. }
  143. Snapshot {
  144. ts: now_ts(),
  145. clusters,
  146. jobs: jobs.to_vec(),
  147. attaches: attaches.to_vec(),
  148. }
  149. }
  150. #[cfg(test)]
  151. mod tests {
  152. use super::*;
  153. use std::collections::HashMap;
  154. fn config() -> Config {
  155. Config {
  156. ssh_user: "u".into(),
  157. ssh_key_path: "/k".into(),
  158. clusters: vec![
  159. crate::config::Cluster {
  160. name: "a".into(),
  161. master_ib_ip: None,
  162. nodes: vec![
  163. crate::config::Node {
  164. name: "n0".into(),
  165. ssh_ip: "10.0.0.1".into(),
  166. ib_ip: None,
  167. rank: 0,
  168. host_key: None,
  169. entity: None,
  170. },
  171. crate::config::Node {
  172. name: "n1".into(),
  173. ssh_ip: "10.0.0.2".into(),
  174. ib_ip: None,
  175. rank: 1,
  176. host_key: None,
  177. entity: None,
  178. },
  179. ],
  180. extra_entities: vec![],
  181. },
  182. crate::config::Cluster {
  183. name: "b".into(),
  184. master_ib_ip: None,
  185. nodes: vec![crate::config::Node {
  186. name: "n2".into(),
  187. ssh_ip: "10.0.0.3".into(),
  188. ib_ip: None,
  189. rank: 0,
  190. host_key: None,
  191. entity: None,
  192. }],
  193. extra_entities: vec![],
  194. },
  195. ],
  196. nccl: std::collections::HashMap::new(),
  197. hf_cache_dir: None,
  198. init_port: 25000,
  199. before_serve: std::collections::HashMap::new(),
  200. exec_aliases: std::collections::HashMap::new(),
  201. home_assistant: None,
  202. global_mods_dir: None,
  203. drop_caches: true,
  204. mount_aliases: std::collections::HashMap::new(),
  205. recipe_sources: Vec::new(),
  206. cluster_sources: Vec::new(),
  207. }
  208. }
  209. #[test]
  210. fn snapshot_reflects_probe_state() {
  211. let config = config();
  212. let mut nodes = HashMap::new();
  213. nodes.insert(
  214. "n0".into(),
  215. crate::web::probe::WebNodeState {
  216. reachable: true,
  217. occupied: true,
  218. containers: vec![ContainerSnap {
  219. name: "sm121_r_0".into(),
  220. status: "Up 3 hours".into(),
  221. cluster: "a".into(),
  222. recipe: "r".into(),
  223. }],
  224. },
  225. );
  226. nodes.insert(
  227. "n1".into(),
  228. crate::web::probe::WebNodeState {
  229. reachable: true,
  230. occupied: false,
  231. containers: vec![],
  232. },
  233. );
  234. // n2 absent → unreachable.
  235. let fleet = crate::web::probe::Fleet {
  236. nodes,
  237. switch_off: std::collections::HashMap::new(),
  238. soft_off: std::collections::HashMap::new(),
  239. };
  240. let snap = build_snapshot(&config, &fleet, &[], &[]);
  241. assert_eq!(snap.clusters.len(), 2);
  242. let a = &snap.clusters[0];
  243. assert!(a.up);
  244. assert!(a.nodes[0].reachable && a.nodes[0].occupied);
  245. assert_eq!(a.nodes[0].containers[0].recipe, "r");
  246. assert!(a.nodes[1].reachable && !a.nodes[1].occupied);
  247. assert_eq!(a.free, 1);
  248. // cluster b has one offline node → down as a whole
  249. let b = &snap.clusters[1];
  250. assert!(!b.up);
  251. assert!(!b.nodes[0].reachable);
  252. assert_eq!(b.free, 0);
  253. assert!(snap.jobs.is_empty());
  254. }
  255. #[test]
  256. fn down_cluster_hides_containers_on_live_nodes() {
  257. // Both nodes of cluster a are live; n1 runs a container but n0 is
  258. // offline → cluster a is down as a whole and the live node's
  259. // container must not be reported.
  260. let config = config();
  261. let mut nodes = HashMap::new();
  262. nodes.insert("n0".into(), crate::web::probe::WebNodeState::default());
  263. nodes.insert(
  264. "n1".into(),
  265. crate::web::probe::WebNodeState {
  266. reachable: true,
  267. occupied: true,
  268. containers: vec![ContainerSnap {
  269. name: "c2".into(),
  270. status: "Up".into(),
  271. cluster: "a".into(),
  272. recipe: "r".into(),
  273. }],
  274. },
  275. );
  276. let fleet = crate::web::probe::Fleet {
  277. nodes,
  278. switch_off: std::collections::HashMap::new(),
  279. soft_off: std::collections::HashMap::new(),
  280. };
  281. let snap = build_snapshot(&config, &fleet, &[], &[]);
  282. let a = &snap.clusters[0];
  283. assert!(!a.up, "one offline node takes the cluster down");
  284. assert_eq!(a.free, 0);
  285. assert!(
  286. a.nodes[1].containers.is_empty(),
  287. "containers hidden on a down cluster"
  288. );
  289. assert!(a.nodes[1].reachable, "reachability is still shown");
  290. }
  291. #[test]
  292. fn switch_off_takes_cluster_down() {
  293. let config = config();
  294. let mut nodes = HashMap::new();
  295. for n in ["n0", "n1", "n2"] {
  296. nodes.insert(
  297. n.into(),
  298. crate::web::probe::WebNodeState {
  299. reachable: true,
  300. occupied: false,
  301. containers: vec![],
  302. },
  303. );
  304. }
  305. let mut so = HashMap::new();
  306. so.insert("a".into(), true);
  307. let fleet = crate::web::probe::Fleet {
  308. nodes,
  309. switch_off: so,
  310. soft_off: std::collections::HashMap::new(),
  311. };
  312. let snap = build_snapshot(&config, &fleet, &[], &[]);
  313. assert!(!snap.clusters[0].up);
  314. assert!(snap.clusters[0].switch_off);
  315. assert_eq!(snap.clusters[0].free, 0);
  316. assert!(snap.clusters[1].up);
  317. assert_eq!(snap.clusters[1].free, 1);
  318. }
  319. #[test]
  320. fn soft_off_does_not_take_cluster_down() {
  321. // A non-blocking entity (fan) is off, all nodes live → cluster
  322. // stays up; the off entity is surfaced for display only.
  323. let config = config();
  324. let mut nodes = HashMap::new();
  325. for n in ["n0", "n1", "n2"] {
  326. nodes.insert(
  327. n.into(),
  328. crate::web::probe::WebNodeState {
  329. reachable: true,
  330. occupied: false,
  331. containers: vec![],
  332. },
  333. );
  334. }
  335. let mut soft = HashMap::new();
  336. soft.insert("a".into(), vec!["switch.fan".to_string()]);
  337. let fleet = crate::web::probe::Fleet {
  338. nodes,
  339. switch_off: HashMap::new(),
  340. soft_off: soft,
  341. };
  342. let snap = build_snapshot(&config, &fleet, &[], &[]);
  343. assert!(
  344. snap.clusters[0].up,
  345. "non-blocking off must not down the cluster"
  346. );
  347. assert_eq!(snap.clusters[0].free, 2);
  348. assert_eq!(snap.clusters[0].soft_off, vec!["switch.fan".to_string()]);
  349. assert!(snap.clusters[1].soft_off.is_empty());
  350. }
  351. #[test]
  352. fn job_event_serializes_with_type_tag() {
  353. let ev = Event::Job {
  354. id: "j1".into(),
  355. kind: "serve",
  356. event: "started".into(),
  357. cluster: "c".into(),
  358. recipe: Some("r".into()),
  359. message: "ok".into(),
  360. };
  361. let json = serde_json::to_value(&ev).unwrap();
  362. assert_eq!(json["type"], "job");
  363. assert_eq!(json["kind"], "serve");
  364. assert_eq!(json["id"], "j1");
  365. assert_eq!(json["event"], "started");
  366. }
  367. }