launcher/src/docker/compose.rs

//! Composed model directories: build a /tmp dir per-node with symlinks into
//! the HF cache (container-side absolute paths) + generated files + hardlinks
//! from node-local files, then bind-mount it into the container. Lets a
//! recipe combine multiple HF repos (e.g. a base checkpoint + an expert
//! overlay), and serve selected files from node-local storage (NVMe) instead
//! of the shared NAS (local_files), without mutating the HF cache.
//!
//! Symlinks point to `/root/.cache/huggingface/...` (the container-side HF
//! cache mount path), NOT the host-side path. They are dangling on the host
//! (the path `/root/.cache/...` doesn't exist there) but resolve correctly
//! inside the container, where the HF cache is mounted at that exact path.
//! This avoids the host↔container mount-root mismatch without requiring
//! relative-path gymnastics across different mount roots.

use super::shq;
use crate::config::{ComposedLocalFile, ComposedModel, ComposedSource, expand_tilde};
use crate::ssh::run_ssh;
use anyhow::{Context, Result};
use base64::Engine;
use base64::engine::general_purpose::STANDARD;
use std::path::Path;
use tracing::info;

/// Build a composed model directory on a node. Creates
/// `/tmp/sm121_<recipe>_<rank>/` with symlinks into the HF cache
/// (container-side absolute paths) + generated files. Returns
/// `(host_path, container_path)` for the bind-mount.
pub(crate) async fn build_composed_dir_on_node(
    config: &crate::config::Config,
    node: &crate::config::Node,
    recipe_name: &str,
    rank: u32,
    composed: &ComposedModel,
    hf_cache: &Path,
    dry_run: bool,
) -> Result<(String, String)> {
    let host_path = format!("/tmp/sm121_{}_{}", recipe_name, rank);
    let container_path = composed.mount.clone();

    if dry_run {
        println!(
            "{}",
            fluent_i18n::t!("composed-dir-dry-run", {
                "host" => host_path,
                "node" => node.name.clone(),
                "mount" => container_path,
                "sources" => composed.sources.len() as u32,
                "files" => composed.files.len() as u32,
                "local" => composed.local_files.len() as u32,
            })
        );
        return Ok((host_path, container_path));
    }

    let cache_str = hf_cache.display().to_string();

    // Create the composed dir.
    run_ssh(config, node, format!("mkdir -p {}", shq(&host_path)), false).await?;

    // Symlink each source's snapshot files into the composed dir.
    for source in &composed.sources {
        symlink_source(config, node, source, &host_path, &cache_str).await?;
    }

    // Hardlink local files (later wins over the symlinks above). These come
    // from node-local storage (e.g. an NVMe copy of a table shard) and must
    // live on the same filesystem as the composed dir (/tmp) to hardlink.
    for lf in &composed.local_files {
        link_local_file(config, node, lf, &host_path).await?;
    }

    // Write generated files (base64-encoded to avoid quoting issues).
    for file in &composed.files {
        let b64 = STANDARD.encode(file.content.as_bytes());
        run_ssh(
            config,
            node,
            format!(
                "printf '%s' '{}' | base64 -d > {}/{}",
                b64,
                shq(&host_path),
                shq(&file.name),
            ),
            false,
        )
        .await?;
    }

    info!(
        "{}",
        fluent_i18n::t!("composed-dir-built", {
            "host" => host_path,
            "node" => node.name.clone(),
            "sources" => composed.sources.len() as u32,
            "files" => composed.files.len() as u32,
            "local" => composed.local_files.len() as u32,
        })
    );

    Ok((host_path, container_path))
}

