Переглянути джерело

feat: add --power-on-if-off and --power-off flags

- Add --power-on-if-off flag to serve command to auto-power-on cluster if switch is off
- Add --power-off flag to stop command to power off cluster after stopping containers
- Power off command now stops containers first with 30s timeout before SSH shutdown
- Extract power_on_cluster_if_off logic to src/cluster_power.rs to keep files under 500 lines
Alexander Korolev 4 місяців тому
батько
коміт
811b8b0848
4 змінених файлів з 184 додано та 12 видалено
  1. 8 0
      src/cli.rs
  2. 115 0
      src/cluster_power.rs
  3. 42 9
      src/docker/mod.rs
  4. 19 3
      src/main.rs

+ 8 - 0
src/cli.rs

@@ -38,6 +38,10 @@ pub enum Commands {
         /// Follow logs after serving (use --no-follow to skip).
         #[arg(long, action = clap::ArgAction::SetFalse, default_value_t = true)]
         follow: bool,
+
+        /// Power on cluster if switch is off (requires HA config).
+        #[arg(long, default_value = "false")]
+        power_on_if_off: bool,
     },
 
     /// Show container logs
@@ -67,6 +71,10 @@ pub enum Commands {
 
         #[arg(long, default_value = "auto")]
         cluster: String,
+
+        /// Power off cluster after stopping containers (requires HA config).
+        #[arg(long, default_value = "false")]
+        power_off: bool,
     },
 
     /// Check cluster health

+ 115 - 0
src/cluster_power.rs

@@ -0,0 +1,115 @@
+//! Cluster power operations - power on cluster if switch is off.
+
+use anyhow::{Context, Result};
+use tracing::info;
+
+use crate::config::{Cluster, Config, NcclConfig, Recipe, resolve_cluster};
+use crate::home_assistant::Client as HaClient;
+use crate::ssh::run_ssh;
+
+/// Power on cluster if switch is off.
+pub async fn power_on_cluster_if_off(
+    config: &Config,
+    recipe: &Recipe,
+    cluster_name: &str,
+    node_name: Option<&str>,
+    dry_run: bool,
+) -> Result<()> {
+    let (cluster_to_check, _nccl_opt) = if let Some(nn) = node_name {
+        let (solo_c, _solo_n) = config
+            .clusters
+            .iter()
+            .find_map(|c| {
+                c.nodes.iter().find(|n| n.name == nn).map(|node| {
+                    let solo = Cluster {
+                        name: format!("solo/{}", node.name),
+                        master_ib_ip: node.ib_ip.clone(),
+                        nodes: vec![node.clone()],
+                        extra_entities: vec![],
+                    };
+                    (solo, None::<NcclConfig>)
+                })
+            })
+            .context(format!("Node '{}' not found in any cluster", nn))?;
+        (solo_c, None)
+    } else {
+        let (c, n) = resolve_cluster(
+            config,
+            recipe,
+            cluster_name,
+            &std::collections::HashSet::new(),
+        )?;
+        (c, Some(n.clone()))
+    };
+
+    let ha_switch_off = if let Some(ha_cfg) = &config.home_assistant {
+        let ha = HaClient::new(ha_cfg).ok();
+        let entities = cluster_to_check.extra_entities.clone();
+        if !entities.is_empty() {
+            if let Some(client) = ha {
+                let mut switch_off = false;
+                for e in &entities {
+                    match client.get_state(e).await {
+                        Some(ref s) if s == "on" => continue,
+                        _ => {
+                            switch_off = true;
+                            break;
+                        }
+                    }
+                }
+                switch_off
+            } else {
+                false
+            }
+        } else {
+            false
+        }
+    } else {
+        false
+    };
+
+    if ha_switch_off {
+        info!("{}", fluent_i18n::t!("cluster-switch-off-detected"));
+        if !dry_run {
+            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 e in &cluster_to_check.extra_entities {
+                    ha.turn_on(e).await?;
+                }
+                info!("{}", fluent_i18n::t!("cluster-switch-powered-on"));
+
+                // Wait for nodes to become available
+                info!("{}", fluent_i18n::t!("cluster-waiting-for-nodes"));
+                let max_wait_secs = 120;
+                let wait_interval_secs = 5;
+                let mut waited_secs = 0;
+                while waited_secs < max_wait_secs {
+                    // Check if all nodes are online
+                    let mut all_online = true;
+                    for node in &cluster_to_check.nodes {
+                        let cmd = "ip link show up 2>/dev/null | head -1 || echo ok";
+                        if run_ssh(config, node, cmd.to_string(), false).await.is_err() {
+                            all_online = false;
+                            break;
+                        }
+                    }
+                    if all_online {
+                        info!("{}", fluent_i18n::t!("cluster-nodes-available"));
+                        break;
+                    }
+                    tokio::time::sleep(std::time::Duration::from_secs(wait_interval_secs)).await;
+                    waited_secs += wait_interval_secs;
+                }
+                if waited_secs >= max_wait_secs {
+                    anyhow::bail!("{}", fluent_i18n::t!("cluster-nodes-not-available-timeout"));
+                }
+            } else {
+                anyhow::bail!("{}", fluent_i18n::t!("ha-config-required-for-power-on"));
+            }
+        } else {
+            info!("{}", fluent_i18n::t!("power-on-dryrun"));
+        }
+    }
+    Ok(())
+}

