launcher/src/staging.rs

//! Declarative data staging — see specs/tech-architecture/staging-feature.md.
//!
//! A recipe declares `[[staging]]` blocks; `serve` executes them on every
//! launch node before containers start, then (for hub-mirrors) serves from
//! the local mirror instead of the shared HF cache. No manual per-node
//! preparation is required — the operator runs plain `sm121 serve`.

pub mod mirror;
pub mod rows;
pub mod view;

use crate::config::{Config, Node, Recipe, StagingKind, expand_tilde};
use anyhow::{Result, bail};
use std::path::{Path, PathBuf};
use tracing::{info, warn};

/// How aggressively to stage (`--stage=auto|skip|force`).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StageMode {
    /// Verify staged content against the ledger, copy what is new (default).
    Auto,
    /// Assume staging is already done; only print the plan (CI / known-good).
    Skip,
    /// Re-verify everything regardless of the ledger.
    Force,
}

impl StageMode {
    pub fn parse(s: &str) -> Result<Self> {
        match s.trim().to_ascii_lowercase().as_str() {
            "" | "auto" => Ok(StageMode::Auto),
            "skip" => Ok(StageMode::Skip),
            "force" => Ok(StageMode::Force),
            other => bail!("unknown --stage value '{other}' (auto|skip|force)"),
        }
    }
}

/// The first `local-view` block's mount: (host view dir, container mount,
/// repo). `serve` appends the mount and rewrites `model` to
/// `<container>/<repo>`, so the engine reads the staged files from local NVMe
/// and everything else through the shared cache's own mount.
pub fn local_view_mount(config: &Config, recipe: &Recipe) -> Option<(PathBuf, String, String)> {
    for (i, s) in recipe.staging.iter().enumerate() {
        if s.kind != StagingKind::LocalView {
            continue;
        }
        let id = s.ledger_id(i);
        let host = PathBuf::from(view::dst_for(s, &id, &config.ssh_user));
        let repo = view::repo_for(s, &recipe.model);
        return Some((host, view::CONTAINER_ROOT.to_string(), repo));
    }
    None
}

/// Free-space preflight for every staging block whose footprint is known.
/// Runs BEFORE any copy and before any container starts: a node that cannot
/// hold the staged files fails the serve with the node, the size and the
/// target directory, instead of dying halfway through a 200 GB copy.
pub async fn preflight(
    config: &Config,
    nodes: &[Node],
    recipe: &Recipe,
    shared_source: &Path,
    dry_run: bool,
) -> Result<()> {
    let mut needs: Vec<view::Need> = Vec::new();
    for (i, block) in recipe.staging.iter().enumerate() {
        if block.kind != StagingKind::LocalView || dry_run {
            continue;
        }
        needs.extend(view::measure(config, nodes, block, i, &recipe.model, shared_source).await?);
    }
    for n in &needs {
        info!("staging space: {}", n.line());
    }
    if let Some(msg) = view::shortfall(&needs) {
        bail!(msg);
    }
    Ok(())
}

/// True when the recipe stages at least one hub-mirror (the HF mount source
/// gets swapped to the mirror's dst in that case).
pub fn has_hub_mirror(recipe: &Recipe) -> bool {
    recipe
        .staging
        .iter()
        .any(|s| s.kind == StagingKind::HubMirror)
}

/// The hub-mirror dst (tilde-expanded) serve must mount instead of the
/// shared HF cache. Multiple hub-mirror blocks must share one dst; blocks
/// without an explicit `dst` use `~/.cache/hf-staging/<id>` — identical to
/// the default `stage_node` picks.
pub fn hub_mirror_dst(config: &Config, recipe: &Recipe) -> Option<PathBuf> {
    let mut out: Option<PathBuf> = None;
    for (i, s) in recipe.staging.iter().enumerate() {
        if s.kind != StagingKind::HubMirror {
            continue;
        }
        let d = s
            .dst
            .clone()
            .unwrap_or_else(|| format!("~/.cache/hf-staging/{}", s.ledger_id(i)));
        let p = PathBuf::from(expand_tilde(&d, &config.ssh_user));
        if let Some(prev) = &out
            && *prev != p
        {
            warn!(
                "multiple hub-mirror staging blocks with different dst — mounting {}",
                prev.display()
            );
        }
        if out.is_none() {
            out = Some(p);
        }
    }
    out
}

