|
|
@@ -0,0 +1,334 @@
|
|
|
+use super::*;
|
|
|
+use crate::cluster_state::NodeState;
|
|
|
+use crate::config::{NcclConfig, Node};
|
|
|
+use std::collections::HashMap;
|
|
|
+use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
+
|
|
|
+// --- Mock ends ---
|
|
|
+
|
|
|
+struct MockProber {
|
|
|
+ state: ClusterState,
|
|
|
+}
|
|
|
+impl NodeProber for MockProber {
|
|
|
+ async fn probe(&self, _config: &Config) -> ClusterState {
|
|
|
+ self.state.clone()
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+struct MockPower {
|
|
|
+ results: HashMap<String, bool>,
|
|
|
+ turned_on: std::sync::Mutex<Vec<String>>,
|
|
|
+}
|
|
|
+impl PowerController for MockPower {
|
|
|
+ async fn entity_is_on(&self, entity: &str) -> Option<bool> {
|
|
|
+ self.results.get(entity).copied()
|
|
|
+ }
|
|
|
+ async fn turn_on(&self, entity: &str) -> Result<()> {
|
|
|
+ self.turned_on.lock().unwrap().push(entity.to_string());
|
|
|
+ Ok(())
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+struct MockPinger {
|
|
|
+ reachable: bool,
|
|
|
+}
|
|
|
+impl NodePinger for MockPinger {
|
|
|
+ async fn is_reachable(&self, _config: &Config, _node: &Node) -> bool {
|
|
|
+ self.reachable
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+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: Some(format!("switch.{}", name)),
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+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,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+fn test_config(clusters: Vec<Cluster>) -> Config {
|
|
|
+ let mut nccl = HashMap::new();
|
|
|
+ for c in &clusters {
|
|
|
+ nccl.insert(c.name.clone(), test_nccl());
|
|
|
+ }
|
|
|
+ Config {
|
|
|
+ ssh_user: "u".to_string(),
|
|
|
+ ssh_key_path: "/k".into(),
|
|
|
+ clusters,
|
|
|
+ nccl,
|
|
|
+ hf_cache_dir: None,
|
|
|
+ init_port: 25000,
|
|
|
+ before_serve: HashMap::new(),
|
|
|
+ exec_aliases: HashMap::new(),
|
|
|
+ home_assistant: None,
|
|
|
+ global_mods_dir: None,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+fn test_recipe(required_nodes: Option<u32>) -> Recipe {
|
|
|
+ toml::from_str(&format!(
|
|
|
+ r#"
|
|
|
+name = "test"
|
|
|
+model = "test/model"
|
|
|
+required_nodes = {}
|
|
|
+"#,
|
|
|
+ required_nodes.map_or("".into(), |n| n.to_string())
|
|
|
+ ))
|
|
|
+ .unwrap()
|
|
|
+}
|
|
|
+
|
|
|
+#[tokio::test]
|
|
|
+async fn cluster_off_auto_power_on_powers_and_selects() {
|
|
|
+ let c = test_cluster("c", &["n1", "n2"]);
|
|
|
+ let config = test_config(vec![c.clone()]);
|
|
|
+ let recipe = test_recipe(Some(2));
|
|
|
+
|
|
|
+ let mut state = ClusterState::default();
|
|
|
+ state.nodes.insert(
|
|
|
+ "n1".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: false,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+ state.nodes.insert(
|
|
|
+ "n2".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: false,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+
|
|
|
+ let mut after = ClusterState::default();
|
|
|
+ after.nodes.insert(
|
|
|
+ "n1".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: true,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+ after.nodes.insert(
|
|
|
+ "n2".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: true,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+
|
|
|
+ let power = MockPower {
|
|
|
+ results: [("switch.n1".into(), false), ("switch.n2".into(), false)]
|
|
|
+ .into_iter()
|
|
|
+ .collect(),
|
|
|
+ turned_on: std::sync::Mutex::new(Vec::new()),
|
|
|
+ };
|
|
|
+
|
|
|
+ let call_count = AtomicUsize::new(0);
|
|
|
+ struct CountProber {
|
|
|
+ first: ClusterState,
|
|
|
+ second: ClusterState,
|
|
|
+ count: AtomicUsize,
|
|
|
+ }
|
|
|
+ impl NodeProber for CountProber {
|
|
|
+ async fn probe(&self, _config: &Config) -> ClusterState {
|
|
|
+ let n = self.count.fetch_add(1, Ordering::SeqCst);
|
|
|
+ if n == 0 {
|
|
|
+ self.first.clone()
|
|
|
+ } else {
|
|
|
+ self.second.clone()
|
|
|
+ }
|
|
|
+ }
|
|
|
+ }
|
|
|
+ let prober = CountProber {
|
|
|
+ first: state,
|
|
|
+ second: after,
|
|
|
+ count: call_count,
|
|
|
+ };
|
|
|
+
|
|
|
+ let pinger = MockPinger { reachable: true };
|
|
|
+
|
|
|
+ let (cluster, _nccl) = prepare_cluster(
|
|
|
+ &config,
|
|
|
+ &recipe,
|
|
|
+ "auto",
|
|
|
+ None,
|
|
|
+ true,
|
|
|
+ &prober,
|
|
|
+ Some(&power),
|
|
|
+ &pinger,
|
|
|
+ )
|
|
|
+ .await
|
|
|
+ .unwrap();
|
|
|
+
|
|
|
+ assert_eq!(cluster.name, "c");
|
|
|
+ assert_eq!(cluster.nodes.len(), 2);
|
|
|
+ let turned_on = power.turned_on.lock().unwrap();
|
|
|
+ assert!(turned_on.contains(&"switch.n1".to_string()));
|
|
|
+ assert!(turned_on.contains(&"switch.n2".to_string()));
|
|
|
+}
|
|
|
+
|
|
|
+#[tokio::test]
|
|
|
+async fn cluster_off_no_auto_power_fails() {
|
|
|
+ let c = test_cluster("c", &["n1", "n2"]);
|
|
|
+ let config = test_config(vec![c.clone()]);
|
|
|
+ let recipe = test_recipe(Some(2));
|
|
|
+
|
|
|
+ let mut state = ClusterState::default();
|
|
|
+ state.nodes.insert(
|
|
|
+ "n1".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: false,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+ state.nodes.insert(
|
|
|
+ "n2".into(),
|
|
|
+ NodeState {
|
|
|
+ reachable: false,
|
|
|
+ occupied: false,
|
|
|
+ },
|
|
|
+ );
|
|
|
+
|
|
|
+ let prober = MockProber { state };
|
|
|
+ let pinger = MockPinger { reachable: false };
|
|
|
+ let power = MockPower {
|
|
|
+ results: HashMap::new(),
|
|
|
+ turned_on: std::sync::Mutex::new(Vec::new()),
|
|
|
+ };
|
|
|
+
|
|
|
+ let result = prepare_cluster(
|
|
|
+ &config,
|
|
|
+ &recipe,
|
|
|
+ "auto",
|
|
|
+ None,
|
|
|
+ false,
|
|
|
+ &prober,
|
|
|
+ Some(&power),
|
|
|
+ &pinger,
|
|
|
+ )
|
|
|
+ .await;
|
|
|
+
|
|
|
+ assert!(result.is_err());
|
|
|
+ let err = result.unwrap_err().to_string();
|
|
|
+ assert!(err.contains("No suitable cluster found") || err.contains("requires"));
|
|
|
+}
|
|
|
+
|
|
|
+#[tokio::test]
|
|
|
+async fn cluster_already_on_selects_directly() {
|
|
|
+ let c = test_cluster("c", &["n1", "n2"]);
|
|
|
+ let config = test_config(vec![c.clone()]);
|
|
|
+ let recipe = test_recipe(Some(2));
|
|
|
+
|
|
|
+ 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 prober = MockProber { state };
|
|
|
+ let pinger = MockPinger { reachable: true };
|
|
|
+ let power = MockPower {
|
|
|
+ results: [("switch.n1".into(), true), ("switch.n2".into(), true)]
|
|
|
+ .into_iter()
|
|
|
+ .collect(),
|
|
|
+ turned_on: std::sync::Mutex::new(Vec::new()),
|
|
|
+ };
|
|
|
+
|
|
|
+ let (cluster, _nccl) = prepare_cluster(
|
|
|
+ &config,
|
|
|
+ &recipe,
|
|
|
+ "auto",
|
|
|
+ None,
|
|
|
+ true,
|
|
|
+ &prober,
|
|
|
+ Some(&power),
|
|
|
+ &pinger,
|
|
|
+ )
|
|
|
+ .await
|
|
|
+ .unwrap();
|
|
|
+
|
|
|
+ assert_eq!(cluster.name, "c");
|
|
|
+ assert_eq!(cluster.nodes.len(), 2);
|
|
|
+ assert!(power.turned_on.lock().unwrap().is_empty());
|
|
|
+}
|
|
|
+
|
|
|
+#[tokio::test]
|
|
|
+async fn occupied_nodes_excluded() {
|
|
|
+ let c = test_cluster("c", &["n1", "n2"]);
|
|
|
+ let config = test_config(vec![c.clone()]);
|
|
|
+ let recipe = test_recipe(Some(1));
|
|
|
+
|
|
|
+ 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 prober = MockProber { state };
|
|
|
+ let pinger = MockPinger { reachable: true };
|
|
|
+ let power = MockPower {
|
|
|
+ results: HashMap::new(),
|
|
|
+ turned_on: std::sync::Mutex::new(Vec::new()),
|
|
|
+ };
|
|
|
+
|
|
|
+ let (cluster, _nccl) = prepare_cluster(
|
|
|
+ &config,
|
|
|
+ &recipe,
|
|
|
+ "auto",
|
|
|
+ None,
|
|
|
+ false,
|
|
|
+ &prober,
|
|
|
+ Some(&power),
|
|
|
+ &pinger,
|
|
|
+ )
|
|
|
+ .await
|
|
|
+ .unwrap();
|
|
|
+
|
|
|
+ assert_eq!(cluster.nodes.len(), 1);
|
|
|
+ assert_eq!(cluster.nodes[0].name, "n2");
|
|
|
+}
|
