From 65632d4b817efda6696caf2d3123913c14549ea2 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 7 Mar 2025 14:08:10 -0500 Subject: [PATCH] downloader code from hf --- Cargo.lock | 4 + crates/desktop/src-tauri/Cargo.toml | 14 +- crates/desktop/src-tauri/src/lib.rs | 1 + .../src/threads/downloader_thread.rs | 11 +- .../src-tauri/src/transfer/download.rs | 268 ++++++++++++++++++ .../desktop/src-tauri/src/transfer/error.rs | 41 +++ crates/desktop/src-tauri/src/transfer/mod.rs | 4 + .../desktop/src-tauri/src/transfer/upload.rs | 204 +++++++++++++ crates/desktop/src-tauri/src/transfer/util.rs | 15 + 9 files changed, 555 insertions(+), 7 deletions(-) create mode 100644 crates/desktop/src-tauri/src/transfer/download.rs create mode 100644 crates/desktop/src-tauri/src/transfer/error.rs create mode 100644 crates/desktop/src-tauri/src/transfer/mod.rs create mode 100644 crates/desktop/src-tauri/src/transfer/upload.rs create mode 100644 crates/desktop/src-tauri/src/transfer/util.rs diff --git a/Cargo.lock b/Cargo.lock index a75a75f549..d1c2589d8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -666,12 +666,14 @@ dependencies = [ "cudarc", "expanduser", "filesys", + "futures 0.3.29", "hf-hub 0.3.2", "image", "libloading 0.8.6", "log", "memory_caching", "once_cell", + "openssl", "rand 0.8.5", "reqwest 0.12.12", "serde", @@ -681,6 +683,8 @@ dependencies = [ "tauri-build", "tauri-plugin-log", "tokenizers 0.15.2", + "tokio", + "tokio-util", ] [[package]] diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 30845e1ae9..7e1ff3aa9f 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -28,24 +28,28 @@ cudarc = "0.13.7" # External anyhow = "1.0" -bytes = "1.10" base64 = "0.22.1" +bytes = "1.10" candle-core = { path = "../../vendor/candle/candle-core", features = ["cuda"] } candle-nn = { path = "../../vendor/candle/candle-nn", features = ["cuda"] } candle-transformers = { path = "../../vendor/candle/candle-transformers", features = ["cuda"] } +expanduser = "1.2.2" +futures = "0.3" hf-hub = "0.3.2" image = "0.25.5" -serde_yaml = "0.9" log = "0.4" once_cell = "1.20.3" -expanduser = "1.2.2" -rand = "0.8" -reqwest = "0.12.12" +openssl = { version = "0.10", features = ["vendored"] } +rand = "0.8.5" +reqwest = { version = "0.12.12", features = ["stream"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +serde_yaml = "0.9" tauri = { version = "2.2.4", features = [] } tauri-plugin-log = "2" tokenizers = "0.15.2" +tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } +tokio-util = { version = "0.7", features = ["codec"] } [features] default = ["cuda"] diff --git a/crates/desktop/src-tauri/src/lib.rs b/crates/desktop/src-tauri/src/lib.rs index 2ee779d49f..8fb77af748 100644 --- a/crates/desktop/src-tauri/src/lib.rs +++ b/crates/desktop/src-tauri/src/lib.rs @@ -3,6 +3,7 @@ pub mod events; pub mod ml; pub mod state; pub mod threads; +pub mod transfer; use crate::endpoints::download_models::download_models; use crate::endpoints::image_endpoint::infer_image; diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 3b9c239928..bef852a3db 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -1,11 +1,18 @@ +use log::error; use tauri::{AppHandle, Emitter}; +use crate::events::notification_event::{NotificationEvent, NotificationModelType}; use crate::state::app_dir::AppDataRoot; pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! { loop { - app.emit("testing", 1); + //let result = app.emit("notification", NotificationEvent::ModelDownloadStarted { + // model_name: "model", + // model_type: NotificationModelType::Unet + //}); + //if let Err(err) = result { + // error!("Could not emit event: {}", err); + //} - println!(">>> Testing: {:?}", app_data_root.path()); std::thread::sleep(std::time::Duration::from_millis(500)); } } diff --git a/crates/desktop/src-tauri/src/transfer/download.rs b/crates/desktop/src-tauri/src/transfer/download.rs new file mode 100644 index 0000000000..becd015b1a --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/download.rs @@ -0,0 +1,268 @@ +//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs + +use crate::transfer::error::Error; +use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME}; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use pyo3::exceptions::PyException; +use pyo3::prelude::*; +use pyo3::types::PyAny; +use rand::{thread_rng, Rng}; +use reqwest::header::{ + HeaderMap, HeaderName, HeaderValue, ToStrError, AUTHORIZATION, CONTENT_LENGTH, CONTENT_RANGE, + RANGE, +}; +use reqwest::Url; +use std::collections::HashMap; +use std::fmt::Display; +use std::fs::remove_file; +use std::io::SeekFrom; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::io::{AsyncReadExt, AsyncSeekExt}; +use tokio::sync::Semaphore; +use tokio::time::sleep; +use tokio_util::codec::{BytesCodec, FramedRead}; + + +/// max_files: Number of open file handles, which determines the maximum number of parallel downloads +/// parallel_failures: Number of maximum failures of different chunks in parallel (cannot exceed max_files) +/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter) +/// +/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in +/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads +#[allow(clippy::too_many_arguments)] +#[pyfunction] +#[pyo3(signature = (url, filename, max_files, chunk_size, parallel_failures=0, max_retries=0, headers=None, callback=None))] +fn download( + url: String, + filename: String, + max_files: usize, + chunk_size: usize, + parallel_failures: usize, + max_retries: usize, + headers: Option>, + callback: Option>, +) -> PyResult<()> { + if parallel_failures > max_files { + return Err(PyException::new_err( + "Error parallel_failures cannot be > max_files".to_string(), + )); + } + if (parallel_failures == 0) != (max_retries == 0) { + return Err(PyException::new_err( + "For retry mechanism you need to set both `parallel_failures` and `max_retries`" + .to_string(), + )); + } + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()? + .block_on(async { + download_async( + url, + filename.clone(), + max_files, + chunk_size, + parallel_failures, + max_retries, + headers, + callback, + ) + .await + }) + .map_err(|err| { + let path = Path::new(&filename); + if path.exists() { + match remove_file(filename) { + Ok(_) => err, + Err(err) => { + PyException::new_err(format!("Error while removing corrupted file: {err}")) + } + } + } else { + err + } + }) +} + + +#[allow(clippy::too_many_arguments)] +async fn download_async( + url: String, + filename: String, + max_files: usize, + chunk_size: usize, + parallel_failures: usize, + max_retries: usize, + input_headers: Option>, + callback: Option>, +) -> PyResult<()> { + let client = reqwest::Client::builder() + // https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526 + .http2_keep_alive_timeout(Duration::from_secs(15)) + .build() + .unwrap(); + + let mut headers = HeaderMap::new(); + let mut auth_token = None; + if let Some(input_headers) = input_headers { + headers.reserve(input_headers.len()); + for (k, v) in input_headers { + let name: HeaderName = k + .try_into() + .map_err(|err| PyException::new_err(format!("Invalid header: {err}")))?; + let value: HeaderValue = AsRef::::as_ref(&v) + .try_into() + .map_err(|err| PyException::new_err(format!("Invalid header value: {err}")))?; + if name == AUTHORIZATION { + auth_token = Some(value); + } else { + headers.insert(name, value); + } + } + }; + + let response = if let Some(token) = auth_token.as_ref() { + client.get(&url).header(AUTHORIZATION, token) + } else { + client.get(&url) + } + .headers(headers.clone()) + .header(RANGE, "bytes=0-0") + .send() + .await + .map_err(|err| PyException::new_err(format!("Error while downloading: {err}")))? + .error_for_status() + .map_err(|err| PyException::new_err(err.to_string()))?; + + // Only call the final redirect URL to avoid overloading the Hub with requests and also + // altering the download count + let redirected_url = response.url(); + if Url::parse(&url) + .map_err(|err| PyException::new_err(format!("failed to parse url: {err}")))? + .host() + == redirected_url.host() + { + if let Some(token) = auth_token { + headers.insert(AUTHORIZATION, token); + } + } + + let content_range = response + .headers() + .get(CONTENT_RANGE) + .ok_or(PyException::new_err("No content length"))? + .to_str() + .map_err(|err| PyException::new_err(format!("Error while downloading: {err}")))?; + + let size: Vec<&str> = content_range.split('/').collect(); + // Content-Range: bytes 0-0/702517648 + // https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Range + let length: usize = size + .last() + .ok_or(PyException::new_err( + "Error while downloading: No size was detected", + ))? + .parse() + .map_err(|err| PyException::new_err(format!("Error while downloading: {err}")))?; + + let mut handles = FuturesUnordered::new(); + let semaphore = Arc::new(Semaphore::new(max_files)); + let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures)); + + for start in (0..length).step_by(chunk_size) { + let url = redirected_url.to_string(); + let filename = filename.clone(); + let client = client.clone(); + let headers = headers.clone(); + + let stop = std::cmp::min(start + chunk_size - 1, length); + let semaphore = semaphore.clone(); + let parallel_failures_semaphore = parallel_failures_semaphore.clone(); + handles.push(tokio::spawn(async move { + let permit = semaphore + .acquire_owned() + .await + .map_err(|err| PyException::new_err(format!("Error while downloading: {err}")))?; + let mut chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await; + let mut i = 0; + if parallel_failures > 0 { + while let Err(dlerr) = chunk { + if i >= max_retries { + return Err(PyException::new_err(format!( + "Failed after too many retries ({max_retries}): {dlerr}" + ))); + } + let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| { + PyException::new_err(format!( + "Failed too many failures in parallel ({parallel_failures}): {dlerr} ({err})" + )) + })?; + + let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME); + sleep(Duration::from_millis(wait_time as u64)).await; + + chunk = download_chunk(&client, &url, &filename, start, stop, headers.clone()).await; + i += 1; + drop(parallel_failure_permit); + } + } + drop(permit); + chunk.map_err(|e| PyException::new_err(format!("Downloading error {e}"))).and(Ok(stop - start)) + })); + } + + // Output the chained result + while let Some(result) = handles.next().await { + match result { + Ok(Ok(size)) => { + if let Some(ref callback) = callback { + callback.call((size,), None)?; + } + } + Ok(Err(py_err)) => { + return Err(py_err); + } + Err(err) => { + return Err(PyException::new_err(format!( + "Error while downloading: {err}" + ))); + } + } + } + Ok(()) +} + +async fn download_chunk( + client: &reqwest::Client, + url: &str, + filename: &str, + start: usize, + stop: usize, + headers: HeaderMap, +) -> Result<(), Error> { + // Process each socket concurrently. + let range = format!("bytes={start}-{stop}"); + let mut file = OpenOptions::new() + .write(true) + .truncate(false) + .create(true) + .open(filename) + .await?; + file.seek(SeekFrom::Start(start as u64)).await?; + let response = client + .get(url) + .headers(headers) + .header(RANGE, range) + .send() + .await? + .error_for_status()?; + let content = response.bytes().await?; + file.write_all(&content).await?; + Ok(()) +} + diff --git a/crates/desktop/src-tauri/src/transfer/error.rs b/crates/desktop/src-tauri/src/transfer/error.rs new file mode 100644 index 0000000000..60c9c173f9 --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/error.rs @@ -0,0 +1,41 @@ +//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs + +use std::fmt::Display; +use reqwest::header::ToStrError; + +#[derive(Debug)] +pub enum Error { + Io(std::io::Error), + Request(reqwest::Error), + ToStrError(ToStrError), +} + +impl From for Error { + fn from(value: std::io::Error) -> Self { + Self::Io(value) + } +} + +impl From for Error { + fn from(value: reqwest::Error) -> Self { + Self::Request(value) + } +} + +impl From for Error { + fn from(value: ToStrError) -> Self { + Self::ToStrError(value) + } +} + +impl Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Io(io) => write!(f, "Io: {io}"), + Self::Request(req) => write!(f, "Request: {req}"), + Self::ToStrError(req) => write!(f, "Response non ascii: {req}"), + } + } +} + +impl std::error::Error for Error {} diff --git a/crates/desktop/src-tauri/src/transfer/mod.rs b/crates/desktop/src-tauri/src/transfer/mod.rs new file mode 100644 index 0000000000..cd6facd1b6 --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/mod.rs @@ -0,0 +1,4 @@ +pub mod download; +pub mod upload; +pub mod util; +pub mod error; \ No newline at end of file diff --git a/crates/desktop/src-tauri/src/transfer/upload.rs b/crates/desktop/src-tauri/src/transfer/upload.rs new file mode 100644 index 0000000000..f790f4ac76 --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/upload.rs @@ -0,0 +1,204 @@ +//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs + +use crate::transfer::error::Error; +use crate::transfer::util::{exponential_backoff, BASE_WAIT_TIME, MAX_WAIT_TIME}; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use pyo3::exceptions::PyException; +use pyo3::prelude::*; +use pyo3::types::PyAny; +use rand::{thread_rng, Rng}; +use reqwest::header::{ + HeaderMap, HeaderName, HeaderValue, ToStrError, AUTHORIZATION, CONTENT_LENGTH, CONTENT_RANGE, + RANGE, +}; +use reqwest::Url; +use std::collections::HashMap; +use std::fmt::Display; +use std::fs::remove_file; +use std::io::SeekFrom; +use std::path::Path; +use std::sync::Arc; +use std::time::Duration; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::io::{AsyncReadExt, AsyncSeekExt}; +use tokio::sync::Semaphore; +use tokio::time::sleep; +use tokio_util::codec::{BytesCodec, FramedRead}; + +/// parts_urls: Dictionary consisting of part numbers as keys and the associated url as values +/// completion_url: The url that should be called when the upload is finished +/// max_files: Number of open file handles, which determines the maximum number of parallel uploads +/// parallel_failures: Number of maximum failures of different chunks in parallel (cannot exceed max_files) +/// max_retries: Number of maximum attempts per chunk. (Retries are exponentially backed off + jitter) +/// +/// The number of threads can be tuned by the environment variable `TOKIO_WORKER_THREADS` as documented in +/// https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.worker_threads +/// +/// See https://docs.aws.amazon.com/AmazonS3/latest/userguide/mpuoverview.html for more information +/// on the multipart upload +#[allow(clippy::too_many_arguments)] +#[pyfunction] +#[pyo3(signature = (file_path, parts_urls, chunk_size, max_files, parallel_failures=0, max_retries=0, callback=None))] +fn multipart_upload( + file_path: String, + parts_urls: Vec, + chunk_size: u64, + max_files: usize, + parallel_failures: usize, + max_retries: usize, + callback: Option>, +) -> PyResult>> { + if parallel_failures > max_files { + return Err(PyException::new_err( + "Error parallel_failures cannot be > max_files".to_string(), + )); + } + if (parallel_failures == 0) != (max_retries == 0) { + return Err(PyException::new_err( + "For retry mechanism you need to set both `parallel_failures` and `max_retries`" + .to_string(), + )); + } + + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()? + .block_on(async { + upload_async( + file_path, + parts_urls, + chunk_size, + max_files, + parallel_failures, + max_retries, + callback, + ) + .await + }) +} + + + +#[allow(clippy::too_many_arguments)] +async fn upload_async( + file_path: String, + parts_urls: Vec, + chunk_size: u64, + max_files: usize, + parallel_failures: usize, + max_retries: usize, + callback: Option>, +) -> PyResult>> { + let client = reqwest::Client::new(); + + let mut handles = FuturesUnordered::new(); + let semaphore = Arc::new(Semaphore::new(max_files)); + let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures)); + + for (part_number, part_url) in parts_urls.iter().enumerate() { + let url = part_url.to_string(); + let path = file_path.to_owned(); + let client = client.clone(); + + let start = (part_number as u64) * chunk_size; + let semaphore = semaphore.clone(); + let parallel_failures_semaphore = parallel_failures_semaphore.clone(); + handles.push(tokio::spawn(async move { + let permit = semaphore + .clone() + .acquire_owned() + .await + .map_err(|err| PyException::new_err(format!("Error acquiring semaphore: {err}")))?; + let mut chunk = upload_chunk(&client, &url, &path, start, chunk_size).await; + let mut i = 0; + if parallel_failures > 0 { + while let Err(ul_err) = chunk { + if i >= max_retries { + return Err(PyException::new_err(format!( + "Failed after too many retries ({max_retries}): {ul_err}" + ))); + } + + let parallel_failure_permit = parallel_failures_semaphore.clone().try_acquire_owned().map_err(|err| { + PyException::new_err(format!( + "Failed too many failures in parallel ({parallel_failures}): {ul_err} ({err})" + )) + })?; + + let wait_time = exponential_backoff(BASE_WAIT_TIME, i, MAX_WAIT_TIME); + sleep(Duration::from_millis(wait_time as u64)).await; + + chunk = upload_chunk(&client, &url, &path, start, chunk_size).await; + i += 1; + drop(parallel_failure_permit); + } + } + drop(permit); + chunk.map_err(|e|{ + match e { + Error::Io(io) => PyException::new_err(format!("Io error {io}")), + Error::Request(req) => PyException::new_err(format!("Error while sending chunk {req}")), + Error::ToStrError(req) => PyException::new_err(format!("Response header contains non ASCII chars: {req}")), + } + }).map(|chunk| (part_number, chunk, chunk_size)) + })); + } + + let mut results: Vec> = vec![HashMap::default(); parts_urls.len()]; + + while let Some(result) = handles.next().await { + match result { + Ok(Ok((part_number, headers, size))) => { + if let Some(ref callback) = callback { + callback.call((size,), None)?; + } + results[part_number] = headers; + } + Ok(Err(py_err)) => { + return Err(py_err); + } + Err(err) => { + return Err(PyException::new_err(format!( + "Error occurred while uploading: {err}" + ))); + } + } + } + + Ok(results) +} + +async fn upload_chunk( + client: &reqwest::Client, + url: &str, + path: &str, + start: u64, + chunk_size: u64, +) -> Result, Error> { + let mut options = OpenOptions::new(); + let mut file = options.read(true).open(path).await?; + let file_size = file.metadata().await?.len(); + let bytes_transferred = std::cmp::min(file_size - start, chunk_size); + + file.seek(SeekFrom::Start(start)).await?; + let chunk = file.take(chunk_size); + + let response = client + .put(url) + .header(CONTENT_LENGTH, bytes_transferred) + .body(reqwest::Body::wrap_stream(FramedRead::new( + chunk, + BytesCodec::new(), + ))) + .send() + .await?; + let response = response.error_for_status()?; + let mut headers = HashMap::new(); + for (name, value) in response.headers().into_iter() { + headers.insert(name.to_string(), value.to_str()?.to_owned()); + } + Ok(headers) +} + diff --git a/crates/desktop/src-tauri/src/transfer/util.rs b/crates/desktop/src-tauri/src/transfer/util.rs new file mode 100644 index 0000000000..70ee4ce908 --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/util.rs @@ -0,0 +1,15 @@ +//! Implemented from https://github.com/huggingface/hf_transfer/blob/main/src/lib.rs + +use rand::{thread_rng, Rng}; +use std::fmt::Display; + +pub const BASE_WAIT_TIME: usize = 300; +pub const MAX_WAIT_TIME: usize = 10_000; + +fn jitter() -> usize { + thread_rng().gen_range(0..=500) +} + +pub fn exponential_backoff(base_wait_time: usize, n: usize, max: usize) -> usize { + (base_wait_time + n.pow(2) + jitter()).min(max) +}