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