launcher/src/config.rs

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

pub use crate::pre_serve::{BeforeServeTable, OnFailure, PreServeCommand};

#[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>,
}

#[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>,
}

/// Top-level config containing cluster definitions and per-cluster
/// pre-serve commands.
#[derive(Clone, Deserialize)]
pub struct Config {
    pub ssh_user: String,
    pub ssh_key_path: String,
    pub clusters: Vec<Cluster>,
    #[serde(default)]
    pub nccl: HashMap<String, NcclConfig>,
    #[serde(default)]
    pub hf_cache_dir: Option<String>,
    /// TCP port used for cluster init (vLLM `--master-port`,
    /// SGLang `--dist-init-addr` port). Default: 25000.
    #[serde(default = "default_init_port")]
    pub init_port: u16,
    /// Per-cluster pre-serve commands.  Looked up by cluster name.
    #[serde(default)]
    pub before_serve: HashMap<String, BeforeServeTable>,
    /// Command aliases for `cargo run exec --alias`.  Keyed by alias
    /// name, value is the shell command to run.
    #[serde(default)]
    pub exec_aliases: HashMap<String, String>,
}

fn default_init_port() -> u16 {
    25000
}

impl Config {
    /// Returns the HF cache directory path for volume mounting.
    /// If set in config, uses that value; otherwise derives from
    /// `ssh_user` to produce `/home/<user>/.cache/huggingface`
    /// (or `/root/.cache/huggingface` for root).
    pub fn resolved_hf_cache_dir(&self) -> String {
        if let Some(ref dir) = self.hf_cache_dir {
            return dir.clone();
        }
        let home = if self.ssh_user == "root" {
            "/root".to_string()
        } else {
            format!("/home/{}", self.ssh_user)
        };
        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)]
    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,
    /// Recipe-level pre-serve commands.  Appended after cluster-level
    /// commands from `[before_serve."<cluster>"]`.
    #[serde(default)]
    pub before_serve: Vec<PreServeCommand>,
}

fn default_engine() -> String {
    "vllm".to_string()
}

fn default_distributed_backend() -> String {
    "mp".to_string()
}

fn default_tensor_parallel_size() -> u32 {
    1
}

/// Selects a cluster and its NCCL config by name, or auto-picks the first
/// cluster that has NCCL config.
pub fn select_cluster<'c>(
    config: &'c Config,
    cluster_name: &str,
) -> Result<(Cluster, &'c NcclConfig)> {
    if cluster_name != "auto" {
        let cluster = config
            .clusters
            .iter()
            .find(|c| c.name == cluster_name)
            .context(format!("Кластер {} не найден", cluster_name))?;
        let nccl = config
            .nccl
            .get(&cluster.name)
            .context("NCCL конфиг не найден")?;
        return Ok((cluster.clone(), nccl));
    }

    for cluster in &config.clusters {
        if let Some(nccl) = config.nccl.get(&cluster.name) {
            info!("✅ Автоподбор кластера: {}", cluster.name);
            return Ok((cluster.clone(), nccl));
        }
    }

    anyhow::bail!("Не удалось подобрать подходящий кластер");
}

/// Looks up a cluster by name without requiring NCCL config.
pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result<Cluster> {
    let cluster = config
        .clusters
        .iter()
        .find(|c| c.name == cluster_name)
        .context(format!("Кластер {} не найден", cluster_name))?;
    Ok(cluster.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)
}

/// Loads a recipe from disk.
pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
    let path = format!("recipes/{}/recipe.toml", recipe_name);
    let content = tokio::fs::read_to_string(&path).await?;
    Ok(toml::from_str(&content)?)
}

/// Returns the cluster and recipe name for a recipe (for stop/logs commands
/// that only receive a recipe name).
pub async fn resolve_cluster_for_recipe(
    config: &Config,
    recipe_name: &str,
) -> Result<(Cluster, Recipe)> {
    let recipe = load_recipe(recipe_name).await?;
    let (cluster, _) = resolve_cluster(config, &recipe, "auto")?;
    Ok((cluster, recipe))
}

#[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,
                },
                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,
                },
            ],
        }
    }

    fn test_nccl() -> NcclConfig {
        NcclConfig {
            ib_hca: "rocep1s0f0".to_string(),
            socket_ifname: "enp1s0f0np0".to_string(),
            ib_gid_index: 3,
            ib_disable: 0,
        }
    }

    // ── Deserialization tests ──────────────────────────────────

    #[test]
    fn test_recipe_deserialize_minimal() {
        let toml_str = r#"
name = "Minimal"
model = "test/model"
"#;
        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, 0);
        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.tensor_parallel_size, 1);
        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"));
        assert_eq!(config.nccl["prod"].ib_hca, "mlx5_0");
    }

    #[test]
    fn test_config_deserialize_with_host_key() {
        let toml_str = r#"
ssh_user = "admin"
ssh_key_path = "/path/to/key"

[[clusters]]
name = "secure"
master_ib_ip = "1.1.1.1"
nodes = [
    { name = "n1", ssh_ip = "2.2.2.2", ib_ip = "1.1.1.1", rank = 0, host_key = "SHA256:abc123" },
]
"#;
        let config: Config = toml::from_str(toml_str).unwrap();
        assert_eq!(
            config.clusters[0].nodes[0].host_key,
            Some("SHA256:abc123".to_string())
        );
    }

    // ── select_cluster tests ───────────────────────────────────

    #[test]
    fn test_select_cluster_by_name() {
        let mut config = Config {
            ssh_user: "user".to_string(),
            ssh_key_path: "/path/to/key".to_string(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        config.nccl.insert("test-cluster".to_string(), test_nccl());

        let (cluster, nccl) = select_cluster(&config, "test-cluster").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".to_string(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        config.nccl.insert("test-cluster".to_string(), test_nccl());

        let (cluster, nccl) = select_cluster(&config, "auto").unwrap();
        assert_eq!(cluster.name, "test-cluster");
        assert_eq!(nccl.ib_hca, "rocep1s0f0");
    }

    #[test]
    fn test_select_cluster_auto_no_match() {
        let config = Config {
            ssh_user: "user".to_string(),
            ssh_key_path: "/path/to/key".to_string(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        let result = select_cluster(&config, "auto");
        assert!(result.is_err());
    }

    #[test]
    fn test_select_cluster_not_found() {
        let config = Config {
            ssh_user: "user".to_string(),
            ssh_key_path: "/path/to/key".to_string(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        let result = select_cluster(&config, "nonexistent");
        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 config = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".to_string(),
            clusters: vec![test_cluster()],
            nccl: nccl_map,
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        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 = Config {
            ssh_user: "u".to_string(),
            ssh_key_path: "/k".to_string(),
            clusters: vec![test_cluster()],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
        };
        let result = select_cluster_by_name(&config, "no-such-cluster");
        assert!(result.is_err());
    }

    // ── Defaults ───────────────────────────────────────────────

    #[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");
    }
}