|
|
@@ -0,0 +1,175 @@
|
|
|
+//! 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),
|
|
|
+ };
|
|
|
+ let opts = CacheOpts {
|
|
|
+ cluster,
|
|
|
+ json,
|
|
|
+ dry_run,
|
|
|
+ };
|
|
|
+ match cmd {
|
|
|
+ C::Status { .. } => status(config, &opts).await,
|
|
|
+ C::Doctor { .. } => doctor(config, &opts).await,
|
|
|
+ }
|
|
|
+}
|
|
|
+
|
|
|
+/// `sm121 cache status` — repos in the cache, their sizes, and the room left.
|
|
|
+pub async fn status(config: &Config, opts: &CacheOpts) -> Result<()> {
|
|
|
+ let node = cache_node(config, opts.cluster.as_deref())?;
|
|
|
+ let home = hf_home(config);
|
|
|
+ let root = crate::docker::shq(&home);
|
|
|
+ let cmd = format!(
|
|
|
+ "echo \"__HOME {home}\"; \
|
|
|
+ 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 free = String::new();
|
|
|
+ let mut repos: Vec<(u64, String)> = Vec::new();
|
|
|
+ for line in out.lines() {
|
|
|
+ if let Some(rest) = line.strip_prefix("__HOME ") {
|
|
|
+ if opts.json {
|
|
|
+ free = rest.to_string();
|
|
|
+ } else {
|
|
|
+ println!("cache: {rest} (on {})", node.name);
|
|
|
+ }
|
|
|
+ continue;
|
|
|
+ }
|
|
|
+ let mut it = line.split_whitespace();
|
|
|
+ if let (Some(bytes), Some(path)) = (it.next(), it.next())
|
|
|
+ && let Ok(b) = bytes.parse::<u64>()
|
|
|
+ {
|
|
|
+ let name = path.rsplit('/').next().unwrap_or(path).to_string();
|
|
|
+ repos.push((b, name));
|
|
|
+ }
|
|
|
+ }
|
|
|
+ repos.sort_by_key(|(b, _)| std::cmp::Reverse(*b));
|
|
|
+ if opts.json {
|
|
|
+ let items: Vec<_> = repos
|
|
|
+ .iter()
|
|
|
+ .map(|(b, n)| json!({"repo": n, "bytes": b, "gib": (*b as f64) / (1u64 << 30) as f64}))
|
|
|
+ .collect();
|
|
|
+ println!(
|
|
|
+ "{}",
|
|
|
+ serde_json::to_string_pretty(&json!({
|
|
|
+ "node": node.name, "hf_home": home, "df": free, "repos": items,
|
|
|
+ }))?
|
|
|
+ );
|
|
|
+ return Ok(());
|
|
|
+ }
|
|
|
+ let total: u64 = repos.iter().map(|(b, _)| b).sum();
|
|
|
+ println!("{:<58} {:>9}", "repo", "GiB");
|
|
|
+ for (b, n) in &repos {
|
|
|
+ println!("{:<58} {:>9.1}", n, *b as f64 / (1u64 << 30) as f64);
|
|
|
+ }
|
|
|
+ println!(
|
|
|
+ "{:<58} {:>9.1}",
|
|
|
+ format!("— {} repos", repos.len()),
|
|
|
+ total as f64 / (1u64 << 30) as f64
|
|
|
+ );
|
|
|
+ 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!(
|
|
|
+ "for t in hfrs uvx flare hf; do printf '%s=%s\\n' \"$t\" \"$(command -v $t || echo MISSING)\"; done; \
|
|
|
+ [ -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(())
|
|
|
+}
|