launcher/src/web/flare_nodes.rs

//! Which of our nodes runs the `flare orbit` CLI.
//!
//! FLARE names its peers itself (`gx10-dgx5`…), so the launcher *finds* the
//! controller: a cluster's rank-0 node is the candidate and the first one
//! whose `orbit status` answers is it. `[flare.clusters.<name>].controller_node`
//! pins it per cluster; the flat `[flare].controller_node` is the
//! single-cluster fallback.
//!
//! Only clusters the config actually **pairs** with an orbit are candidates:
//! several of our clusters share nodes (`dgx1` is rank 0 of `4nodes-switch`,
//! `6nodes-switch`, `2nodes-dac`, `qwen38`, `1node-1`), so probing an unpaired
//! cluster with the flat `sparks` default made another cluster's lane look
//! like its own.

use crate::config::{Cluster, Config, Node};

/// The pinned controller of ONE cluster: a per-cluster override wins, then the
/// flat `[flare]` field — and either must actually be a node of that cluster.
pub fn pinned_for<'a>(config: &'a Config, cluster: &'a Cluster) -> Option<&'a Node> {
    let f = config.flare.as_ref()?;
    let want = f
        .clusters
        .get(&cluster.name)
        .and_then(|c| c.controller_node.clone())
        .or_else(|| f.controller_node.clone())?;
    cluster.nodes.iter().find(|n| n.name == want)
}

/// Rank-0 candidates of ONE cluster (pinned controller first) — the fleet
/// probe must ask each cluster with its own orbit name, so `prod4` groups are
/// not looked for in the `sparks` orbit (that silently showed nothing).
pub fn candidates_for<'a>(
    config: &'a Config,
    cluster: &'a crate::config::Cluster,
) -> Vec<&'a Node> {
    let mut out: Vec<&Node> = Vec::new();
    if let Some(n) = pinned_for(config, cluster) {
        out.push(n);
    }
    for n in &cluster.nodes {
        if n.rank == 0 && !out.iter().any(|x| x.ssh_ip == n.ssh_ip) {
            out.push(n);
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    fn config() -> Config {
        toml::from_str(
            r#"
            ssh_user = "u"
            ssh_key_path = "/k"

            [flare]
            orbit_cluster = "sparks"
            controller_node = "n5"

            [flare.clusters.prod4]
            orbit_cluster = "prod4"
            controller_node = "n1"

            [[clusters]]
            name = "prod4"
            nodes = [
              { name = "n1", ssh_ip = "10.0.0.1", rank = 0 },
              { name = "n2", ssh_ip = "10.0.0.2", rank = 1 },
            ]

            [[clusters]]
            name = "six"
            nodes = [
              { name = "n1", ssh_ip = "10.0.0.1", rank = 0 },
              { name = "n2", ssh_ip = "10.0.0.2", rank = 1 },
              { name = "n5", ssh_ip = "10.0.0.5", rank = 2 },
            ]
            "#,
        )
        .expect("config toml")
    }

    #[test]
    fn pinned_is_per_cluster_and_in_that_cluster() {
        let cfg = config();
        let prod4 = cfg.clusters.iter().find(|c| c.name == "prod4").unwrap();
        let six = cfg.clusters.iter().find(|c| c.name == "six").unwrap();
        assert_eq!(pinned_for(&cfg, prod4).map(|n| n.name.as_str()), Some("n1"));
        // `six` has no per-cluster pin; the flat one (n5) is one of its nodes
        assert_eq!(pinned_for(&cfg, six).map(|n| n.name.as_str()), Some("n5"));
    }

    #[test]
    fn pairing_only_names_clusters_the_config_mentions() {
        use crate::native::flare::orbit_cluster_paired;
        let cfg = config();
        assert_eq!(
            orbit_cluster_paired(&cfg, "prod4").as_deref(),
            Some("prod4")
        );
        // dgx5 is six's rank 2, not its controller — no pairing
        assert_eq!(orbit_cluster_paired(&cfg, "six"), None);
    }
}