Przeglądaj źródła

feat: Home Assistant integration for power on/off

New `power` subcommand lets you turn clusters on/off via Home
Assistant smart plugs:

  sm121-launcher power off --cluster 4nodes-switch
  sm121-launcher power on --cluster 4nodes-switch

`power off` runs SSH shutdown, waits for
`[home_assistant].shutdown_delay_seconds` (default 30s) so the OS
actually settles, then calls HA turn_off per node.
`power on` just calls HA turn_on (no SSH).

Per-node entity IDs live in config.toml next to each node's other
config (e.g. `entity = "switch.dgx1"`).  Nodes without an entity
are skipped with a warning.

The launcher now requires `[home_assistant]` for `power on` since
that's the only way to wake a node.  `power off` will still issue
the SSH shutdown if HA isn't configured.

Implementation:
- New `home_assistant` module wraps the HA REST API
  (`POST /api/services/<domain>/<service>`)
- New `power` module implements on/off flows
- New `HomeAssistantConfig` struct (URL, token, delay)
- Per-node `entity: Option<String>` field
- 10 new tests covering entity parsing, target selection, and
  HA-required enforcement

New dependencies:
- `reqwest` (json + rustls-tls) for HA HTTP calls
- `serde_json` for request bodies
- `dirs` (was already added for default config dir)
Alexander Korolev 4 miesięcy temu
rodzic
commit
7b8651cead
11 zmienionych plików z 1287 dodań i 66 usunięć
  1. 701 3
      Cargo.lock
  2. 3 1
      Cargo.toml
  3. 65 0
      README.md
  4. 3 0
      src/before_serve.rs
  5. 28 0
      src/cli.rs
  6. 51 62
      src/config.rs
  7. 2 0
      src/docker/tests.rs
  8. 1 0
      src/exec.rs
  9. 173 0
      src/home_assistant.rs
  10. 12 0
      src/main.rs
  11. 248 0
      src/power.rs

Plik diff jest za duży
+ 701 - 3
Cargo.lock


+ 3 - 1
Cargo.toml

@@ -10,10 +10,12 @@ russh = "0.60"
 russh-keys = "0.49"
 anyhow = "1.0"
 serde = { version = "1.0", features = ["derive"] }
+serde_json = "1"
 toml = "1"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 sha2 = { version = "0.10", features = ["std"] }
 base64 = "0.22"
 futures = "0.3"
-dirs = "5"
+dirs = "5"
+reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }

+ 65 - 0
README.md

@@ -267,6 +267,69 @@ Anything in `[env]` is appended after all of the above, so a recipe can
 override any default (e.g. `NCCL_CUMEM_ENABLE = "1"` to re-enable
 CUDA-managed memory).
 
