609 lines
20 KiB
Rust
609 lines
20 KiB
Rust
use std::collections::HashMap;
|
|
use std::io::{Read, Seek, SeekFrom, Write};
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
use anyhow::{Context, Result, bail};
|
|
use serde::Serialize;
|
|
|
|
use wg_app_link::private;
|
|
|
|
use crate::session::transport::{Launch, Transport};
|
|
|
|
const USER_AGENT: &str = concat!("ai-server/", env!("CARGO_PKG_VERSION"));
|
|
|
|
const CHUNK: usize = 256 * 1024;
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct LocalModel {
|
|
pub key: String,
|
|
pub repo: String,
|
|
pub file: String,
|
|
pub bytes: u64,
|
|
}
|
|
|
|
/// What a run is doing, or did. Flat rather than a tagged enum carrying its
|
|
/// message, because the phone switches on this and a string it can compare is
|
|
/// easier to render than a variant it has to destructure.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum DownloadState {
|
|
Running,
|
|
/// Reading the finished file back to check it against the hash HuggingFace
|
|
/// publishes. Its own state because it takes real time on a multi-gigabyte
|
|
/// file and "still working" is the honest thing to show, rather than a bar
|
|
/// sitting at 100% for half a minute.
|
|
Verifying,
|
|
Finished,
|
|
Failed,
|
|
Cancelled,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct DownloadStatus {
|
|
pub key: String,
|
|
pub run: u64,
|
|
pub repo: String,
|
|
pub file: String,
|
|
pub state: DownloadState,
|
|
pub done: u64,
|
|
/// What `Content-Length` said, or absent when the server did not say.
|
|
/// Absent means "unknown", never "zero".
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub total: Option<u64>,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub error: Option<String>,
|
|
pub started: f64,
|
|
#[serde(skip_serializing_if = "Option::is_none")]
|
|
pub finished: Option<f64>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct Progress {
|
|
state: DownloadState,
|
|
done: u64,
|
|
total: Option<u64>,
|
|
error: Option<String>,
|
|
started: f64,
|
|
finished: Option<f64>,
|
|
}
|
|
|
|
struct Run {
|
|
id: u64,
|
|
key: String,
|
|
repo: String,
|
|
file: String,
|
|
progress: Mutex<Progress>,
|
|
cancel: AtomicBool,
|
|
}
|
|
|
|
impl Run {
|
|
fn status(&self) -> DownloadStatus {
|
|
let p = self.progress.lock().unwrap();
|
|
DownloadStatus {
|
|
key: self.key.clone(),
|
|
run: self.id,
|
|
repo: self.repo.clone(),
|
|
file: self.file.clone(),
|
|
state: p.state,
|
|
done: p.done,
|
|
total: p.total,
|
|
error: p.error.clone(),
|
|
started: p.started,
|
|
finished: p.finished,
|
|
}
|
|
}
|
|
|
|
fn finish(&self, state: DownloadState, error: Option<String>) {
|
|
let mut p = self.progress.lock().unwrap();
|
|
p.state = state;
|
|
p.error = error;
|
|
p.finished = Some(crate::session::now());
|
|
}
|
|
}
|
|
|
|
pub struct ModelStore {
|
|
dir: PathBuf,
|
|
runs: Mutex<HashMap<String, Arc<Run>>>,
|
|
next_run: AtomicU64,
|
|
}
|
|
|
|
impl ModelStore {
|
|
pub fn new(dir: PathBuf) -> Self {
|
|
Self {
|
|
dir,
|
|
runs: Mutex::new(HashMap::new()),
|
|
next_run: AtomicU64::new(1),
|
|
}
|
|
}
|
|
|
|
/// The repo and file come from a phone, so they are treated as hostile:
|
|
/// every component must be an ordinary name. Rejecting rather than
|
|
/// sanitising, since a silently rewritten path would download the right
|
|
/// bytes to the wrong place.
|
|
fn path_for(&self, repo: &str, file: &str) -> Result<PathBuf> {
|
|
let mut path = self.dir.clone();
|
|
for part in repo.split('/').chain(file.split('/')) {
|
|
if part.is_empty() || part == "." || part == ".." || part.contains('\\') {
|
|
bail!("\"{repo}/{file}\" is not a name this can store: \"{part}\"");
|
|
}
|
|
path.push(part);
|
|
}
|
|
Ok(path)
|
|
}
|
|
|
|
pub fn key_for(repo: &str, file: &str) -> String {
|
|
format!("{repo}/{file}")
|
|
}
|
|
|
|
pub fn list(&self) -> Vec<LocalModel> {
|
|
let mut found = Vec::new();
|
|
collect(&self.dir, &self.dir, &mut found);
|
|
found.sort_by(|a, b| a.key.cmp(&b.key));
|
|
found
|
|
}
|
|
|
|
pub fn downloads(&self) -> Vec<DownloadStatus> {
|
|
let runs = self.runs.lock().unwrap();
|
|
let mut all: Vec<_> = runs.values().map(|run| run.status()).collect();
|
|
all.sort_by_key(|status| std::cmp::Reverse(status.run));
|
|
all
|
|
}
|
|
|
|
/// Starts fetching `file` from `repo`, or returns the run already doing so.
|
|
/// Idempotent on purpose: a phone that lost its connection will press the
|
|
/// button again, and that must join the existing run rather than start a
|
|
/// second one writing the same file.
|
|
pub fn start(self: &Arc<Self>, repo: &str, file: &str) -> Result<DownloadStatus> {
|
|
let key = Self::key_for(repo, file);
|
|
let target = self.path_for(repo, file)?;
|
|
if target.is_file() {
|
|
bail!("{key} is already downloaded");
|
|
}
|
|
|
|
let mut runs = self.runs.lock().unwrap();
|
|
if let Some(existing) = runs.get(&key)
|
|
&& existing.progress.lock().unwrap().state == DownloadState::Running
|
|
{
|
|
return Ok(existing.status());
|
|
}
|
|
|
|
let run = Arc::new(Run {
|
|
id: self.next_run.fetch_add(1, Ordering::Relaxed),
|
|
key: key.clone(),
|
|
repo: repo.to_string(),
|
|
file: file.to_string(),
|
|
progress: Mutex::new(Progress {
|
|
state: DownloadState::Running,
|
|
done: 0,
|
|
total: None,
|
|
error: None,
|
|
started: crate::session::now(),
|
|
finished: None,
|
|
}),
|
|
cancel: AtomicBool::new(false),
|
|
});
|
|
runs.insert(key, Arc::clone(&run));
|
|
let status = run.status();
|
|
drop(runs);
|
|
|
|
let store = Arc::clone(self);
|
|
std::thread::spawn(move || {
|
|
let outcome = store.fetch(&run, &target);
|
|
match outcome {
|
|
Ok(()) if run.cancel.load(Ordering::Relaxed) => {
|
|
run.finish(DownloadState::Cancelled, None);
|
|
tracing::info!("download {} cancelled", run.key);
|
|
}
|
|
Ok(()) => {
|
|
run.finish(DownloadState::Finished, None);
|
|
tracing::info!("download {} finished", run.key);
|
|
}
|
|
Err(err) => {
|
|
let message = format!("{err:#}");
|
|
tracing::warn!("download {} failed: {message}", run.key);
|
|
run.finish(DownloadState::Failed, Some(message));
|
|
}
|
|
}
|
|
});
|
|
Ok(status)
|
|
}
|
|
|
|
pub fn cancel(&self, key: &str) -> Result<DownloadStatus> {
|
|
let runs = self.runs.lock().unwrap();
|
|
let Some(run) = runs.get(key) else {
|
|
bail!("no download for {key}");
|
|
};
|
|
run.cancel.store(true, Ordering::Relaxed);
|
|
Ok(run.status())
|
|
}
|
|
|
|
pub fn delete(&self, key: &str) -> Result<()> {
|
|
let (repo, file) = key.rsplit_once('/').context("a key is repo/file")?;
|
|
let target = self.path_for(repo, file)?;
|
|
let partial = partial_of(&target);
|
|
if !target.is_file() && !partial.is_file() {
|
|
bail!("{key} is not downloaded");
|
|
}
|
|
for path in [&target, &partial] {
|
|
if path.is_file() {
|
|
std::fs::remove_file(path).with_context(|| format!("remove {}", path.display()))?;
|
|
}
|
|
}
|
|
self.runs.lock().unwrap().remove(key);
|
|
Ok(())
|
|
}
|
|
|
|
fn fetch(&self, run: &Run, target: &Path) -> Result<()> {
|
|
let partial = partial_of(target);
|
|
let identity = identity_of(target);
|
|
if let Some(parent) = target.parent() {
|
|
private::create_dir(parent)?;
|
|
}
|
|
|
|
let known = std::fs::read_to_string(&identity)
|
|
.ok()
|
|
.map(|s| s.trim().to_string());
|
|
let have = match known {
|
|
Some(_) => partial.metadata().map(|m| m.len()).unwrap_or(0),
|
|
None => 0,
|
|
};
|
|
|
|
let url = format!(
|
|
"https://huggingface.co/{}/resolve/main/{}",
|
|
run.repo,
|
|
run.file.replace(' ', "%20")
|
|
);
|
|
let (mut response, mut resumed) = request(&url, have)?;
|
|
let mut etag = etag_of(&response);
|
|
|
|
if resumed && etag.is_some() && etag != known {
|
|
tracing::info!(
|
|
"{} changed upstream since the partial was written -- starting again",
|
|
run.key,
|
|
);
|
|
let (fresh, fresh_resumed) = request(&url, 0)?;
|
|
response = fresh;
|
|
resumed = fresh_resumed;
|
|
etag = etag_of(&response);
|
|
}
|
|
|
|
let total: Option<u64> = if resumed {
|
|
response
|
|
.headers()
|
|
.get("content-range")
|
|
.and_then(|v| v.to_str().ok())
|
|
.and_then(|v| {
|
|
v.rsplit_once('/')
|
|
.map(|(_, whole)| whole.trim().to_string())
|
|
})
|
|
.and_then(|whole| whole.parse().ok())
|
|
} else {
|
|
response
|
|
.headers()
|
|
.get("content-length")
|
|
.and_then(|v| v.to_str().ok()?.parse().ok())
|
|
};
|
|
let mut done = if resumed { have } else { 0 };
|
|
{
|
|
let mut p = run.progress.lock().unwrap();
|
|
p.done = done;
|
|
p.total = total;
|
|
}
|
|
|
|
let mut file = std::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.truncate(false)
|
|
.open(&partial)
|
|
.with_context(|| format!("open {}", partial.display()))?;
|
|
if resumed {
|
|
file.seek(SeekFrom::Start(have))
|
|
.context("seek to resume point")?;
|
|
} else {
|
|
file.set_len(0)
|
|
.context("truncate a partial we cannot resume onto")?;
|
|
}
|
|
if let Some(etag) = &etag {
|
|
std::fs::write(&identity, etag).ok();
|
|
}
|
|
|
|
let mut reader = response.body_mut().as_reader();
|
|
let mut buffer = vec![0u8; CHUNK];
|
|
loop {
|
|
if run.cancel.load(Ordering::Relaxed) {
|
|
file.flush().ok();
|
|
return Ok(());
|
|
}
|
|
let read = reader
|
|
.read(&mut buffer)
|
|
.context("reading from HuggingFace")?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
file.write_all(&buffer[..read])
|
|
.context("writing the model file")?;
|
|
done += read as u64;
|
|
run.progress.lock().unwrap().done = done;
|
|
}
|
|
file.flush().context("flushing the model file")?;
|
|
drop(file);
|
|
|
|
// Checked before the rename, so a file that fails never gets the real
|
|
// name and `list` never offers it. With the identity check above this
|
|
// should not fire; it is here because a wrong model is the kind of
|
|
// failure that surfaces as bad output rather than as an error.
|
|
if let Some(expected) = published_sha256(&run.repo, &run.file) {
|
|
run.progress.lock().unwrap().state = DownloadState::Verifying;
|
|
let actual = sha256_of(&partial)?;
|
|
if actual != expected {
|
|
std::fs::remove_file(&partial).ok();
|
|
std::fs::remove_file(&identity).ok();
|
|
bail!(
|
|
"{} arrived corrupted -- HuggingFace publishes sha256 {expected}, what \
|
|
arrived hashes to {actual}. It has been deleted; downloading again \
|
|
starts clean.",
|
|
run.key,
|
|
);
|
|
}
|
|
}
|
|
|
|
std::fs::rename(&partial, target)
|
|
.with_context(|| format!("finish {}", target.display()))?;
|
|
std::fs::remove_file(&identity).ok();
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn sha256_of(path: &Path) -> Result<String> {
|
|
use sha2::{Digest, Sha256};
|
|
let mut file =
|
|
std::fs::File::open(path).with_context(|| format!("reopen {}", path.display()))?;
|
|
let mut hasher = Sha256::new();
|
|
let mut buffer = vec![0u8; CHUNK];
|
|
loop {
|
|
let read = file.read(&mut buffer).context("reading back to verify")?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
hasher.update(&buffer[..read]);
|
|
}
|
|
Ok(hasher
|
|
.finalize()
|
|
.iter()
|
|
.map(|byte| format!("{byte:02x}"))
|
|
.collect())
|
|
}
|
|
|
|
fn request(url: &str, from: u64) -> Result<(ureq::http::Response<ureq::Body>, bool)> {
|
|
let mut get = ureq::get(url).header("User-Agent", USER_AGENT);
|
|
if from > 0 {
|
|
get = get.header("Range", &format!("bytes={from}-"));
|
|
}
|
|
let response = get.call().with_context(|| format!("GET {url}"))?;
|
|
let resumed = response.status() == 206;
|
|
Ok((response, resumed))
|
|
}
|
|
|
|
fn etag_of(response: &ureq::http::Response<ureq::Body>) -> Option<String> {
|
|
Some(
|
|
response
|
|
.headers()
|
|
.get("etag")?
|
|
.to_str()
|
|
.ok()?
|
|
.trim()
|
|
.to_string(),
|
|
)
|
|
}
|
|
|
|
fn identity_of(target: &Path) -> PathBuf {
|
|
let mut name = target.as_os_str().to_os_string();
|
|
name.push(".part.etag");
|
|
PathBuf::from(name)
|
|
}
|
|
|
|
fn partial_of(target: &Path) -> PathBuf {
|
|
let mut name = target.as_os_str().to_os_string();
|
|
name.push(".part");
|
|
PathBuf::from(name)
|
|
}
|
|
|
|
fn collect(root: &Path, dir: &Path, found: &mut Vec<LocalModel>) {
|
|
let Ok(entries) = std::fs::read_dir(dir) else {
|
|
return;
|
|
};
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
if path.is_dir() {
|
|
collect(root, &path, found);
|
|
continue;
|
|
}
|
|
if path.extension().is_none_or(|e| e != "gguf") {
|
|
continue;
|
|
}
|
|
let Ok(relative) = path.strip_prefix(root) else {
|
|
continue;
|
|
};
|
|
let key = relative.to_string_lossy().replace('\\', "/");
|
|
let Some((repo, file)) = key.rsplit_once('/') else {
|
|
continue;
|
|
};
|
|
found.push(LocalModel {
|
|
key: key.clone(),
|
|
repo: repo.to_string(),
|
|
file: file.to_string(),
|
|
bytes: entry.metadata().map(|m| m.len()).unwrap_or(0),
|
|
});
|
|
}
|
|
}
|
|
|
|
const FAR_MODELS_DIR: &str = "~/.local/share/ai-app/models";
|
|
|
|
/// One answer, because two things ask: the list a spawn screen offers,
|
|
/// and the path a session hands `llama-server`. A machine that listed one
|
|
/// directory and served from another would offer models that then failed
|
|
/// to load, which reads as the model being broken.
|
|
pub fn dir_on(transport: &Transport, local: &Path) -> String {
|
|
match transport {
|
|
Transport::Here => local.to_string_lossy().into_owned(),
|
|
Transport::Ssh { ssh, .. } => ssh
|
|
.models_dir
|
|
.as_ref()
|
|
.map_or(FAR_MODELS_DIR.to_string(), |dir| {
|
|
dir.to_string_lossy().into_owned()
|
|
}),
|
|
}
|
|
}
|
|
|
|
/// The local half of this is [`ModelStore::list`], reading the same shape
|
|
/// off this machine's disk; a caller picks by transport, since a setup
|
|
/// with no ssh *is* this machine and asking a shell about it would be a
|
|
/// slower way to the same answer. What must not happen is offering this
|
|
/// backend's downloads for a session on another machine: the file has to
|
|
/// be where `llama-server` runs, and a list that says otherwise is a
|
|
/// claim about the wrong filesystem.
|
|
///
|
|
/// `dir` is that machine's models directory, `~` included -- expanded on
|
|
/// the far side, which is the only place that knows what it is. A
|
|
/// directory that is not there is an empty list rather than a failure: a
|
|
/// machine that has never had a model put on it is an ordinary state, and
|
|
/// the same one as a machine whose directory exists and is empty.
|
|
pub async fn on_machine(transport: &Transport, dir: &str) -> Result<Vec<LocalModel>> {
|
|
let script = "p=$1; case $p in \"~\") p=$HOME;; \"~/\"*) p=$HOME/${p#\"~/\"};; esac; \
|
|
[ -d \"$p\" ] || exit 0; \
|
|
find \"$p\" -type f -name '*.gguf' -printf '%s\\t%P\\0'";
|
|
let launch = Launch::new(
|
|
"sh",
|
|
vec![
|
|
"-c".to_string(),
|
|
script.to_string(),
|
|
"sh".to_string(),
|
|
dir.to_string(),
|
|
],
|
|
None,
|
|
);
|
|
let out = transport.capture(&launch).await?;
|
|
let mut found: Vec<LocalModel> = out
|
|
.split('\0')
|
|
.filter(|record| !record.is_empty())
|
|
.filter_map(|record| record.split_once('\t'))
|
|
.filter_map(|(bytes, key)| {
|
|
let (repo, file) = key.rsplit_once('/')?;
|
|
Some(LocalModel {
|
|
key: key.to_string(),
|
|
repo: repo.to_string(),
|
|
file: file.to_string(),
|
|
bytes: bytes.trim().parse().unwrap_or(0),
|
|
})
|
|
})
|
|
.collect();
|
|
found.sort_by(|a, b| a.key.cmp(&b.key));
|
|
Ok(found)
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct RemoteRepo {
|
|
pub id: String,
|
|
pub downloads: u64,
|
|
pub likes: u64,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Serialize)]
|
|
#[serde(rename_all = "camelCase")]
|
|
pub struct RemoteFile {
|
|
pub path: String,
|
|
pub bytes: u64,
|
|
pub have: bool,
|
|
}
|
|
|
|
pub fn search(query: &str) -> Result<Vec<RemoteRepo>> {
|
|
let url = format!(
|
|
"https://huggingface.co/api/models?search={}&filter=gguf&limit=25&sort=downloads&direction=-1",
|
|
urlencode(query)
|
|
);
|
|
let body = get_json(&url)?;
|
|
let list = body
|
|
.as_array()
|
|
.context("HuggingFace returned something that is not a list")?;
|
|
Ok(list
|
|
.iter()
|
|
.filter_map(|m| {
|
|
Some(RemoteRepo {
|
|
id: m.get("id")?.as_str()?.to_string(),
|
|
downloads: m
|
|
.get("downloads")
|
|
.and_then(serde_json::Value::as_u64)
|
|
.unwrap_or(0),
|
|
likes: m
|
|
.get("likes")
|
|
.and_then(serde_json::Value::as_u64)
|
|
.unwrap_or(0),
|
|
})
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
fn published_sha256(repo: &str, file: &str) -> Option<String> {
|
|
let url = format!("https://huggingface.co/api/models/{repo}/tree/main?expand=true");
|
|
let body = get_json(&url).ok()?;
|
|
body.as_array()?.iter().find_map(|f| {
|
|
(f.get("path")?.as_str()? == file)
|
|
.then(|| f.get("lfs")?.get("oid")?.as_str().map(str::to_string))?
|
|
})
|
|
}
|
|
|
|
pub fn files(repo: &str, store: &ModelStore) -> Result<Vec<RemoteFile>> {
|
|
let url = format!("https://huggingface.co/api/models/{repo}/tree/main");
|
|
let body = get_json(&url)?;
|
|
let list = body
|
|
.as_array()
|
|
.context("HuggingFace returned something that is not a list")?;
|
|
let have: std::collections::HashSet<String> = store.list().into_iter().map(|m| m.key).collect();
|
|
let mut files: Vec<RemoteFile> = list
|
|
.iter()
|
|
.filter_map(|f| {
|
|
let path = f.get("path")?.as_str()?.to_string();
|
|
if !path.ends_with(".gguf") {
|
|
return None;
|
|
}
|
|
Some(RemoteFile {
|
|
bytes: f
|
|
.get("size")
|
|
.and_then(serde_json::Value::as_u64)
|
|
.unwrap_or(0),
|
|
have: have.contains(&ModelStore::key_for(repo, &path)),
|
|
path,
|
|
})
|
|
})
|
|
.collect();
|
|
files.sort_by_key(|f| f.bytes);
|
|
Ok(files)
|
|
}
|
|
|
|
fn get_json(url: &str) -> Result<serde_json::Value> {
|
|
let text = ureq::get(url)
|
|
.header("User-Agent", USER_AGENT)
|
|
.call()
|
|
.and_then(|mut r| r.body_mut().read_to_string())
|
|
.with_context(|| format!("GET {url}"))?;
|
|
serde_json::from_str(&text).with_context(|| format!("{url} did not return JSON"))
|
|
}
|
|
|
|
fn urlencode(value: &str) -> String {
|
|
value
|
|
.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(),
|
|
other => format!("%{other:02X}"),
|
|
})
|
|
.collect()
|
|
}
|