launcher/src/config/types.rs

use super::ExtraEntity;
use crate::pre_serve::BeforeServeTable;
use serde::Deserialize;
use std::collections::HashMap;

#[derive(Clone, Debug, Deserialize)]
pub struct Node {
    pub name: String,
    pub ssh_ip: String,
    /// InfiniBand IP; optional (falls back to `ssh_ip`).
    #[serde(default)]
    pub ib_ip: Option<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,
    /// RoCE GID index: absent = auto-detect per node (`NCCL_IB_GID_INDEX`
    /// resolved from the RoCE-v2 IPv4 GID at serve time); Some(n) = pinned.
    #[serde(default)]
    pub ib_gid_index: Option<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,
    /// Master IB IP; optional (falls back to rank-0 node's effective IP).
    #[serde(default)]
    pub master_ib_ip: Option<String>,
    pub nodes: Vec<Node>,
    #[serde(default)]
    pub extra_entities: Vec<ExtraEntity>,
}

#[derive(Clone, Deserialize)]
pub struct Config {
    /// SSH user; may be empty — a `cluster_sources` file can provide it.
    #[serde(default)]
    pub ssh_user: String,
    #[serde(default)]
    pub ssh_key_path: std::path::PathBuf,
    #[serde(default)]
    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>,
    /// RouterOS devices for graceful `power off` (see [`MikrotikConfig`]).
    #[serde(default)]
    pub mikrotik: Option<MikrotikConfig>,
    #[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,
    /// Mount-alias registry: recipe `[[extra_mounts]]` `alias` resolves here; the
    /// alias wins over the recipe's `host` default. Keeps recipes portable.
    #[serde(default)]
    pub mount_aliases: HashMap<String, String>,
    /// Extra recipe dirs (each `<name>/recipe.toml`) searched after `./recipes` (`~` expanded).
    #[serde(default)]
    pub recipe_sources: Vec<String>,
    /// Ordered git recipe repos (first entry = highest priority on a
    /// recipe-name collision). Cloned/pulled into a local cache by
    /// `crate::recipe_repos` and appended to the recipe search path in the
    /// same order; archive state lives in the separate `repos.json`
    /// registry, not here.
    #[serde(default)]
    pub recipe_repos: Vec<RecipeRepo>,
    /// Extra image-build dirs searched after `./docker-images`: repo roots
    /// (with `docker-images/`) or dirs of image dirs; `~` expanded.
    #[serde(default)]
    pub docker_image_sources: Vec<String>,
    /// Extra cluster def files (`[[clusters]]`, `[nccl.<name>]`); this file wins.
    #[serde(default)]
    pub cluster_sources: Vec<String>,
    /// Optional image-registry push target: after a successful local build
    /// the image is pushed as `{base}/{image-name}:{tag}` on the build node.
    #[serde(default)]
    pub registry: Option<RegistryConfig>,
}

/// `[registry]` — where freshly built images are pushed (the build node
/// must already be authenticated: `docker login <host>` once on the node).
#[derive(Clone, Debug, Deserialize)]
pub struct RegistryConfig {
    /// Repository base, e.g. `ghcr.io/kilork` → push ref
    /// `ghcr.io/kilork/<image-name>:<tag>`.
    pub base: String,
}

/// One `[[recipe_repos]]` entry: a git repo holding `recipes/` (and
/// `docker-images/` build contexts), cloned/pulled by the launcher into a
/// local cache and merged into the recipe set. List order is priority:
/// a recipe name present in an earlier repo shadows the later one.
#[derive(Clone, Debug, Deserialize)]
pub struct RecipeRepo {
    /// Git URL (https/ssh/`file://`).
    pub url: String,
    /// Optional private-repo credential (deploy key / PAT); embedded in the
    /// clone URL only inside the local cache, never committed.
    #[serde(default)]
    pub token: Option<String>,
}

#[derive(Clone, Deserialize)]
pub struct HomeAssistantConfig {
    pub url: String,
    pub token: String,
    #[serde(default = "default_shutdown_delay")]
    pub shutdown_delay_seconds: u64,
}

/// `[mikrotik]` — RouterOS devices shut down gracefully over the RouterOS
/// API (default port 8728) during `power off`. A device's `name` is
/// referenced verbatim from a cluster's `extra_entities` (or a node's
/// `entity`); names not present here keep going through Home Assistant.
#[derive(Clone, Deserialize)]
pub struct MikrotikConfig {
    pub devices: Vec<MikrotikDevice>,
}

