launcher/src/cache.rs

//! Shared Hugging Face cache management: `sm121 cache status|doctor`.
//!
//! The cache lives on the NAS and is reached over SSH through the cluster (an
//! add-on on a Pi has no NAS mount, so this is the only path that works
//! everywhere). A node that holds the mount runs the actual inspection —
//! `[cache].node` pins it, otherwise the rank-0 node of `--cluster` is used.
//!
//! Read-only here: `status` reports what is cached and how much room is left,
//! `doctor` reports whether the tooling (`hfrs`, `uvx`, `flare`) that the
//! update job needs is actually present. Destructive steps (`update`, `prune`)
//! are separate subcommands so they never run by accident.

use anyhow::Result;
use serde_json::json;

use crate::config::Config;

/// Env/CLI knobs for the cache commands.
pub struct CacheOpts {
    pub cluster: Option<String>,
    pub json: bool,
    pub dry_run: bool,
}

/// The node that can see the cache: `[cache].node`, else rank 0 of the chosen
/// cluster, else rank 0 of the first cluster.
fn cache_node(config: &Config, cluster: Option<&str>) -> Result<crate::config::Node> {
    if let Some(name) = config.cache.as_ref().and_then(|c| c.node.clone()) {
        return config
            .clusters
            .iter()
            .flat_map(|c| c.nodes.iter())
            .find(|n| n.name == name)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("[cache].node '{name}' is not a configured node"));
    }
    let cl = match cluster {
        Some(name) => crate::config::select_cluster_by_name(config, name)?,
        None => config
            .clusters
            .first()
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("no clusters configured"))?,
    };
    cl.nodes
        .iter()
        .find(|n| n.rank == 0)
        .or_else(|| cl.nodes.first())
        .cloned()
        .ok_or_else(|| anyhow::anyhow!("cluster '{}' has no nodes", cl.name))
}

/// HF cache root (config `[cache].hf_home`, else the ssh user's default).
pub fn hf_home(config: &Config) -> String {
    config
        .cache
        .as_ref()
        .and_then(|c| c.hf_home.clone())
        .unwrap_or_else(|| config.resolved_hf_cache_dir().display().to_string())
}

/// CLI dispatch (`sm121 cache <subcommand>`).
pub async fn dispatch(
    config: &Config,
    cmd: &crate::cli::CacheCommands,
    dry_run: bool,
) -> Result<()> {
    use crate::cli::CacheCommands as C;
    let (cluster, json) = match cmd {
        C::Status { cluster, json } | C::Doctor { cluster, json } => (cluster.clone(), *json),
        C::Update { cluster, json, .. } => (cluster.clone(), *json),
    };
    let opts = CacheOpts {
        cluster,
        json,
        dry_run,
    };
    match cmd {
        C::Status { .. } => status(config, &opts).await,
        C::Doctor { .. } => doctor(config, &opts).await,
        C::Update { repos, apply, .. } => update(config, &opts, repos, *apply).await,
    }
}

/// One cached repo (name + size), shared by the CLI table and the API.
#[derive(Clone, serde::Serialize)]
pub struct CacheRepo {
    pub repo: String,
    pub bytes: u64,
    pub gib: f64,
}

/// What the cache looks like right now: the node it was read from, the mount
/// point, the filesystem line and every repo with its size.
#[derive(Clone, serde::Serialize)]
pub struct CacheStatus {
    pub node: String,
    pub hf_home: String,
    pub df: String,
    pub repos: Vec<CacheRepo>,
    pub total_gib: f64,
}

