ソースを参照

feat: optional cluster-level extra_entities for power commands

Some setups have a network switch (or PDU) that backs multiple nodes
in a cluster.  Add `extra_entities` to `Cluster` so these can be
toggled in addition to per-node entities:

  [[clusters]]
  name = "4nodes-switch"
  extra_entities = ["switch.network_switch_4nodes"]

Toggled during `power on` and `power off` in parallel with per-node
entities, after the SSH-shutdown + delay phase (for off).

Doubles as a place for any other shared infrastructure switch (e.g.
rack PDU, master KVM).

README and tests updated.  No new dependencies.
Alexander Korolev 4 ヶ月 前
親
コミット
7b9f963036
5 ファイル変更、93 行追加、23 行削除
  1. 7 0
      README.md
  2. 1 0
      src/before_serve.rs
  3. 6 0
      src/config.rs
  4. 1 0
      src/docker/tests.rs
  5. 78 23
      src/power.rs

+ 7 - 0
README.md

@@ -289,6 +289,12 @@ controls it:
 [[clusters]]
 name = "4nodes-switch"
 master_ib_ip = "192.168.200.15"
+
+# Optional: extra switches to toggle alongside the nodes.  Useful for
+# a network switch that powers the cluster, or a PDU that backs the
+# rack.  Toggled in addition to any per-node entities.
+extra_entities = ["switch.network_switch_4nodes"]
+
 nodes = [
     { name = "dgx1", entity = "switch.dgx1", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
     # ...
@@ -297,6 +303,7 @@ nodes = [
 
 Without an `entity` the node is skipped (with a warning) — the SSH
 shutdown still happens, but Home Assistant is not consulted.
+Cluster-level `extra_entities` are always toggled (no opt-in).
 
 ### `power off` flow
 

+ 1 - 0
src/before_serve.rs

@@ -168,6 +168,7 @@ mod tests {
                     entity: None,
                 },
             ],
+            extra_entities: vec![],
         }
     }
 

+ 6 - 0
src/config.rs

@@ -30,6 +30,11 @@ pub struct Cluster {
     pub name: String,
     pub master_ib_ip: String,
     pub nodes: Vec<Node>,
+    /// Extra Home Assistant switches to toggle alongside the nodes
+    /// during `power on`/`power off`.  E.g. a network switch that powers
+    /// the cluster.  Toggled together with per-node entities.
+    #[serde(default)]
+    pub extra_entities: Vec<String>,
 }
 
 /// Top-level config (clusters + per-cluster pre-serve commands).
@@ -243,6 +248,7 @@ mod tests {
                     entity: None,
                 },
             ],
+            extra_entities: vec![],
         }
     }
 

+ 1 - 0
src/docker/tests.rs

@@ -23,6 +23,7 @@ fn tcls() -> Cluster {
                 entity: None,
             },
         ],
+        extra_entities: vec![],
     }
 }
 

+ 78 - 23
src/power.rs

@@ -2,9 +2,11 @@
 //!
 //! Power-off: SSH `sudo shutdown now` to each node → wait
 //! `[home_assistant].shutdown_delay_seconds` (default 30s) → POST
-//! `turn_off` to Home Assistant for each node with an `entity`.
+//! `turn_off` to Home Assistant for each per-node `entity` and each
+//! cluster-level `extra_entities`.
 //!
-//! Power-on: just POST `turn_on` to Home Assistant (the smart plug
+//! Power-on: just POST `turn_on` to Home Assistant for each per-node
+//! `entity` and each cluster-level `extra_entities` (the smart plug
 //! handles waking the machine).
 
 use anyhow::Result;
@@ -36,6 +38,9 @@ pub async fn power_off(
                 println!("would: HA turn_off {}", entity);
             }
         }
+        for extra in &cluster.extra_entities {
+            println!("would: HA turn_off {} (cluster-level)", extra);
+        }
         return Ok(());
     }
 
