refactor: separate cluster state from SSH probing for testability

- Add src/cluster_state.rs with ClusterState struct and probe()
- occupied_nodes() now delegates to ClusterState::probe()
- Unreachable nodes (SSH failure) are marked unavailable for selection
- 7 new unit tests covering free/occupied/unreachable/missing nodes
  and cluster-auto-selection logic, all without real SSH

src/cluster_state.rs

@@ -0,0 +1,400 @@
+//! 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;
+
+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>,
+}
+
+#[allow(dead_code)]
+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,
+                        },
+                    );
+                }
+                _ => {}
+            }
+        }
+        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)
+    }
+
+    /// Set of occupied node names (reachable + has containers).
+    /// Used to filter occupied nodes from deployment targets.
+    pub fn occupied_names(&self) -> HashSet<&str> {
+        self.nodes
+            .iter()
+            .filter(|(_, s)| s.occupied)
+            .map(|(n, _)| n.as_str())
+            .collect()
+    }
+
+    /// Set of unreachable node names (SSH failure).
+    pub fn unreachable_names(&self) -> HashSet<&str> {
+        self.nodes
+            .iter()
+            .filter(|(_, s)| !s.reachable)
+            .map(|(n, _)| n.as_str())
+            .collect()
+    }
+}
+
+/// Select a cluster given the current cluster state.
+#[allow(dead_code)]
+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,
+            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,
+            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,
+            global_mods_dir: None,
+        };
+        assert!(select_cluster(&cfg, "auto", Some(1), &state).is_err());
+    }
+}

src/config.rs

@@ -133,7 +133,7 @@ pub struct Recipe {
     pub mods: ModsTable,
 }
 /// Resolve cluster.
-fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
+pub(crate) fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
     if let Some(c) = clusters.iter().find(|c| c.name == name) {
         return Ok(c);
     }

src/docker/mod.rs

@@ -237,38 +237,14 @@ pub async fn resolve_cluster_by_container_label<'c>(
 /// occupied node names.  A node is occupied if any sm121 container is
 /// running on it.
 pub async fn occupied_nodes(config: &Config) -> Result<HashSet<String>> {
-    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 cmd = "docker ps -q --filter label=sm121.cluster --format '{{.Names}}' | head -1"
-                .to_string();
-            let config = config.clone();
-            let node = node.clone();
-            let name = node.name.clone();
-            handles.push(tokio::spawn(async move {
-                (name, run_ssh(&config, &node, cmd, false).await)
-            }));
-        }
-    }
-    let mut occupied = HashSet::new();
-    for h in futures::future::join_all(handles).await {
-        match h {
-            Ok((name, Ok(output))) if !output.trim().is_empty() => {
-                occupied.insert(name);
-            }
-            Ok((name, Err(_))) => {
-                // Node is unreachable — mark as unavailable
-                occupied.insert(name);
-            }
-            _ => {}
+    let state = crate::cluster_state::ClusterState::probe(config).await;
+    let mut result: HashSet<String> = HashSet::new();
+    for (name, ns) in &state.nodes {
+        if !ns.reachable || ns.occupied {
+            result.insert(name.clone());
         }
     }
-    Ok(occupied)
+    Ok(result)
 }
 /// Find the cluster that has running containers for a recipe by reading
 /// the `sm121.cluster` label from live containers.  Used by stop/logs

src/main.rs

@@ -8,6 +8,7 @@ fluent_i18n::i18n!("locales", fallback = "en");
 mod before_serve;
 mod cli;
 mod cluster_power;
+mod cluster_state;
 mod config;
 mod docker;
 mod doctor;