/// Gather the cache status (`status` prints it, the dashboard serves it).
pub async fn collect(config: &Config, opts: &CacheOpts) -> Result<CacheStatus> {
    let node = cache_node(config, opts.cluster.as_deref())?;
    let home = hf_home(config);
    let root = crate::docker::shq(&home);
    let cmd = format!(
        "df -BG --output=size,used,avail {root} 2>/dev/null | tail -1; \
         du -s --block-size=1 {root}/hub/*--* 2>/dev/null | awk '{{print $1, $2}}'"
    );
    let out = crate::ssh::run_ssh(config, &node, cmd, opts.dry_run).await?;
    let mut df = String::new();
    let mut repos: Vec<CacheRepo> = Vec::new();
    for line in out.lines() {
        let toks: Vec<&str> = line.split_whitespace().collect();
        match toks.as_slice() {
            // `df -BG`: "931G 700G 231G" — size/used/avail, no path column
            [size, used, avail] if size.parse::<u64>().is_err() => {
                df = format!("{size} total, {used} used, {avail} avail");
            }
            // `du -s --block-size=1`: "<bytes> <path>"
            [bytes, path] => {
                if let Ok(b) = bytes.parse::<u64>() {
                    let name = path.rsplit('/').next().unwrap_or(path).to_string();
                    repos.push(CacheRepo {
                        repo: name,
                        bytes: b,
                        gib: b as f64 / (1u64 << 30) as f64,
                    });
                }
            }
            _ => {}
        }
    }
    repos.sort_by_key(|r| std::cmp::Reverse(r.bytes));
    let total_gib = repos.iter().map(|r| r.gib).sum();
    Ok(CacheStatus {
        node: node.name,
        hf_home: home,
        df,
        repos,
        total_gib,
    })
}

/// `sm121 cache status` — repos in the cache, their sizes, and the room left.
pub async fn status(config: &Config, opts: &CacheOpts) -> Result<()> {
    let st = collect(config, opts).await?;
    if opts.json {
        println!("{}", serde_json::to_string_pretty(&st)?);
        return Ok(());
    }
    println!("cache: {} (on {}, df {})", st.hf_home, st.node, st.df);
    println!("{:<58} {:>9}", "repo", "GiB");
    for r in &st.repos {
        println!("{:<58} {:>9.1}", r.repo, r.gib);
    }
    println!(
        "{:<58} {:>9.1}",
        format!("— {} repos", st.repos.len()),
        st.total_gib
    );
    Ok(())
}

/// `sm121 cache doctor` — is the tooling the maintenance job needs present?
pub async fn doctor(config: &Config, opts: &CacheOpts) -> Result<()> {
    let node = cache_node(config, opts.cluster.as_deref())?;
    let home = hf_home(config);
    let cmd = format!(
        // A non-login shell has a bare PATH: also look where the tools land.
        "PATH=\"$HOME/.local/bin:$HOME/.cargo/bin:$PATH\"; \
         for t in hfrs uvx hf; do printf '%s=%s\\n' \"$t\" \"$(command -v $t || echo MISSING)\"; done; \
         if [ -x \"$HOME/flare-bin/flare\" ]; then echo \"flare=$HOME/flare-bin/flare\"; \
         elif command -v flare >/dev/null; then echo \"flare=$(command -v flare)\"; else echo flare=MISSING; fi; \
         [ -d {root} ] && echo \"hf_home=ok\" || echo \"hf_home=MISSING\"; \
         touch {root}/.sm121-write-test 2>/dev/null && echo write=ok && rm -f {root}/.sm121-write-test || echo write=FAILED; \
         df -BG --output=avail {root} 2>/dev/null | tail -1 | sed 's/^/free=/'",
        root = crate::docker::shq(&home)
    );
    let out = crate::ssh::run_ssh(config, &node, cmd, opts.dry_run).await?;
    if opts.json {
        let map: serde_json::Map<String, serde_json::Value> = out
            .lines()
            .filter_map(|l| l.split_once('='))
            .map(|(k, v)| (k.trim().to_string(), json!(v.trim())))
            .collect();
        println!(
            "{}",
            serde_json::to_string_pretty(&json!({"node": node.name, "checks": map}))?
        );
        return Ok(());
    }
    println!(
        "cache doctor on {} ({}), keep_revs={:?}:",
        node.name,
        home,
        config.cache.as_ref().and_then(|c| c.keep_revs)
    );
    for line in out.lines() {
        let ok = !line.contains("MISSING") && !line.contains("FAILED");
        println!("  {} {}", if ok { "✓" } else { "✗" }, line);
    }
    Ok(())
}

