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::{Context, 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 -q --filter label=sm121.cluster --format '{{.Names}}' | head -1"
                        .to_string();
                handles.push(tokio::spawn(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.
pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
    required_nodes: Option<u32>,
    state: &ClusterState,
) -> Result<(Cluster, &'c NcclConfig)> {
    if cluster_name != "auto" {
        let cluster = crate::config::resolve_cluster_name(&config.clusters, cluster_name)?;
        let nccl = config
            .nccl
            .get(&cluster.name)
            .context("no NCCL config for cluster")?;
        return Ok((cluster.clone(), nccl));
    }
    let mut candidates: Vec<(&Cluster, &NcclConfig)> = config
        .clusters
        .iter()
        .filter_map(|c| Some((c, config.nccl.get(&c.name)?)))
        .filter(|(c, _)| required_nodes.is_none_or(|r| state.free_count(c) >= r))
        .collect();
    candidates.sort_by_key(|(c, _)| c.nodes.len());
    if let Some((cluster, nccl)) = candidates.into_iter().next() {
        let avail = state.free_count(cluster);
        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: "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: "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: 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());
    }
}