/// Resolve a source repo's snapshot on the node and symlink its files into
/// the composed dir. Symlink targets are container-side absolute paths
/// (`/root/.cache/huggingface/...`) so they resolve inside the container.
///
/// With no `glob`, the ENTIRE snapshot hierarchy is symlinked recursively
/// (every file in every subdir, real dirs) — so a later source can overlay
/// individual files at any path. With a `glob`, only matching top-level (or
/// explicit-subdir) entries are symlinked.
async fn symlink_source(
    config: &crate::config::Config,
    node: &crate::config::Node,
    source: &ComposedSource,
    host_path: &str,
    cache_str: &str,
) -> Result<()> {
    let model_dir = crate::config::model_cache_dir_name(&source.repo);
    let refs_path = format!("{}/hub/{}/refs/main", cache_str, model_dir);
    let glob = source.glob.as_deref();

    // Validate the glob: only allow glob-safe characters to prevent shell
    // injection (the glob is inserted raw into the shell `for` statement).
    if let Some(g) = glob
        && !g
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || "*?._-/".contains(c))
    {
        anyhow::bail!(
            "composed_model source '{}': glob '{}' contains invalid characters \
             (allowed: alphanumeric, * ? . _ - /)",
            source.repo,
            g
        );
    }

    // Build a shell `case` clause to skip excluded paths. Each entry is
    // matched against the file's path relative to the snapshot root (e.g.
    // `L5.safetensors`, or `subdir/file.bin`). Same character set as glob.
    let exclude_clause = match &source.exclude {
        None => String::new(),
        Some(entries) if entries.is_empty() => String::new(),
        Some(entries) => {
            for e in entries {
                if e.is_empty()
                    || !e
                        .chars()
                        .all(|c| c.is_ascii_alphanumeric() || "*?._-/".contains(c))
                {
                    anyhow::bail!(
                        "composed_model source '{}': exclude entry '{}' is invalid \
                         (allowed: alphanumeric, * ? . _ - /)",
                        source.repo,
                        e
                    );
                }
            }
            format!("case \"$rel\" in {}) continue ;; esac; ", entries.join("|"))
        }
    };

    // Enumerate: no glob → recursive `find` (whole hierarchy, every file in
    // every subdir); glob → shell glob match (one level, `/` allowed for a
    // specific subdir). Both feed the same per-file body below.
    let enumerate: String = match glob {
        None => String::from(
            "find \"$snap\" -mindepth 1 \\( -type l -o -type f \\) | while IFS= read -r f; do",
        ),
        Some(g) => format!("for f in \"$snap\"/{}; do", g),
    };

    // One SSH round-trip: resolve refs/main, then symlink each matching file.
    // The symlink TARGET is the CONTAINER-side path (/root/.cache/huggingface/...)
    // not the host-side path — the link is dangling on the host but resolves
    // inside the container where the HF cache is mounted at that exact path.
    // `ln -snf`: -s symlink, -n no-deref (treat existing symlink as a file,
    // don't dive into the dir it points to), -f force (overwrite earlier).
    //
    // Paths are assembled via `printf` (not embedded in double-quoted
    // assignments) so shq's single-quote wrapping doesn't become literal
    // characters in the path — same pattern as resolve_path_ref_on_node.
    //
    // Relative paths are preserved: a glob like `grids/*.json` creates
    // `$host/grids/L10.json` (with the subdir), and the symlink target is
    // `$csnap/grids/L10.json` (matching structure). With no glob, `find`
    // recurses so `mtp-draft/foo.bin` etc. are symlinked individually.
    // `mkdir -p` creates parent dirs as needed (real dirs, not symlinks —
    // so later sources can overlay files at any path).
    let cmd = format!(
        "ref=\"$(cat {refs_q} 2>/dev/null)\" && test -n \"$ref\" && \
         snap=\"$(printf '%s/hub/%s/snapshots/%s' {cache_q} {model_q} \"$ref\")\" && \
         test -d \"$snap\" && \
         csnap=\"$(printf '/root/.cache/huggingface/hub/%s/snapshots/%s' {model_q} \"$ref\")\" && \
         {enumerate} \
           [ -e \"$f\" ] || [ -L \"$f\" ] || continue; \
           rel=\"${{f#\"$snap\"/}}\"; \
           {exclude_clause}\
           parent=\"$(dirname {host_q}/\"$rel\")\"; \
           mkdir -p \"$parent\" && \
           ln -snf \"$csnap/$rel\" {host_q}/\"$rel\"; \
         done",
        refs_q = shq(&refs_path),
        cache_q = shq(cache_str),
        model_q = shq(&model_dir),
        host_q = shq(host_path),
        enumerate = enumerate,
        exclude_clause = exclude_clause,
    );

    run_ssh(config, node, cmd, false).await.with_context(|| {
        format!(
            "composed_model source '{}' not found in HF cache on '{}' \
             — run `huggingface-cli download {}` on that node",
            source.repo, node.name, source.repo
        )
    })?;
    Ok(())
}

/// Hardlink a node-local file into the composed dir. Runs after the repo
/// symlink pass, so a local file overwrites the repo's symlink at the same
/// name (local wins — that is the point: e.g. an NVMe table shard instead of
/// the shared-NAS version). Hardlinked, NOT copied — a sparse local copy
/// stays sparse and costs zero extra disk — so the host path must be on the
/// same filesystem as the composed dir (/tmp).
async fn link_local_file(
    config: &crate::config::Config,
    node: &crate::config::Node,
    lf: &ComposedLocalFile,
    host_path: &str,
) -> Result<()> {
    let host = expand_tilde(&lf.host, &config.ssh_user);
    let name = lf
        .name
        .as_deref()
        .filter(|n| !n.is_empty())
        .or_else(|| host.rsplit('/').next())
        .filter(|s| !s.is_empty())
        .ok_or_else(|| anyhow::anyhow!("composed_model local_file with empty host path"))?;
    if name.starts_with('/')
        || name.contains("..")
        || !name
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || "._-".contains(c))
    {
        anyhow::bail!(
            "composed_model local_file '{}': name '{}' is invalid \
             (allowed: alphanumeric . _ - ; flat name, no directories)",
            host,
            name
        );
    }
    // Existence check + hardlink in one round-trip; `ln -f` replaces an
    // earlier symlink at the same path.
    let cmd = format!(
        "test -e {h} && ln -f {h} {d}/{n}",
        h = shq(&host),
        d = shq(host_path),
        n = shq(name),
    );
    run_ssh(config, node, cmd, false).await.with_context(|| {
        format!(
            "composed_model local_file: '{}' not found on '{}' — a per-node \
             local path is required (and it must share a filesystem with /tmp to hardlink)",
            host, node.name
        )
    })?;
    Ok(())
}