/// Expected size of `repo@rev` from the HF API (bytes), for the progress line.
/// Sums the file sizes; a failure only costs the percentage (progress falls
/// back to raw GB), so it never blocks the update.
async fn expected_bytes(repo: &str, rev: &str) -> Option<u64> {
    let url = format!("https://huggingface.co/api/models/{repo}/tree/{rev}?recursive=1");
    let client = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(20))
        .build()
        .ok()?;
    let body: serde_json::Value = client.get(url).send().await.ok()?.json().await.ok()?;
    let mut total = 0u64;
    for f in body.as_array()? {
        if f.get("type").and_then(|t| t.as_str()) != Some("file") {
            continue;
        }
        let size = f
            .get("size")
            .and_then(serde_json::Value::as_u64)
            .unwrap_or(0);
        let path = f.get("path").and_then(|p| p.as_str()).unwrap_or("");
        // .gitattributes etc. are noise, not weights
        if !path.starts_with('.') {
            total += size;
        }
    }
    (total > 0).then_some(total)
}

/// Live state of a running `cache update` (the dashboard renders it, the CLI
/// renders the same fields as a progress line).
#[derive(Clone, Default, serde::Serialize)]
pub struct UpdateProgress {
    pub running: bool,
    pub stage: String,
    pub repo: Option<String>,
    pub index: usize,
    pub total: usize,
    pub bytes: u64,
    pub expected: Option<u64>,
    pub percent: Option<f64>,
    pub failed: Vec<String>,
    pub node: String,
    pub message: String,
}

/// `sm121 cache update` — refresh every cached repo to its refs (or the ones
/// selected with `--repo`), the same loop `hf-cache-maintain.sh` runs, but over
/// SSH so the machine driving it needs no mount.
///
/// Progress: the remote shell samples the repo's blob directory every 10 s
/// while `hfrs` runs and prints `PROGRESS <bytes>`; the launcher turns that into
/// a percentage of the API-reported size. One SSH for the whole download, and
/// the sampler touches only a directory listing, so the download itself is not
/// slowed.
pub async fn update(
    config: &Config,
    opts: &CacheOpts,
    repos: &[String],
    apply: bool,
) -> Result<()> {
    let mut last = String::new();
    update_with(config, opts, repos, apply, |p| {
        if p.stage == "download" {
            let text = p.message.clone();
            if text != last {
                use std::io::Write;
                print!("\r  {text:<70}");
                let _ = std::io::stdout().flush();
                last = text;
            }
        } else if p.stage == "repo-done" {
            if !last.is_empty() {
                println!();
                last.clear();
            }
            println!("  {}", p.message);
        } else if p.stage == "dry-run" {
            println!("  {}", p.message);
        }
    })
    .await
}

