launcher/src/config.rs

pub use crate::pre_serve::{BeforeServeTable, ModsTable, OnFailure, PreServeCommand};
use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::HashMap;
use tracing::info;
#[derive(Clone, Deserialize)]
pub struct Node {
    pub name: String,
    pub ssh_ip: String,
    pub ib_ip: String,
    pub rank: u32,
    #[serde(default)]
    pub host_key: Option<String>,
    #[serde(default)]
    pub entity: Option<String>,
}
#[derive(Clone, Deserialize)]
pub struct NcclConfig {
    pub ib_hca: String,
    pub socket_ifname: String,
    pub ib_gid_index: u32,
    pub ib_disable: u32,
}
#[derive(Clone, Deserialize)]
pub struct Cluster {
    pub name: String,
    pub master_ib_ip: String,
    pub nodes: Vec<Node>,
    /// Extra Home Assistant switches to toggle alongside the nodes during power on/off.
    #[serde(default)]
    pub extra_entities: Vec<String>,
}
/// Top-level config (clusters + per-cluster pre-serve commands).
#[derive(Clone, Deserialize)]
pub struct Config {
    pub ssh_user: String,
    pub ssh_key_path: std::path::PathBuf,
    pub clusters: Vec<Cluster>,
    #[serde(default)]
    pub nccl: HashMap<String, NcclConfig>,
    #[serde(default)]
    pub hf_cache_dir: Option<String>,
    /// TCP port for cluster init (vLLM/SGLang)
    #[serde(default = "default_init_port")]
    pub init_port: u16,
    /// Per-cluster pre-serve commands, command aliases, HA config
    #[serde(default)]
    pub before_serve: HashMap<String, BeforeServeTable>, // <-- move here
    #[serde(default)]
    pub exec_aliases: HashMap<String, String>,
    #[serde(default)]
    pub home_assistant: Option<HomeAssistantConfig>,
    /// Global mods directory — script files merged with recipe mods/ into
    /// the container's /mods/.  Recipe-level [mods] commands explicitly
    /// reference files from either source via /mods/<file>.
    #[serde(default)]
    pub global_mods_dir: Option<std::path::PathBuf>,
}
/// Home Assistant config for `power on`/`power off`.
#[derive(Clone, Deserialize)]
pub struct HomeAssistantConfig {
    pub url: String,
    pub token: String,
    #[serde(default = "default_shutdown_delay")]
    pub shutdown_delay_seconds: u64,
}
fn default_shutdown_delay() -> u64 {
    30
}
fn default_init_port() -> u16 {
    25000
}
fn default_port() -> u16 {
    30000
}
impl Config {
    pub fn resolved_hf_cache_dir(&self) -> std::path::PathBuf {
        if let Some(ref dir) = self.hf_cache_dir {
            return std::path::PathBuf::from(dir);
        }
        let home = if self.ssh_user == "root" {
            "/root"
        } else {
            // Safe to return a temporary String; PathBuf::from copies it
            return std::path::PathBuf::from(format!("/home/{}/.cache/huggingface", self.ssh_user));
        };
        std::path::PathBuf::from(format!("{}/.cache/huggingface", home))
    }
}
#[derive(Clone, Deserialize)]
pub struct Recipe {
    #[serde(default = "default_engine")]
    pub engine: String,
    pub name: String,
    #[serde(default)]
    pub docker_image: String,
    #[serde(default)]
    pub model: String,
    #[serde(default = "default_port")]
    pub port: u16,
    #[serde(default)]
    pub required_nodes: Option<u32>,
    #[serde(default)]
    pub pipeline_parallel_size: Option<u32>,
    #[serde(default = "default_tensor_parallel_size")]
    pub tensor_parallel_size: u32,
    #[serde(default)]
    pub args: Vec<String>,
    #[serde(default)]
    pub env: HashMap<String, String>,
    #[serde(default)]
    pub cluster: Option<String>,
    #[serde(default)]
    pub entrypoint: Vec<String>,
    #[serde(default = "default_distributed_backend")]
    pub distributed_backend: String,
    /// Shared memory size for the container (e.g. "64gb"). Omitted when empty.
    #[serde(default = "default_shm_size")]
    pub shm_size: String,
    /// Recipe-level pre-serve commands (appended after cluster-level).
    #[serde(default)]
    pub before_serve: Vec<PreServeCommand>,
    /// Mods: commands that run *inside* the container via entrypoint wrapping
    /// before the engine process starts. Requires `entrypoint` to be set.
    #[serde(default)]
    pub mods: ModsTable,
}
/// Resolves cluster by exact name, then unique prefix. Errors if none or ambiguous.
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(", ")
        ),
    }
}
fn default_engine() -> String {
    "vllm".to_string()
}
fn default_distributed_backend() -> String {
    "mp".to_string()
}
fn default_shm_size() -> String {
    String::new()
}
fn default_tensor_parallel_size() -> u32 {
    1
}
/// Selects a cluster + NCCL config by name or auto-picks the smallest
/// cluster with NCCL config satisfying `required_nodes`.
pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
    required_nodes: Option<u32>,
) -> Result<(Cluster, &'c NcclConfig)> {
    if cluster_name != "auto" {
        let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
        let nccl = config
            .nccl
            .get(&cluster.name)
            .context(fluent_i18n::t!("nccl-config-missing"))?;
        return Ok((cluster.clone(), nccl));
    }
    let mut candidates: Vec<(&Cluster, &NcclConfig)> = config
        .clusters
        .iter()
        .filter_map(|c| Some((c, config.nccl.get(&c.name)?)))
        .filter(|(c, _)| required_nodes.is_none_or(|r| (c.nodes.len() as u32) == r))
        .collect();
    candidates.sort_by_key(|(c, _)| c.nodes.len());
    if let Some((cluster, nccl)) = candidates.into_iter().next() {
        info!(
            "{}",
            fluent_i18n::t!("cluster-selected", {
                "name" => cluster.name.to_string(),
                "nodes" => cluster.nodes.len(),
            })
        );
        return Ok((cluster.clone(), nccl));
    }
    anyhow::bail!(fluent_i18n::t!("cluster-auto-fail"));
}
/// Looks up a cluster by name (or unique prefix) without requiring NCCL config.
pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
    Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
}
/// Resolves the cluster for a recipe, applying recipe-level cluster pinning.
pub fn resolve_cluster<'c>(
    config: &'c Config,
    recipe: &Recipe,
    cluster_name: &str,
) -> Result<(Cluster, &'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)
}
/// Resolve the path to a recipe's `mods/` directory, if it exists.
/// Checks cwd first, then the default config dir.
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)
}
/// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
/// falling back to cwd for the error).
pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
    let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
    if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
        return Ok(toml::from_str(&content)?);
    }
    if let Some(dir) = default_config_dir() {
        let in_default = dir.join("recipes").join(recipe_name).join("recipe.toml");
        if let Ok(content) = tokio::fs::read_to_string(&in_default).await {
            return Ok(toml::from_str(&content)?);
        }
    }
    // Neither found — let the cwd path produce a clean error.
    let content = tokio::fs::read_to_string(&cwd_path).await?;
    Ok(toml::from_str(&content)?)
}
/// OS-appropriate default config dir (e.g. `~/.config/sm121-launcher`),
/// or `None` for sandboxes with no known config dir.
pub fn default_config_dir() -> Option<std::path::PathBuf> {
    dirs::config_dir().map(|p| p.join("sm121-launcher"))
}
#[cfg(test)]
mod tests {
    use super::*;
    fn test_cluster() -> Cluster {
        Cluster {
            name: "test-cluster".to_string(),
            master_ib_ip: "192.168.1.1".to_string(),
            nodes: vec![
                Node {
                    name: "node0".to_string(),
                    ssh_ip: "10.0.0.1".to_string(),
                    ib_ip: "192.168.1.1".to_string(),
                    rank: 0,
                    host_key: None,
                    entity: None,
                },
                Node {
                    name: "node1".to_string(),
                    ssh_ip: "10.0.0.2".to_string(),
                    ib_ip: "192.168.1.2".to_string(),
                    rank: 1,
                    host_key: None,
                    entity: None,
                },
            ],
            extra_entities: vec![],
        }
    }
    fn test_config() -> Config {
        Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".into(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            global_mods_dir: None,
        }
    }
    fn test_nccl() -> NcclConfig {
        NcclConfig {
            ib_hca: "rocep1s0f0".to_string(),
            socket_ifname: "enp1s0f0np0".to_string(),
            ib_gid_index: 3,
            ib_disable: 0,
        }
    }
    #[test]
    fn test_recipe_deserialize_minimal() {
        let toml_str = "\nname = \"Minimal\"\nmodel = \"test/model\"\n";
        let recipe: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(recipe.name, "Minimal");
        assert_eq!(recipe.model, "test/model");
        assert_eq!(recipe.engine, "vllm");
        assert_eq!(recipe.port, 30000);
        assert!(recipe.args.is_empty());
        assert!(recipe.env.is_empty());
        assert_eq!(recipe.distributed_backend, "mp");
        assert_eq!(recipe.tensor_parallel_size, 1);
    }
    #[test]
    fn test_recipe_deserialize_full() {
        let toml_str = r#"
engine = "sglang"
name = "Full Recipe"
docker_image = "org/image:tag"
model = "org/model"
port = 30000
required_nodes = 4
pipeline_parallel_size = 2
distributed_backend = "ray"
cluster = "my-cluster"
args = ["--arg1", "val1", "--arg2"]
[env]
KEY = "value"
OTHER = "123"
"#;
        let recipe: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(recipe.engine, "sglang");
        assert_eq!(recipe.name, "Full Recipe");
        assert_eq!(recipe.docker_image, "org/image:tag");
        assert_eq!(recipe.model, "org/model");
        assert_eq!(recipe.port, 30000);
        assert_eq!(recipe.required_nodes, Some(4));
        assert_eq!(recipe.pipeline_parallel_size, Some(2));
        assert_eq!(recipe.distributed_backend, "ray");
        assert_eq!(recipe.cluster, Some("my-cluster".to_string()));
        assert_eq!(recipe.args, vec!["--arg1", "val1", "--arg2"]);
        assert_eq!(recipe.env.get("KEY").unwrap(), "value");
        assert_eq!(recipe.env.get("OTHER").unwrap(), "123");
    }
    #[test]
    fn test_recipe_deserialize_no_required_nodes() {
        let toml_str = r#"
name = "Flexible"
model = "test/model"
"#;
        let recipe: Recipe = toml::from_str(toml_str).unwrap();
        assert_eq!(recipe.required_nodes, None);
        assert_eq!(recipe.pipeline_parallel_size, None);
    }
    #[test]
    fn test_config_deserialize() {
        let toml_str = r#"
ssh_user = "admin"
ssh_key_path = "/home/admin/.ssh/id_rsa"
[[clusters]]
name = "prod"
master_ib_ip = "10.0.0.1"
nodes = [
    { name = "node1", ssh_ip = "10.0.0.10", ib_ip = "10.0.0.1", rank = 0 },
    { name = "node2", ssh_ip = "10.0.0.11", ib_ip = "10.0.0.2", rank = 1 },
]
[nccl."prod"]
ib_hca = "mlx5_0"
socket_ifname = "eth0"
ib_gid_index = 3
ib_disable = 0
"#;
        let config: Config = toml::from_str(toml_str).unwrap();
        assert_eq!(config.ssh_user, "admin");
        assert_eq!(config.clusters.len(), 1);
        assert_eq!(config.clusters[0].nodes.len(), 2);
        assert_eq!(config.clusters[0].nodes[0].host_key, None);
        assert!(config.nccl.contains_key("prod") && config.nccl["prod"].ib_hca == "mlx5_0");
    }
    #[test]
    fn test_config_deserialize_with_host_key() {
        let toml_str = "\nssh_user = \"admin\"\nssh_key_path = \"/path/to/key\"\n[[clusters]]\nname = \"secure\"\nmaster_ib_ip = \"1.1.1.1\"\nnodes = [\n    { name = \"n1\", ssh_ip = \"2.2.2.2\", ib_ip = \"1.1.1.1\", rank = 0, host_key = \"SHA256:abc123\" },\n]\n";
        let config: Config = toml::from_str(toml_str).unwrap();
        assert_eq!(
            config.clusters[0].nodes[0].host_key,
            Some("SHA256:abc123".into())
        );
    }
    #[test]
    fn test_select_cluster_by_name() {
        let mut config = Config {
            ssh_user: "user".to_string(),
            ssh_key_path: "/path/to/key".into(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            global_mods_dir: None,
        };
        config.nccl.insert("test-cluster".to_string(), test_nccl());
        let (cluster, nccl) = select_cluster(&config, "test-cluster", None).unwrap();
        assert_eq!(cluster.name, "test-cluster");
        assert_eq!(nccl.ib_hca, "rocep1s0f0");
    }
    #[test]
    fn test_select_cluster_auto() {
        let mut config = Config {
            ssh_user: "user".to_string(),
            ssh_key_path: "/path/to/key".into(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            global_mods_dir: None,
        };
        config.nccl.insert("test-cluster".to_string(), test_nccl());
        let (cluster, nccl) = select_cluster(&config, "auto", None).unwrap();
        assert_eq!(cluster.name, "test-cluster");
        assert_eq!(nccl.ib_hca, "rocep1s0f0");
    }
    #[test]
    fn test_select_cluster_auto_no_match() {
        let config = test_config();
        let result = select_cluster(&config, "auto", None);
        assert!(result.is_err());
    }
    #[test]
    fn test_select_cluster_not_found() {
        let mut config = test_config();
        config.clusters = vec![];
        let result = select_cluster(&config, "nonexistent", None);
        assert!(result.is_err());
    }
    #[test]
    fn test_select_cluster_by_name_found() {
        let mut nccl_map = HashMap::new();
        nccl_map.insert("test-cluster".to_string(), test_nccl());
        let mut config = test_config();
        config.nccl = nccl_map;
        let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
        assert_eq!(cluster.name, "test-cluster");
        assert_eq!(cluster.nodes.len(), 2);
    }
    #[test]
    fn test_select_cluster_by_name_not_found() {
        let config = test_config();
        let result = select_cluster_by_name(&config, "no-such-cluster");
        assert!(result.is_err());
    }
    #[test]
    fn test_select_cluster_by_prefix() {
        let mut config = test_config();
        config.nccl.insert("test-cluster".to_string(), test_nccl());
        // Unique prefix works
        let cluster = select_cluster_by_name(&config, "test-c").unwrap();
        assert_eq!(cluster.name, "test-cluster");
        // No match errors
        let result = select_cluster_by_name(&config, "no-such");
        assert!(result.is_err());
    }
    #[test]
    fn test_russian_pluralization() {
        let _prev = fluent_i18n::set_locale(Some("ru")).ok();
        let s1 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 1u64 });
        assert!(s1.contains("1 узел"), "got: {}", s1);
        let s2 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 2u64 });
        assert!(s2.contains("2 узла"), "got: {}", s2);
        let s5 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 5u64 });
        assert!(s5.contains("5 узлов"), "got: {}", s5);
    }
    #[test]
    fn test_default_engine_is_vllm() {
        let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
        assert_eq!(recipe.engine, "vllm");
    }
    #[test]
    fn test_default_distributed_backend_is_mp() {
        let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
        assert_eq!(recipe.distributed_backend, "mp");
    }
    #[test]
    fn default_config_dir_appends_sm121_launcher() {
        let Some(dir) = default_config_dir() else {
            return;
        };
        assert!(dir.ends_with("sm121-launcher") && dir.parent().is_some());
    }
}