From a4f876daa96de46292ccfc591a336dc94b31ce3e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 10 Mar 2025 21:51:07 -0400 Subject: [PATCH] Multithreaded downloader + notifications (#447) --- Cargo.lock | 85 +++++- crates/desktop/src-tauri/Cargo.toml | 16 +- .../src/events/notification_event.rs | 8 + crates/desktop/src-tauri/src/lib.rs | 1 + .../src/ml/downloads/download_all_models.rs | 10 +- crates/desktop/src-tauri/src/ml/model_type.rs | 23 +- crates/desktop/src-tauri/src/state/app_dir.rs | 52 +++- .../src/threads/downloader_thread.rs | 168 +++++++++++- .../src-tauri/src/transfer/download.rs | 241 ++++++++++++++++++ .../desktop/src-tauri/src/transfer/error.rs | 49 ++++ crates/desktop/src-tauri/src/transfer/mod.rs | 4 + .../desktop/src-tauri/src/transfer/upload.rs | 189 ++++++++++++++ crates/desktop/src-tauri/src/transfer/util.rs | 15 ++ .../RenderingPrimitives/RealTimeDrawEngine.ts | 2 +- 14 files changed, 830 insertions(+), 33 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..e3a41f9e10 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", @@ -680,7 +682,11 @@ dependencies = [ "tauri", "tauri-build", "tauri-plugin-log", + "tempdir", + "tempfile", "tokenizers 0.15.2", + "tokio", + "tokio-util", ] [[package]] @@ -3792,12 +3798,12 @@ dependencies = [ [[package]] name = "errno" -version = "0.3.9" +version = "0.3.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "534c5cf6194dfab3db3242765c03bbe257cf92f22b38f6bc0c58d59108a820ba" +checksum = "33d852cb9b869c2a9b3df2f71a3074817f01e1844f839a144f5fcef059a4eb5d" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.59.0", ] [[package]] @@ -3927,9 +3933,9 @@ dependencies = [ [[package]] name = "fastrand" -version = "2.1.0" +version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fc0510504f03c51ada170672ac806f1f105a88aa97a5281117e1ddc3368e51a" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" [[package]] name = "fdeflate" @@ -4657,6 +4663,18 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "getrandom" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "wasi 0.13.3+wasi-0.2.2", + "windows-targets 0.52.6", +] + [[package]] name = "getset" version = "0.1.4" @@ -6011,7 +6029,7 @@ dependencies = [ "chumsky", "email-encoding", "email_address", - "fastrand 2.1.0", + "fastrand 2.3.0", "futures-util", "hostname 0.3.1", "httpdate", @@ -6229,6 +6247,12 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" +[[package]] +name = "linux-raw-sys" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db9c683daf087dc577b7506e9695b3d556a9f3849903fa28186283afd6809e9" + [[package]] name = "local-channel" version = "0.1.5" @@ -9664,10 +9688,23 @@ dependencies = [ "bitflags 2.5.0", "errno", "libc", - "linux-raw-sys", + "linux-raw-sys 0.4.14", "windows-sys 0.52.0", ] +[[package]] +name = "rustix" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7178faa4b75a30e269c71e61c353ce2748cf3d76f0c44c393f4e60abf49b825" +dependencies = [ + "bitflags 2.5.0", + "errno", + "libc", + "linux-raw-sys 0.9.2", + "windows-sys 0.59.0", +] + [[package]] name = "rustls" version = "0.21.12" @@ -11899,14 +11936,16 @@ dependencies = [ [[package]] name = "tempfile" -version = "3.10.1" +version = "3.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85b77fafb263dd9d05cbeac119526425676db3784113aa9295c88498cbf8bff1" +checksum = "2c317e0a526ee6120d8dabad239c8dadca62b24b6f168914bbbc8e2fb1f0e567" dependencies = [ "cfg-if 1.0.0", - "fastrand 2.1.0", - "rustix", - "windows-sys 0.52.0", + "fastrand 2.3.0", + "getrandom 0.3.1", + "once_cell", + "rustix 1.0.2", + "windows-sys 0.59.0", ] [[package]] @@ -13279,6 +13318,15 @@ version = "0.11.0+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" +[[package]] +name = "wasi" +version = "0.13.3+wasi-0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26816d2e1a4a36a2940b96c5296ce403917633dff8f3440e9b236ed6f6bacad2" +dependencies = [ + "wit-bindgen-rt", +] + [[package]] name = "wasite" version = "0.1.0" @@ -14048,6 +14096,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "wit-bindgen-rt" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3268f3d866458b787f390cf61f4bbb563b922d091359f9608842999eaee3943c" +dependencies = [ + "bitflags 2.5.0", +] + [[package]] name = "wry" version = "0.48.1" @@ -14138,8 +14195,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e105d177a3871454f754b33bb0ee637ecaaac997446375fd3e5d43a2ed00c909" dependencies = [ "libc", - "linux-raw-sys", - "rustix", + "linux-raw-sys 0.4.14", + "rustix 0.38.34", ] [[package]] diff --git a/crates/desktop/src-tauri/Cargo.toml b/crates/desktop/src-tauri/Cargo.toml index 30845e1ae9..f0650c2353 100644 --- a/crates/desktop/src-tauri/Cargo.toml +++ b/crates/desktop/src-tauri/Cargo.toml @@ -28,24 +28,30 @@ 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" +tempdir = "0.3.7" +tempfile = "3.18.0" 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/events/notification_event.rs b/crates/desktop/src-tauri/src/events/notification_event.rs index 153ccfca3c..6ce4d9c182 100644 --- a/crates/desktop/src-tauri/src/events/notification_event.rs +++ b/crates/desktop/src-tauri/src/events/notification_event.rs @@ -6,6 +6,8 @@ use serde::Serialize; pub enum NotificationModelType { Unet, Vae, + Json, + ClipEncoder, } #[derive(Clone, Serialize)] @@ -15,6 +17,12 @@ pub enum NotificationEvent<'a> { model_name: &'a str, model_type: NotificationModelType, }, + ModelDownloadProgress { + model_name: &'a str, + model_type: NotificationModelType, + currently_downloaded_bytes: usize, + total_file_bytes: usize, + }, ModelDownloadComplete { model_name: &'a str, model_type: NotificationModelType, 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/ml/downloads/download_all_models.rs b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs index 9b97957447..b83ab08495 100644 --- a/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs +++ b/crates/desktop/src-tauri/src/ml/downloads/download_all_models.rs @@ -5,11 +5,11 @@ use crate::state::app_dir::AppDataRoot; pub async fn download_all_models(app_data_root: &AppDataRoot) -> anyhow::Result<()> { let weights_dir = app_data_root.weights_dir(); - download_model_registry_file(ModelType::ClipJson, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?; - download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?; + //download_model_registry_file(ModelType::ClipJson, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboUnet, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboVae, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboClipEncoder, weights_dir).await?; + //download_model_registry_file(ModelType::SdxlTurboClipEncoder2, weights_dir).await?; Ok(()) } diff --git a/crates/desktop/src-tauri/src/ml/model_type.rs b/crates/desktop/src-tauri/src/ml/model_type.rs index 597acc3ffc..69b95acb09 100644 --- a/crates/desktop/src-tauri/src/ml/model_type.rs +++ b/crates/desktop/src-tauri/src/ml/model_type.rs @@ -1,5 +1,6 @@ -use std::path::PathBuf; +use crate::events::notification_event::NotificationModelType; use crate::state::app_dir::AppWeightsDir; +use std::path::PathBuf; #[derive(Copy, Clone, Debug)] pub enum ModelType { @@ -11,6 +12,26 @@ pub enum ModelType { } impl ModelType { + pub fn get_name(&self) -> &'static str { + match self { + Self::ClipJson => "CLIP (JSON)", + Self::SdxlTurboUnet => "SDXL Turbo UNET", + Self::SdxlTurboVae => "SDXL Turbo VAE", + Self::SdxlTurboClipEncoder => "Clip Encoder #1", + Self::SdxlTurboClipEncoder2 => "Clip Encoder #2", + } + } + + pub fn get_notification_type(&self) -> NotificationModelType { + match self { + Self::ClipJson => NotificationModelType::Json, + Self::SdxlTurboUnet => NotificationModelType::Unet, + Self::SdxlTurboVae => NotificationModelType::Vae, + Self::SdxlTurboClipEncoder => NotificationModelType::ClipEncoder, + Self::SdxlTurboClipEncoder2 => NotificationModelType::ClipEncoder, + } + } + pub fn get_hf_id(&self) -> Option<&'static str> { match self { Self::ClipJson => Some("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k"), diff --git a/crates/desktop/src-tauri/src/state/app_dir.rs b/crates/desktop/src-tauri/src/state/app_dir.rs index e321ccee0e..24a690dbec 100644 --- a/crates/desktop/src-tauri/src/state/app_dir.rs +++ b/crates/desktop/src-tauri/src/state/app_dir.rs @@ -1,17 +1,22 @@ use crate::state::os_platform::OsPlatform; use expanduser::expanduser; use std::path::{Path, PathBuf}; +use tempdir::TempDir; +use tempfile::{Builder, NamedTempFile}; use crate::ml::model_type::ModelType; const ASSETS_SUBDIRECTORY : &str = "assets"; const WEIGHTS_SUBDIRECTORY : &str = "weights"; +const TEMPORARY_SUBDIRECTORY : &str = "temp"; + /// The path to the application data directory, which includes "asset" and "weights" data. #[derive(Clone)] pub struct AppDataRoot { path: PathBuf, assets_dir: AppAssetsDir, weights_dir: AppWeightsDir, + temp_dir: TemporaryDir, } #[derive(Clone)] @@ -24,6 +29,11 @@ pub struct AppWeightsDir { path: PathBuf, } +#[derive(Clone)] +pub struct TemporaryDir { + path: PathBuf, +} + impl AppDataRoot { pub fn create_existing>(dir: P) -> anyhow::Result { let mut dir = dir.as_ref().to_path_buf(); @@ -51,11 +61,13 @@ impl AppDataRoot { let assets_dir = AppAssetsDir::create_existing(dir.join(ASSETS_SUBDIRECTORY))?; let weights_dir = AppWeightsDir::create_existing(dir.join(WEIGHTS_SUBDIRECTORY))?; - + let temp_dir = TemporaryDir::create_existing(dir.join(TEMPORARY_SUBDIRECTORY))?; + Ok(Self { path: dir, assets_dir, weights_dir, + temp_dir, }) } @@ -66,7 +78,11 @@ impl AppDataRoot { pub fn weights_dir(&self) -> &AppWeightsDir { &self.weights_dir } - + + pub fn temp_dir(&self) -> &TemporaryDir { + &self.temp_dir + } + pub fn path(&self) -> &Path { &self.path } @@ -121,3 +137,35 @@ impl AppWeightsDir { self.path.join(model_registry.get_filename()) } } + +impl TemporaryDir { + pub fn create_existing>(dir: P) -> anyhow::Result { + let mut dir = dir.as_ref().to_path_buf(); + match dir.canonicalize() { + Ok(d) => dir = d, + Err(err) => { + println!("Error canonicalizing {:?}: {}", dir, err); + } + } + if !dir.exists() { + println!("Creating directory {:?}", dir); + std::fs::create_dir(&dir)?; + } + Ok(Self { + path: dir, + }) + } + + pub fn with_prefix(&self, prefix: &str) -> anyhow::Result { + let tempdir = TempDir::new_in(&self.path, prefix)?; + Ok(tempdir) + } + + pub fn new_named_temp_file(&self) -> anyhow::Result { + let tempfile = Builder::new() + .prefix("temp_") + .suffix(".bin") + .tempfile_in(&self.path)?; + Ok(tempfile) + } +} diff --git a/crates/desktop/src-tauri/src/threads/downloader_thread.rs b/crates/desktop/src-tauri/src/threads/downloader_thread.rs index 3b9c239928..490588820d 100644 --- a/crates/desktop/src-tauri/src/threads/downloader_thread.rs +++ b/crates/desktop/src-tauri/src/threads/downloader_thread.rs @@ -1,11 +1,169 @@ -use tauri::{AppHandle, Emitter}; +use crate::events::notification_event::{NotificationEvent, NotificationModelType}; +use crate::ml::model_type::ModelType; use crate::state::app_dir::AppDataRoot; +use crate::transfer::download::{download_async, ProgressUpdate}; +use futures::stream::FuturesUnordered; +use futures::StreamExt; +use log::{debug, error, info}; +use std::collections::VecDeque; +use std::path::PathBuf; +use std::sync::mpsc::Receiver; +use std::sync::{Arc, Mutex}; +use tauri::{AppHandle, Emitter}; +use tempfile::NamedTempFile; +use tokio::task::JoinHandle; + +const NOTIFICATION_CHANNEL_NAME : &str = "notification"; +const MAX_FILES : usize = 8; +const CHUNK_SIZE : usize = 1024 * 1024; +const PARALLEL_FAILURES : usize = 8; +const MAX_RETRIES : usize = 30; + +struct ModelDownloadStatusChannel { + pub model_name: String, + pub model_type: NotificationModelType, + pub final_file_path: PathBuf, + pub receiver: Receiver, +} pub async fn downloader_thread(app_data_root: AppDataRoot, app: AppHandle) -> ! { - loop { - app.emit("testing", 1); + + let mut status_check_queue : Arc>> = Arc::new(Mutex::new(Vec::new())); + + { + let status_check_queue2 = status_check_queue.clone(); + let app2 = app.clone(); + + tokio::spawn(download_notify_thread(status_check_queue2, app2)); + } + + let mut download_queue = VecDeque::new(); + + // TODO: Automatic enqueue of known important models + enqueue new models on-demand + download_queue.push_back(ModelType::ClipJson); + download_queue.push_back(ModelType::SdxlTurboUnet); + download_queue.push_back(ModelType::SdxlTurboVae); + download_queue.push_back(ModelType::SdxlTurboClipEncoder); + download_queue.push_back(ModelType::SdxlTurboClipEncoder2); + + let mut handles = FuturesUnordered::new(); + + while let Some(model) = download_queue.pop_front() { + let url = model.get_download_url().to_string(); + + let filename = app_data_root.weights_dir().model_path(&model); + + if filename.exists() { + info!("File already exists: {:?}", &filename); + continue; + } + + let headers = None; - println!(">>> Testing: {:?}", app_data_root.path()); - std::thread::sleep(std::time::Duration::from_millis(500)); + let (tx, rx) = std::sync::mpsc::channel(); + + info!("Download: {:?}", filename); + + let temp_file = match app_data_root.temp_dir().new_named_temp_file() { + Ok(file) => file, + Err(err) => { + error!("Couldn't create a temporary file: {:?}", err); + continue; + } + }; + + let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadStarted { + model_name: model.get_name(), + model_type: model.get_notification_type(), + }); + + if let Err(err) = result { + error!("Could not emit event: {}", err); + } + + let task = tokio::spawn(download_async( + url, + temp_file, + filename.clone(), + MAX_FILES, + CHUNK_SIZE, + PARALLEL_FAILURES, + MAX_RETRIES, + headers, + Some(tx), + )); + + match status_check_queue.lock() { + Err(err) => error!("Could not unlock: {:?}", err), + Ok(mut lock) => { + lock.push(ModelDownloadStatusChannel { + receiver: rx, + final_file_path: filename.clone(), + model_name: model.get_name().to_string(), + model_type: model.get_notification_type(), + }) + } + } + + handles.push(task); + } + + while let Some(result) = handles.next().await { + match result { + Err(err) => error!("Error downloading: {:?}", err), + Ok(_) => {}, + } + } + + loop { + // TODO(bt,2025-03-10): Queue management for new files entering the queue. + tokio::time::sleep(std::time::Duration::from_millis(60_000)).await; + } +} + +pub async fn download_notify_thread(queue: Arc>>, app: AppHandle) -> ! { + loop { + match queue.lock() { + Err(err) => error!("Could not unlock: {:?}", err), + Ok(mut lock) => { + debug!(">>>>> Notifier has {} items", lock.len()); + + let mut reenqueue = Vec::with_capacity(lock.len()); + + for item in (*lock).drain(..) { + if let Ok(update) = item.receiver.try_recv() { + let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadProgress { + model_name: &item.model_name, + model_type: item.model_type, + currently_downloaded_bytes: update.complete, + total_file_bytes: update.total_length, + }); + + if let Err(err) = result { + error!("Could not emit event: {}", err); + } + } + + if item.final_file_path.exists() { + debug!("No longer sending notifications for {:?}", &item.final_file_path); + + let result = app.emit(NOTIFICATION_CHANNEL_NAME, NotificationEvent::ModelDownloadComplete { + model_name: &item.model_name, + model_type: item.model_type, + }); + + if let Err(err) = result { + error!("Could not emit event: {}", err); + } + } else { + reenqueue.push(item); + } + } + + lock.extend(reenqueue); + } + } + + tokio::time::sleep(std::time::Duration::from_millis(500)).await; } } 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..9dd2f1f52e --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/download.rs @@ -0,0 +1,241 @@ +//! 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 reqwest::header::{ + HeaderMap, HeaderName, HeaderValue, AUTHORIZATION, 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::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::Sender; +use std::time::Duration; +use log::info; +use tempfile::NamedTempFile; +use tokio::fs::OpenOptions; +use tokio::io::AsyncWriteExt; +use tokio::io::AsyncSeekExt; +use tokio::sync::Semaphore; +use tokio::time::sleep; + +pub struct ProgressUpdate { + pub complete: usize, + pub total_length: usize, +} + +/// 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)] +pub async fn download_async>( + url: String, + filename: NamedTempFile, + final_filename: P, + max_files: usize, + chunk_size: usize, + parallel_failures: usize, + max_retries: usize, + input_headers: Option>, + maybe_progress_sender: Option>, +) -> Result<(), Error> { + info!(">>> Downloading: {}", filename.as_ref().display()); + + let client = reqwest::Client::builder() + // https://github.com/hyperium/hyper/issues/2136#issuecomment-589488526 + .http2_keep_alive_timeout(Duration::from_secs(15)) + .build()?; + + 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| Error::new_err(format!("Invalid header: {err}")))?; + let value: HeaderValue = AsRef::::as_ref(&v) + .try_into() + .map_err(|err| Error::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| Error::new_err(format!("Error while downloading: {err}")))? + .error_for_status() + .map_err(|err| Error::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| Error::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(Error::new_err("No content length".to_string()))? + .to_str() + .map_err(|err| Error::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(Error::new_err( + "Error while downloading: No size was detected".to_string(), + ))? + .parse() + .map_err(|err| Error::new_err(format!("Error while downloading: {err}")))?; + + + info!("File {:?} is {:?} length", filename.as_ref(), length); + + let mut handles = FuturesUnordered::new(); + let semaphore = Arc::new(Semaphore::new(max_files)); + let parallel_failures_semaphore = Arc::new(Semaphore::new(parallel_failures)); + let max_chunk_complete = Arc::new(AtomicUsize::new(0)); + + for start in (0..length).step_by(chunk_size) { + let url = redirected_url.to_string(); + let filename = filename.as_ref().to_owned(); + 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(); + let maybe_progress = maybe_progress_sender.clone(); + let max_chunk_complete = max_chunk_complete.clone(); + handles.push(tokio::spawn(async move { + let permit = semaphore + .acquire_owned() + .await + .map_err(|err| Error::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(Error::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| { + Error::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); + + if let Some(sender) = maybe_progress.as_ref() { + // NB: Calling fetch_max twice. First call will return the lower value. + let _old = max_chunk_complete.fetch_max(stop, Ordering::SeqCst); + let max_chunk = max_chunk_complete.fetch_max(stop, Ordering::SeqCst); + let progress_update = ProgressUpdate { + complete: max_chunk, + total_length: length, + }; + if let Err(_err) = sender.send(progress_update) { + // Fail silently, not essential to issue progress updates + } + } + } + } + drop(permit); + chunk.map_err(|e| Error::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(Error::new_err(format!( + "Error while downloading: {err}" + ))); + } + } + } + + let file = filename.persist(final_filename) + .map_err(|e| Error::new_err(format!("Could not rename to final filename: {:?}", e)))?; + + file.sync_all()?; + + Ok(()) +} + +async fn download_chunk>( + client: &reqwest::Client, + url: &str, + filename: P, + 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..fcad35635c --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/error.rs @@ -0,0 +1,49 @@ +//! 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), + Misc(String), +} + +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}"), + Self::Misc(err) => write!(f, "Misc error: {err}"), + } + } +} + +impl std::error::Error for Error {} + +impl Error { + pub fn new_err(reason: String) -> Self { + Self::Misc(reason) + } +} 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..acd6fd7bb9 --- /dev/null +++ b/crates/desktop/src-tauri/src/transfer/upload.rs @@ -0,0 +1,189 @@ +//! 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 reqwest::header::CONTENT_LENGTH; +use std::collections::HashMap; +use std::fmt::Display; +use std::io::SeekFrom; +use std::sync::Arc; +use std::time::Duration; +use tokio::fs::OpenOptions; +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)] +fn multipart_upload( + file_path: String, + parts_urls: Vec, + chunk_size: u64, + max_files: usize, + parallel_failures: usize, + max_retries: usize, +) -> Result>, Error> { + if parallel_failures > max_files { + return Err(Error::new_err( + "Error parallel_failures cannot be > max_files".to_string(), + )); + } + if (parallel_failures == 0) != (max_retries == 0) { + return Err(Error::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, + ) + .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, +) -> Result>, Error> { + 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| Error::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(Error::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| { + Error::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) => Error::new_err(format!("Io error {io}")), + // Error::Request(req) => Error::new_err(format!("Error while sending chunk {req}")), + // Error::ToStrError(req) => Error::new_err(format!("Response header contains non ASCII chars: {req}")), + //} + e + }).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(Error::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) +} diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts index fcb5311271..a1f151596c 100644 --- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts +++ b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts @@ -214,7 +214,7 @@ export class RealTimeDrawEngine { // Add mouse events for preview canvas copying this.previewCopyListener(); - //this.listenToServerEvents(); + this.listenToServerEvents(); }