+ 42 - 9
src/docker/mod.rs

@@ -15,6 +15,15 @@ use std::collections::HashSet;
 use tokio::time::{Duration, sleep};
 use tracing::info;
 
+/// Options for the serve operation.
+#[derive(Clone, Debug)]
+pub struct ServeOpts {
+    pub keep: bool,
+    pub skip_before_serve: bool,
+    pub dry_run: bool,
+    pub power_on_if_off: bool,
+}
+
 pub(crate) fn shq(s: &str) -> String {
     let mut out = String::with_capacity(s.len() + 2);
     out.push('\'');
@@ -34,9 +43,7 @@ pub async fn serve(
     recipe_name: &str,
     cluster_name: &str,
     node_name: Option<&str>,
-    keep: bool,
-    skip_before_serve: bool,
-    dry_run: bool,
+    opts: ServeOpts,
 ) -> Result<Cluster> {
     let recipe = load_recipe(recipe_name).await?;
     let rn = recipe.name.clone();
@@ -45,6 +52,19 @@ pub async fn serve(
         "{}",
         fluent_i18n::t!("recipe-starting", { "name" => rn, "engine" => re, "backend" => recipe.distributed_backend.map_or("none".to_string(), |b| b.to_string()) })
     );
+
+    // Check if cluster switch is off and power on if requested
+    if opts.power_on_if_off {
+        crate::cluster_power::power_on_cluster_if_off(
+            config,
+            &recipe,
+            cluster_name,
+            node_name,
+            opts.dry_run,
+        )
+        .await?;
+    }
+
     let (cluster, nccl): (Cluster, NcclConfig) = if let Some(nn) = node_name {
         let (solo_c, solo_n) = config
             .clusters
@@ -96,11 +116,24 @@ pub async fn serve(
             );
         }
     }
-    crate::before_serve::run(config, &cluster, &recipe, skip_before_serve, dry_run).await?;
+    crate::before_serve::run(
+        config,
+        &cluster,
+        &recipe,
+        opts.skip_before_serve,
+        opts.dry_run,
+    )
+    .await?;
 
-    let mods_remote_path =
-        mods::upload_mods_to_cluster(config, &cluster, &recipe, recipe_name, node_name, dry_run)
-            .await?;
+    let mods_remote_path = mods::upload_mods_to_cluster(
+        config,
+        &cluster,
+        &recipe,
+        recipe_name,
+        node_name,
+        opts.dry_run,
+    )
+    .await?;
 
     let image = try_repaired_image(config, &cluster, &recipe).await?;
     let mut recipe = recipe;
@@ -116,10 +149,10 @@ pub async fn serve(
         &recipe,
         StartClusterOpts {
             cache_dir,
-            keep,
+            keep: opts.keep,
             node_name: node_name.map(String::from),
             mods_remote_path,
-            dry_run,
+            dry_run: opts.dry_run,
         },
     )
     .await?;

+ 19 - 3
src/main.rs

@@ -7,6 +7,7 @@ fluent_i18n::i18n!("locales", fallback = "en");
 
 mod before_serve;
 mod cli;
+mod cluster_power;
 mod config;
 mod docker;
 mod doctor;
@@ -84,15 +85,19 @@ async fn main() -> Result<()> {
             keep,
             skip_before_serve,
             follow,
+            power_on_if_off,
         } => {
             let cluster = docker::serve(
                 &config,
                 &recipe,
                 &cluster,
                 node.as_deref(),
-                keep,
-                skip_before_serve,
-                dry_run,
+                docker::ServeOpts {
+                    keep,
+                    skip_before_serve,
+                    dry_run,
+                    power_on_if_off,
+                },
             )
             .await?;
             if follow && !dry_run {
@@ -122,6 +127,7 @@ async fn main() -> Result<()> {
             recipe: recipe_name,
             cluster,
             node,
+            power_off,
         } => {
             if !node.is_empty() {
                 // User specified specific nodes — stop only those.
@@ -161,6 +167,16 @@ async fn main() -> Result<()> {
                     anyhow::bail!(msg);
                 }
             }
+            // Handle power_off flag
+            if power_off {
+                let cluster = resolve_cluster_for_op(&config, &recipe_name, &cluster).await?;
+                let node_str = if node.is_empty() {
+                    None
+                } else {
+                    Some(node[0].as_str())
+                };
+                crate::power::power_off(&config, &cluster, node_str, dry_run).await?;
+            }
         }
         Commands::Health {
             cluster: cluster_name,