From cdf760ca240c1b3a20281a45bae75a21be25ea69 Mon Sep 17 00:00:00 2001 From: boxbeam Date: Wed, 27 Mar 2024 15:42:44 -0400 Subject: [PATCH 1/5] feat(download): Add support for segmented models --- crates/tabby-common/src/registry.rs | 3 ++ crates/tabby-download/src/lib.rs | 50 +++++++++++++++++++++++++++-- 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/crates/tabby-common/src/registry.rs b/crates/tabby-common/src/registry.rs index 24e9c2d0352d..40e11fbf15b7 100644 --- a/crates/tabby-common/src/registry.rs +++ b/crates/tabby-common/src/registry.rs @@ -12,7 +12,10 @@ pub struct ModelInfo { pub prompt_template: Option, #[serde(skip_serializing_if = "Option::is_none")] pub chat_template: Option, + #[serde(skip_serializing_if = "Vec::is_empty")] pub urls: Vec, + #[serde(skip_serializing_if = "Vec::is_empty")] + pub segmented_urls: Vec, pub sha256: String, } diff --git a/crates/tabby-download/src/lib.rs b/crates/tabby-download/src/lib.rs index a746fdd45fd8..7da8fb6cacba 100644 --- a/crates/tabby-download/src/lib.rs +++ b/crates/tabby-download/src/lib.rs @@ -1,9 +1,13 @@ //! Responsible for downloading ML models for use with tabby. -use std::{fs, path::Path}; +use std::{ + fs::{self, File}, + io::{BufRead, BufReader, Write}, + path::Path, +}; use aim_downloader::{bar::WrappedBar, error::DownloadError, https}; use anyhow::{anyhow, bail, Result}; -use tabby_common::registry::{parse_model_id, ModelRegistry}; +use tabby_common::registry::{parse_model_id, ModelInfo, ModelRegistry}; use tokio_retry::{ strategy::{jitter, ExponentialBackoff}, Retry, @@ -37,6 +41,10 @@ async fn download_model_impl( } } + if !model_info.segmented_urls.is_empty() { + return download_split_model(&model_info, &model_path).await; + } + let registry = std::env::var("TABBY_DOWNLOAD_HOST").unwrap_or("huggingface.co".to_owned()); let Some(model_url) = model_info.urls.iter().find(|x| x.contains(®istry)) else { return Err(anyhow!( @@ -52,6 +60,44 @@ async fn download_model_impl( Ok(()) } +async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Result<()> { + if !model_info.urls.is_empty() { + return Err(anyhow!( + "{}: Cannot specify both `urls` and `segmented_urls`", + model_info.name + )); + } + let mut paths = vec![]; + for (index, url) in model_info.segmented_urls.iter().enumerate() { + let ext = format!( + "{}.{}", + model_path.extension().unwrap_or_default().to_string_lossy(), + index.to_string() + ); + let path = model_path.with_extension(ext); + let strategy = ExponentialBackoff::from_millis(100).map(jitter).take(2); + let download_job = Retry::spawn(strategy, || download_file(url, &path)); + download_job.await?; + paths.push(path); + } + info!("Merging split model files..."); + let mut file = File::open(model_path)?; + for path in paths { + let mut reader = BufReader::new(File::open(&path)?); + loop { + let buffer = reader.fill_buf()?; + file.write_all(buffer)?; + let len = buffer.len(); + reader.consume(len); + if len == 0 { + break; + } + } + std::fs::remove_file(path)?; + } + Ok(()) +} + async fn download_file(url: &str, path: &Path) -> Result<()> { let dir = path .parent() From ab324d28f83d08879bc10edfb333b183b3712bef Mon Sep 17 00:00:00 2001 From: boxbeam Date: Thu, 28 Mar 2024 17:44:39 -0400 Subject: [PATCH 2/5] Fix segmented downloads --- crates/tabby-common/src/registry.rs | 8 ++++---- crates/tabby-download/src/lib.rs | 29 ++++++++++++++++++++++------- 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/crates/tabby-common/src/registry.rs b/crates/tabby-common/src/registry.rs index 40e11fbf15b7..590df4d92536 100644 --- a/crates/tabby-common/src/registry.rs +++ b/crates/tabby-common/src/registry.rs @@ -12,10 +12,10 @@ pub struct ModelInfo { pub prompt_template: Option, #[serde(skip_serializing_if = "Option::is_none")] pub chat_template: Option, - #[serde(skip_serializing_if = "Vec::is_empty")] - pub urls: Vec, - #[serde(skip_serializing_if = "Vec::is_empty")] - pub segmented_urls: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub urls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub segmented_urls: Option>, pub sha256: String, } diff --git a/crates/tabby-download/src/lib.rs b/crates/tabby-download/src/lib.rs index 7da8fb6cacba..5b40d5ad96dc 100644 --- a/crates/tabby-download/src/lib.rs +++ b/crates/tabby-download/src/lib.rs @@ -1,6 +1,6 @@ //! Responsible for downloading ML models for use with tabby. use std::{ - fs::{self, File}, + fs::{self, File, OpenOptions}, io::{BufRead, BufReader, Write}, path::Path, }; @@ -41,12 +41,17 @@ async fn download_model_impl( } } - if !model_info.segmented_urls.is_empty() { + if !model_info.segmented_urls.is_none() { return download_split_model(&model_info, &model_path).await; } let registry = std::env::var("TABBY_DOWNLOAD_HOST").unwrap_or("huggingface.co".to_owned()); - let Some(model_url) = model_info.urls.iter().find(|x| x.contains(®istry)) else { + let Some(model_url) = model_info + .urls + .iter() + .flatten() + .find(|x| x.contains(®istry)) + else { return Err(anyhow!( "Invalid mirror <{}> for model urls: {:?}", registry, @@ -55,33 +60,43 @@ async fn download_model_impl( }; let strategy = ExponentialBackoff::from_millis(100).map(jitter).take(2); - let download_job = Retry::spawn(strategy, || download_file(model_url, model_path.as_path())); + let download_job = Retry::spawn(strategy, || download_file(&model_url, model_path.as_path())); download_job.await?; Ok(()) } async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Result<()> { - if !model_info.urls.is_empty() { + if !model_info.urls.is_none() { return Err(anyhow!( "{}: Cannot specify both `urls` and `segmented_urls`", model_info.name )); } let mut paths = vec![]; - for (index, url) in model_info.segmented_urls.iter().enumerate() { + let segmented_urls = model_info.segmented_urls.clone().unwrap_or_default(); + for (index, url) in segmented_urls.iter().enumerate() { let ext = format!( "{}.{}", model_path.extension().unwrap_or_default().to_string_lossy(), index.to_string() ); let path = model_path.with_extension(ext); + info!( + "Downloading {path:?} ({index} / {total})", + index = index + 1, + total = segmented_urls.len() + ); let strategy = ExponentialBackoff::from_millis(100).map(jitter).take(2); let download_job = Retry::spawn(strategy, || download_file(url, &path)); download_job.await?; paths.push(path); } info!("Merging split model files..."); - let mut file = File::open(model_path)?; + println!("{model_path:?}"); + let mut file = OpenOptions::new() + .append(true) + .create(true) + .open(model_path)?; for path in paths { let mut reader = BufReader::new(File::open(&path)?); loop { From 9978225468596c3833370b9d7256d217a6cb3de5 Mon Sep 17 00:00:00 2001 From: boxbeam Date: Thu, 28 Mar 2024 17:45:55 -0400 Subject: [PATCH 3/5] Apply suggestion --- crates/tabby-common/src/registry.rs | 2 +- crates/tabby-download/src/lib.rs | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/tabby-common/src/registry.rs b/crates/tabby-common/src/registry.rs index 590df4d92536..f82388c7cf99 100644 --- a/crates/tabby-common/src/registry.rs +++ b/crates/tabby-common/src/registry.rs @@ -15,7 +15,7 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub urls: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub segmented_urls: Option>, + pub partition_urls: Option>, pub sha256: String, } diff --git a/crates/tabby-download/src/lib.rs b/crates/tabby-download/src/lib.rs index 5b40d5ad96dc..85b3553d6878 100644 --- a/crates/tabby-download/src/lib.rs +++ b/crates/tabby-download/src/lib.rs @@ -41,7 +41,7 @@ async fn download_model_impl( } } - if !model_info.segmented_urls.is_none() { + if !model_info.partition_urls.is_none() { return download_split_model(&model_info, &model_path).await; } @@ -68,13 +68,13 @@ async fn download_model_impl( async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Result<()> { if !model_info.urls.is_none() { return Err(anyhow!( - "{}: Cannot specify both `urls` and `segmented_urls`", + "{}: Cannot specify both `urls` and `partition_urls`", model_info.name )); } let mut paths = vec![]; - let segmented_urls = model_info.segmented_urls.clone().unwrap_or_default(); - for (index, url) in segmented_urls.iter().enumerate() { + let partition_urls = model_info.partition_urls.clone().unwrap_or_default(); + for (index, url) in partition_urls.iter().enumerate() { let ext = format!( "{}.{}", model_path.extension().unwrap_or_default().to_string_lossy(), @@ -84,7 +84,7 @@ async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Resu info!( "Downloading {path:?} ({index} / {total})", index = index + 1, - total = segmented_urls.len() + total = partition_urls.len() ); let strategy = ExponentialBackoff::from_millis(100).map(jitter).take(2); let download_job = Retry::spawn(strategy, || download_file(url, &path)); From 848b28ec2213ddbf08f6fb17a236364d55ef07da Mon Sep 17 00:00:00 2001 From: boxbeam Date: Thu, 28 Mar 2024 17:50:49 -0400 Subject: [PATCH 4/5] Allow specifying multiple mirrors for partitioned models --- crates/tabby-common/src/registry.rs | 2 +- crates/tabby-download/src/lib.rs | 21 +++++++++++++++++++-- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/crates/tabby-common/src/registry.rs b/crates/tabby-common/src/registry.rs index f82388c7cf99..dd0fe59c8c08 100644 --- a/crates/tabby-common/src/registry.rs +++ b/crates/tabby-common/src/registry.rs @@ -15,7 +15,7 @@ pub struct ModelInfo { #[serde(skip_serializing_if = "Option::is_none")] pub urls: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub partition_urls: Option>, + pub partition_urls: Option>>, pub sha256: String, } diff --git a/crates/tabby-download/src/lib.rs b/crates/tabby-download/src/lib.rs index 85b3553d6878..502ac857e708 100644 --- a/crates/tabby-download/src/lib.rs +++ b/crates/tabby-download/src/lib.rs @@ -14,6 +14,10 @@ use tokio_retry::{ }; use tracing::{info, warn}; +fn download_host() -> String { + std::env::var("TABBY_DOWNLOAD_HOST").unwrap_or("huggingface.co".to_owned()) +} + async fn download_model_impl( registry: &ModelRegistry, name: &str, @@ -45,7 +49,7 @@ async fn download_model_impl( return download_split_model(&model_info, &model_path).await; } - let registry = std::env::var("TABBY_DOWNLOAD_HOST").unwrap_or("huggingface.co".to_owned()); + let registry = download_host(); let Some(model_url) = model_info .urls .iter() @@ -74,7 +78,20 @@ async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Resu } let mut paths = vec![]; let partition_urls = model_info.partition_urls.clone().unwrap_or_default(); - for (index, url) in partition_urls.iter().enumerate() { + let mirror = download_host(); + + let Some(urls) = partition_urls + .iter() + .find(|urls| urls.iter().all(|url| url.contains(&mirror))) + else { + return Err(anyhow!( + "Invalid mirror <{}> for model urls: {:?}", + mirror, + partition_urls + )); + }; + + for (index, url) in urls.iter().enumerate() { let ext = format!( "{}.{}", model_path.extension().unwrap_or_default().to_string_lossy(), From 2f17e2cc00486e43d1c2fec2a2e2266e853f0ccd Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Thu, 28 Mar 2024 21:58:09 +0000 Subject: [PATCH 5/5] [autofix.ci] apply automated fixes --- crates/tabby-download/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/tabby-download/src/lib.rs b/crates/tabby-download/src/lib.rs index 502ac857e708..731d980d3eb3 100644 --- a/crates/tabby-download/src/lib.rs +++ b/crates/tabby-download/src/lib.rs @@ -45,8 +45,8 @@ async fn download_model_impl( } } - if !model_info.partition_urls.is_none() { - return download_split_model(&model_info, &model_path).await; + if model_info.partition_urls.is_some() { + return download_split_model(model_info, &model_path).await; } let registry = download_host(); @@ -64,13 +64,13 @@ async fn download_model_impl( }; let strategy = ExponentialBackoff::from_millis(100).map(jitter).take(2); - let download_job = Retry::spawn(strategy, || download_file(&model_url, model_path.as_path())); + let download_job = Retry::spawn(strategy, || download_file(model_url, model_path.as_path())); download_job.await?; Ok(()) } async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Result<()> { - if !model_info.urls.is_none() { + if model_info.urls.is_some() { return Err(anyhow!( "{}: Cannot specify both `urls` and `partition_urls`", model_info.name @@ -95,7 +95,7 @@ async fn download_split_model(model_info: &ModelInfo, model_path: &Path) -> Resu let ext = format!( "{}.{}", model_path.extension().unwrap_or_default().to_string_lossy(), - index.to_string() + index ); let path = model_path.with_extension(ext); info!(