aria_engine/
download.rs

1//! Regional public-hub model auto-download for the Rust SDK.
2//!
3//! Matches `aria-engine download`: `.com` → Hugging Face, `.cn` → ModelScope.
4//! Dashboard zip meta is not used. A Dashboard `sk-` / `bfvk-` token is ignored
5//! for hub auth.
6
7use std::io::{Read, Write};
8use std::path::{Path, PathBuf};
9use std::time::Duration;
10use thiserror::Error;
11
12#[derive(Debug, Error)]
13pub enum DownloadError {
14    #[error("invalid model name: {0}")]
15    InvalidModelName(String),
16    #[error("{0}")]
17    Request(String),
18    #[error("download stream failed: {0}")]
19    Stream(String),
20    #[error("invalid bundle after download: {0}")]
21    InvalidBundle(String),
22    #[error("io error: {0}")]
23    Io(#[from] std::io::Error),
24}
25
26const DEFAULT_SITE: &str = "https://ariacompute.com";
27const DEFAULT_SDK: &str = "v1.0";
28const HUB_REQUIRED: &[&str] = &["config.json", "weight.bin"];
29const HUB_OPTIONAL: &[&str] = &[
30    "tokenizer.json",
31    "tokenizer.model",
32    "tokenizer_config.json",
33    "special_tokens_map.json",
34    "vocab.json",
35    "merges.txt",
36];
37
38fn aria_home() -> Result<PathBuf, DownloadError> {
39    if let Ok(override_home) = std::env::var("ARIA_COMPUTE_HOME") {
40        if !override_home.is_empty() {
41            return Ok(PathBuf::from(override_home));
42        }
43    }
44    let home = if cfg!(windows) {
45        std::env::var("USERPROFILE").map_err(|_| {
46            DownloadError::Io(std::io::Error::new(
47                std::io::ErrorKind::NotFound,
48                "could not resolve home directory",
49            ))
50        })?
51    } else {
52        std::env::var("HOME").map_err(|_| {
53            DownloadError::Io(std::io::Error::new(
54                std::io::ErrorKind::NotFound,
55                "could not resolve home directory",
56            ))
57        })?
58    };
59    Ok(PathBuf::from(home).join(".ariacompute"))
60}
61
62fn models_dir() -> Result<PathBuf, DownloadError> {
63    Ok(aria_home()?.join("models"))
64}
65
66/// Parse a model name such as `gemma-4-e2b-it_q4` into `(slug, quant)`.
67/// Quant follows the `_q4`/`_q8`/`_q326`/`_q3.26` suffix; defaults to `int4`.
68/// Optional codebook-share `_channel` / `_group` (e.g. `*_q326_channel`) is ignored.
69fn parse_bundle_name(model: &str) -> Result<(String, String), DownloadError> {
70    if model.is_empty() || model.contains('/') || model.contains('\\') {
71        return Err(DownloadError::InvalidModelName(model.to_string()));
72    }
73    let (slug, quant) = if let Some(idx) = model.rfind("_q") {
74        let (slug, suffix) = model.split_at(idx);
75        let mut suffix = &suffix[2..];
76        if let Some(core) = suffix.strip_suffix("_channel") {
77            suffix = core;
78        } else if let Some(core) = suffix.strip_suffix("_group") {
79            suffix = core;
80        }
81        let quant = match suffix {
82            "4" => "int4",
83            "8" => "int8",
84            "326" | "3.26" => "int326",
85            other => {
86                return Err(DownloadError::InvalidModelName(format!(
87                    "unknown quant suffix _q{other}"
88                )))
89            }
90        };
91        (slug.to_string(), quant.to_string())
92    } else {
93        (model.to_string(), "int4".to_string())
94    };
95    if slug.is_empty() {
96        return Err(DownloadError::InvalidModelName(model.to_string()));
97    }
98    Ok((slug, quant))
99}
100
101fn preferred_public_hub(site: Option<&str>) -> &'static str {
102    if site
103        .unwrap_or("")
104        .to_ascii_lowercase()
105        .contains("ariacompute.cn")
106    {
107        "modelscope"
108    } else {
109        "huggingface"
110    }
111}
112
113fn hub_bearer(token: &str) -> Option<&str> {
114    let t = token.trim();
115    if t.is_empty() {
116        return None;
117    }
118    let low = t.to_ascii_lowercase();
119    if low.starts_with("sk-") || low.starts_with("bfvk-") {
120        None
121    } else {
122        Some(t)
123    }
124}
125
126fn unquote_yaml(v: &str) -> String {
127    let t = v.trim();
128    let b = t.as_bytes();
129    if b.len() >= 2
130        && ((b[0] == b'"' && *b.last().unwrap() == b'"')
131            || (b[0] == b'\'' && *b.last().unwrap() == b'\''))
132    {
133        t[1..t.len() - 1].to_string()
134    } else {
135        t.to_string()
136    }
137}
138
139fn config_yml_scalar(key: &str) -> Option<String> {
140    let path = aria_home().ok()?.join("config.yml");
141    let raw = std::fs::read_to_string(path).ok()?;
142    for line in raw.lines() {
143        if line.starts_with(' ') || line.starts_with('\t') {
144            continue;
145        }
146        let s = line.trim();
147        if s.is_empty() || s.starts_with('#') {
148            continue;
149        }
150        let Some((k, v)) = s.split_once(':') else {
151            continue;
152        };
153        if k.trim() != key {
154            continue;
155        }
156        let val = unquote_yaml(v);
157        if val.is_empty() {
158            return None;
159        }
160        return Some(val);
161    }
162    None
163}
164
165fn hub_token_field(source: &str) -> &'static str {
166    if source == "modelscope" {
167        "modelscope_api_token"
168    } else {
169        "hf_token"
170    }
171}
172
173fn resolve_hub_token(
174    source: &str,
175    token: &str,
176    hf_token: Option<&str>,
177    modelscope_api_token: Option<&str>,
178) -> Option<String> {
179    let named = if source == "modelscope" {
180        modelscope_api_token.unwrap_or("")
181    } else {
182        hf_token.unwrap_or("")
183    };
184    let from_cfg = config_yml_scalar(hub_token_field(source)).unwrap_or_default();
185    for cand in [named, token, from_cfg.as_str()] {
186        if let Some(b) = hub_bearer(cand) {
187            return Some(b.to_string());
188        }
189    }
190    None
191}
192
193fn hub_path_names(model: &str) -> Vec<String> {
194    let mut names = vec![model.to_string()];
195    let mut lower = model.to_ascii_lowercase();
196    let mut core = model.to_string();
197    for suf in ["_channel", "_group"] {
198        if lower.ends_with(suf) {
199            core = model[..model.len() - suf.len()].to_string();
200            lower = core.to_ascii_lowercase();
201            break;
202        }
203    }
204    let mut stems = vec![core.clone()];
205    if lower.ends_with("_q326") {
206        stems.push(format!("{}q3.26", &core[..core.len() - 5]));
207    } else if lower.ends_with("_q3.26") {
208        stems.push(format!("{}q326", &core[..core.len() - 6]));
209    }
210    for stem in stems {
211        for share in ["", "_channel", "_group"] {
212            let cand = format!("{stem}{share}");
213            if !names.iter().any(|n| n == &cand) {
214                names.push(cand);
215            }
216        }
217    }
218    names
219}
220
221fn hub_file_urls(source: &str, model: &str, file: &str) -> Vec<String> {
222    let mut urls = Vec::new();
223    for name in hub_path_names(model) {
224        if source == "modelscope" {
225            for repo in [format!("AriaCompute/{name}"), "AriaCompute/model".into()] {
226                urls.push(format!(
227                    "https://www.modelscope.cn/models/{repo}/resolve/master/{DEFAULT_SDK}/{name}/{file}"
228                ));
229                urls.push(format!(
230                    "https://modelscope.cn/models/{repo}/resolve/master/{DEFAULT_SDK}/{name}/{file}"
231                ));
232            }
233        } else {
234            for repo in [format!("ariacompute/{name}"), "ariacompute/model".into()] {
235                urls.push(format!(
236                    "https://huggingface.co/{repo}/resolve/main/{DEFAULT_SDK}/{name}/{file}"
237                ));
238            }
239        }
240    }
241    urls
242}
243
244fn is_valid_bundle(dir: &Path) -> bool {
245    let weight = dir.join("weight.bin");
246    let config = dir.join("config.json");
247    if !weight.is_file() || !config.is_file() {
248        return false;
249    }
250    let Ok(raw) = std::fs::read_to_string(&config) else {
251        return false;
252    };
253    let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) else {
254        return false;
255    };
256    v.get("format")
257        .and_then(|x| x.as_str())
258        .map(|f| f == "aria-quant-bundle")
259        .unwrap_or(false)
260}
261
262fn atomic_replace(src: &Path, dst: &Path) -> std::io::Result<()> {
263    if dst.exists() {
264        std::fs::remove_dir_all(dst)?;
265    }
266    std::fs::rename(src, dst)
267}
268
269fn auth_error(source: &str, code: u16) -> DownloadError {
270    let field = if source == "modelscope" {
271        "modelscope_api_token"
272    } else {
273        "hf_token"
274    };
275    DownloadError::Request(format!(
276        "auth failed HTTP {code}; set {field} via aria-engine auth (do not pass a Dashboard sk-/bfvk- key as the hub token)"
277    ))
278}
279
280fn fetch_url_to_file(url: &str, dest: &Path, token: Option<&str>) -> Result<(), DownloadError> {
281    let agent = ureq::AgentBuilder::new()
282        .timeout(Duration::from_secs(600))
283        .build();
284    let mut req = agent.get(url);
285    if let Some(t) = token {
286        req = req.set("Authorization", &format!("Bearer {t}"));
287    }
288    let resp = match req.call() {
289        Ok(r) => r,
290        Err(ureq::Error::Status(code, _)) if code == 401 || code == 403 => {
291            return Err(DownloadError::Request(format!("HTTP {code}")));
292        }
293        Err(ureq::Error::Status(code, _)) => {
294            return Err(DownloadError::Request(format!("HTTP {code}")));
295        }
296        Err(e) => return Err(DownloadError::Request(e.to_string())),
297    };
298    if let Some(parent) = dest.parent() {
299        std::fs::create_dir_all(parent)?;
300    }
301    let mut reader = resp.into_reader();
302    let mut out = std::fs::File::create(dest)?;
303    let mut buf = [0u8; 1024 * 1024];
304    loop {
305        let n = reader
306            .read(&mut buf)
307            .map_err(|e| DownloadError::Stream(e.to_string()))?;
308        if n == 0 {
309            break;
310        }
311        out.write_all(&buf[..n])?;
312    }
313    Ok(())
314}
315
316fn fetch_hub_file(
317    source: &str,
318    model: &str,
319    file: &str,
320    dest: &Path,
321    token: Option<&str>,
322    required: bool,
323) -> Result<bool, DownloadError> {
324    let mut last: Option<DownloadError> = None;
325    for url in hub_file_urls(source, model, file) {
326        match fetch_url_to_file(&url, dest, token) {
327            Ok(()) => return Ok(true),
328            Err(DownloadError::Request(msg))
329                if msg.contains("HTTP 401") || msg.contains("HTTP 403") =>
330            {
331                let code = if msg.contains("401") { 401 } else { 403 };
332                return Err(auth_error(source, code));
333            }
334            Err(e) => last = Some(e),
335        }
336    }
337    if required {
338        Err(DownloadError::Request(format!(
339            "{source}: missing {file}{}",
340            last.map(|e| format!(": {e}")).unwrap_or_default()
341        )))
342    } else {
343        Ok(false)
344    }
345}
346
347/// Download `model` from the regional public hub into
348/// `~/.ariacompute/models/{model}`, then return that directory.
349///
350/// If a valid bundle already exists at the cache path, the download is skipped.
351/// Hub auth: explicit `hf_token` / `modelscope_api_token`, then generic `token`,
352/// then `~/.ariacompute/config.yml` (same keys as `aria-engine auth`).
353/// Dashboard `sk-` / `bfvk-` keys are not sent to the hub.
354pub fn download_model(
355    model: &str,
356    token: &str,
357    site: Option<&str>,
358) -> Result<PathBuf, DownloadError> {
359    download_model_auth(model, token, site, None, None)
360}
361
362/// Like [`download_model`], with named hub tokens matching `aria-engine auth`.
363pub fn download_model_auth(
364    model: &str,
365    token: &str,
366    site: Option<&str>,
367    hf_token: Option<&str>,
368    modelscope_api_token: Option<&str>,
369) -> Result<PathBuf, DownloadError> {
370    parse_bundle_name(model)?;
371    let site = site.unwrap_or(DEFAULT_SITE);
372    let source = preferred_public_hub(Some(site));
373    let hub_owned = resolve_hub_token(source, token, hf_token, modelscope_api_token);
374    let hub_token = hub_owned.as_deref();
375    let cache = models_dir()?.join(model);
376
377    if cache.exists() && is_valid_bundle(&cache) {
378        return Ok(cache);
379    }
380
381    let staging = models_dir()?.join(format!(".{}.partial", model));
382    if staging.exists() {
383        std::fs::remove_dir_all(&staging)?;
384    }
385    std::fs::create_dir_all(&staging)?;
386    let result = (|| {
387        for file in HUB_REQUIRED {
388            fetch_hub_file(source, model, file, &staging.join(file), hub_token, true)?;
389        }
390        for extra in HUB_OPTIONAL {
391            let _ = fetch_hub_file(source, model, extra, &staging.join(extra), hub_token, false);
392        }
393        if !is_valid_bundle(&staging) {
394            return Err(DownloadError::InvalidBundle(
395                "need weight.bin + aria-quant-bundle config.json".into(),
396            ));
397        }
398        atomic_replace(&staging, &cache)?;
399        Ok(cache.clone())
400    })();
401    if result.is_err() && staging.exists() {
402        let _ = std::fs::remove_dir_all(&staging);
403    }
404    result
405}
406
407#[cfg(test)]
408mod tests {
409    use super::*;
410    use std::sync::Mutex;
411
412    static ENV_LOCK: Mutex<()> = Mutex::new(());
413
414    #[test]
415    fn preferred_hub_follows_site_tld() {
416        assert_eq!(
417            preferred_public_hub(Some("https://ariacompute.com")),
418            "huggingface"
419        );
420        assert_eq!(
421            preferred_public_hub(Some("https://ariacompute.cn")),
422            "modelscope"
423        );
424        assert_eq!(preferred_public_hub(None), "huggingface");
425    }
426
427    #[test]
428    fn dashboard_token_not_sent_to_hub() {
429        assert!(hub_bearer("sk-bf-95076ed1-8c1a-4efa-b33c-f52c1d7f9f24").is_none());
430        assert!(hub_bearer("bfvk-test").is_none());
431        assert_eq!(hub_bearer("hf_abc"), Some("hf_abc"));
432    }
433
434    #[test]
435    fn hub_urls_follow_upload_layout() {
436        let hf = hub_file_urls("huggingface", "gemma-4-e2b-it_q4", "config.json");
437        assert!(hf.iter().any(|u| u.contains(
438            "/ariacompute/gemma-4-e2b-it_q4/resolve/main/v1.0/gemma-4-e2b-it_q4/config.json"
439        )));
440        let ms = hub_file_urls("modelscope", "gemma-4-e2b-it_q4", "weight.bin");
441        assert!(ms
442            .iter()
443            .any(|u| u.contains("/v1.0/gemma-4-e2b-it_q4/weight.bin")));
444        assert!(hf.iter().chain(ms.iter()).all(|u| !u.contains("/api/dashboard/")));
445    }
446
447    #[test]
448    fn parse_channel_suffix() {
449        let (slug, quant) = parse_bundle_name("foo_q326_channel").unwrap();
450        assert_eq!(slug, "foo");
451        assert_eq!(quant, "int326");
452    }
453
454    #[test]
455    fn cached_bundle_skips_download() {
456        let _guard = ENV_LOCK.lock().unwrap();
457        let tmp = tempfile::tempdir().unwrap();
458        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
459        let cache = tmp.path().join("models").join("foo_q4");
460        std::fs::create_dir_all(&cache).unwrap();
461        std::fs::write(cache.join("weight.bin"), b"x").unwrap();
462        std::fs::write(
463            cache.join("config.json"),
464            br#"{"format":"aria-quant-bundle"}"#,
465        )
466        .unwrap();
467        let got = download_model("foo_q4", "", None).unwrap();
468        assert_eq!(got, cache);
469        std::env::remove_var("ARIA_COMPUTE_HOME");
470    }
471
472    #[test]
473    fn resolve_named_and_config_yml() {
474        let _guard = ENV_LOCK.lock().unwrap();
475        let tmp = tempfile::tempdir().unwrap();
476        std::env::set_var("ARIA_COMPUTE_HOME", tmp.path());
477        std::fs::write(
478            tmp.path().join("config.yml"),
479            "hf_token: hf_from_yml\nmodelscope_api_token: \"ms_from_yml\"\n",
480        )
481        .unwrap();
482        assert_eq!(
483            resolve_hub_token("huggingface", "hf_generic", Some("hf_named"), None).as_deref(),
484            Some("hf_named")
485        );
486        assert_eq!(
487            resolve_hub_token("modelscope", "", None, Some("ms_named")).as_deref(),
488            Some("ms_named")
489        );
490        assert_eq!(
491            resolve_hub_token("huggingface", "", None, None).as_deref(),
492            Some("hf_from_yml")
493        );
494        assert_eq!(
495            resolve_hub_token("modelscope", "", None, None).as_deref(),
496            Some("ms_from_yml")
497        );
498        assert_eq!(
499            resolve_hub_token("huggingface", "sk-bf-not-hub", None, None).as_deref(),
500            Some("hf_from_yml")
501        );
502        std::env::remove_var("ARIA_COMPUTE_HOME");
503    }
504}
aria_engine/download.rs

aria_engine/
download.rs
