|
@@ -68,6 +68,7 @@ pub async fn dispatch(
|
|
|
use crate::cli::CacheCommands as C;
|
|
use crate::cli::CacheCommands as C;
|
|
|
let (cluster, json) = match cmd {
|
|
let (cluster, json) = match cmd {
|
|
|
C::Status { cluster, json } | C::Doctor { cluster, json } => (cluster.clone(), *json),
|
|
C::Status { cluster, json } | C::Doctor { cluster, json } => (cluster.clone(), *json),
|
|
|
|
|
+ C::Update { cluster, json, .. } => (cluster.clone(), *json),
|
|
|
};
|
|
};
|
|
|
let opts = CacheOpts {
|
|
let opts = CacheOpts {
|
|
|
cluster,
|
|
cluster,
|
|
@@ -77,6 +78,7 @@ pub async fn dispatch(
|
|
|
match cmd {
|
|
match cmd {
|
|
|
C::Status { .. } => status(config, &opts).await,
|
|
C::Status { .. } => status(config, &opts).await,
|
|
|
C::Doctor { .. } => doctor(config, &opts).await,
|
|
C::Doctor { .. } => doctor(config, &opts).await,
|
|
|
|
|
+ C::Update { repos, apply, .. } => update(config, &opts, repos, *apply).await,
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
|
|
@@ -177,3 +179,131 @@ pub async fn doctor(config: &Config, opts: &CacheOpts) -> Result<()> {
|
|
|
}
|
|
}
|
|
|
Ok(())
|
|
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)
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// `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 node = cache_node(config, opts.cluster.as_deref())?;
|
|
|
|
|
+ let home = hf_home(config);
|
|
|
|
|
+ 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() {
|
|
|
|
|
+ println!("nothing to update on {}", node.name);
|
|
|
|
|
+ return Ok(());
|
|
|
|
|
+ }
|
|
|
|
|
+ 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;
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+ if failures > 0 {
|
|
|
|
|
+ anyhow::bail!("{failures} repo(s) failed to update on {}", node.name);
|
|
|
|
|
+ }
|
|
|
|
|
+ Ok(())
|
|
|
|
|
+}
|