huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-20 06:12:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(launcher): pass weights cache override to the download process (#274)

Browse Source 
closes #273
This commit is contained in:
OlivierDehaene 2023-05-03 23:39:35 +02:00 committed by GitHub
parent 85aa7e2e7b
commit b67908e0cf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
launcher/src/main.rs
View File

@ -542,7 +542,7 @@ fn download_convert_model(
// Copy current process env // Copy current process env
let mut env: Vec<(OsString, OsString)> = env::vars_os().collect(); let mut env: Vec<(OsString, OsString)> = env::vars_os().collect();
// If huggingface_hub_cache is set, pass it to the shard // If huggingface_hub_cache is set, pass it to the download process
// Useful when running inside a docker container // Useful when running inside a docker container
if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache { if let Some(ref huggingface_hub_cache) = args.huggingface_hub_cache {
env.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into())); env.push(("HUGGINGFACE_HUB_CACHE".into(), huggingface_hub_cache.into()));
@ -560,6 +560,15 @@ fn download_convert_model(
env.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into())) env.push(("HUGGING_FACE_HUB_TOKEN".into(), api_token.into()))
}; };
// If args.weights_cache_override is some, pass it to the download process
// Useful when running inside a HuggingFace Inference Endpoint
if let Some(weights_cache_override) = &args.weights_cache_override {
env.push((
"WEIGHTS_CACHE_OVERRIDE".into(),
weights_cache_override.into(),
));
};
// Start process // Start process
tracing::info!("Starting download process."); tracing::info!("Starting download process.");
let mut download_process = match Popen::create( let mut download_process = match Popen::create(