launcher/src/config.rs

use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::{HashMap, HashSet};
use tracing::info;

/// A cluster-level HA entity (smart plug, switch, fan, …).
///
/// Two TOML forms:
/// ```toml
/// # blocking (default) — "off" takes the whole cluster down
/// extra_entities = ["switch.network"]
/// # non-blocking — tracked & toggled by `power`, but "off" does NOT
/// # make the cluster offline (e.g. a fan that may be off while serving)
/// extra_entities = [{ entity = "switch.fan", blocks = false }]
/// ```
#[derive(Clone, Debug, Deserialize, PartialEq)]
#[serde(untagged)]
pub enum ExtraEntity {
    /// Bare string — always blocking.
    Simple(String),
    /// Table `{ entity = "...", blocks = <bool> }`; `blocks` defaults to true.
    Detailed {
        entity: String,
        #[serde(default = "default_true")]
        blocks: bool,
    },
}

fn default_true() -> bool {
    true
}

impl ExtraEntity {
    /// The HA entity id (`domain.name`).
    pub fn entity(&self) -> &str {
        match self {
            ExtraEntity::Simple(s) => s,
            ExtraEntity::Detailed { entity, .. } => entity,
        }
    }
    /// Whether an "off" state counts against the cluster's up/down.
    /// Bare strings are always blocking; `blocks = false` opts out.
    pub fn blocks(&self) -> bool {
        match self {
            ExtraEntity::Simple(_) => true,
            ExtraEntity::Detailed { blocks, .. } => *blocks,
        }
    }
}

mod types;
pub use types::*;

impl Node {
    /// Host IP for engine exports; falls back to `ssh_ip` when `ib_ip` unset
    /// (single-node runs need no IB addressing).
    pub fn effective_ib_ip(&self) -> &str {
        self.ib_ip.as_deref().unwrap_or(&self.ssh_ip)
    }
}

impl Cluster {
    /// Effective master IP: explicit `master_ib_ip`, else the rank-0 node's
    /// effective (ib→ssh fallback) IP, else the first node's ssh IP, else empty.
    pub fn effective_master_ip(&self) -> String {
        if let Some(ip) = &self.master_ib_ip {
            return ip.clone();
        }
        self.nodes
            .iter()
            .find(|n| n.rank == 0)
            .or_else(|| self.nodes.first())
            .map(|n| n.effective_ib_ip().to_string())
            .unwrap_or_default()
    }
}

pub 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);
    }
    let matches: Vec<&Cluster> = clusters
        .iter()
        .filter(|c| c.name.starts_with(name))
        .collect();
    match matches.len() {
        0 => anyhow::bail!(fluent_i18n::t!("cluster-not-found", { "name" => name.to_string() })),
        1 => Ok(matches[0]),
        n => anyhow::bail!(
            "Pref '{}' matches {} clusters: {}",
            name,
            n,
            matches
                .iter()
                .map(|c| c.name.as_str())
                .collect::<Vec<_>>()
                .join(", ")
        ),
    }
}

/// A `cluster_sources` file: extra `[[clusters]]` entries, optionally with
/// per-cluster `[nccl.<name>]` / `[before_serve.<name>]` tables. Merged into
/// the main [`Config`] by [`merge_cluster_sources`].
#[derive(Deserialize, Default)]
pub struct ClusterSource {
    /// SSH credentials for this source's clusters; merged into the main
    /// config only when it leaves them empty (first source wins).
    #[serde(default)]
    pub ssh_user: Option<String>,
    #[serde(default)]
    pub ssh_key_path: Option<std::path::PathBuf>,
    #[serde(default)]
    pub clusters: Vec<Cluster>,
    #[serde(default)]
    pub nccl: HashMap<String, NcclConfig>,
    #[serde(default)]
    pub before_serve: HashMap<String, BeforeServeTable>,
    /// HA endpoint for `power on`/`power off` from other cwds; merged only
    /// when the main config has none.
    #[serde(default)]
    pub home_assistant: Option<HomeAssistantConfig>,
}

/// Merge each `Config.cluster_sources` file into `config`: extra clusters
/// (main-config entries win on name conflict; among sources, first wins),
/// plus per-cluster `nccl` / `before_serve` tables.
pub async fn merge_cluster_sources(config: &mut Config) -> Result<()> {
    for raw in &config.cluster_sources {
        let path = local_path(raw);
        let content = tokio::fs::read_to_string(&path)
            .await
            .with_context(|| format!("cluster source '{}'", path.display()))?;
        let src: ClusterSource = toml::from_str(&content)
            .with_context(|| format!("invalid cluster source '{}'", path.display()))?;
        if config.ssh_user.is_empty() {
            config.ssh_user = src.ssh_user.clone().unwrap_or_default();
        }
        if config.ssh_key_path.as_os_str().is_empty() {
            config.ssh_key_path = src.ssh_key_path.clone().unwrap_or_default();
        }
        if config.home_assistant.is_none() {
            config.home_assistant = src.home_assistant.clone();
        }
        for c in src.clusters {
            if !config.clusters.iter().any(|x| x.name == c.name) {
                config.clusters.push(c);
            }
        }
        for (name, v) in src.nccl {
            config.nccl.entry(name).or_insert(v);
        }
        for (name, v) in src.before_serve {
            config.before_serve.entry(name).or_insert(v);
        }
    }
    Ok(())
}

/// Expand `~/...` against the local user's home (unlike `expand_tilde`,
/// which targets the ssh user's home *on the node*). Absolute paths pass
/// through.
pub fn local_path(s: &str) -> std::path::PathBuf {
    if let (Some(rest), Some(home)) = (s.strip_prefix("~/"), dirs::home_dir()) {
        return home.join(rest);
    }
    std::path::PathBuf::from(s)
}

