Эх сурвалжийг харах

cache update + hf search/card

'cache update' runs the same hfrs loop as hf-cache-maintain.sh but over SSH, so
the driver needs no NAS mount (the add-on has none) — it needs the tooling on the
node, which doctor now verifies (hfrs is installed on all eight nodes).

Progress without slowing the download: the remote shell samples the repo's blob
directory every 10 s while hfrs runs and prints 'PROGRESS <bytes>' — one SSH for
the whole job, a directory listing per tick, and the launcher renders a
percentage against the size the HF API reports (falls back to GB when the API
gives no total, e.g. a docs-only revision).

'sm121 hf search <q>' lists models by downloads with likes/updated; 'sm121 hf
card <repo> [--revision R]' prints licence/base model/tags, every branch with
its resolved commit and, with --revision, the file list and total size — enough
to pick 'org/name@revision' for a recipe or a FLARE package.
Alexander Korolev 1 өдөр өмнө
parent
commit
396e02e61a
5 өөрчлөгдсөн 374 нэмэгдсэн , 6 устгасан
  1. 130 0
      src/cache.rs
  2. 21 0
      src/cli.rs
  3. 24 0
      src/cli/hf.rs
  4. 195 0
      src/hf.rs
  5. 4 6
      src/main.rs

+ 130 - 0
src/cache.rs

