launcher/src/config/types.rs

use serde::Deserialize;
use std::collections::HashMap;

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

#[derive(Clone, Debug, 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, Debug, Deserialize)]
pub struct NcclConfig {
    pub ib_hca: String,
    pub socket_ifname: String,
    pub ib_gid_index: u32,
    pub ib_disable: u32,
    #[serde(default)]
    pub gloo_socket_ifname: Option<String>,
    #[serde(default)]
    pub ucx_net_devices: Option<String>,
    #[serde(default)]
    pub nccl_cross_nic: Option<u32>,
    #[serde(default)]
    pub nccl_proto: Option<String>,
    #[serde(default)]
    pub nccl_p2p_level: Option<String>,
}

#[derive(Clone, Debug, Deserialize)]
pub struct Cluster {
    pub name: String,
    pub master_ib_ip: String,
    pub nodes: Vec<Node>,
    #[serde(default)]
    pub extra_entities: Vec<String>,
}

#[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>,
    #[serde(default = "default_init_port")]
    pub init_port: u16,
    #[serde(default)]
    pub before_serve: HashMap<String, BeforeServeTable>,
    #[serde(default)]
    pub exec_aliases: HashMap<String, String>,
    #[serde(default)]
    pub home_assistant: Option<HomeAssistantConfig>,
    #[serde(default)]
    pub global_mods_dir: Option<std::path::PathBuf>,
    /// Drop page cache before starting containers.
    #[serde(default = "default_drop_caches")]
    pub drop_caches: bool,
    /// Named mount-alias registry — host paths referenced by name from recipe
    /// `[[extra_mounts]]` entries. Lets recipes stay portable (reference an
    /// alias like `jit_cache`) while the operator sets the real path here.
    /// A recipe may also carry a `host` default; the alias, when present,
    /// overrides it.
    #[serde(default)]
    pub mount_aliases: HashMap<String, String>,
}

#[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
}
fn default_engine() -> String {
    "vllm".to_string()
}
fn default_shm_size() -> String {
    String::new()
}
fn default_tensor_parallel_size() -> u32 {
    1
}
fn default_sync_image() -> bool {
    true
}
fn default_drop_caches() -> bool {
    true
}

#[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: RecipeArgs,
    #[serde(default)]
    pub env: HashMap<String, String>,
    #[serde(default)]
    pub cluster: Option<String>,
    #[serde(default)]
    pub entrypoint: Vec<String>,
    pub distributed_backend: Option<DistributedBackend>,
    #[serde(default = "default_shm_size")]
    pub shm_size: String,
    #[serde(default)]
    pub before_serve: Vec<PreServeCommand>,
    /// Commands to run on each node **after** stopping containers
    /// (`sm121 stop`). Same shape as [`before_serve`]: raw shell commands
    /// (`Bare`/`Configured`) or built-in script aliases (`Script`). Runs
    /// per-node in parallel, sequentially within a node. Typical use:
    /// `after_stop = [{ script = "restore_dgx_services" }]` to bring back
    /// services that `before_serve` stopped.
    #[serde(default)]
    pub after_stop: Vec<PreServeCommand>,
    /// Whether this recipe should always be served single-node on the first
    /// free node found (solo mode). When set, `sm121 serve` ignores
    /// `required_nodes` / multi-node cluster selection and launches exactly
    /// one container with no distributed flags. The CLI `--solo`/`--no-solo`
    /// flag overrides this per-invocation.
    #[serde(default)]
    pub solo: Option<bool>,
    #[serde(default)]
    pub mods: ModsTable,
    #[serde(default = "default_sync_image")]
    pub sync_image: bool,
    /// Build directory for docker-images/ (default: derived from recipe name).
    #[serde(default)]
    pub docker_image_dir: Option<String>,
    /// Recipe-driven overrides for the image's `versions.env`. Each entry is
    /// merged over the image's pins at build time (recipe wins) and folded
    /// into the build checksum, so changing a pin here triggers a rebuild
    /// without editing `docker-images/<dir>/versions.env`. Keys must match
    /// the `versions.env` names the image's `Dockerfile` declares as `ARG`.
    #[serde(default)]
    pub build: RecipeBuild,
    /// When set, the launcher builds a /tmp composed-model dir per-node
    /// (symlinks into the HF cache + generated files), bind-mounts it at
    /// `mount`, and rewrites `model` to that path. Lets a recipe combine
    /// multiple HF repos (e.g. a base checkpoint + an expert overlay)
    /// without mutating the HF cache. See [`ComposedModel`].
    #[serde(default)]
    pub composed_model: Option<ComposedModel>,
    /// Extra bind-mounts to add to `docker run` (e.g. persistent JIT cache).
    /// Each entry resolves its host path from `Config.mount_aliases` when
    /// `alias` is set and present there (alias wins); otherwise falls back to
    /// `host` (with `~` expansion). At least one of `alias`/`host` is required.
    #[serde(default)]
    pub extra_mounts: Vec<ExtraMount>,
}

