launcher/src/cluster_state.rs

//! Cluster state — tracks which nodes are reachable and occupied,
//! separate from the SSH probing that collects it.  This split lets
//! the cluster-selection logic be tested without real hardware.
//!
//! State collection (the "control" side) lives in [`ClusterState::probe`].
//! Selection logic (the "state" side) lives on [`ClusterState`] methods.

use anyhow::Result;
use std::collections::{HashMap, HashSet};
use tracing::{info, warn};

use crate::config::{Cluster, Config, NcclConfig};
use crate::ssh::run_ssh;

/// Per-node state gathered during a probe.
#[derive(Clone, Debug, Default)]
pub struct NodeState {
    /// SSH connection succeeded (node is at least booted and reachable).
    pub reachable: bool,
    /// Node has at least one running sm121 container.
    pub occupied: bool,
}

/// Aggregate view of the cluster fleet — built by [`probe`].
#[derive(Clone, Debug, Default)]
pub struct ClusterState {
    /// Node name → state.  Absence means the node was not probed.
    pub nodes: HashMap<String, NodeState>,
}

/// Cluster-level health verdict.
///
/// A cluster is **UP** only when *all* of its nodes are reachable AND its
/// HA switch entities (`extra_entities`, when HA is configured) report
/// "on". One offline node — or a switch off — takes the whole cluster
/// down: containers on surviving nodes are ignored, `free` is 0.
#[derive(Clone, Debug, PartialEq, serde::Serialize)]
pub struct ClusterVerdict {
    pub up: bool,
    /// HA extra entities report anything other than "on".
    pub switch_off: bool,
    /// Free nodes; 0 when the cluster is down.
    pub free: u32,
}

/// Derive the verdict from per-node `(reachable, occupied)` pairs plus the
/// HA switch state. The single source of truth for cluster health — both
/// the `health` command and the web dashboard route through this.
pub fn verdict(reachables: &[(bool, bool)], switch_off: bool) -> ClusterVerdict {
    let all_reachable = reachables.iter().all(|(r, _)| *r);
    let up = all_reachable && !switch_off;
    let free = if up {
        reachables.iter().filter(|(_, o)| !o).count() as u32
    } else {
        0
    };
    ClusterVerdict {
        up,
        switch_off,
        free,
    }
}

impl ClusterState {
    /// Probe all unique nodes across the config and collect their state.
    /// This is the only function that touches SSH (the "control" side).
    pub async fn probe(config: &Config) -> Self {
        let mut seen = HashSet::new();
        let mut handles = Vec::new();
        for cluster in &config.clusters {
            for node in &cluster.nodes {
                if seen.contains(&node.ssh_ip) {
                    continue;
                }
                seen.insert(node.ssh_ip.clone());
                let config = config.clone();
                let node = node.clone();
                let name = node.name.clone();
                let cmd = "docker ps --filter label=sm121.cluster --format '{{.Names}}' | head -1"
                    .to_string();
                handles.push(crate::locale::spawn_with_locale(async move {
                    (name, run_ssh(&config, &node, cmd, false).await)
                }));
            }
        }

        let mut nodes: HashMap<String, NodeState> = HashMap::new();
        for h in futures::future::join_all(handles).await {
            match h {
                Ok((name, Ok(output))) => {
                    let occupied = !output.trim().is_empty();
                    nodes.insert(
                        name,
                        NodeState {
                            reachable: true,
                            occupied,
                        },
                    );
                }
                Ok((name, Err(e))) => {
                    // SSH failure → unreachable, cannot be deployed on
                    nodes.insert(
                        name.clone(),
                        NodeState {
                            reachable: false,
                            occupied: false,
                        },
                    );
                    warn!("node probe failed: {name}: {e:#}");
                }
                Err(e) => {
                    // A probe task panicked (JoinError) — log it so a node
                    // silently dropping out of the state map is diagnosable.
                    warn!("Cluster probe task panicked: {}", e);
                }
            }
        }
        Self { nodes }
    }

    /// Nodes in `cluster` that are reachable AND not occupied.
    pub fn free_nodes<'c>(&self, cluster: &'c Cluster) -> Vec<&'c crate::config::Node> {
        cluster.nodes.iter().filter(|n| self.free(n)).collect()
    }

    /// Number of free nodes in `cluster`.
    pub fn free_count(&self, cluster: &Cluster) -> u32 {
        self.free_nodes(cluster).len() as u32
    }

    /// Whether a single node is free (reachable + not occupied).
    fn free(&self, node: &crate::config::Node) -> bool {
        self.nodes
            .get(&node.name)
            .is_some_and(|s| s.reachable && !s.occupied)
    }
}

/// Select a cluster given the current cluster state.
///
/// NCCL config is only required for multi-node recipes
/// (`required_nodes > 1`). For single-node recipes, clusters without
/// NCCL config are allowed — the config is returned as `None`.
pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
    required_nodes: Option<u32>,
    state: &ClusterState,
) -> Result<(Cluster, Option<&'c NcclConfig>)> {
    if cluster_name != "auto" {
        let cluster = crate::config::resolve_cluster_name(&config.clusters, cluster_name)?;
        let nccl = config.nccl.get(&cluster.name);
        if required_nodes.is_some_and(|n| n > 1) && nccl.is_none() {
            anyhow::bail!("no NCCL config for cluster '{}'", cluster.name);
        }
        return Ok((cluster.clone(), nccl));
    }
    // For single-node recipes, NCCL config is not required — allow clusters
    // without it. For multi-node, require it.
    let candidates: Vec<&Cluster> = if required_nodes.is_some_and(|n| n > 1) {
        config
            .clusters
            .iter()
            .filter_map(|c| {
                let _nccl = config.nccl.get(&c.name)?;
                Some(c)
            })
            .filter(|c| required_nodes.is_none_or(|r| state.free_count(c) >= r))
            .collect()
    } else {
        config
            .clusters
            .iter()
            .filter(|c| required_nodes.is_none_or(|r| state.free_count(c) >= r))
            .collect()
    };
    let mut candidates = candidates;
    candidates.sort_by_key(|c| c.nodes.len());
    if let Some(cluster) = candidates.into_iter().next() {
        let avail = state.free_count(cluster);
        let nccl = config.nccl.get(&cluster.name);
        info!(
            "{}",
            fluent_i18n::t!("cluster-selected", {
                "name" => cluster.name.to_string(),
                "nodes" => avail,
            })
        );
        return Ok((cluster.clone(), nccl));
    }
    anyhow::bail!(fluent_i18n::t!("cluster-auto-fail"));
}

#[cfg(test)]
mod tests;