#[derive(Clone, Deserialize)]
pub struct MikrotikDevice {
    /// Entity-style name used in cluster `extra_entities`.
    pub name: String,
    pub address: String,
    pub username: String,
    pub password: String,
    #[serde(default = "default_mikrotik_port")]
    pub api_port: u16,
    /// Optional HA relay entity that powers the router on (the API cannot
    /// wake an off device); without it, `power on` tells the operator to
    /// plug it in manually.
    #[serde(default)]
    pub on_entity: Option<String>,
}

fn default_mikrotik_port() -> u16 {
    8728
}

fn default_shutdown_delay() -> u64 {
    30
}
fn default_init_port() -> u16 {
    25000
}
pub(crate) fn default_port() -> u16 {
    30000
}
pub(crate) fn default_engine() -> String {
    "vllm".to_string()
}
pub(crate) fn default_shm_size() -> String {
    String::new()
}
pub(crate) fn default_tensor_parallel_size() -> u32 {
    1
}
pub(crate) fn default_sync_image() -> bool {
    true
}
fn default_drop_caches() -> bool {
    true
}

/// Recipe-level build config: `versions` — a map of `versions.env` pin
/// overrides merged over the image's pins before build (after the profile).
#[derive(Clone, Debug, Default, Deserialize)]
pub struct RecipeBuild {
    /// Pins merged over the image's `versions.env` (recipe wins). Keys are
    /// `versions.env` variable names; values the exact version/commit.
    #[serde(default)]
    pub versions: HashMap<String, String>,
}

/// A composed model directory: the launcher builds a `/tmp` dir per-node
/// with symlinks into the container-side HF cache (`/root/.cache/...` —
/// dangling on the host, resolved by the HF-cache mount in-container) +
/// generated files, then bind-mounts it — combining multiple HF repos (a
/// base checkpoint + an expert overlay) without mutating the HF cache.
#[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>,
    /// Real files HARDLINKED from a per-node local host path into the
    /// composed dir (e.g. a local NVMe copy of a table shard instead of the
    /// shared-NAS version the repo symlinks would point at). Later local
    /// files win over earlier sources. Must be on the node's local disk
    /// (same filesystem as /tmp) — hardlinked, not copied.
    #[serde(default)]
    pub local_files: Vec<ComposedLocalFile>,
}

/// A local file hardlinked into a [`ComposedModel`] dir (per node).
/// Used to serve a shard from node-local storage (NVMe) while the rest of
/// the model stays symlinked into the shared HF cache.
#[derive(Clone, Debug, Default, Deserialize)]
pub struct ComposedLocalFile {
    /// Host path on each node (leading `~`/`~/` expanded to the SSH user's
    /// home). Must exist on every launch node.
    pub host: String,
    /// Optional path inside the composed dir; defaults to the host basename.
    #[serde(default)]
    pub name: Option<String>,
}

/// A source repo for [`ComposedModel`]. The repo's HF cache snapshot is
/// symlinked into the composed dir, optionally filtered by `glob` and
/// `exclude`.
#[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>,
    /// Optional list of path patterns to skip (matched against the file's
    /// path relative to the snapshot root, via shell `case`). Lets a later
    /// source overlay skip specific files so an earlier source's version
    /// survives — e.g. `exclude = ["L5.safetensors"]` keeps the base
    /// checkpoint's `L5.safetensors` instead of the overlay's. Same allowed
    /// characters as `glob` (alphanumeric + `*?._-/`).
    #[serde(default)]
    pub exclude: Option<Vec<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`); else the
    /// fallback `host` path; both mutually exclusive with `repo`.
    #[serde(default)]
    pub alias: Option<String>,
    #[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))
}

/// Deployment target — how the recipe is launched on each node.
/// `docker` (default): `docker run` with the recipe's `entrypoint`.
/// `tmux`: a detached `tmux` session running the engine directly (no Docker).
/// Docker-only features (mods, composed_model, image sync) are rejected
/// when `deploy = "tmux"`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum DeployTarget {
    #[default]
    Docker,
    Tmux,
    /// Bare SSH command on the head node (rank 0) driving an engine's own
    /// cluster manager (DGPP's `dgpp-cluster up/down`). Head-coordinated:
    /// the manager brings up the peers itself.
    Native,
}

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