/// Same as `update`, but reports every step to `report` — the dashboard's
/// background job and the CLI progress line are two renderings of one loop.
pub async fn update_with<F: FnMut(&UpdateProgress)>(
    config: &Config,
    opts: &CacheOpts,
    repos: &[String],
    apply: bool,
    mut report: F,
) -> Result<()> {
    let node = cache_node(config, opts.cluster.as_deref())?;
    let home = hf_home(config);
    let mut prog = UpdateProgress {
        running: true,
        stage: "listing".into(),
        node: node.name.clone(),
        message: "listing cached repos".into(),
        ..Default::default()
    };
    report(&prog);
    let list_cmd = format!(
        "ls -d {root}/hub/*--* 2>/dev/null",
        root = crate::docker::shq(&home)
    );
    let listing = crate::ssh::run_ssh(config, &node, list_cmd, opts.dry_run).await?;
    let mut targets: Vec<(String, String)> = Vec::new();
    for dir in listing.lines() {
        let Some(base) = dir.rsplit('/').next() else {
            continue;
        };
        let Some((kind, rest)) = base.split_once("--") else {
            continue;
        };
        if kind != "models" {
            continue;
        }
        let repo = rest.replacen("--", "/", 1);
        if !repos.is_empty() && !repos.iter().any(|r| r == &repo) {
            continue;
        }
        targets.push((repo, dir.to_string()));
    }
    if targets.is_empty() {
        prog.running = false;
        prog.stage = "done".into();
        prog.message = format!("nothing to update on {}", node.name);
        report(&prog);
        return Ok(());
    }
    prog.total = targets.len();
    let mut failures = 0usize;
    for (i, (repo, dir)) in targets.iter().enumerate() {
        let rev = "main";
        prog.index = i + 1;
        prog.repo = Some(repo.clone());
        prog.bytes = 0;
        prog.percent = None;
        prog.expected = expected_bytes(repo, rev).await;
        prog.stage = "repo-start".into();
        prog.message = format!("{repo} @ {rev}");
        report(&prog);
        if !apply || opts.dry_run {
            prog.stage = "dry-run".into();
            prog.message = format!("[dry-run] hfrs download {repo} --revision {rev}");
            report(&prog);
            continue;
        }
        let quoted_dir = crate::docker::shq(dir);
        let cmd = format!(
            "export PATH=\"$HOME/.local/bin:$HOME/.cargo/bin:$PATH\"; \
             ( hfrs download {repo} --type model --revision {rev} \
                 --disable-progress-bars --quiet >/dev/null 2>&1 & p=$!; \
               while kill -0 $p 2>/dev/null; do \
                 b=$(du -sb {quoted_dir}/blobs 2>/dev/null | cut -f1); \
                 echo \"PROGRESS ${{b:-0}}\"; sleep 10; done; wait $p; echo \"EXIT $?\" )"
        );
        prog.stage = "download".into();
        report(&prog);
        let out = crate::ssh::run_ssh(config, &node, cmd, false).await?;
        let mut ok = true;
        for line in out.lines() {
            if let Some(bytes) = line.strip_prefix("PROGRESS ") {
                prog.bytes = bytes.trim().parse().unwrap_or(0);
                prog.percent = prog
                    .expected
                    .filter(|t| *t > 1_000_000)
                    .map(|t| (prog.bytes as f64 / t as f64 * 100.0).min(100.0));
                prog.message = match prog.expected {
                    Some(t) if t > 1_000_000 => format!(
                        "{repo} {:.1}% ({:.1} / {:.1} GB)",
                        prog.percent.unwrap_or(0.0),
                        prog.bytes as f64 / 1e9,
                        t as f64 / 1e9
                    ),
                    _ => format!("{repo} {:.1} GB", prog.bytes as f64 / 1e9),
                };
                report(&prog);
            } else if let Some(code) = line.strip_prefix("EXIT ") {
                ok = code.trim() == "0";
            }
        }
        prog.stage = "repo-done".into();
        if ok {
            prog.message = format!("[ok]   {repo} @ {rev}");
        } else {
            prog.message = format!("[FAIL] {repo} @ {rev}");
            prog.failed.push(repo.clone());
            failures += 1;
        }
        report(&prog);
    }
    prog.running = false;
    prog.stage = "done".into();
    prog.repo = None;
    prog.message = if failures > 0 {
        format!("{failures} repo(s) failed to update on {}", node.name)
    } else {
        format!("updated {} repo(s) on {}", targets.len(), node.name)
    };
    report(&prog);
    if failures > 0 {
        anyhow::bail!("{}", prog.message);
    }
    Ok(())
}