@@ -72,23 +77,18 @@ pub async fn power_off(
     info!("⏳ waiting {}s for nodes to settle", delay);
     tokio::time::sleep(Duration::from_secs(delay)).await;
 
-    // 3) HA turn_off for every node that has an entity.
+    // 3) HA turn_off: per-node entities, then cluster-level extras.
     if let Some(ha_cfg) = &config.home_assistant {
         let ha = HaClient::new(ha_cfg)?;
         crate::home_assistant::warn_if_token_looks_short(&ha_cfg.token);
-        for node in &targets {
-            if let Some(entity) = &node.entity {
-                if let Err(e) = split_entity(entity) {
-                    warn!("{}: invalid entity '{}': {}", node.name, entity, e);
-                    continue;
-                }
-                if let Err(e) = ha.turn_off(entity).await {
-                    warn!("HA turn_off {}: {}", entity, e);
-                }
-            } else {
-                warn!("⚠ {}: no `entity` configured, skipping HA", node.name);
-            }
-        }
+        toggle_entities(
+            &ha,
+            "turn_off",
+            &targets,
+            &cluster.extra_entities,
+            "cluster",
+        )
+        .await;
     } else {
         info!("no [home_assistant] config; OS shutdown was performed above");
     }
@@ -99,8 +99,9 @@ pub async fn power_off(
     Ok(())
 }
 
-/// Flow for `power on`: HA `turn_on` for each node with an `entity`.
-/// No SSH needed — the smart plug handles waking the machine.
+/// Flow for `power on`: HA `turn_on` for each node with an `entity`
+/// and each cluster-level `extra_entities`.  No SSH needed — the smart
+/// plug handles waking the machine.
 pub async fn power_on(
     config: &Config,
     cluster: &Cluster,
@@ -127,20 +128,59 @@ pub async fn power_on(
                 println!("would skip: {} (no entity)", node.name);
             }
         }
+        for extra in &cluster.extra_entities {
+            println!("would: HA turn_on {} (cluster-level)", extra);
+        }
         return Ok(());
     }
 
     info!("💡 power on: {} node(s)", targets.len());
-    for node in &targets {
+    toggle_entities(&ha, "turn_on", &targets, &cluster.extra_entities, "cluster").await;
+    Ok(())
+}
+
+/// Call HA `service` for each per-node entity, then each cluster-level
+/// `extra_entities`.  Validation is the same in both cases (entity IDs
+/// must be `domain.name`).
+async fn toggle_entities(
+    ha: &HaClient,
+    service: &'static str,
+    targets: &[Node],
+    extras: &[String],
+    extras_label: &'static str,
+) {
+    for node in targets {
         if let Some(entity) = &node.entity {
-            if let Err(e) = ha.turn_on(entity).await {
-                warn!("HA turn_on {}: {}", entity, e);
+            if let Err(e) = split_entity(entity) {
+                warn!("{}: invalid entity '{}': {}", node.name, entity, e);
+                continue;
+            }
+            if let Err(e) = call(ha, entity, service).await {
+                warn!("HA {} {}: {}", service, entity, e);
             }
         } else {
-            warn!("⚠ {}: no `entity` configured, skipping", node.name);
+            warn!("⚠ {}: no `entity` configured, skipping HA", node.name);
         }
     }
-    Ok(())
+    for extra in extras {
+        if let Err(e) = split_entity(extra) {
+            warn!("{} ({}): invalid entity '{}'", extras_label, extra, e);
+            continue;
+        }
+        if let Err(e) = call(ha, extra, service).await {
+            warn!("HA {} {} ({}): {}", service, extra, extras_label, e);
+        }
+    }
+}
+
+/// Dispatch to the right HA service.  Inlined because each call only
+/// differs by `service`, and inlining keeps the trait surface small.
+async fn call(ha: &HaClient, entity: &str, service: &'static str) -> Result<()> {
+    match service {
+        "turn_off" => ha.turn_off(entity).await,
+        "turn_on" => ha.turn_on(entity).await,
+        other => anyhow::bail!("unsupported HA service '{}'", other),
+    }
 }
 
 fn select_targets(cluster: &Cluster, node_name: Option<&str>, op: &str) -> Result<Vec<Node>> {
@@ -182,6 +222,7 @@ nodes = [
     { name = "n1", ssh_ip = "10.0.0.1", ib_ip = "10.0.0.1", rank = 0, entity = "switch.n1" },
     { name = "n2", ssh_ip = "10.0.0.2", ib_ip = "10.0.0.2", rank = 1 },
 ]
+extra_entities = ["switch.network_x"]
 [nccl."x"]
 ib_hca = "h"
 socket_ifname = "s"
@@ -245,4 +286,18 @@ ib_disable = 0
         let cfg = cfg_no_ha();
         assert!(cfg.home_assistant.is_none());
     }
+
+    #[test]
+    fn cluster_extra_entities_parsed() {
+        let cfg = cfg_with_ha();
+        let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
+        assert_eq!(c.extra_entities, vec!["switch.network_x"]);
+    }
+
+    #[test]
+    fn cluster_without_extra_entities_is_empty() {
+        let cfg = cfg_no_ha();
+        let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
+        assert!(c.extra_entities.is_empty());
+    }
 }