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 (or `None` for
/// single-node recipes, where NCCL config is not required).
///
/// Timing for [`wait_for_all_nodes`] after a power-on: how long to wait
/// for the machines to boot before the first poll, the poll interval, and
/// the max total poll window. Injectable so tests don't sleep real seconds
/// (a previous hardcoded 60s boot wait made the whole suite take 60s).
#[derive(Clone, Debug)]
pub struct BootTiming {
    pub boot_wait_secs: u64,
    pub poll_interval_secs: u64,
    pub max_poll_secs: u64,
}

impl Default for BootTiming {
    fn default() -> Self {
        Self {
            boot_wait_secs: 60,
            poll_interval_secs: 10,
            max_poll_secs: 120,
        }
    }
}

/// Bundles the injectable end-effectors + power-on flag so
/// [`prepare_cluster`] stays under clippy's `too_many_arguments` threshold
/// without an `#[allow]`.
pub struct PrepareRequest<'a, P: NodeProber, C: PowerController, N: NodePinger> {
    pub prober: &'a P,
    pub power: Option<&'a C>,
    pub pinger: &'a N,
    pub auto_power_on: bool,
    pub boot_timing: BootTiming,
    /// `--force`: skip the occupancy filter during cluster selection.
    /// Nodes with another container running are kept in the available set
    /// (reachable is still required — can't SSH to an unreachable node).
    pub force_occupied: bool,
}

pub async fn prepare_cluster<P, C, N>(
    config: &Config,
    recipe: &Recipe,
    cluster_name: &str,
    node_name: Option<&str>,
    req: &PrepareRequest<'_, P, C, N>,
) -> Result<(Cluster, Option<NcclConfig>)>
where
    P: NodeProber,
    C: PowerController,
    N: NodePinger,
{
    let state = req.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 req.auto_power_on && state_has_unreachable(&state, config) {
            let (c, _nccl) = resolve_cluster_simple(config, recipe, cluster_name)?;
            if let Some(pwr) = req.power {
                power_on_if_off(&c, pwr).await?;
            }
            // Wait for nodes to become reachable.
            wait_for_all_nodes(config, &c.nodes, req.pinger, &req.boot_timing).await?;
            // Re-probe after power-on.
            req.prober.probe(config).await
        } else {
            state
        };

        // `--force`: clear occupied flags so select_cluster counts reachable
        // nodes (ignoring occupancy) and the downstream filter keeps them.
        let mut state = state;
        if req.force_occupied {
            for ns in state.nodes.values_mut() {
                if ns.reachable {
                    ns.occupied = false;
                }
            }
        }

        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 && (req.force_occupied || !ns.occupied))
            })
            .collect();
        let c = Cluster { nodes: avail, ..c };
        (c, n.cloned())
    };

    // 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.
///
/// Returns `None` for the NCCL config when the cluster has no NCCL entry
/// — this is valid for single-node deployments where NCCL env vars are
/// never emitted.
fn resolve_solo_node(config: &Config, node_name: &str) -> Result<(Cluster, Option<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();
                (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, so being optimistic is safe:
/// a full probe follows immediately after power-on completes.
fn resolve_cluster_simple<'c>(
    config: &'c Config,
    recipe: &Recipe,
    cluster_name: &str,
) -> Result<(Cluster, Option<&'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)
}

/// Dry-run cluster resolution: treat every configured node as reachable and
/// free, then select. Skips all SSH probing/power/pinging so a dry-run works
/// even when the cluster is off or unreachable (the whole point of a dry-run
/// is to preview the generated command, not to verify hardware).
///
/// Like [`prepare_cluster`], returns the cluster with **all** its nodes (dry-run
/// does not filter for free nodes, since none are actually inspected).
pub fn prepare_cluster_dry_run(
    config: &Config,
    recipe: &Recipe,
    cluster_name: &str,
    node_name: Option<&str>,
) -> Result<(Cluster, Option<NcclConfig>)> {
    if let Some(nn) = node_name {
        let (c, n) = resolve_solo_node(config, nn)?;
        return Ok((c, n));
    }
    let (c, n) = resolve_cluster_simple(config, recipe, cluster_name)?;
    Ok((c, n.cloned()))
}

/// 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(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,
    timing: &BootTiming,
) -> Result<()> {
    info!("{}", fluent_i18n::t!("cluster-waiting-for-nodes"));

    if timing.boot_wait_secs > 0 {
        info!(
            "{}",
            fluent_i18n::t!("boot-wait", { "secs" => timing.boot_wait_secs })
        );
        tokio::time::sleep(Duration::from_secs(timing.boot_wait_secs)).await;
    }

    let mut waited = 0;
    while waited < timing.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(timing.poll_interval_secs)).await;
        // Advance by at least 1s so a zero poll-interval (tests) can't spin
        // forever when the pinger reports unreachable.
        waited += timing.poll_interval_secs.max(1);
    }
    anyhow::bail!("{}", fluent_i18n::t!("cluster-nodes-not-available-timeout"));
}
#[cfg(test)]
#[path = "tests.rs"]
mod tests;