aria_engine/
download.rs

1//! Dashboard-only model auto-download for the Rust SDK.
2//!
3//! Mirrors the `dashboard` branch of `openai/src/download.rs`: resolve the
4//! `slug`/`quant` from the model name, request the meta URL with a bearer
5//! token, stream the zip, validate the zip magic, extract (flattening a single
6//! top-level subdir), and verify the resulting bundle.
7
8use std::io::{Read, Write};
9use std::path::{Path, PathBuf};
10use thiserror::Error;
11
12#[derive(Debug, Error)]
13pub enum DownloadError {
14    #[error("invalid model name: {0}")]
15    InvalidModelName(String),
16    #[error("dashboard request failed: {0}")]
17    Request(String),
18    #[error("download stream failed: {0}")]
19    Stream(String),
20    #[error("invalid zip archive: {0}")]
21    BadZip(String),
22    #[error("extraction failed: {0}")]
23    Extract(String),
24    #[error("invalid bundle after download: {0}")]
25    InvalidBundle(String),
26    #[error("io error: {0}")]
27    Io(#[from] std::io::Error),
28}
29
30const DEFAULT_SITE: &str = "https://ariacompute.com";
31
32fn aria_home() -> Result<PathBuf, DownloadError> {
33    if let Ok(override_home) = std::env::var("ARIA_COMPUTE_HOME") {
34        if !override_home.is_empty() {
35            return Ok(PathBuf::from(override_home));
36        }
37    }
38    let home = if cfg!(windows) {
39        std::env::var("USERPROFILE").map_err(|_| {
40            DownloadError::Io(std::io::Error::new(
41                std::io::ErrorKind::NotFound,
42                "could not resolve home directory",
43            ))
44        })?
45    } else {
46        std::env::var("HOME").map_err(|_| {
47            DownloadError::Io(std::io::Error::new(
48                std::io::ErrorKind::NotFound,
49                "could not resolve home directory",
50            ))
51        })?
52    };
53    Ok(PathBuf::from(home).join(".ariacompute"))
54}
55
56fn models_dir() -> Result<PathBuf, DownloadError> {
57    Ok(aria_home()?.join("models"))
58}
59
60/// Parse a model name such as `gemma-4-e2b-it_q4` into `(slug, quant)`.
61/// Quant follows the `_q4`/`_q8`/`_q326`/`_q3.26` suffix; defaults to `int4`.
62/// Optional codebook-share `_channel` / `_group` (e.g. `*_q326_channel`) is ignored.
63fn parse_bundle_name(model: &str) -> Result<(String, String), DownloadError> {
64    if model.is_empty() || model.contains('/') || model.contains('\\') {
65        return Err(DownloadError::InvalidModelName(model.to_string()));
66    }
67    let (slug, quant) = if let Some(idx) = model.rfind("_q") {
68        let (slug, suffix) = model.split_at(idx);
69        let mut suffix = &suffix[2..];
70        if let Some(core) = suffix.strip_suffix("_channel") {
71            suffix = core;
72        } else if let Some(core) = suffix.strip_suffix("_group") {
73            suffix = core;
74        }
75        let quant = match suffix {
76            "4" => "int4",
77            "8" => "int8",
78            "326" | "3.26" => "int326",
79            other => {
80                return Err(DownloadError::InvalidModelName(format!(
81                    "unknown quant suffix _q{other}"
82                )))
83            }
84        };
85        (slug.to_string(), quant.to_string())
86    } else {
87        (model.to_string(), "int4".to_string())
88    };
89    if slug.is_empty() {
90        return Err(DownloadError::InvalidModelName(model.to_string()));
91    }
92    Ok((slug, quant))
93}
94
95fn url_encode(s: &str) -> String {
96    let mut out = String::with_capacity(s.len());
97    for b in s.bytes() {
98        match b {
99            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
100                out.push(b as char)
101            }
102            _ => out.push_str(&format!("%{:02X}", b)),
103        }
104    }
105    out
106}
107
108#[derive(serde::Deserialize)]
109struct DashboardMeta {
110    url: String,
111}
112
113fn meta_url(site: &str, slug: &str, quant: &str) -> String {
114    let sdk = "v1.0";
115    format!(
116        "{}/api/dashboard/models/{}/download?quant={}&sdk={}&format=json",
117        site.trim_end_matches('/'),
118        url_encode(slug),
119        url_encode(quant),
120        url_encode(sdk),
121    )
122}
123
124fn is_valid_bundle(dir: &Path) -> bool {
125    let weight = dir.join("weight.bin");
126    let config = dir.join("config.json");
127    if !weight.is_file() || !config.is_file() {
128        return false;
129    }
130    let Ok(raw) = std::fs::read_to_string(&config) else {
131        return false;
132    };
133    let Ok(v) = serde_json::from_str::<serde_json::Value>(&raw) else {
134        return false;
135    };
136    v.get("format")
137        .and_then(|x| x.as_str())
138        .map(|f| f == "aria-quant-bundle")
139        .unwrap_or(false)
140}
141
142/// Flatten a single top-level subdir (when config.json sits inside it).
143fn flatten_single_subdir(dir: &Path) -> std::io::Result<()> {
144    let mut entries: Vec<_> = std::fs::read_dir(dir)?
145        .filter_map(|e| e.ok())
146        .map(|e| e.path())
147        .collect();
148    if entries.len() != 1 {
149        return Ok(());
150    }
151    let only = &entries[0];
152    if !only.is_dir() {
153        return Ok(());
154    }
155    if !only.join("config.json").is_file() {
156        return Ok(());
157    }
158    let tmp = dir.join(format!(".flatten_{}", std::process::id()));
159    std::fs::rename(only, &tmp)?;
160    for entry in std::fs::read_dir(&tmp)? {
161        let entry = entry?;
162        std::fs::rename(entry.path(), dir.join(entry.file_name()))?;
163    }
164    std::fs::remove_dir_all(&tmp)?;
165    entries.clear();
166    Ok(())
167}
168
169fn extract_zip(data: &[u8], dest: &Path) -> Result<(), DownloadError> {
170    if data.len() < 4 || &data[0..2] != b"PK" {
171        return Err(DownloadError::BadZip("missing PK magic".into()));
172    }
173    std::fs::create_dir_all(dest)?;
174    let cursor = std::io::Cursor::new(data);
175    let mut archive = zip::ZipArchive::new(cursor)
176        .map_err(|e| DownloadError::BadZip(e.to_string()))?;
177    for i in 0..archive.len() {
178        let mut file = archive
179            .by_index(i)
180            .map_err(|e| DownloadError::Extract(e.to_string()))?;
181        let name = match file.enclosed_name() {
182            Some(n) => n.to_path_buf(),
183            None => continue,
184        };
185        let out_path = dest.join(&name);
186        if file.is_dir() {
187            std::fs::create_dir_all(&out_path)?;
188        } else {
189            if let Some(parent) = out_path.parent() {
190                std::fs::create_dir_all(parent)?;
191            }
192            let mut buf = Vec::with_capacity(file.size() as usize);
193            file.read_to_end(&mut buf)?;
194            let mut out = std::fs::File::create(&out_path)?;
195            out.write_all(&buf)?;
196        }
197    }
198    flatten_single_subdir(dest).map_err(DownloadError::Io)?;
199    Ok(())
200}
201
202fn atomic_replace(src: &Path, dst: &Path) -> std::io::Result<()> {
203    if dst.exists() {
204        std::fs::remove_dir_all(dst)?;
205    }
206    std::fs::rename(src, dst)
207}
208
209/// Download `model` from the Dashboard private source into
210/// `~/.ariacompute/models/{model}`, then return that directory.
211///
212/// If a valid bundle already exists at the cache path, the download is skipped.
213pub fn download_model(
214    model: &str,
215    token: &str,
216    site: Option<&str>,
217) -> Result<PathBuf, DownloadError> {
218    let (slug, quant) = parse_bundle_name(model)?;
219    let site = site.unwrap_or(DEFAULT_SITE);
220    let cache = models_dir()?.join(model);
221
222    if cache.exists() && is_valid_bundle(&cache) {
223        return Ok(cache);
224    }
225
226    let url = meta_url(site, &slug, &quant);
227    let agent = ureq::AgentBuilder::new().build();
228    let meta: DashboardMeta = agent
229        .get(&url)
230        .set("Authorization", &format!("Bearer {token}"))
231        .call()
232        .map_err(|e| DownloadError::Request(e.to_string()))?
233        .into_json::<DashboardMeta>()
234        .map_err(|e| DownloadError::Request(e.to_string()))?;
235    if meta.url.is_empty() {
236        return Err(DownloadError::Request(
237            "dashboard meta returned empty url".into(),
238        ));
239    }
240
241    let mut reader = agent
242        .get(&meta.url)
243        .set("Authorization", &format!("Bearer {token}"))
244        .call()
245        .map_err(|e| DownloadError::Stream(e.to_string()))?
246        .into_reader();
247    let mut data = Vec::new();
248    reader
249        .read_to_end(&mut data)
250        .map_err(|e| DownloadError::Stream(e.to_string()))?;
251
252    let staging = models_dir()?.join(format!(".{}.partial", model));
253    if staging.exists() {
254        std::fs::remove_dir_all(&staging)?;
255    }
256    extract_zip(&data, &staging)?;
257    if !is_valid_bundle(&staging) {
258        let _ = std::fs::remove_dir_all(&staging);
259        return Err(DownloadError::InvalidBundle(
260            "downloaded archive did not contain a valid aria-quant-bundle".into(),
261        ));
262    }
263    atomic_replace(&staging, &cache)?;
264    Ok(cache)
265}
aria_engine/download.rs

aria_engine/
download.rs
