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>,
}

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(_))) => {
                    // SSH failure → unreachable, cannot be deployed on
                    nodes.insert(
                        name,
                        NodeState {
                            reachable: false,
                            occupied: false,
                        },
                    );
                }
                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 {
    use super::*;
    use crate::config::{Config, NcclConfig, Node};

    fn test_node(name: &str) -> Node {
        Node {
            name: name.to_string(),
            ssh_ip: format!("10.0.0.{}", name.chars().last().unwrap_or('0')),
            ib_ip: Some("10.0.1.1".to_string()),
            rank: 0,
            host_key: None,
            entity: None,
        }
    }

    fn test_cluster(name: &str, node_names: &[&str]) -> Cluster {
        Cluster {
            name: name.to_string(),
            master_ib_ip: Some("10.0.1.1".to_string()),
            nodes: node_names.iter().map(|n| test_node(n)).collect(),
            extra_entities: vec![],
        }
    }

    fn test_nccl() -> NcclConfig {
        NcclConfig {
            ib_hca: "r".to_string(),
            socket_ifname: "e".to_string(),
            ib_gid_index: Some(3),
            ib_disable: 0,
            gloo_socket_ifname: None,
            ucx_net_devices: None,
            nccl_cross_nic: None,
            nccl_proto: None,
            nccl_p2p_level: None,
        }
    }

    #[test]
    fn free_count_all_reachable() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        state.nodes.insert(
            "n2".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        let c = test_cluster("c", &["n1", "n2"]);
        assert_eq!(state.free_count(&c), 2);
    }

    #[test]
    fn free_count_occupied_excluded() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: true,
            },
        );
        state.nodes.insert(
            "n2".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        let c = test_cluster("c", &["n1", "n2"]);
        assert_eq!(state.free_count(&c), 1);
    }

    #[test]
    fn free_count_unreachable_excluded() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: false,
                occupied: false,
            },
        );
        state.nodes.insert(
            "n2".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        let c = test_cluster("c", &["n1", "n2"]);
        assert_eq!(state.free_count(&c), 1);
    }

    #[test]
    fn free_count_missing_node_not_free() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        // n2 not probed — absent is treated as not free
        let c = test_cluster("c", &["n1", "n2"]);
        assert_eq!(state.free_count(&c), 1);
    }

    #[test]
    fn select_smallest_with_enough_free() {
        let mut state = ClusterState::default();
        // big cluster: 4 nodes, 2 free
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: true,
            },
        );
        state.nodes.insert(
            "n2".into(),
            NodeState {
                reachable: true,
                occupied: true,
            },
        );
        state.nodes.insert(
            "n3".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        state.nodes.insert(
            "n4".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        // small cluster: 2 nodes, 1 free
        state.nodes.insert(
            "n5".into(),
            NodeState {
                reachable: true,
                occupied: true,
            },
        );
        state.nodes.insert(
            "n6".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );

        let cfg = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![
                test_cluster("big", &["n1", "n2", "n3", "n4"]),
                test_cluster("small", &["n5", "n6"]),
            ],
            nccl: {
                let mut m = HashMap::new();
                m.insert("big".into(), test_nccl());
                m.insert("small".into(), test_nccl());
                m
            },
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            global_mods_dir: None,
        };

        // Request 1 node — small cluster is smallest with free >= 1
        let (c, _) = select_cluster(&cfg, "auto", Some(1), &state).unwrap();
        assert_eq!(c.name, "small");

        // Request 2 nodes — only big cluster has 2 free
        let (c, _) = select_cluster(&cfg, "auto", Some(2), &state).unwrap();
        assert_eq!(c.name, "big");
    }

    #[test]
    fn select_cluster_not_found_errors() {
        let state = ClusterState::default();
        let cfg = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            global_mods_dir: None,
        };
        assert!(select_cluster(&cfg, "nonexistent", None, &state).is_err());
    }

    #[test]
    fn no_cluster_has_enough_free_errors() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: true,
            },
        );
        let cfg = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![test_cluster("c", &["n1"])],
            nccl: {
                let mut m = HashMap::new();
                m.insert("c".into(), test_nccl());
                m
            },
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            global_mods_dir: None,
        };
        assert!(select_cluster(&cfg, "auto", Some(1), &state).is_err());
    }

    #[test]
    fn single_node_allows_cluster_without_nccl() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        let cfg = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![test_cluster("c", &["n1"])],
            nccl: HashMap::new(), // no NCCL config
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            global_mods_dir: None,
        };
        // Single-node recipe on a cluster without NCCL config → OK
        let (c, nccl) = select_cluster(&cfg, "auto", Some(1), &state).unwrap();
        assert_eq!(c.name, "c");
        assert!(nccl.is_none());
    }

    #[test]
    fn multi_node_requires_nccl() {
        let mut state = ClusterState::default();
        state.nodes.insert(
            "n1".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        state.nodes.insert(
            "n2".into(),
            NodeState {
                reachable: true,
                occupied: false,
            },
        );
        let cfg = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![test_cluster("c", &["n1", "n2"])],
            nccl: HashMap::new(), // no NCCL config
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            global_mods_dir: None,
        };
        // Multi-node recipe on a cluster without NCCL config → error
        assert!(select_cluster(&cfg, "auto", Some(2), &state).is_err());
    }
}