|
|
@@ -82,59 +82,85 @@ pub async fn dispatch(
|
|
|
}
|
|
|
}
|
|
|
|
|
|
-/// `sm121 cache status` — repos in the cache, their sizes, and the room left.
|
|
|
-pub async fn status(config: &Config, opts: &CacheOpts) -> Result<()> {
|
|
|
+/// 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!(
|
|
|
- "echo \"__HOME {home}\"; \
|
|
|
- df -BG --output=size,used,avail {root} 2>/dev/null | tail -1; \
|
|
|
+ "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();
|
|
|
+ let mut df = String::new();
|
|
|
+ let mut repos: Vec<CacheRepo> = 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);
|
|
|
+ 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");
|
|
|
}
|
|
|
- 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));
|
|
|
+ // `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(|(b, _)| std::cmp::Reverse(*b));
|
|
|
+ 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 {
|
|
|
- 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,
|
|
|
- }))?
|
|
|
- );
|
|
|
+ println!("{}", serde_json::to_string_pretty(&st)?);
|
|
|
return Ok(());
|
|
|
}
|
|
|
- let total: u64 = repos.iter().map(|(b, _)| b).sum();
|
|
|
+ println!("cache: {} (on {}, df {})", st.hf_home, st.node, st.df);
|
|
|
println!("{:<58} {:>9}", "repo", "GiB");
|
|
|
- for (b, n) in &repos {
|
|
|
- println!("{:<58} {:>9.1}", n, *b as f64 / (1u64 << 30) as f64);
|
|
|
+ for r in &st.repos {
|
|
|
+ println!("{:<58} {:>9.1}", r.repo, r.gib);
|
|
|
}
|
|
|
println!(
|
|
|
"{:<58} {:>9.1}",
|
|
|
- format!("— {} repos", repos.len()),
|
|
|
- total as f64 / (1u64 << 30) as f64
|
|
|
+ format!("— {} repos", st.repos.len()),
|
|
|
+ st.total_gib
|
|
|
);
|
|
|
Ok(())
|
|
|
}
|
|
|
@@ -208,6 +234,23 @@ async fn expected_bytes(repo: &str, rev: &str) -> Option<u64> {
|
|
|
(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.
|
|
|
@@ -222,9 +265,49 @@ pub async fn update(
|
|
|
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)
|
|
|
@@ -248,62 +331,85 @@ pub async fn update(
|
|
|
targets.push((repo, dir.to_string()));
|
|
|
}
|
|
|
if targets.is_empty() {
|
|
|
- println!("nothing to update on {}", node.name);
|
|
|
+ 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 (repo, dir) in &targets {
|
|
|
- for rev in ["main"] {
|
|
|
- let total = expected_bytes(repo, rev).await;
|
|
|
- if !apply || opts.dry_run {
|
|
|
- println!(" [dry-run] hfrs download {repo} --revision {rev}");
|
|
|
- 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 $?\" )"
|
|
|
- );
|
|
|
- let out = crate::ssh::run_ssh(config, &node, cmd, false).await?;
|
|
|
- let mut last_line_len = 0usize;
|
|
|
- let mut ok = true;
|
|
|
- for line in out.lines() {
|
|
|
- if let Some(bytes) = line.strip_prefix("PROGRESS ") {
|
|
|
- let cur: u64 = bytes.trim().parse().unwrap_or(0);
|
|
|
- let text = match total {
|
|
|
- Some(t) if t > 1_000_000 => format!(
|
|
|
- "{repo} {:.1}% ({:.1} / {:.1} GB)",
|
|
|
- (cur as f64 / t as f64 * 100.0).min(100.0),
|
|
|
- cur as f64 / 1e9,
|
|
|
- t as f64 / 1e9
|
|
|
- ),
|
|
|
- _ => format!("{repo} {:.1} GB", cur as f64 / 1e9),
|
|
|
- };
|
|
|
- print!("\r {text:<70}");
|
|
|
- last_line_len = text.len();
|
|
|
- use std::io::Write;
|
|
|
- let _ = std::io::stdout().flush();
|
|
|
- } else if let Some(code) = line.strip_prefix("EXIT ") {
|
|
|
- ok = code.trim() == "0";
|
|
|
- }
|
|
|
- }
|
|
|
- if last_line_len > 0 {
|
|
|
- println!();
|
|
|
- }
|
|
|
- if ok {
|
|
|
- println!(" [ok] {repo} @ {rev}");
|
|
|
- } else {
|
|
|
- println!(" [FAIL] {repo} @ {rev}");
|
|
|
- failures += 1;
|
|
|
+ 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!("{failures} repo(s) failed to update on {}", node.name);
|
|
|
+ anyhow::bail!("{}", prog.message);
|
|
|
}
|
|
|
Ok(())
|
|
|
}
|
