launcher/src/staging/view.rs

//! `local-view` staging block executor.
//!
//! Some checkpoints stream a large side table at every token (the DSV4.1
//! Engram tables read ~1 MB of rows per token). Serving those rows from the
//! shared NAS costs an order of magnitude more than from local NVMe, while the
//! bulk of the weights belongs on shared storage. A `local-view` block
//! materializes, per launch node, a directory holding **real copies** of the
//! files matching `local` and symlinks (into the shared cache's container
//! path) for everything else; `serve` mounts it and points `model` at it.
//!
//! The node-side work is `scripts/local_view.py`; this module only builds the
//! command lines, decides the container mount and computes the bytes a node
//! must have free (used by the pre-start free-space check).

use crate::config::{Config, Node, Staging, expand_tilde};
use crate::ssh::{run_ssh_long_running, upload_script};
use anyhow::{Result, bail};
use std::path::Path;
use tracing::info;

pub const SCRIPT_BODY: &str = include_str!("scripts/local_view.py");

/// Container-side path the view is mounted at (the model references it).
pub const CONTAINER_ROOT: &str = "/root/.cache/hf-view";

fn script_path(id: &str) -> String {
    format!("/tmp/sm121-local-view-{id}.py")
}

/// The repo a block views: explicit `repo`, else the recipe's model.
pub fn repo_for(block: &Staging, model: &str) -> String {
    block.repo.clone().unwrap_or_else(|| model.to_string())
}

/// Node-side view directory (tilde-expanded).
pub fn dst_for(block: &Staging, id: &str, ssh_user: &str) -> String {
    block
        .dst
        .as_deref()
        .map(|d| expand_tilde(d, ssh_user))
        .unwrap_or_else(|| format!("~/.cache/hf-view/{id}"))
}

/// `measure` (bytes a node needs) or `apply` (build the view).
pub struct ViewSpec<'a> {
    pub repo: &'a str,
    pub local: &'a [String],
    pub verify: &'a str,
    pub min_free: Option<&'a str>,
    pub force: bool,
}

/// Command line for `python3 <script> apply --src … --repo … --dst … --local …`.
pub fn command_for(id: &str, mode: &str, src: &str, dst: &str, spec: &ViewSpec<'_>) -> String {
    let mut cmd = format!(
        "python3 {} {} --src {} --repo {} --dst {} --container-root {} --verify {}",
        crate::docker::shq(&script_path(id)),
        mode,
        crate::docker::shq(src),
        crate::docker::shq(spec.repo),
        crate::docker::shq(dst),
        crate::docker::shq(CONTAINER_ROOT),
        spec.verify,
    );
    for g in spec.local {
        cmd.push_str(&format!(" --local {}", crate::docker::shq(g)));
    }
    if let Some(m) = spec.min_free {
        cmd.push_str(&format!(" --min-free {}", crate::docker::shq(m)));
    }
    if spec.force {
        cmd.push_str(" --force");
    }
    cmd
}

/// What one node needs free for a view, plus what it has (bytes).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Need {
    pub node: String,
    pub repo: String,
    pub need: u64,
    pub free: u64,
    pub dst: String,
}

impl Need {
    pub fn short(&self) -> bool {
        self.free < self.need
    }

    pub fn line(&self) -> String {
        format!(
            "{}: {} → needs {:.1} GiB, has {:.1} GiB{}",
            self.node,
            self.repo,
            self.need as f64 / (1u64 << 30) as f64,
            self.free as f64 / (1u64 << 30) as f64,
            if self.short() { "  ⚠ SHORT" } else { "" }
        )
    }
}

/// Parse the `measure` output (`{"need": …, "free": …}`) of the node script.
pub fn parse_measure(node: &str, repo: &str, dst: &str, stdout: &str) -> Result<Need> {
    let v: serde_json::Value = serde_json::from_str(stdout.trim()).map_err(|e| {
        anyhow::anyhow!("local-view measure on {node}: unparsable output {stdout:?}: {e}")
    })?;
    let num = |k: &str| v.get(k).and_then(serde_json::Value::as_u64).unwrap_or(0);
    Ok(Need {
        node: node.to_string(),
        repo: repo.to_string(),
        need: num("need"),
        free: num("free"),
        dst: dst.to_string(),
    })
}

/// The aggregated, human-facing shortfall message (`None` when every node fits).
pub fn shortfall(needs: &[Need]) -> Option<String> {
    let short: Vec<&Need> = needs.iter().filter(|n| n.short()).collect();
    if short.is_empty() {
        return None;
    }
    let mut msg =
        String::from("not enough free space on the stage target(s) — nothing was started:\n");
    for n in &short {
        msg.push_str(&format!(
            "  {}: {} needs {:.1} GiB in {}, has {:.1} GiB\n",
            n.node,
            n.repo,
            n.need as f64 / (1u64 << 30) as f64,
            n.dst,
            n.free as f64 / (1u64 << 30) as f64
        ));
    }
    msg.push_str("free space (docker rmi / stale staging dirs) or point dst at a bigger volume.");
    Some(msg)
}

