Mirrors
/
hf_text-generation-inference
mirror of https://github.com/huggingface/text-generation-inference.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat(launcher): Pass CUDA_VISIBLE_DEVICES to the shard

This commit is contained in:
OlivierDehaene 2022-11-04 18:31:08 +01:00
parent 427d7cc444
commit cea6051eff
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
launcher/src/main.rs
View File

@ -306,6 +306,14 @@ fn shard_manager(
)); ));
}; };
// If the CUDA_VISIBLE_DEVICES env var is set, pass it to the shard
if let Ok(cuda_visible_devices) = env::var("CUDA_VISIBLE_DEVICES") {
env.push((
"CUDA_VISIBLE_DEVICES".parse().unwrap(),
cuda_visible_devices.parse().unwrap(),
));
};
// Start process // Start process
tracing::info!("Starting shard {}", rank); tracing::info!("Starting shard {}", rank);
let mut p = match Popen::create( let mut p = match Popen::create(