@@ -68,6 +68,7 @@ pub async fn dispatch(
     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,
@@ -77,6 +78,7 @@ pub async fn dispatch(
     match cmd {
         C::Status { .. } => status(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(())
 }
+
+/// 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(())
+}

+ 21 - 0
src/cli.rs

@@ -1,5 +1,7 @@
 use crate::runlock::RunlockCommands;
+pub(crate) mod hf;
 use clap::{Parser, Subcommand, ValueEnum};
+pub use hf::HfCommands;
 
 pub use crate::cli_sub::{
     ConfigCommands, PowerAction, PublishArgs, RegistryCommands, ReposArgs, ReposCommands,
@@ -309,6 +311,12 @@ pub enum Commands {
     #[command(subcommand)]
     Runlock(RunlockCommands),
 
+    /// Hugging Face catalogue: search models, read a card, list revisions
+    Hf {
+        #[command(subcommand)]
+        cmd: HfCommands,
+    },
+
     /// Shared Hugging Face cache on the NAS: `status` / `doctor` (read-only
     /// here; update/prune are separate subcommands).
     Cache {
@@ -459,6 +467,19 @@ pub enum CacheCommands {
         #[arg(long)]
         json: bool,
     },
+    /// Refresh every cached repo to its refs (dry-run unless --apply)
+    Update {
+        #[arg(long)]
+        cluster: Option<String>,
+        #[arg(long)]
+        json: bool,
+        /// Only these repos (repeatable, org/name)
+        #[arg(long = "repo")]
+        repos: Vec<String>,
+        /// Really download (default: preview)
+        #[arg(long)]
+        apply: bool,
+    },
     /// Is the tooling the maintenance job needs present on that node?
     Doctor {
         #[arg(long)]

+ 24 - 0
src/cli/hf.rs

@@ -0,0 +1,24 @@
+//! `sm121 hf …` subcommands.
+
+use clap::Subcommand;
+
+/// `sm121 hf …`
+#[derive(Subcommand)]
+pub enum HfCommands {
+    /// Search models by name (newest/most downloaded first)
+    Search {
+        query: String,
+        #[arg(long, default_value_t = 20)]
+        limit: usize,
+        #[arg(long)]
+        json: bool,
+    },
+    /// Model card: licence, base model, tags, revisions; `--revision` lists files
+    Card {
+        repo: String,
+        #[arg(long)]
+        revision: Option<String>,
+        #[arg(long)]
+        json: bool,
+    },
+}

+ 195 - 0
src/hf.rs

@@ -0,0 +1,195 @@
+//! Hugging Face catalogue helpers: `sm121 hf search|card`.
+//!
+//! Read-only, straight against the public HF API — the point is to pick a
+//! model **and a revision** for a recipe without opening a browser: search by
+//! name, see the card (tags, size, gated/licence, last update) and list the
+//! revisions with their resolved commits, then take that `org/name@revision`
+//! into the recipe / `[flare].revision`.
+
+use anyhow::{Result, bail};
+
+fn client() -> Result<reqwest::Client> {
+    Ok(reqwest::Client::builder()
+        .timeout(std::time::Duration::from_secs(25))
+        .user_agent("sm121")
+        .build()?)
+}
+
+/// Token from the environment, so gated repos resolve.
+fn token() -> Option<String> {
+    std::env::var("HF_TOKEN").ok().filter(|t| !t.is_empty())
+}
+
+async fn get_json(url: &str) -> Result<serde_json::Value> {
+    let mut req = client()?.get(url);
+    if let Some(t) = token() {
+        req = req.bearer_auth(t);
+    }
+    let res = req.send().await?;
+    if !res.status().is_success() {
+        bail!("{url} → HTTP {}", res.status());
+    }
+    Ok(res.json().await?)
+}
+
+/// `sm121 hf search <query>`
+pub async fn search(query: &str, limit: usize, json: bool) -> Result<()> {
+    let url = format!(
+        "https://huggingface.co/api/models?search={}&limit={limit}&sort=downloads&direction=-1&full=true",
+        urlencode(query)
+    );
+    let body = get_json(&url).await?;
+    let items = body.as_array().cloned().unwrap_or_default();
+    if json {
+        println!("{}", serde_json::to_string_pretty(&body)?);
+        return Ok(());
+    }
+    if items.is_empty() {
+        println!("no models match '{query}'");
+        return Ok(());
+    }
+    println!("{:<52} {:>9} {:>6}  updated", "model", "downloads", "likes");
+    for m in &items {
+        let id = m.get("id").and_then(|v| v.as_str()).unwrap_or("?");
+        let dl = m.get("downloads").and_then(|v| v.as_u64()).unwrap_or(0);
+        let likes = m.get("likes").and_then(|v| v.as_u64()).unwrap_or(0);
+        let upd = m
+            .get("lastModified")
+            .and_then(|v| v.as_str())
+            .map(|s| s.get(0..10).unwrap_or(s))
+            .unwrap_or("");
+        println!("{id:<52} {dl:>9} {likes:>6}  {upd}");
+    }
+    Ok(())
+}
+
+/// `sm121 hf card <repo> [--revision R]` — the card summary plus every
+/// revision (branch/tag) with its commit, so a recipe can pin one.
+pub async fn card(repo: &str, revision: Option<&str>, json: bool) -> Result<()> {
+    let info = get_json(&format!(
+        "https://huggingface.co/api/models/{repo}?full=true"
+    ))
+    .await?;
+    let refs = get_json(&format!("https://huggingface.co/api/models/{repo}/refs")).await?;
+    let files = match revision {
+        Some(rev) => Some(
+            get_json(&format!(
+                "https://huggingface.co/api/models/{repo}/tree/{rev}?recursive=1"
+            ))
+            .await?,
+        ),
+        None => None,
+    };
+    if json {
+        println!(
+            "{}",
+            serde_json::to_string_pretty(&serde_json::json!({
+                "model": info, "refs": refs, "files": files,
+            }))?
+        );
+        return Ok(());
+    }
+    let s = |k: &str| info.get(k).and_then(|v| v.as_str()).unwrap_or("");
+    let n = |k: &str| info.get(k).and_then(|v| v.as_u64());
+    println!(
+        "{}",
+        info.get("id").and_then(|v| v.as_str()).unwrap_or(repo)
+    );
+    if let Some(c) = info.get("cardData") {
+        let lic = c.get("license").and_then(|v| v.as_str()).unwrap_or("—");
+        let base = c
+            .get("base_model")
+            .and_then(|v| {
+                v.as_str().map(str::to_string).or_else(|| {
+                    v.as_array().map(|a| {
+                        a.iter()
+                            .filter_map(|x| x.as_str())
+                            .collect::<Vec<_>>()
+                            .join(", ")
+                    })
+                })
+            })
+            .unwrap_or_else(|| "—".into());
+        println!("  licence: {lic}   base: {base}");
+    }
+    if !s("pipeline_tag").is_empty() {
+        println!("  task: {}", s("pipeline_tag"));
+    }
+    println!(
+        "  updated: {}   downloads: {}",
+        s("lastModified"),
+        n("downloads").unwrap_or(0)
+    );
+    if info.get("gated").and_then(|v| v.as_bool()) == Some(true) {
+        println!("  ⚠ gated repo — needs HF_TOKEN");
+    }
+    if let Some(tags) = info.get("tags").and_then(|v| v.as_array()) {
+        let list: Vec<&str> = tags.iter().filter_map(|t| t.as_str()).take(12).collect();
+        println!("  tags: {}", list.join(", "));
+    }
+    if let Some(siblings) = info.get("siblings").and_then(|v| v.as_array()) {
+        println!("  files: {} (use --revision to list)", siblings.len());
+    }
+    if let Some(branches) = refs.get("branches").and_then(|v| v.as_array()) {
+        println!("\nrevisions (branches):");
+        for b in branches {
+            println!(
+                "  {:<24} {}",
+                b.get("name").and_then(|v| v.as_str()).unwrap_or("?"),
+                b.get("targetCommit")
+                    .and_then(|v| v.as_str())
+                    .map(|c| c.get(0..12).unwrap_or(c))
+                    .unwrap_or("")
+            );
+        }
+    }
+    if let Some(tags) = refs.get("tags").and_then(|v| v.as_array()) {
+        let named: Vec<String> = tags
+            .iter()
+            .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(str::to_string))
+            .collect();
+        if !named.is_empty() {
+            println!("tags: {}", named.join(", "));
+        }
+    }
+    if let Some(files) = files.and_then(|f| f.as_array().cloned()) {
+        let mut total = 0u64;
+        println!("\nfiles at {}:", revision.unwrap_or("main"));
+        for f in files.iter().take(40) {
+            let size = f.get("size").and_then(|v| v.as_u64()).unwrap_or(0);
+            total += size;
+            println!(
+                "  {:>10.2} GB  {}",
+                size as f64 / 1e9,
+                f.get("path").and_then(|v| v.as_str()).unwrap_or("?")
+            );
+        }
+        println!("  total ≈ {:.1} GB", total as f64 / 1e9);
+    }
+    Ok(())
+}
+
+fn urlencode(s: &str) -> String {
+    s.bytes()
+        .map(|b| match b {
+            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
+                (b as char).to_string()
+            }
+            b' ' => "+".to_string(),
+            _ => format!("%{b:02X}"),
+        })
+        .collect()
+}
+
+/// CLI dispatch.
+pub async fn dispatch(cmd: &crate::cli::HfCommands) -> Result<()> {
+    use crate::cli::HfCommands as H;
+    match cmd {
+        H::Search { query, limit, json } => search(query, *limit, *json).await,
+        H::Card {
+            repo,
+            revision,
+            json,
+        } => card(repo, revision.as_deref(), *json).await,
+    }
+}

+ 4 - 6
src/main.rs

@@ -44,6 +44,7 @@ mod exec;
 mod gconfig;
 mod handlers;
 mod health;
+mod hf;
 mod home_assistant;
 mod host_keys;
 mod locale;
@@ -380,6 +381,7 @@ async fn main() -> Result<()> {
             .await?;
         }
         Commands::Runlock(cmd) => crate::runlock::run(&config, &cmd)?,
+        Commands::Hf { cmd } => crate::hf::dispatch(&cmd).await?,
         Commands::Cache { cmd } => crate::cache::dispatch(&config, &cmd, dry_run).await?,
         Commands::Repos(a) => {
             let cmd = a.command.clone().unwrap_or(crate::cli::ReposCommands::List);
@@ -420,12 +422,8 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
-/// Handle `serve` command: launch inference containers on a cluster.
-///
-/// Dispatches to `docker::serve` or `tmux::serve` based on the recipe's
-/// `deploy` field. Docker-only CLI opts (sync_image, build_node,
-/// force_build, keep, auto_power_on) are ignored with a warning when
-/// `deploy = "tmux"`.
+/// Handle `serve`: dispatch on the recipe's `deploy` (docker / native / tmux).
+/// Docker-only CLI opts are ignored with a warning on the tmux path.
 pub(crate) async fn handle_serve(
     config: &Config,
     recipe: &str,