Skip to content

GitLab

"examples/backends/vscode:/vscode.git/clone" did not exist on "ece08dc926679909dd7a0d980f073cdee47414c2"
Unverified Commit 95ce83d5 authored by KavinKrishnan's avatar KavinKrishnan Committed by GitHub
Browse files

feat: Integrate Model Express Client into Dynamo Model Downloads (#2574) 


Signed-off-by: default avatarKavin Krishnan <kavink@nvidia.com>
Co-authored-by: default avatarKavinKrishnan <kavin.krishnan@nvidia.com>
parent 9c4e9950
Show whitespace changes
Inline Side-by-side
Showing with 292 additions and 107 deletions
Cargo.lock
View file @ 95ce83d5
...@@ -711,9 +711,9 @@ checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" ...@@ -711,9 +711,9 @@ checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
[[package]] [[package]]
name = "bit_field" name = "bit_field"
version = "0.10.2" version = "0.10.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc827186963e592360843fb5ba4b973e145841266c1357f7180c43526f2e5b61" checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6"
[[package]] [[package]]
name = "bitflags" name = "bitflags"
...@@ -996,7 +996,7 @@ version = "0.27.0" ...@@ -996,7 +996,7 @@ version = "0.27.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3fce8dd7fcfcbf3a0a87d8f515194b49d6135acab73e18bd380d1d93bb1a15eb" checksum = "3fce8dd7fcfcbf3a0a87d8f515194b49d6135acab73e18bd380d1d93bb1a15eb"
dependencies = [ dependencies = [
"clap 4.5.45", "clap 4.5.46",
"heck 0.4.1", "heck 0.4.1",
"indexmap 2.11.0", "indexmap 2.11.0",
"log", "log",
...@@ -1110,9 +1110,9 @@ dependencies = [ ...@@ -1110,9 +1110,9 @@ dependencies = [
[[package]] [[package]]
name = "clap" name = "clap"
version = "4.5.45" version = "4.5.46"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fc0e74a703892159f5ae7d3aac52c8e6c392f5ae5f359c70b5881d60aaac318" checksum = "2c5e4fcf9c21d2e544ca1ee9d8552de13019a42aa7dbf32747fa7aaf1df76e57"
dependencies = [ dependencies = [
"clap_builder", "clap_builder",
"clap_derive", "clap_derive",
...@@ -1120,9 +1120,9 @@ dependencies = [ ...@@ -1120,9 +1120,9 @@ dependencies = [
[[package]] [[package]]
name = "clap_builder" name = "clap_builder"
version = "4.5.44" version = "4.5.46"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3e7f4214277f3c7aa526a59dd3fbe306a370daee1f8b7b8c987069cd8e888a8" checksum = "fecb53a0e6fcfb055f686001bc2e2592fa527efaf38dbe81a6a9563562e57d41"
dependencies = [ dependencies = [
"anstream", "anstream",
"anstyle", "anstyle",
...@@ -2059,7 +2059,7 @@ dependencies = [ ...@@ -2059,7 +2059,7 @@ dependencies = [
"anyhow", "anyhow",
"async-stream", "async-stream",
"async-trait", "async-trait",
"clap 4.5.45", "clap 4.5.46",
"dynamo-async-openai", "dynamo-async-openai",
"dynamo-engine-llamacpp", "dynamo-engine-llamacpp",
"dynamo-engine-mistralrs", "dynamo-engine-mistralrs",
...@@ -3011,7 +3011,7 @@ dependencies = [ ...@@ -3011,7 +3011,7 @@ dependencies = [
"js-sys", "js-sys",
"libc", "libc",
"r-efi", "r-efi",
"wasi 0.14.2+wasi-0.2.4", "wasi 0.14.3+wasi-0.2.4",
"wasm-bindgen", "wasm-bindgen",
] ]
...@@ -3631,9 +3631,9 @@ dependencies = [ ...@@ -3631,9 +3631,9 @@ dependencies = [
[[package]] [[package]]
name = "image-webp" name = "image-webp"
version = "0.2.3" version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f6970fe7a5300b4b42e62c52efa0187540a5bef546c60edaf554ef595d2e6f0b" checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3"
dependencies = [ dependencies = [
"byteorder-lite", "byteorder-lite",
"quick-error 2.0.1", "quick-error 2.0.1",
...@@ -3890,7 +3890,7 @@ dependencies = [ ...@@ -3890,7 +3890,7 @@ dependencies = [
"anyhow", "anyhow",
"base64 0.21.7", "base64 0.21.7",
"bytecount", "bytecount",
"clap 4.5.45", "clap 4.5.46",
"fancy-regex 0.11.0", "fancy-regex 0.11.0",
"fraction", "fraction",
"getrandom 0.2.16", "getrandom 0.2.16",
...@@ -4300,7 +4300,7 @@ name = "metrics" ...@@ -4300,7 +4300,7 @@ name = "metrics"
version = "0.4.1" version = "0.4.1"
dependencies = [ dependencies = [
"axum 0.8.4", "axum 0.8.4",
"clap 4.5.45", "clap 4.5.46",
"dynamo-llm", "dynamo-llm",
"dynamo-runtime", "dynamo-runtime",
"futures", "futures",
...@@ -4430,7 +4430,7 @@ dependencies = [ ...@@ -4430,7 +4430,7 @@ dependencies = [
"anyhow", "anyhow",
"candle-core 0.9.1 (git+https://github.com/EricLBuehler/candle.git?rev=95d713f9)", "candle-core 0.9.1 (git+https://github.com/EricLBuehler/candle.git?rev=95d713f9)",
"candle-nn", "candle-nn",
"clap 4.5.45", "clap 4.5.46",
"either", "either",
"futures", "futures",
"image", "image",
...@@ -4475,7 +4475,7 @@ dependencies = [ ...@@ -4475,7 +4475,7 @@ dependencies = [
"candle-nn", "candle-nn",
"cfgrammar", "cfgrammar",
"chrono", "chrono",
"clap 4.5.45", "clap 4.5.46",
"csv", "csv",
"derive-new", "derive-new",
"derive_more 2.0.1", "derive_more 2.0.1",
...@@ -5115,7 +5115,7 @@ dependencies = [ ...@@ -5115,7 +5115,7 @@ dependencies = [
"anyhow", "anyhow",
"base64 0.22.1", "base64 0.22.1",
"bstr", "bstr",
"clap 4.5.45", "clap 4.5.46",
"fancy-regex 0.13.0", "fancy-regex 0.13.0",
"futures", "futures",
"image", "image",
...@@ -5803,9 +5803,9 @@ dependencies = [ ...@@ -5803,9 +5803,9 @@ dependencies = [
[[package]] [[package]]
name = "quinn" name = "quinn"
version = "0.11.8" version = "0.11.9"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "626214629cda6781b6dc1d316ba307189c85ba657213ce642d9c77670f8202c8" checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20"
dependencies = [ dependencies = [
"bytes", "bytes",
"cfg_aliases", "cfg_aliases",
...@@ -5814,7 +5814,7 @@ dependencies = [ ...@@ -5814,7 +5814,7 @@ dependencies = [
"quinn-udp", "quinn-udp",
"rustc-hash 2.1.1", "rustc-hash 2.1.1",
"rustls", "rustls",
"socket2 0.5.10", "socket2 0.6.0",
"thiserror 2.0.16", "thiserror 2.0.16",
"tokio", "tokio",
"tracing", "tracing",
...@@ -5823,9 +5823,9 @@ dependencies = [ ...@@ -5823,9 +5823,9 @@ dependencies = [
[[package]] [[package]]
name = "quinn-proto" name = "quinn-proto"
version = "0.11.12" version = "0.11.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49df843a9161c85bb8aae55f101bc0bac8bcafd637a620d9122fd7e0b2f7422e" checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31"
dependencies = [ dependencies = [
"bytes", "bytes",
"getrandom 0.3.3", "getrandom 0.3.3",
...@@ -5844,16 +5844,16 @@ dependencies = [ ...@@ -5844,16 +5844,16 @@ dependencies = [
[[package]] [[package]]
name = "quinn-udp" name = "quinn-udp"
version = "0.5.13" version = "0.5.14"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcebb1209ee276352ef14ff8732e24cc2b02bbac986cd74a4c81bcb2f9881970" checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd"
dependencies = [ dependencies = [
"cfg_aliases", "cfg_aliases",
"libc", "libc",
"once_cell", "once_cell",
"socket2 0.5.10", "socket2 0.6.0",
"tracing", "tracing",
"windows-sys 0.59.0", "windows-sys 0.60.2",
] ]
[[package]] [[package]]
...@@ -8802,11 +8802,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" ...@@ -8802,11 +8802,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]] [[package]]
name = "wasi" name = "wasi"
version = "0.14.2+wasi-0.2.4" version = "0.14.3+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" checksum = "6a51ae83037bdd272a9e28ce236db8c07016dd0d50c27038b3f407533c030c95"
dependencies = [ dependencies = [
"wit-bindgen-rt", "wit-bindgen",
] ]
[[package]] [[package]]
...@@ -9351,13 +9351,10 @@ dependencies = [ ...@@ -9351,13 +9351,10 @@ dependencies = [
] ]
[[package]] [[package]]
name = "wit-bindgen-rt" name = "wit-bindgen"
version = "0.39.0" version = "0.45.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" checksum = "052283831dbae3d879dc7f51f3d92703a316ca49f91540417d38591826127814"
dependencies = [
"bitflags 2.9.3",
]
[[package]] [[package]]
name = "writeable" name = "writeable"
......
lib/bindings/python/Cargo.lock
View file @ 95ce83d5
...@@ -575,9 +575,9 @@ checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" ...@@ -575,9 +575,9 @@ checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7"
[[package]] [[package]]
name = "bit_field" name = "bit_field"
version = "0.10.2" version = "0.10.3"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "dc827186963e592360843fb5ba4b973e145841266c1357f7180c43526f2e5b61" checksum = "1e4b40c7323adcfc0a41c4b88143ed58346ff65a288fc144329c5c45e05d70c6"
[[package]] [[package]]
name = "bitflags" name = "bitflags"
...@@ -799,9 +799,9 @@ dependencies = [ ...@@ -799,9 +799,9 @@ dependencies = [
[[package]] [[package]]
name = "clap" name = "clap"
version = "4.5.45" version = "4.5.46"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1fc0e74a703892159f5ae7d3aac52c8e6c392f5ae5f359c70b5881d60aaac318" checksum = "2c5e4fcf9c21d2e544ca1ee9d8552de13019a42aa7dbf32747fa7aaf1df76e57"
dependencies = [ dependencies = [
"clap_builder", "clap_builder",
"clap_derive", "clap_derive",
...@@ -809,9 +809,9 @@ dependencies = [ ...@@ -809,9 +809,9 @@ dependencies = [
[[package]] [[package]]
name = "clap_builder" name = "clap_builder"
version = "4.5.44" version = "4.5.46"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b3e7f4214277f3c7aa526a59dd3fbe306a370daee1f8b7b8c987069cd8e888a8" checksum = "fecb53a0e6fcfb055f686001bc2e2592fa527efaf38dbe81a6a9563562e57d41"
dependencies = [ dependencies = [
"anstream", "anstream",
"anstyle", "anstyle",
...@@ -2228,7 +2228,7 @@ dependencies = [ ...@@ -2228,7 +2228,7 @@ dependencies = [
"js-sys", "js-sys",
"libc", "libc",
"r-efi", "r-efi",
"wasi 0.14.2+wasi-0.2.4", "wasi 0.14.3+wasi-0.2.4",
"wasm-bindgen", "wasm-bindgen",
] ]
...@@ -2673,9 +2673,9 @@ dependencies = [ ...@@ -2673,9 +2673,9 @@ dependencies = [
[[package]] [[package]]
name = "image-webp" name = "image-webp"
version = "0.2.3" version = "0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f6970fe7a5300b4b42e62c52efa0187540a5bef546c60edaf554ef595d2e6f0b" checksum = "525e9ff3e1a4be2fbea1fdf0e98686a6d98b4d8f937e1bf7402245af1909e8c3"
dependencies = [ dependencies = [
"byteorder-lite", "byteorder-lite",
"quick-error", "quick-error",
...@@ -4153,9 +4153,9 @@ dependencies = [ ...@@ -4153,9 +4153,9 @@ dependencies = [
[[package]] [[package]]
name = "quinn" name = "quinn"
version = "0.11.8" version = "0.11.9"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "626214629cda6781b6dc1d316ba307189c85ba657213ce642d9c77670f8202c8" checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20"
dependencies = [ dependencies = [
"bytes", "bytes",
"cfg_aliases", "cfg_aliases",
...@@ -4164,7 +4164,7 @@ dependencies = [ ...@@ -4164,7 +4164,7 @@ dependencies = [
"quinn-udp", "quinn-udp",
"rustc-hash 2.1.1", "rustc-hash 2.1.1",
"rustls", "rustls",
"socket2 0.5.10", "socket2 0.6.0",
"thiserror 2.0.16", "thiserror 2.0.16",
"tokio", "tokio",
"tracing", "tracing",
...@@ -4173,9 +4173,9 @@ dependencies = [ ...@@ -4173,9 +4173,9 @@ dependencies = [
[[package]] [[package]]
name = "quinn-proto" name = "quinn-proto"
version = "0.11.12" version = "0.11.13"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "49df843a9161c85bb8aae55f101bc0bac8bcafd637a620d9122fd7e0b2f7422e" checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31"
dependencies = [ dependencies = [
"bytes", "bytes",
"getrandom 0.3.3", "getrandom 0.3.3",
...@@ -4194,16 +4194,16 @@ dependencies = [ ...@@ -4194,16 +4194,16 @@ dependencies = [
[[package]] [[package]]
name = "quinn-udp" name = "quinn-udp"
version = "0.5.13" version = "0.5.14"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "fcebb1209ee276352ef14ff8732e24cc2b02bbac986cd74a4c81bcb2f9881970" checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd"
dependencies = [ dependencies = [
"cfg_aliases", "cfg_aliases",
"libc", "libc",
"once_cell", "once_cell",
"socket2 0.5.10", "socket2 0.6.0",
"tracing", "tracing",
"windows-sys 0.59.0", "windows-sys 0.60.2",
] ]
[[package]] [[package]]
...@@ -6081,11 +6081,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" ...@@ -6081,11 +6081,11 @@ checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]] [[package]]
name = "wasi" name = "wasi"
version = "0.14.2+wasi-0.2.4" version = "0.14.3+wasi-0.2.4"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" checksum = "6a51ae83037bdd272a9e28ce236db8c07016dd0d50c27038b3f407533c030c95"
dependencies = [ dependencies = [
"wit-bindgen-rt", "wit-bindgen",
] ]
[[package]] [[package]]
...@@ -6507,13 +6507,10 @@ dependencies = [ ...@@ -6507,13 +6507,10 @@ dependencies = [
] ]
[[package]] [[package]]
name = "wit-bindgen-rt" name = "wit-bindgen"
version = "0.39.0" version = "0.45.0"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" checksum = "052283831dbae3d879dc7f51f3d92703a316ca49f91540417d38591826127814"
dependencies = [
"bitflags 2.9.3",
]
[[package]] [[package]]
name = "writeable" name = "writeable"
......
lib/llm/Cargo.toml
View file @ 95ce83d5
...@@ -36,6 +36,8 @@ testing-etcd = [] ...@@ -36,6 +36,8 @@ testing-etcd = []
block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:ndarray", "dep:nix"] block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:ndarray", "dep:nix"]
cuda = ["dep:cudarc"] cuda = ["dep:cudarc"]
integration = [] integration = []
# NOTE: This feature will be enabled once ModelExpress packages are published
# model-express = ["dep:model_express_client", "dep:model_express_common"]
[[bench]] [[bench]]
name = "tokenizer" name = "tokenizer"
...@@ -81,6 +83,11 @@ url = { workspace = true } ...@@ -81,6 +83,11 @@ url = { workspace = true }
uuid = { workspace = true } uuid = { workspace = true }
xxhash-rust = { workspace = true } xxhash-rust = { workspace = true }
# ModelExpress client for model downloading
# NOTE: These will be uncommented once the packages are published to crates.io
# model_express_client = { version = "0.1.0", optional = true }
# model_express_common = { version = "0.1.0", optional = true }
akin = "0.4.0" akin = "0.4.0"
blake3 = "1" blake3 = "1"
bytemuck = "1.22" bytemuck = "1.22"
......
lib/llm/src/hub.rs
View file @ 95ce83d5
...@@ -13,18 +13,20 @@ ...@@ -13,18 +13,20 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#![allow(unexpected_cfgs)]
use hf_hub::api::tokio::ApiBuilder; use hf_hub::api::tokio::ApiBuilder;
use std::env; use std::env;
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
const IGNORED: [&str; 5] = [ #[cfg(feature = "model-express")]
".gitattributes", use model_express_client::{
"LICENSE", Client as MxClient, ClientConfig as MxClientConfig, ModelProvider as MxModelProvider,
"LICENSE.txt", };
"README.md", #[cfg(feature = "model-express")]
"USE_POLICY.md", use model_express_common::download as mx;
];
const MODEL_EXPRESS_ENDPOINT_ENV_VAR: &str = "MODEL_EXPRESS_URL";
const HF_TOKEN_ENV_VAR: &str = "HF_TOKEN"; const HF_TOKEN_ENV_VAR: &str = "HF_TOKEN";
/// Checks if a file is a model weight file /// Checks if a file is a model weight file
...@@ -36,63 +38,142 @@ fn is_weight_file(filename: &str) -> bool { ...@@ -36,63 +38,142 @@ fn is_weight_file(filename: &str) -> bool {
|| filename.ends_with(".ckpt.index") || filename.ends_with(".ckpt.index")
} }
/// Attempt to download a model from Hugging Face /// Attempt to download a model from Hugging Face using ModelExpress client
/// Only called when model-express feature is enabled, otherwise it will fall back to homonymous hf-hub function
/// Returns the directory it is in
/// If ignore_weights is true, model weight files will be skipped
#[cfg(feature = "model-express")]
pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Result<PathBuf> {
let name = name.as_ref();
let model_name = name.display().to_string();
// Only use ModelExpress if the environment variable is explicitly set
if let Ok(endpoint) = env::var(MODEL_EXPRESS_ENDPOINT_ENV_VAR) {
tracing::info!(
"ModelExpress endpoint configured, attempting to use ModelExpress for model: {model_name}"
);
let config: MxClientConfig = MxClientConfig::default().with_endpoint(endpoint.clone());
let result = match MxClient::new(config.clone()).await {
Ok(mut client) => {
tracing::info!("Successfully connected to ModelExpress server");
match client
.request_model_with_provider_and_fallback(
&model_name,
MxModelProvider::HuggingFace,
)
.await
{
Ok(()) => {
tracing::info!("Server download succeeded for model: {model_name}");
get_mx_model_path_from_cache(&model_name)
}
Err(e) => {
tracing::warn!(
"Server download failed for model '{model_name}': {e}. Falling back to direct download."
);
mx_download_direct(&model_name).await
}
}
}
Err(e) => {
tracing::warn!(
"Cannot connect to ModelExpress server: {e}. Using direct download."
);
mx_download_direct(&model_name).await
}
};
match result {
Ok(path) => {
tracing::info!(
"ModelExpress download completed successfully for model: {model_name}"
);
return Ok(path);
}
Err(e) => {
tracing::warn!(
"ModelExpress download failed for model '{model_name}': {e}. Falling back to hf-hub."
);
}
}
}
tracing::info!("Using hf-hub for model: {model_name}");
download_with_hf_hub(&model_name, ignore_weights).await
}
/// Attempt to download a model from Hugging Face using hf-hub directly
/// Called when model-express feature is not enabled
/// Returns the directory it is in /// Returns the directory it is in
/// If ignore_weights is true, model weight files will be skipped /// If ignore_weights is true, model weight files will be skipped
#[cfg(not(feature = "model-express"))]
pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Result<PathBuf> { pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Result<PathBuf> {
let name = name.as_ref(); let name = name.as_ref();
let model_name = name.display().to_string();
if env::var(MODEL_EXPRESS_ENDPOINT_ENV_VAR).is_ok() {
tracing::warn!(
"ModelExpress endpoint configured but model-express feature not enabled. Using hf-hub."
);
}
tracing::info!("Using hf-hub for model: {model_name}");
download_with_hf_hub(&model_name, ignore_weights).await
}
// Direct download using the ModelExpress client.
#[cfg(feature = "model-express")]
async fn mx_download_direct(model_name: &str) -> anyhow::Result<PathBuf> {
let cache_dir = get_model_express_cache_dir();
mx::download_model(model_name, MxModelProvider::HuggingFace, Some(cache_dir)).await
}
/// Attempt to download a model from Hugging Face with hf-hub
/// Returns the directory it is in
/// If ignore_weights is true, model weight files will be skipped
async fn download_with_hf_hub(model_name: &str, ignore_weights: bool) -> anyhow::Result<PathBuf> {
let token = env::var(HF_TOKEN_ENV_VAR).ok(); let token = env::var(HF_TOKEN_ENV_VAR).ok();
let api = ApiBuilder::from_env() let api = ApiBuilder::from_env()
.with_progress(true) .with_progress(true)
.with_token(token) .with_token(token)
.high() .high()
.build()?; .build()?;
let model_name = name.display().to_string();
let repo = api.model(model_name.clone()); let repo = api.model(model_name.to_string());
let info = match repo.info().await { let info = repo.info().await
Ok(info) => info, .map_err(|e| anyhow::anyhow!("Failed to fetch model '{model_name}' from HuggingFace: {e}. Is this a valid HuggingFace ID?"))?;
Err(e) => {
return Err(anyhow::anyhow!(
"Failed to fetch model '{}' from HuggingFace: {}. Is this a valid HuggingFace ID?",
model_name,
e
));
}
};
if info.siblings.is_empty() { if info.siblings.is_empty() {
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!(
"Model '{}' exists but contains no downloadable files.", "Model '{model_name}' exists but contains no downloadable files."
model_name
)); ));
} }
let mut p = PathBuf::new(); let mut model_path = PathBuf::new();
let mut files_downloaded = false; let mut files_downloaded = false;
for sib in info.siblings { for sibling in info.siblings {
if IGNORED.contains(&sib.rfilename.as_str()) || is_image(&sib.rfilename) { if is_ignored_file(&sibling.rfilename) || is_image_file(&sibling.rfilename) {
continue; continue;
} }
// If ignore_weights is true, skip weight files if ignore_weights && is_weight_file(&sibling.rfilename) {
if ignore_weights && is_weight_file(&sib.rfilename) {
continue; continue;
} }
match repo.get(&sib.rfilename).await { match repo.get(&sibling.rfilename).await {
Ok(path) => { Ok(path) => {
p = path; model_path = path;
files_downloaded = true; files_downloaded = true;
} }
Err(e) => { Err(e) => {
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!(
"Failed to download file '{}' from model '{}': {}", "Failed to download file '{}' from model '{model_name}': {e}",
sib.rfilename, sibling.rfilename
model_name,
e
)); ));
} }
} }
...@@ -105,23 +186,126 @@ pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Re ...@@ -105,23 +186,126 @@ pub async fn from_hf(name: impl AsRef<Path>, ignore_weights: bool) -> anyhow::Re
"valid" "valid"
}; };
return Err(anyhow::anyhow!( return Err(anyhow::anyhow!(
"No {} files found for model '{}'.", "No {file_type} files found for model '{model_name}'."
file_type, ));
model_name }
match model_path.parent() {
Some(path) => Ok(path.to_path_buf()),
None => Err(anyhow::anyhow!(
"Invalid HF cache path: {}",
model_path.display()
)),
}
}
fn is_ignored_file(filename: &str) -> bool {
const IGNORED_FILES: [&str; 5] = [
".gitattributes",
"LICENSE",
"LICENSE.txt",
"README.md",
"USE_POLICY.md",
];
IGNORED_FILES.contains(&filename)
}
fn is_image_file(filename: &str) -> bool {
filename.ends_with(".png")
|| filename.ends_with("PNG")
|| filename.ends_with(".jpg")
|| filename.ends_with("JPG")
|| filename.ends_with(".jpeg")
|| filename.ends_with("JPEG")
}
#[cfg(feature = "model-express")]
fn get_mx_model_path_from_cache(model_name: &str) -> anyhow::Result<PathBuf> {
let cache_dir = get_model_express_cache_dir();
let model_dir = cache_dir.join(model_name);
if !model_dir.exists() {
return Err(anyhow::anyhow!(
"Model '{model_name}' was downloaded but directory not found at expected location: {}",
model_dir.display()
)); ));
} }
match p.parent() { Ok(model_dir)
Some(p) => Ok(p.to_path_buf()), }
None => Err(anyhow::anyhow!("Invalid HF cache path: {}", p.display())),
#[cfg(feature = "model-express")]
fn get_model_express_cache_dir() -> PathBuf {
if let Ok(cache_path) = env::var("HF_HUB_CACHE") {
return PathBuf::from(cache_path);
} }
if let Ok(cache_path) = env::var("MODEL_EXPRESS_PATH") {
return PathBuf::from(cache_path);
}
let home = env::var("HOME")
.or_else(|_| env::var("USERPROFILE"))
.unwrap_or_else(|_| ".".to_string());
PathBuf::from(home).join(".cache/huggingface/hub")
} }
fn is_image(s: &str) -> bool { #[cfg(test)]
s.ends_with(".png") mod tests {
|| s.ends_with("PNG") use super::*;
|| s.ends_with(".jpg")
|| s.ends_with("JPG") #[tokio::test]
|| s.ends_with(".jpeg") async fn test_from_hf_with_model_express() {
|| s.ends_with("JPEG") let test_path = PathBuf::from("test-model");
let _result: anyhow::Result<PathBuf> = from_hf(test_path, false).await;
}
#[cfg(feature = "model-express")]
#[test]
fn test_get_model_express_cache_dir() {
let cache_dir = get_model_express_cache_dir();
assert!(!cache_dir.to_string_lossy().is_empty());
assert!(cache_dir.is_absolute() || cache_dir.starts_with("."));
}
#[test]
fn test_is_ignored_file() {
assert!(is_ignored_file(".gitattributes"));
assert!(is_ignored_file("LICENSE"));
assert!(is_ignored_file("LICENSE.txt"));
assert!(is_ignored_file("README.md"));
assert!(is_ignored_file("USE_POLICY.md"));
assert!(!is_ignored_file("model.bin"));
assert!(!is_ignored_file("tokenizer.json"));
assert!(!is_ignored_file("config.json"));
}
#[test]
fn test_is_weight_file() {
assert!(is_weight_file("model.bin"));
assert!(is_weight_file("model.safetensors"));
assert!(is_weight_file("model.h5"));
assert!(is_weight_file("model.msgpack"));
assert!(is_weight_file("model.ckpt.index"));
assert!(!is_weight_file("tokenizer.json"));
assert!(!is_weight_file("config.json"));
assert!(!is_weight_file("README.md"));
}
#[test]
fn test_is_image_file() {
assert!(is_image_file("image.png"));
assert!(is_image_file("image.PNG"));
assert!(is_image_file("photo.jpg"));
assert!(is_image_file("photo.JPG"));
assert!(is_image_file("picture.jpeg"));
assert!(is_image_file("picture.JPEG"));
assert!(!is_image_file("model.bin"));
assert!(!is_image_file("tokenizer.json"));
assert!(!is_image_file("config.json"));
assert!(!is_image_file("README.md"));
}
} }
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment