launcher/src/cluster_power.rs

//! 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_with_retry;

/// 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"));

                // Initial boot delay — nodes need time for POST → bootloader → OS → SSH
                let boot_wait_secs = 60;
                info!("⏳ Waiting {}s for nodes to boot...", boot_wait_secs);
                tokio::time::sleep(std::time::Duration::from_secs(boot_wait_secs)).await;

                let max_poll_secs = 120;
                let poll_interval_secs = 10;
                let mut waited_secs = 0;
                while waited_secs < max_poll_secs {
                    let mut all_online = true;
                    for node in &cluster_to_check.nodes {
                        let cmd = "echo ok".to_string();
                        if run_ssh_with_retry(config, node, cmd, false, 2)
                            .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(poll_interval_secs)).await;
                    waited_secs += poll_interval_secs;
                }
                if waited_secs >= max_poll_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(())
}