/// Serve-side staging pass: validate the CLI knobs, run the recipe's
/// `[[staging]]` blocks on every launch node, and return the HF mount
/// source for this serve: `(mirror, cache_dir)`. A hub-mirror block takes
/// over the mount (`mirror=None`, `cache_dir=dst`); otherwise the standard
/// `mirror::resolve` applies (`--mirror` / `--hf-cache`).
pub async fn hf_mount(
    config: &Config,
    nodes: &[Node],
    recipe: &Recipe,
    mirror_cli: Option<&str>,
    hf_cache_cli: Option<&str>,
    stage_cli: Option<&str>,
    dry_run: bool,
) -> Result<(Option<PathBuf>, PathBuf)> {
    if mirror_cli.is_some() && hf_cache_cli.is_some() {
        anyhow::bail!(
            "--hf-cache and --mirror are mutually exclusive (both pick the HF mount source)"
        );
    }
    let staged = has_hub_mirror(recipe);
    if staged && mirror_cli.is_some() {
        anyhow::bail!(
            "--mirror and a recipe [[staging]] hub-mirror cannot be combined — declare the mirror in the recipe"
        );
    }
    if staged && hf_cache_cli.is_some() {
        anyhow::bail!(
            "--hf-cache is ignored for a hub-mirror staging recipe (dst wins); drop one of them"
        );
    }
    if !recipe.staging.is_empty() {
        let mode = stage_cli
            .map(StageMode::parse)
            .transpose()?
            .unwrap_or(StageMode::Auto);
        let shared_source = match hf_cache_cli {
            Some(p) => PathBuf::from(expand_tilde(p, &config.ssh_user)),
            None => config.resolved_hf_cache_dir(),
        };
        preflight(config, nodes, recipe, &shared_source, dry_run).await?;
        run(config, nodes, recipe, &shared_source, mode, dry_run).await?;
    }
    if staged {
        let dst = hub_mirror_dst(config, recipe)
            .ok_or_else(|| anyhow::anyhow!("hub-mirror staging block has no dst"))?;
        return Ok((None, dst));
    }
    let (mirror, cache_dir) = crate::docker::mirror::resolve(config, mirror_cli, hf_cache_cli);
    if let Some(ref m) = mirror {
        info!(
            "serving from local HF mirror {} (models copied per node on first use)",
            m.display()
        );
    }
    Ok((mirror, cache_dir))
}

