From e63a21eb4da9db60c973412820dd4c8644883579 Mon Sep 17 00:00:00 2001 From: OlivierDehaene Date: Sun, 9 Apr 2023 20:00:05 +0200 Subject: [PATCH] feat(launcher): allow disabling hf_transfer (#161) --- launcher/src/main.rs | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/launcher/src/main.rs b/launcher/src/main.rs index 80466fe..a598a8b 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -200,7 +200,16 @@ fn main() -> ExitCode { }; // Enable hf transfer for insane download speeds - env.push(("HF_HUB_ENABLE_HF_TRANSFER".into(), "1".into())); + let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string()); + env.push(( + "HF_HUB_ENABLE_HF_TRANSFER".into(), + enable_hf_transfer.into(), + )); + + // Parse Inference API token + if let Ok(api_token) = env::var("HF_API_TOKEN") { + env.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into())) + }; // Start process tracing::info!("Starting download process."); @@ -555,7 +564,16 @@ fn shard_manager( env.push(("SAFETENSORS_FAST_GPU".into(), "1".into())); // Enable hf transfer for insane download speeds - env.push(("HF_HUB_ENABLE_HF_TRANSFER".into(), "1".into())); + let enable_hf_transfer = env::var("HF_HUB_ENABLE_HF_TRANSFER").unwrap_or("1".to_string()); + env.push(( + "HF_HUB_ENABLE_HF_TRANSFER".into(), + enable_hf_transfer.into(), + )); + + // Parse Inference API token + if let Ok(api_token) = env::var("HF_API_TOKEN") { + env.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into())) + }; // If huggingface_hub_cache is some, pass it to the shard // Useful when running inside a docker container