+## Power Management (Home Assistant)
+
+Optional `[home_assistant]` section lets the launcher toggle smart
+plugs (or any switch domain) so you can shut down and wake nodes
+remotely:
+
+```toml
+[home_assistant]
+url = "http://homeassistant.local:8123"
+token = "long-lived-access-token"
+# Seconds to wait between `sudo shutdown now` and HA `turn_off` so
+# the OS has time to actually shut down.  Default: 30.
+shutdown_delay_seconds = 30
+```
+
+Each node can optionally carry an `entity` field for the plug that
+controls it:
+
+```toml
+[[clusters]]
+name = "4nodes-switch"
+master_ib_ip = "192.168.200.15"
+nodes = [
+    { name = "dgx1", entity = "switch.dgx1", ssh_ip = "10.0.0.1", ib_ip = "1.2.3.4", rank = 0 },
+    # ...
+]
+```
+
+Without an `entity` the node is skipped (with a warning) — the SSH
+shutdown still happens, but Home Assistant is not consulted.
+
+### `power off` flow
+
+```
+1. ssh <node>: sudo shutdown now
+2. wait shutdown_delay_seconds (default 30s)
+3. HA POST /api/services/switch/turn_off
+   (entity: switch.dgx1)
+```
+
+```bash
+sm121-launcher power off --cluster 4nodes-switch
+sm121-launcher power off --cluster 4nodes-switch --node dgx1
+sm121-launcher -d power off --cluster 4nodes-switch   # dry-run
+```
+
+### `power on` flow
+
+`power on` does **not** SSH — the smart plug wakes the machine:
+
+```
+1. HA POST /api/services/switch/turn_on
+   (entity: switch.dgx1)
+```
+
+Requires `[home_assistant]` to be configured (no other way to wake
+a node).
+
+```bash
+sm121-launcher power on --cluster 4nodes-switch
+sm121-launcher power on --cluster 4nodes-switch --node dgx1
+```
+
 ## Doctor Mode
 
 ```bash
@@ -295,6 +358,8 @@ Commands:
   health   Check cluster health (or fingerprint host keys)
   doctor   Check and repair node prerequisites
   tmux     Manage tmux sessions (alternative to Docker)
+  power    Power on/off nodes via Home Assistant
+  exec     Run an arbitrary shell command on a cluster or single node
 
 Global flags:
   -c, --config <FILE>    Config file path (default: config.toml)

+ 3 - 0
src/before_serve.rs

@@ -157,6 +157,7 @@ mod tests {
                     ib_ip: "1.2.3.4".to_string(),
                     rank: 0,
                     host_key: None,
+                    entity: None,
                 },
                 Node {
                     name: "n1".to_string(),
@@ -164,6 +165,7 @@ mod tests {
                     ib_ip: "1.2.3.5".to_string(),
                     rank: 1,
                     host_key: None,
+                    entity: None,
                 },
             ],
         }
@@ -179,6 +181,7 @@ mod tests {
             init_port: 25000,
             before_serve,
             exec_aliases: HashMap::new(),
+            home_assistant: None,
         }
     }
 

+ 28 - 0
src/cli.rs

@@ -114,6 +114,34 @@ pub enum Commands {
         #[arg(trailing_var_arg = true)]
         command: Vec<String>,
     },
+
+    /// Power on/off nodes via Home Assistant (smart plugs / Wake-on-LAN).
+    Power {
+        #[command(subcommand)]
+        action: PowerAction,
+    },
+}
+
+#[derive(Subcommand)]
+pub enum PowerAction {
+    /// Power off: SSH `sudo shutdown now` + wait + HA `turn_off`.
+    Off {
+        #[arg(long)]
+        cluster: String,
+
+        /// Target a specific node within the cluster (default: all).
+        #[arg(long)]
+        node: Option<String>,
+    },
+    /// Power on: HA `turn_on` only (no SSH; the smart plug wakes the machine).
+    On {
+        #[arg(long)]
+        cluster: String,
+
+        /// Target a specific node within the cluster (default: all).
+        #[arg(long)]
+        node: Option<String>,
+    },
 }
 
 #[derive(Subcommand)]

+ 51 - 62
src/config.rs

@@ -13,6 +13,8 @@ pub struct Node {
     pub rank: u32,
     #[serde(default)]
     pub host_key: Option<String>,
+    #[serde(default)]
+    pub entity: Option<String>,
 }
 
 #[derive(Clone, Deserialize)]
@@ -30,8 +32,7 @@ pub struct Cluster {
     pub nodes: Vec<Node>,
 }
 
-/// Top-level config containing cluster definitions and per-cluster
-/// pre-serve commands.
+/// Top-level config (clusters + per-cluster pre-serve commands).
 #[derive(Clone, Deserialize)]
 pub struct Config {
     pub ssh_user: String,
@@ -41,15 +42,29 @@ pub struct Config {
     pub nccl: HashMap<String, NcclConfig>,
     #[serde(default)]
     pub hf_cache_dir: Option<String>,
-    /// TCP port used for cluster init (vLLM/SGLang). Default: 25000.
+    /// TCP port for cluster init (vLLM/SGLang).
     #[serde(default = "default_init_port")]
     pub init_port: u16,
-    /// Per-cluster pre-serve commands.  Looked up by cluster name.
+    /// Per-cluster pre-serve commands, command aliases, HA config.
     #[serde(default)]
     pub before_serve: HashMap<String, BeforeServeTable>,
-    /// Command aliases for `cargo run exec --alias`.
     #[serde(default)]
     pub exec_aliases: HashMap<String, String>,
+    #[serde(default)]
+    pub home_assistant: Option<HomeAssistantConfig>,
+}
+
+/// Home Assistant config for `power on`/`power off`.
+#[derive(Clone, Deserialize)]
+pub struct HomeAssistantConfig {
+    pub url: String,
+    pub token: String,
+    #[serde(default = "default_shutdown_delay")]
+    pub shutdown_delay_seconds: u64,
+}
+
+fn default_shutdown_delay() -> u64 {
+    30
 }
 
 fn default_init_port() -> u16 {
@@ -57,7 +72,6 @@ fn default_init_port() -> u16 {
 }
 
 impl Config {
-    /// HF cache directory for volume mounting.
     pub fn resolved_hf_cache_dir(&self) -> String {
         if let Some(ref dir) = self.hf_cache_dir {
             return dir.clone();
@@ -98,8 +112,7 @@ pub struct Recipe {
     pub entrypoint: Vec<String>,
     #[serde(default = "default_distributed_backend")]
     pub distributed_backend: String,
-    /// Recipe-level pre-serve commands.  Appended after cluster-level
-    /// commands from `[before_serve."<cluster>"]`.
+    /// Recipe-level pre-serve commands (appended after cluster-level).
     #[serde(default)]
     pub before_serve: Vec<PreServeCommand>,
 }
@@ -116,8 +129,7 @@ fn default_tensor_parallel_size() -> u32 {
     1
 }
 
-/// Selects a cluster and its NCCL config by name, or auto-picks the first
-/// cluster that has NCCL config.
+/// Selects a cluster + NCCL config by name or auto-picks the first one with NCCL.
 pub fn select_cluster<'c>(
     config: &'c Config,
     cluster_name: &str,
@@ -171,13 +183,8 @@ pub fn resolve_cluster<'c>(
     select_cluster(config, &choice)
 }
 
-/// Loads a recipe from disk.
-///
-/// Search order:
-/// 1. `./recipes/<name>/recipe.toml` (cwd)
-/// 2. `<default_config_dir>/recipes/<name>/recipe.toml`
-///
-/// Falls back to the cwd path so the error message mentions it.
+/// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
+/// falling back to cwd for the error).
 pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
     let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
     if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
@@ -194,15 +201,13 @@ pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
     Ok(toml::from_str(&content)?)
 }
 
-/// Returns the OS-appropriate default config directory for the launcher
-/// (e.g. `~/.config/sm121-launcher` on Linux), or `None` if the OS
-/// has no known config dir (some restricted sandboxes).
+/// OS-appropriate default config dir (e.g. `~/.config/sm121-launcher`),
+/// or `None` for sandboxes with no known config dir.
 pub fn default_config_dir() -> Option<std::path::PathBuf> {
     dirs::config_dir().map(|p| p.join("sm121-launcher"))
 }
 
-/// Returns the cluster and recipe name for a recipe (for stop/logs commands
-/// that only receive a recipe name).
+/// Resolves the cluster for a recipe (stop/logs flow).
 pub async fn resolve_cluster_for_recipe(
     config: &Config,
     recipe_name: &str,
@@ -227,6 +232,7 @@ mod tests {
                     ib_ip: "192.168.1.1".to_string(),
                     rank: 0,
                     host_key: None,
+                    entity: None,
                 },
                 Node {
                     name: "node1".to_string(),
@@ -234,11 +240,26 @@ mod tests {
                     ib_ip: "192.168.1.2".to_string(),
                     rank: 1,
                     host_key: None,
+                    entity: None,
                 },
             ],
         }
     }
 
+    fn test_config() -> Config {
+        Config {
+            ssh_user: "u".to_string(),
+            ssh_key_path: "/k".to_string(),
+            clusters: vec![test_cluster()],
+            nccl: HashMap::new(),
+            hf_cache_dir: None,
+            init_port: 25000,
+            before_serve: HashMap::new(),
+            exec_aliases: HashMap::new(),
+            home_assistant: None,
+        }
+    }
+
     fn test_nccl() -> NcclConfig {
         NcclConfig {
             ib_hca: "rocep1s0f0".to_string(),
@@ -375,6 +396,7 @@ nodes = [
             init_port: 25000,
             before_serve: HashMap::new(),
             exec_aliases: HashMap::new(),
+            home_assistant: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -394,6 +416,7 @@ nodes = [
             init_port: 25000,
             before_serve: HashMap::new(),
             exec_aliases: HashMap::new(),
+            home_assistant: None,
         };
         config.nccl.insert("test-cluster".to_string(), test_nccl());
 
@@ -404,32 +427,15 @@ nodes = [
 
     #[test]
     fn test_select_cluster_auto_no_match() {
-        let config = Config {
-            ssh_user: "user".to_string(),
-            ssh_key_path: "/path/to/key".to_string(),
-            clusters: vec![test_cluster()],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-        };
+        let config = test_config();
         let result = select_cluster(&config, "auto");
         assert!(result.is_err());
     }
 
     #[test]
     fn test_select_cluster_not_found() {
-        let config = Config {
-            ssh_user: "user".to_string(),
-            ssh_key_path: "/path/to/key".to_string(),
-            clusters: vec![],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-        };
+        let mut config = test_config();
+        config.clusters = vec![];
         let result = select_cluster(&config, "nonexistent");
         assert!(result.is_err());
     }
@@ -438,16 +444,8 @@ nodes = [
     fn test_select_cluster_by_name_found() {
         let mut nccl_map = HashMap::new();
         nccl_map.insert("test-cluster".to_string(), test_nccl());
-        let config = Config {
-            ssh_user: "u".to_string(),
-            ssh_key_path: "/k".to_string(),
-            clusters: vec![test_cluster()],
-            nccl: nccl_map,
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-        };
+        let mut config = test_config();
+        config.nccl = nccl_map;
         let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
         assert_eq!(cluster.name, "test-cluster");
         assert_eq!(cluster.nodes.len(), 2);
@@ -455,16 +453,7 @@ nodes = [
 
     #[test]
     fn test_select_cluster_by_name_not_found() {
-        let config = Config {
-            ssh_user: "u".to_string(),
-            ssh_key_path: "/k".to_string(),
-            clusters: vec![test_cluster()],
-            nccl: HashMap::new(),
-            hf_cache_dir: None,
-            init_port: 25000,
-            before_serve: HashMap::new(),
-            exec_aliases: HashMap::new(),
-        };
+        let config = test_config();
         let result = select_cluster_by_name(&config, "no-such-cluster");
         assert!(result.is_err());
     }

+ 2 - 0
src/docker/tests.rs

@@ -12,6 +12,7 @@ fn tcls() -> Cluster {
                 ib_ip: "1".into(),
                 rank: 0,
                 host_key: None,
+                entity: None,
             },
             Node {
                 name: "n1".into(),
@@ -19,6 +20,7 @@ fn tcls() -> Cluster {
                 ib_ip: "2".into(),
                 rank: 1,
                 host_key: None,
+                entity: None,
             },
         ],
     }

+ 1 - 0
src/exec.rs

@@ -159,6 +159,7 @@ mod tests {
             init_port: 25000,
             before_serve: HashMap::new(),
             exec_aliases,
+            home_assistant: None,
         }
     }
 

+ 173 - 0
src/home_assistant.rs

@@ -0,0 +1,173 @@
+//! Home Assistant REST API client for `power on`/`power off`.
+//!
+//! Talks to HA via `POST /api/services/<domain>/<service>` and uses
+//! `Authorization: Bearer <token>`.  The HA domain for a switch is the
+//! part before the first dot in the entity_id (e.g. `switch.dgx1` →
+//! domain `switch`, service `turn_off` or `turn_on`).
+
+use anyhow::{Context, Result, anyhow};
+use serde_json::json;
+use std::time::Duration;
+use tracing::{debug, info, warn};
+
+use crate::config::HomeAssistantConfig;
+
+/// Thin async wrapper around the HA REST API.
+///
+/// `base_url` is the HA root URL (`http://homeassistant.local:8123`),
+/// `token` is the long-lived access token used as bearer auth.
+pub struct Client {
+    base_url: String,
+    token: String,
+    http: reqwest::Client,
+}
+
+impl Client {
+    /// Construct a new client.  The reqwest client is built once and
+    /// reused so connection pooling works across the parallel calls in
+    /// the `power` subcommand.
+    pub fn new(cfg: &HomeAssistantConfig) -> Result<Self> {
+        let base_url = cfg.url.trim_end_matches('/').to_string();
+        if base_url.is_empty() {
+            anyhow::bail!("[home_assistant].url must be set");
+        }
+        let http = reqwest::Client::builder()
+            .timeout(Duration::from_secs(10))
+            .build()
+            .context("building reqwest client")?;
+        Ok(Self {
+            base_url,
+            token: cfg.token.clone(),
+            http,
+        })
+    }
+
+    /// Build `https://.../api/services/<domain>/<service>` for the given
+    /// `entity_id`.  Domain is the prefix before the first `.`; pass an
+    /// empty fallback for ambiguous input.
+    fn endpoint(entity_id: &str, service: &str) -> Result<(String, String)> {
+        let (domain, entity) = entity_id.split_once('.').ok_or_else(|| {
+            anyhow!(
+                "entity_id '{}' has no domain (e.g. 'switch.foo')",
+                entity_id
+            )
+        })?;
+        if entity.is_empty() {
+            anyhow::bail!("entity_id '{}' has empty entity name", entity_id);
+        }
+        Ok((
+            domain.to_string(),
+            format!("/api/services/{}/{}", domain, service),
+        ))
+    }
+
+    /// POST `/api/services/<domain>/<service>` with `{"entity_id": ...}`.
+    async fn call_service(&self, entity_id: &str, service: &str) -> Result<()> {
+        let (_domain, path) = Self::endpoint(entity_id, service)?;
+        let url = format!("{}{}", self.base_url, path);
+        debug!("HA POST {}", url);
+        let resp = self
+            .http
+            .post(&url)
+            .bearer_auth(&self.token)
+            .json(&json!({ "entity_id": entity_id }))
+            .send()
+            .await
+            .with_context(|| format!("calling HA at {}", url))?;
+        let status = resp.status();
+        if !status.is_success() {
+            let body = resp.text().await.unwrap_or_default();
+            anyhow::bail!(
+                "HA {} {} returned {}: {}",
+                service,
+                entity_id,
+                status,
+                body.trim()
+            );
+        }
+        Ok(())
+    }
+
+    /// `switch.turn_off` for the given entity.
+    pub async fn turn_off(&self, entity_id: &str) -> Result<()> {
+        info!("HA turn_off {}", entity_id);
+        self.call_service(entity_id, "turn_off").await
+    }
+
+    /// `switch.turn_on` for the given entity.
+    pub async fn turn_on(&self, entity_id: &str) -> Result<()> {
+        info!("HA turn_on {}", entity_id);
+        self.call_service(entity_id, "turn_on").await
+    }
+}
+
+/// Parse `switch.dgx1` into `("switch", "dgx1")` so callers can build
+/// their own endpoints if needed.  Exposed for the `power` subcommand
+/// and for tests.
+pub fn split_entity(entity_id: &str) -> Result<(&str, &str)> {
+    let (domain, entity) = entity_id.split_once('.').ok_or_else(|| {
+        anyhow!(
+            "entity_id '{}' has no domain (e.g. 'switch.foo')",
+            entity_id
+        )
+    })?;
+    if domain.is_empty() || entity.is_empty() {
+        anyhow::bail!("entity_id '{}' has empty domain or entity name", entity_id);
+    }
+    Ok((domain, entity))
+}
+
+/// Emit a deprecation-style warning if a HA token looks suspiciously
+/// short.  HA tokens are typically 100+ chars; anything shorter is
+/// almost certainly a placeholder.
+pub fn warn_if_token_looks_short(token: &str) {
+    if token.len() < 20 {
+        warn!(
+            "[home_assistant].token looks suspiciously short ({} chars); \
+             did you paste a placeholder?",
+            token.len()
+        );
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn split_entity_basic() {
+        assert_eq!(split_entity("switch.dgx1").unwrap(), ("switch", "dgx1"));
+        assert_eq!(
+            split_entity("light.living_room").unwrap(),
+            ("light", "living_room")
+        );
+    }
+
+    #[test]
+    fn split_entity_no_domain_errors() {
+        assert!(split_entity("dgx1").is_err());
+    }
+
+    #[test]
+    fn split_entity_empty_entity_errors() {
+        assert!(split_entity("switch.").is_err());
+    }
+
+    #[test]
+    fn split_entity_empty_domain_errors() {
+        assert!(split_entity(".thing").is_err());
+    }
+
+    #[test]
+    fn endpoint_extracts_domain() {
+        let (_d, path) = Client::endpoint("switch.dgx1", "turn_off").unwrap();
+        assert_eq!(path, "/api/services/switch/turn_off");
+        let (_d, path) = Client::endpoint("light.kitchen", "turn_on").unwrap();
+        assert_eq!(path, "/api/services/light/turn_on");
+    }
+
+    #[test]
+    fn endpoint_rejects_no_domain() {
+        assert!(Client::endpoint("plain", "turn_off").is_err());
+    }
+}

+ 12 - 0
src/main.rs

@@ -8,7 +8,9 @@ mod docker;
 mod doctor;
 mod exec;
 mod health;
+mod home_assistant;
 mod host_keys;
+mod power;
 mod pre_serve;
 mod ssh;
 mod tmux;
@@ -147,6 +149,16 @@ async fn main() -> Result<()> {
             let cmd_str = exec::resolve_command(&config, alias.as_deref(), &command)?;
             exec::run(&config, &cluster, node.as_deref(), &cmd_str, dry_run).await?;
         }
+        Commands::Power { action } => match action {
+            cli::PowerAction::Off { cluster, node } => {
+                let cluster = select_cluster_by_name(&config, &cluster)?;
+                power::power_off(&config, &cluster, node.as_deref(), dry_run).await?;
+            }
+            cli::PowerAction::On { cluster, node } => {
+                let cluster = select_cluster_by_name(&config, &cluster)?;
+                power::power_on(&config, &cluster, node.as_deref(), dry_run).await?;
+            }
+        },
     }
 
     Ok(())