/// Recipe-level build configuration. Currently only `versions` (a map of
/// `versions.env` pin overrides) is supported; the launcher merges these
/// over the image's own `versions.env` before building.
#[derive(Clone, Debug, Default, Deserialize)]
pub struct RecipeBuild {
    /// Pins merged over the image's `versions.env` (recipe wins). Keys are
    /// `versions.env` variable names (e.g. `VLLM_COMMIT`); values are the
    /// exact version/tag/commit. See [`Recipe::build`].
    #[serde(default)]
    pub versions: HashMap<String, String>,
}

/// A composed model directory: the launcher builds a `/tmp` dir per-node
/// with symlinks into the HF cache (container-side absolute paths) +
/// generated files, then bind-mounts it into the container. Lets a recipe
/// combine multiple HF repos (e.g. a base checkpoint + an expert overlay)
/// without mutating the HF cache.
///
/// Symlinks point to `/root/.cache/huggingface/...` (the container-side HF
/// cache mount), not the host-side path. They are dangling on the host but
/// resolve correctly inside the container where the HF cache is mounted at
/// that exact path — avoiding the host↔container mount-root mismatch.
#[derive(Clone, Debug, Default, Deserialize)]
pub struct ComposedModel {
    /// Container-side mount target for the composed dir (e.g. `/root/model`).
    /// `recipe.model` is rewritten to this path at serve time.
    pub mount: String,
    /// Source repos to symlink into the composed dir. Later entries win
    /// (overwrite earlier symlinks), so order matters for overlays.
    #[serde(default)]
    pub sources: Vec<ComposedSource>,
    /// Real files (not symlinks) to write into the composed dir.
    #[serde(default)]
    pub files: Vec<ComposedFile>,
}

/// A source repo for [`ComposedModel`]. The repo's HF cache snapshot is
/// symlinked into the composed dir, optionally filtered by `glob`.
#[derive(Clone, Debug, Default, Deserialize)]
pub struct ComposedSource {
    /// HF Hub repo id (e.g. `aidendle94/GLM-5.2-MXFP4-Experts-GPTQ`).
    pub repo: String,
    /// Optional glob to filter which snapshot files are symlinked (e.g.
    /// `L*.safetensors` — only expert weight files). Empty/omitted = all
    /// files. Only glob-safe characters are allowed (alphanumeric + `*?._-`).
    #[serde(default)]
    pub glob: Option<String>,
}

/// A generated file written into the composed dir (not a symlink).
#[derive(Clone, Debug, Default, Deserialize)]
pub struct ComposedFile {
    /// Filename in the composed dir (e.g. `nf4_layers.json`).
    pub name: String,
    /// Raw file content (written verbatim via base64).
    pub content: String,
}

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".to_string()
        } else {
            format!("/home/{}", self.ssh_user)
        };
        std::path::PathBuf::from(format!("{home}/.cache/huggingface"))
    }
}

/// A recipe-level extra bind-mount for `docker run`.
///
/// Two mutually exclusive kinds:
/// - **Static** (`alias`/`host`): host path is known at recipe-write time.
///   `alias` resolves against `Config.mount_aliases` and, when present there,
///   overrides `host`. If the alias is absent from config, `host` is used as
///   the default. `~` (and `~/`) at the start of the resolved path is expanded
///   to the remote SSH user's home. At least one of `alias` or `host` must be
///   set.
/// - **Model path reference** (`repo` + optional `subfolder`): NOT a bind-mount.
///   The path is resolved **at serve time, per node**, from the node's HF
///   cache: `<hf_cache>/hub/<model_cache_dir_name(repo)>/snapshots/<refs/main>/<subfolder>`.
///   `refs/main` always points at the latest commit, so the path survives
///   `huggingface-cli download` updates — no fragile symlink. The recipe
///   references this path via a `{{<container>}}` placeholder in args/env
///   (e.g. `--speculative-config`), which the launcher substitutes per-node
///   before building the docker command. The full HF cache is already mounted
///   at `/root/.cache/huggingface`, so the resolved path's symlinks resolve
///   correctly (a separate bind-mount would break them).
///   If the subfolder is absent on a node, serve aborts with a hint to
///   `huggingface-cli download <repo>` (the launcher never downloads).
#[derive(Clone, Debug, Default, Deserialize)]
pub struct ExtraMount {
    /// Name resolved via `Config.mount_aliases`; overrides `host` when present.
    /// Mutually exclusive with `repo`.
    #[serde(default)]
    pub alias: Option<String>,
    /// Fallback host path (used when the alias is not in `mount_aliases`).
    /// Mutually exclusive with `repo`.
    #[serde(default)]
    pub host: Option<String>,
    /// HF Hub repo id, e.g. `aidendle94/GLM-5.2-Hybrid-FP8-MXFP4`. When set,
    /// the path is resolved at serve time from the node's HF cache and
    /// substituted into args/env via `{{<container>}}` (so `alias`/`host`
    /// must be empty).
    #[serde(default)]
    pub repo: Option<String>,
    /// Subfolder within the repo snapshot (e.g. `mtp-draft`).
    /// Empty/omitted = the whole snapshot dir.
    #[serde(default)]
    pub subfolder: Option<String>,
    /// For static mounts: the container-side mount target. For repo-based
    /// path references: the placeholder name (e.g. `mtp_draft`) that gets
    /// substituted as `{{mtp_draft}}` in args/env.
    pub container: String,
    /// Mount read-only (`:ro`). Only for static mounts.
    #[serde(default)]
    pub read_only: bool,
    /// If set, inject `-e <env>=<container path>` so recipes don't hardcode
    /// the path in two places (mount + env). Only for static mounts.
    #[serde(default)]
    pub env: Option<String>,
}