/// Run every `[[staging]]` block on every launch node (blocks per node run
/// in sequence, nodes in parallel). `shared_source` is the host-side HF cache
/// a hub-mirror defaults to copy FROM (the `--hf-cache` / config dir).
pub async fn run(
    config: &Config,
    nodes: &[Node],
    recipe: &Recipe,
    shared_source: &Path,
    mode: StageMode,
    dry_run: bool,
) -> Result<()> {
    if recipe.staging.is_empty() {
        return Ok(());
    }
    if mode == StageMode::Skip {
        info!("staging: --stage=skip — assuming staged");
        return Ok(());
    }
    // Validate ids before touching any node.
    let mut seen: Vec<String> = Vec::new();
    for (i, s) in recipe.staging.iter().enumerate() {
        let id = s.ledger_id(i);
        if seen.contains(&id) {
            bail!("duplicate [[staging]] id '{id}' in recipe");
        }
        seen.push(id);
    }

    let mut handles: Vec<tokio::task::JoinHandle<Result<()>>> = Vec::new();
    // sparse-rows splits the table across the engine's tensor-parallel
    // ranks: recipe TP wins, else one rank per launch node.
    let tp = if recipe.tensor_parallel_size > 0 {
        recipe.tensor_parallel_size
    } else {
        nodes.len() as u32
    };
    for node in nodes {
        let config = config.clone();
        let node = node.clone();
        let recipe = recipe.clone();
        let source = shared_source.to_path_buf();
        handles.push(crate::locale::spawn_with_locale(async move {
            let ctx = mirror::Ctx {
                config: &config,
                recipe: &recipe,
                shared_source: &source,
                mode,
                dry_run,
            };
            for (i, block) in recipe.staging.iter().enumerate() {
                match block.kind {
                    StagingKind::HubMirror => mirror::stage_node(&ctx, &node, block, i).await?,
                    StagingKind::LocalView => {
                        let vctx = view::StageCtx {
                            config: &config,
                            recipe: &recipe,
                            shared_source: &source,
                            force: mode == StageMode::Force,
                            dry_run,
                        };
                        view::stage_node(&vctx, &node, block, i).await?
                    }
                    StagingKind::SparseRows => {
                        let rctx = rows::Ctx {
                            config: &config,
                            shared_source: &source,
                            mode,
                            dry_run,
                            rank: node.rank,
                            tp,
                        };
                        rows::stage_node(&rctx, &node, block, i).await?
                    }
                }
            }
            Ok(())
        }));
    }
    for h in handles {
        h.await??;
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::Node;
    use crate::config::{DeployTarget, DistributedBackend, Staging};
    use crate::pre_serve::{RecipeArgs, VenvSetup};
    use std::collections::HashMap;

    fn cfg() -> Config {
        Config {
            ssh_user: "u".into(),
            ssh_key_path: "/k".into(),
            clusters: vec![],
            nccl: HashMap::new(),
            hf_cache_dir: None,
            init_port: 25000,
            before_serve: HashMap::new(),
            exec_aliases: HashMap::new(),
            home_assistant: None,
            mikrotik: None,
            global_mods_dir: None,
            drop_caches: true,
            mount_aliases: HashMap::new(),
            recipe_sources: Vec::new(),
            recipe_repos: Vec::new(),
            docker_image_sources: Vec::new(),
            cluster_sources: Vec::new(),
            registry: None,
        }
    }

    fn trec() -> Recipe {
        Recipe {
            build_profile: None,
            parent_model: None,
            archived: false,
            engine: "vllm".into(),
            name: "T".into(),
            docker_image: "img".into(),
            model: "org/main".into(),
            port: 8000,
            required_nodes: Some(1),
            pipeline_parallel_size: Some(1),
            tensor_parallel_size: 1,
            args: RecipeArgs::List(vec![]),
            cold_args: RecipeArgs::default(),
            env: HashMap::new(),
            cluster: None,
            entrypoint: vec![],
            distributed_backend: Some(DistributedBackend::Mp),
            before_serve: vec![],
            after_stop: vec![],
            solo: None,
            deploy: DeployTarget::Docker,
            venv: None,
            venv_setup: VenvSetup::default(),
            venv_setup_script: None,
            mods: Default::default(),
            shm_size: String::new(),
            devices: vec![],
            memory_limit: None,
            cpus: None,
            sync_image: true,
            docker_image_dir: None,
            extra_mounts: vec![],
            extra_mounts_manifest: None,
            build: Default::default(),
            composed_model: None,
            staging: vec![],
            dgpp: None,
        }
    }

    #[test]
    fn stage_mode_parse() {
        assert_eq!(StageMode::parse("").unwrap(), StageMode::Auto);
        assert_eq!(StageMode::parse("auto").unwrap(), StageMode::Auto);
        assert_eq!(StageMode::parse("SKIP").unwrap(), StageMode::Skip);
        assert_eq!(StageMode::parse("force").unwrap(), StageMode::Force);
        assert!(StageMode::parse("nope").is_err());
    }

    #[tokio::test]
    async fn sparse_rows_dry_run_dispatches() {
        // dry-run touches no SSH: both nodes just print their plan (rank
        // 0 and 1 of a 2-rank split) and the pass succeeds — this exercises
        // the sparse-rows dispatch arm and the tp derivation.
        let mut r = trec();
        r.staging.push(Staging {
            kind: crate::config::StagingKind::SparseRows,
            src: Some("{{hf}}/hub/snap".into()),
            dst: Some("~/.cache/engram".into()),
            ..Default::default()
        });
        let nodes = vec![
            Node {
                name: "n0".into(),
                ssh_ip: "10.0.0.1".into(),
                ib_ip: None,
                rank: 0,
                host_key: None,
                entity: None,
            },
            Node {
                name: "n1".into(),
                ssh_ip: "10.0.0.2".into(),
                ib_ip: None,
                rank: 1,
                host_key: None,
                entity: None,
            },
        ];
        crate::staging::run(
            &cfg(),
            &nodes,
            &r,
            std::path::Path::new("/mnt/nas/hf"),
            StageMode::Auto,
            true,
        )
        .await
        .unwrap();
    }

    #[test]
    fn hub_mirror_dst_tilde_and_detection() {
        let mut r = trec();
        assert!(!has_hub_mirror(&r));
        assert!(hub_mirror_dst(&cfg(), &r).is_none());
        r.staging.push(Staging {
            kind: StagingKind::HubMirror,
            dst: Some("~/.cache/hf-local".into()),
            ..Default::default()
        });
        assert!(has_hub_mirror(&r));
        assert_eq!(
            hub_mirror_dst(&cfg(), &r),
            Some(PathBuf::from("/home/u/.cache/hf-local"))
        );
    }
}