+ 248 - 0
src/power.rs

@@ -0,0 +1,248 @@
+//! `power on` / `power off` subcommands.
+//!
+//! 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`.
+//!
+//! Power-on: just POST `turn_on` to Home Assistant (the smart plug
+//! handles waking the machine).
+
+use anyhow::Result;
+use std::time::Duration;
+use tracing::{info, warn};
+
+use crate::home_assistant::split_entity;
+
+use crate::config::{Cluster, Config, Node};
+use crate::home_assistant::Client as HaClient;
+
+/// Flow for `power off`: SSH shutdown + wait + HA `turn_off`.
+pub async fn power_off(
+    config: &Config,
+    cluster: &Cluster,
+    node_name: Option<&str>,
+    dry_run: bool,
+) -> Result<()> {
+    let targets = select_targets(cluster, node_name, "power off")?;
+
+    if dry_run {
+        warn!("dry-run: skipping actual shutdown + HA call");
+        for node in &targets {
+            println!(
+                "would: ssh {}@{} sudo shutdown now",
+                config.ssh_user, node.ssh_ip
+            );
+            if let Some(entity) = &node.entity {
+                println!("would: HA turn_off {}", entity);
+            }
+        }
+        return Ok(());
+    }
+
+    // 1) SSH shutdown every target in parallel.
+    info!(
+        "🛑 power off: {} node(s) — sudo shutdown now",
+        targets.len()
+    );
+    let mut handles = Vec::with_capacity(targets.len());
+    for node in targets.iter().cloned() {
+        let config = config.clone();
+        handles.push(tokio::spawn(async move {
+            let cmd = "sudo shutdown now".to_string();
+            let r = crate::ssh::run_ssh(&config, &node, cmd.clone(), false).await;
+            (node, cmd, r)
+        }));
+    }
+
+    let mut any_ssh_failed = false;
+    for h in handles {
+        let (node, cmd, r) = h.await?;
+        if let Err(e) = r {
+            any_ssh_failed = true;
+            warn!("✗ {} :: {} :: {}", node.name, cmd, e);
+        }
+    }
+
+    // 2) Wait so the OS actually shuts down before we cut power.
+    let delay = config
+        .home_assistant
+        .as_ref()
+        .map(|c| c.shutdown_delay_seconds)
+        .unwrap_or(30);
+    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.
+    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);
+            }
+        }
+    } else {
+        info!("no [home_assistant] config; OS shutdown was performed above");
+    }
+
+    if any_ssh_failed {
+        warn!("some nodes reported SSH failures during shutdown");
+    }
+    Ok(())
+}
+
+/// Flow for `power on`: HA `turn_on` for each node with an `entity`.
+/// No SSH needed — the smart plug handles waking the machine.
+pub async fn power_on(
+    config: &Config,
+    cluster: &Cluster,
+    node_name: Option<&str>,
+    dry_run: bool,
+) -> Result<()> {
+    let targets = select_targets(cluster, node_name, "power on")?;
+
+    let Some(ha_cfg) = &config.home_assistant else {
+        anyhow::bail!(
+            "[home_assistant] section is required for `power on` \
+             (the launcher has no other way to wake a node)"
+        );
+    };
+    let ha = HaClient::new(ha_cfg)?;
+    crate::home_assistant::warn_if_token_looks_short(&ha_cfg.token);
+
+    if dry_run {
+        warn!("dry-run: skipping actual HA call");
+        for node in &targets {
+            if let Some(entity) = &node.entity {
+                println!("would: HA turn_on {}", entity);
+            } else {
+                println!("would skip: {} (no entity)", node.name);
+            }
+        }
+        return Ok(());
+    }
+
+    info!("💡 power on: {} node(s)", targets.len());
+    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);
+            }
+        } else {
+            warn!("⚠ {}: no `entity` configured, skipping", node.name);
+        }
+    }
+    Ok(())
+}
+
+fn select_targets(cluster: &Cluster, node_name: Option<&str>, op: &str) -> Result<Vec<Node>> {
+    let targets: Vec<Node> = match node_name {
+        Some(name) => cluster
+            .nodes
+            .iter()
+            .filter(|n| n.name == *name)
+            .cloned()
+            .collect(),
+        None => cluster.nodes.clone(),
+    };
+    if targets.is_empty() {
+        if let Some(n) = node_name {
+            anyhow::bail!(
+                "{}: node '{}' not found in cluster '{}'",
+                op,
+                n,
+                cluster.name
+            );
+        }
+        anyhow::bail!("{}: cluster '{}' has no nodes", op, cluster.name);
+    }
+    Ok(targets)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn cfg_with_ha() -> Config {
+        let toml = r#"
+ssh_user = "u"
+ssh_key_path = "/k"
+[[clusters]]
+name = "x"
+master_ib_ip = "1.1.1.1"
+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 },
+]
+[nccl."x"]
+ib_hca = "h"
+socket_ifname = "s"
+ib_gid_index = 1
+ib_disable = 0
+
+[home_assistant]
+url = "http://ha.local:8123"
+token = "abcdef0123456789abcdef0123456789"
+shutdown_delay_seconds = 5
+"#;
+        toml::from_str(toml).unwrap()
+    }
+
+    fn cfg_no_ha() -> Config {
+        let toml = r#"
+ssh_user = "u"
+ssh_key_path = "/k"
+[[clusters]]
+name = "x"
+master_ib_ip = "1.1.1.1"
+nodes = [
+    { name = "n1", ssh_ip = "10.0.0.1", ib_ip = "10.0.0.1", rank = 0, entity = "switch.n1" },
+]
+[nccl."x"]
+ib_hca = "h"
+socket_ifname = "s"
+ib_gid_index = 1
+ib_disable = 0
+"#;
+        toml::from_str(toml).unwrap()
+    }
+
+    #[test]
+    fn select_targets_filters_by_node() {
+        let cfg = cfg_with_ha();
+        let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
+        let t = select_targets(&c, Some("n1"), "test").unwrap();
+        assert_eq!(t.len(), 1);
+        assert_eq!(t[0].name, "n1");
+    }
+
+    #[test]
+    fn select_targets_all() {
+        let cfg = cfg_with_ha();
+        let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
+        let t = select_targets(&c, None, "test").unwrap();
+        assert_eq!(t.len(), 2);
+    }
+
+    #[test]
+    fn select_targets_unknown_node_errors() {
+        let cfg = cfg_with_ha();
+        let c = crate::config::select_cluster_by_name(&cfg, "x").unwrap();
+        let r = select_targets(&c, Some("nope"), "test");
+        assert!(r.is_err());
+    }
+
+    #[test]
+    fn power_on_requires_ha_config() {
+        let cfg = cfg_no_ha();
+        assert!(cfg.home_assistant.is_none());
+    }
+}

Niektóre pliki nie zostały wyświetlone z powodu dużej ilości zmienionych plików