/// Directory name a repo id maps to inside `<hf_cache>/hub/`.
/// `org/name` -> `models--org--name` (HF convention: prefix `models--`,
/// replace `/` with `--`). Pure & deterministic; used by the model-subfolder
/// mount resolver.
pub fn model_cache_dir_name(repo: &str) -> String {
    format!("models--{}", repo.replace('/', "--"))
}

/// Expand a leading `~` / `~/` to the remote SSH user's home directory.
/// Paths without a leading `~` are returned unchanged. Used for host paths
/// before they are sent to `docker run -v`.
pub fn expand_tilde(path: &str, ssh_user: &str) -> String {
    let home = if ssh_user == "root" {
        "/root".to_string()
    } else {
        format!("/home/{}", ssh_user)
    };
    if let Some(rest) = path.strip_prefix("~/") {
        format!("{}/{}", home, rest)
    } else if path == "~" {
        home
    } else {
        path.to_string()
    }
}

/// Resolve a recipe `ExtraMount`'s host path: the config alias wins over the
/// recipe `host` default; if neither is set, error. `~` is expanded.
pub fn resolve_mount_host(
    alias: Option<&str>,
    host: Option<&str>,
    mount_aliases: &HashMap<String, String>,
    ssh_user: &str,
    container: &str,
) -> anyhow::Result<String> {
    let raw = alias
        .and_then(|a| {
            if a.trim().is_empty() {
                None
            } else {
                mount_aliases.get(a).map(String::as_str)
            }
        })
        .or_else(|| host.filter(|h| !h.trim().is_empty()))
        .ok_or_else(|| {
            anyhow::anyhow!(
                "extra_mount '{}' has neither a resolvable alias nor a host path",
                container
            )
        })?;
    Ok(expand_tilde(raw, ssh_user))
}

#[cfg(test)]
mod tests {
    use super::{expand_tilde, model_cache_dir_name, resolve_mount_host};
    use std::collections::HashMap;

    #[test]
    fn model_cache_dir_name_maps_repo_id() {
        assert_eq!(
            model_cache_dir_name("aidendle94/GLM-5.2-Hybrid-FP8-MXFP4"),
            "models--aidendle94--GLM-5.2-Hybrid-FP8-MXFP4"
        );
        assert_eq!(model_cache_dir_name("org/name"), "models--org--name");
        assert_eq!(model_cache_dir_name("bare"), "models--bare");
    }

    #[test]
    fn expand_tilde_user_and_root() {
        assert_eq!(expand_tilde("~/.cache", "kilork"), "/home/kilork/.cache");
        assert_eq!(expand_tilde("~/", "kilork"), "/home/kilork/");
        assert_eq!(expand_tilde("~", "root"), "/root");
        assert_eq!(expand_tilde("/abs/path", "kilork"), "/abs/path");
        assert_eq!(expand_tilde("rel/path", "kilork"), "rel/path");
    }

    #[test]
    fn resolve_mount_alias_wins_over_host() {
        let aliases = HashMap::from([("jit_cache".to_string(), "/mnt/nas/jit".to_string())]);
        // alias present in config → alias wins, host ignored
        let r = resolve_mount_host(Some("jit_cache"), Some("~/.cache/jit"), &aliases, "u", "/c");
        assert_eq!(r.unwrap(), "/mnt/nas/jit");
    }

    #[test]
    fn resolve_mount_falls_back_to_host_with_tilde() {
        let aliases = HashMap::new();
        // alias not in config → host default used, with ~ expansion
        let r = resolve_mount_host(
            Some("jit_cache"),
            Some("~/.cache/jit"),
            &aliases,
            "kilork",
            "/c",
        );
        assert_eq!(r.unwrap(), "/home/kilork/.cache/jit");
    }

    #[test]
    fn resolve_mount_host_only() {
        let aliases = HashMap::new();
        let r = resolve_mount_host(None, Some("/var/cache"), &aliases, "u", "/c");
        assert_eq!(r.unwrap(), "/var/cache");
    }

    #[test]
    fn resolve_mount_neither_errors() {
        let aliases = HashMap::new();
        let r = resolve_mount_host(None, None, &aliases, "u", "/c");
        assert!(r.is_err());
    }
}