/// Ask every launch node how much space a block's `local` copies need.
pub async fn measure(
    config: &Config,
    nodes: &[Node],
    block: &Staging,
    index: usize,
    model: &str,
    shared_source: &Path,
) -> Result<Vec<Need>> {
    let id = block.ledger_id(index);
    let repo = repo_for(block, model);
    let spec = ViewSpec {
        repo: &repo,
        local: &block.local,
        verify: block.verify.as_str(),
        min_free: block.min_free.as_deref(),
        force: false,
    };
    let mut out = Vec::new();
    for node in nodes {
        let dst = dst_for(block, &id, &config.ssh_user);
        upload_script(config, node, SCRIPT_BODY, &script_path(&id), false).await?;
        let cmd = command_for(
            &id,
            "measure",
            &shared_source.display().to_string(),
            &dst,
            &spec,
        );
        let stdout = crate::ssh::run_ssh(config, node, cmd, false).await?;
        out.push(parse_measure(&node.name, &repo, &dst, &stdout)?);
    }
    Ok(out)
}

/// Everything one node's `local-view` pass needs.
pub struct StageCtx<'a> {
    pub config: &'a Config,
    pub recipe: &'a crate::config::Recipe,
    pub shared_source: &'a Path,
    pub force: bool,
    pub dry_run: bool,
}

/// Stage one `local-view` block on one node.
pub async fn stage_node(
    ctx: &StageCtx<'_>,
    node: &Node,
    block: &Staging,
    index: usize,
) -> Result<()> {
    let (config, model, shared_source, force, dry_run) = (
        ctx.config,
        ctx.recipe.model.as_str(),
        ctx.shared_source,
        ctx.force,
        ctx.dry_run,
    );
    if block.local.is_empty() {
        bail!(
            "[[staging]] local-view block '{}' needs at least one `local` glob (the files that must be real node-local copies)",
            block.ledger_id(index)
        );
    }
    let id = block.ledger_id(index);
    let repo = repo_for(block, model);
    let dst = dst_for(block, &id, &config.ssh_user);
    let spec = ViewSpec {
        repo: &repo,
        local: &block.local,
        verify: block.verify.as_str(),
        min_free: block.min_free.as_deref(),
        force,
    };
    let src = shared_source.display().to_string();
    if dry_run {
        info!(
            "dry-run staging[{}] on {}: local-view {} → {} (local: {:?}, min_free {:?})",
            id, node.name, repo, dst, block.local, block.min_free
        );
        return Ok(());
    }
    upload_script(config, node, SCRIPT_BODY, &script_path(&id), false).await?;
    let cmd = command_for(&id, "apply", &src, &dst, &spec);
    run_ssh_long_running(config, node, cmd, false).await?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn need(node: &str, need: u64, free: u64) -> Need {
        Need {
            node: node.into(),
            repo: "org/model".into(),
            need,
            free,
            dst: "/v".into(),
        }
    }

    #[test]
    fn measure_output_parses() {
        let n = parse_measure(
            "dgx5",
            "org/m",
            "/v",
            r#"{"repo":"org/m","need":203,"free":400}"#,
        )
        .unwrap();
        assert_eq!(n.need, 203);
        assert_eq!(n.free, 400);
        assert!(!n.short());
    }

    #[test]
    fn shortfall_lists_only_the_short_nodes() {
        let needs = vec![need("a", 100, 200), need("b", 300, 250)];
        let msg = shortfall(&needs).expect("b is short");
        assert!(msg.contains("b:"), "{msg}");
        assert!(!msg.contains("a:"), "{msg}");
        assert!(msg.contains("nothing was started"), "{msg}");
        assert!(shortfall(&[need("a", 1, 2)]).is_none());
    }

    #[test]
    fn command_carries_globs_and_container_root() {
        let globs = vec!["model-00047-of-00048.safetensors".to_string()];
        let spec = ViewSpec {
            repo: "org/model",
            local: &globs,
            verify: "sample-hash",
            min_free: Some("10G"),
            force: true,
        };
        let cmd = command_for("s0", "apply", "/src", "/dst", &spec);
        // shq() shell-quotes arguments, so the glob may come back quoted
        assert!(cmd.contains("model-00047-of-00048.safetensors"), "{cmd}");
        assert!(cmd.contains("--local"), "{cmd}");
        assert!(cmd.contains(CONTAINER_ROOT), "{cmd}");
        assert!(cmd.contains("--min-free '10G'"), "{cmd}");
        assert!(cmd.ends_with("--force"), "{cmd}");
    }
}