launcher/src/orchestrate/mod.rs

//! Serve orchestrator — combines probing, power control, and node-pinging
//! into a single `prepare_cluster` step that `serve` calls.
//!
//! Every external interaction goes through the trait bounds `P: NodeProber`,
//! `C: PowerController`, `N: NodePinger`, making the control flow testable
//! with mock implementations.

use std::time::Duration;

use anyhow::{Context, Result};
use tracing::info;

use crate::cluster_state::{ClusterState, select_cluster};
use crate::config::{Cluster, Config, NcclConfig, Node, Recipe};
use crate::ends::{NodePinger, NodeProber, PowerController, cluster_entities};

/// Prepare a cluster for serving: probe, optionally power on, wait, then
/// select a cluster with its free nodes.
///
/// Returns a `Cluster` whose `nodes` field contains only free (reachable
/// and not occupied) nodes, plus the matching NCCL config.
#[allow(clippy::too_many_arguments)]
pub async fn prepare_cluster<P, C, N>(
    config: &Config,
    recipe: &Recipe,
    cluster_name: &str,
    node_name: Option<&str>,
    auto_power_on: bool,
    prober: &P,
    power: Option<&C>,
    pinger: &N,
) -> Result<(Cluster, NcclConfig)>
where
    P: NodeProber,
    C: PowerController,
    N: NodePinger,
{
    let state = prober.probe(config).await;

    let resolved = if let Some(nn) = node_name {
        resolve_solo_node(config, nn)?
    } else {
        // Auto-power-on: if any nodes are unreachable, try powering on.
        let state = if auto_power_on && state_has_unreachable(&state, config) {
            let (c, _nccl) = resolve_cluster_simple(config, recipe, cluster_name)?;
            if let Some(pwr) = power {
                power_on_if_off(&state, &c, pwr).await?;
            }
            // Wait for nodes to become reachable.
            wait_for_all_nodes(config, &c.nodes, pinger).await?;
            // Re-probe after power-on.
            prober.probe(config).await
        } else {
            state
        };

        let (c, n) = select_cluster(config, cluster_name, recipe.required_nodes, &state)?;
        let avail: Vec<Node> = c
            .nodes
            .into_iter()
            .filter(|n| {
                state
                    .nodes
                    .get(&n.name)
                    .is_some_and(|ns| ns.reachable && !ns.occupied)
            })
            .collect();
        let c = Cluster { nodes: avail, ..c };
        (c, n.clone())
    };

    // Validate node count.
    if let Some(required) = recipe.required_nodes {
        let actual = resolved.0.nodes.len() as u32;
        if actual < required {
            anyhow::bail!(
                "Recipe '{}' requires {} nodes, but cluster '{}' has only {} nodes",
                recipe.name,
                required,
                resolved.0.name,
                actual
            );
        }
    }

    Ok(resolved)
}

/// Collect a single node from any cluster.
fn resolve_solo_node(config: &Config, node_name: &str) -> Result<(Cluster, NcclConfig)> {
    let (solo_c, solo_n) = config
        .clusters
        .iter()
        .find_map(|c| {
            c.nodes.iter().find(|n| n.name == node_name).map(|node| {
                let solo = Cluster {
                    name: c.name.clone(),
                    master_ib_ip: node.ib_ip.clone(),
                    nodes: vec![node.clone()],
                    extra_entities: vec![],
                };
                let nc = config
                    .nccl
                    .get(&c.name)
                    .cloned()
                    .unwrap_or_else(|| NcclConfig {
                        ib_hca: String::new(),
                        socket_ifname: "lo".to_string(),
                        ib_gid_index: 3,
                        ib_disable: 1,
                        gloo_socket_ifname: None,
                        ucx_net_devices: None,
                        nccl_cross_nic: None,
                        nccl_proto: None,
                        nccl_p2p_level: None,
                    });
                (solo, nc)
            })
        })
        .context(format!("Node '{}' not found in any cluster", node_name))?;
    Ok((solo_c, solo_n))
}

/// Resolve cluster by name or auto treating every config node as free
/// (reachable, not occupied). Only used for the power-on pre-check
/// before any probe has been made.
fn resolve_cluster_simple<'c>(
    config: &'c Config,
    recipe: &Recipe,
    cluster_name: &str,
) -> Result<(Cluster, &'c NcclConfig)> {
    use crate::cluster_state::NodeState;
    let mut state = ClusterState::default();
    for cluster in &config.clusters {
        for node in &cluster.nodes {
            state.nodes.entry(node.name.clone()).or_insert(NodeState {
                reachable: true,
                occupied: false,
            });
        }
    }
    select_cluster(config, cluster_name, recipe.required_nodes, &state)
}

/// Does the state contain any unreachable nodes?
fn state_has_unreachable(state: &ClusterState, config: &Config) -> bool {
    for cluster in &config.clusters {
        for node in &cluster.nodes {
            if state.nodes.get(&node.name).is_some_and(|ns| !ns.reachable) {
                return true;
            }
        }
    }
    false
}

/// Power on HA entities that are not already on.
async fn power_on_if_off(
    _state: &ClusterState,
    cluster: &Cluster,
    power: &impl PowerController,
) -> Result<bool> {
    let entities = cluster_entities(cluster);
    if entities.is_empty() {
        return Ok(false);
    }

    let any_off = {
        let mut off = false;
        for e in &entities {
            match power.entity_is_on(e).await {
                Some(true) => {}
                _ => off = true,
            }
        }
        off
    };

    if !any_off {
        return Ok(false);
    }

    info!("{}", fluent_i18n::t!("cluster-switch-off-detected"));
    for e in &entities {
        power.turn_on(e).await?;
    }
    info!("{}", fluent_i18n::t!("cluster-switch-powered-on"));
    Ok(true)
}

/// Wait until every node in the list accepts an SSH ping, with a timeout.
async fn wait_for_all_nodes(
    config: &Config,
    nodes: &[Node],
    pinger: &impl NodePinger,
) -> Result<()> {
    info!("{}", fluent_i18n::t!("cluster-waiting-for-nodes"));

    let boot_wait_secs = 60;
    info!("⏳ Waiting {}s for nodes to boot...", boot_wait_secs);
    tokio::time::sleep(Duration::from_secs(boot_wait_secs)).await;

    let max_poll_secs = 120;
    let poll_interval_secs = 10;
    let mut waited = 0;
    while waited < max_poll_secs {
        let mut all_online = true;
        for node in nodes {
            if !pinger.is_reachable(config, node).await {
                all_online = false;
                break;
            }
        }
        if all_online {
            info!("{}", fluent_i18n::t!("cluster-nodes-available"));
            return Ok(());
        }
        tokio::time::sleep(Duration::from_secs(poll_interval_secs)).await;
        waited += poll_interval_secs;
    }
    anyhow::bail!("{}", fluent_i18n::t!("cluster-nodes-not-available-timeout"));
}
#[cfg(test)]
#[path = "tests.rs"]
mod tests;