/// Whether NCCL configuration is required for cluster selection.
///
/// NCCL config is only needed for multi-node distributed runs — for
/// single-node recipes (`required_nodes == Some(1)`), no NCCL env vars
/// are emitted, so the config is not required.
fn nccl_required(required_nodes: Option<u32>) -> bool {
    required_nodes.is_some_and(|n| n > 1)
}

pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
    required_nodes: Option<u32>,
    occupied: &HashSet<String>,
) -> Result<(Cluster, Option<&'c NcclConfig>)> {
    if cluster_name != "auto" {
        let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
        let nccl = config.nccl.get(&cluster.name);
        if nccl_required(required_nodes) && nccl.is_none() {
            anyhow::bail!("{}", fluent_i18n::t!("nccl-config-missing"));
        }
        return Ok((cluster.clone(), nccl));
    }
    let free = |c: &Cluster| -> u32 {
        c.nodes
            .iter()
            .filter(|n| !occupied.contains(&n.name))
            .count() as u32
    };
    // For single-node recipes, NCCL config is not required — allow clusters
    // without it. For multi-node, require it.
    let candidates: Vec<&Cluster> = if nccl_required(required_nodes) {
        config
            .clusters
            .iter()
            .filter_map(|c| {
                let nccl = config.nccl.get(&c.name)?;
                Some((c, nccl))
            })
            .filter(|(c, _)| required_nodes.is_none_or(|r| free(c) >= r))
            .map(|(c, _)| c)
            .collect()
    } else {
        config
            .clusters
            .iter()
            .filter(|c| required_nodes.is_none_or(|r| free(c) >= r))
            .collect()
    };
    let mut candidates = candidates;
    candidates.sort_by_key(|c| c.nodes.len());
    if let Some(cluster) = candidates.into_iter().next() {
        let avail = free(cluster);
        let nccl = config.nccl.get(&cluster.name);
        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"));
}

pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
    Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
}

/// Resolve a cluster for a **build**.
///
/// A build runs the Docker image build on a single node, so NCCL config is
/// never required here, even when the recipe's `required_nodes > 1` (the
/// resulting image is later pushed/synced to the serving cluster). For
/// `cluster_name != "auto"` the named cluster is used as-is; for `"auto"` the
/// smallest cluster containing `build_node` is selected.
pub fn resolve_build_cluster(
    config: &Config,
    cluster_name: &str,
    build_node: &str,
) -> Result<Cluster> {
    if cluster_name != "auto" {
        return select_cluster_by_name(config, cluster_name);
    }
    config
        .clusters
        .iter()
        .filter(|c| c.nodes.iter().any(|n| n.name == build_node))
        .min_by_key(|c| c.nodes.len())
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("no cluster contains build node '{}'", build_node))
}

pub fn resolve_cluster<'c>(
    config: &'c Config,
    recipe: &Recipe,
    cluster_name: &str,
    occupied: &HashSet<String>,
) -> Result<(Cluster, Option<&'c NcclConfig>)> {
    let choice = if cluster_name != "auto" {
        cluster_name.to_string()
    } else if let Some(ref rc) = recipe.cluster {
        rc.clone()
    } else {
        "auto".to_string()
    };
    select_cluster(config, &choice, recipe.required_nodes, occupied)
}

pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
    let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods"));
    if cwd.is_dir() {
        return Some(cwd);
    }
    let cfg = default_config_dir()?
        .join("recipes")
        .join(name)
        .join("mods");
    cfg.is_dir().then_some(cfg)
}

/// All recipe dirs in search order: `./recipes`, `.` (published-bundle
/// clones are flat), each `recipe_sources` entry (`~` expanded), then the
/// global config dir's `recipes/`.
pub fn recipe_dirs(config: &Config) -> Vec<std::path::PathBuf> {
    let mut v = vec![
        std::path::PathBuf::from("recipes"),
        std::path::PathBuf::from("."),
    ];
    for s in &config.recipe_sources {
        v.push(local_path(s));
    }
    if let Some(dir) = default_config_dir() {
        v.push(dir.join("recipes"));
    }
    v
}

/// Load a recipe by name. Search order: `./recipes/<name>/recipe.toml` (cwd),
/// then each `sources` entry (a dir containing `<name>/recipe.toml`, `~`
/// expanded), then `<default_config_dir>/recipes/<name>/recipe.toml`.
pub async fn load_recipe(recipe_name: &str, sources: &[String]) -> Result<Recipe> {
    // Search order: `recipes/<name>` (launcher workspace), then `<name>`
    // (published bundle layout — a clone of the recipe registry is flat),
    // then each source, then the global config dir.
    let mut paths: Vec<std::path::PathBuf> = vec![
        std::path::PathBuf::from(format!("recipes/{}/recipe.toml", recipe_name)),
        std::path::PathBuf::from(format!("{}/recipe.toml", recipe_name)),
    ];
    for src in sources {
        paths.push(local_path(src).join(recipe_name).join("recipe.toml"));
    }
    if let Some(dir) = default_config_dir() {
        paths.push(dir.join("recipes").join(recipe_name).join("recipe.toml"));
    }
    for p in &paths {
        if let Ok(content) = tokio::fs::read_to_string(p).await {
            return Ok(toml::from_str(&content)?);
        }
    }
    anyhow::bail!(
        "recipe '{}' not found; searched: {}",
        recipe_name,
        paths
            .iter()
            .map(|p| p.to_string_lossy())
            .collect::<Vec<_>>()
            .join(", ")
    )
}

pub fn default_config_dir() -> Option<std::path::PathBuf> {
    dirs::config_dir().map(|p| p.join("sm121"))
}

#[cfg(test)]
mod tests;