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

fix(launcher): parse num cuda devices from CUDA_VISIBLE_DEVICES and NVIDIA_VISIBLE_DEVICES

Browse Source
This commit is contained in:
OlivierDehaene 2023-05-30 12:52:18 +02:00
parent 5fde8d9991
commit 146e72c3be
1 changed files with 8 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
launcher/src/main.rs
View File

@ -455,11 +455,11 @@ fn shutdown_shards(shutdown: Arc<Mutex<bool>>, shutdown_receiver: &mpsc::Receive
} }
fn num_cuda_devices() -> Option<usize> { fn num_cuda_devices() -> Option<usize> {
if let Ok(cuda_visible_devices) = env::var("CUDA_VISIBLE_DEVICES") { let devices = env::var("CUDA_VISIBLE_DEVICES")
let n_devices = cuda_visible_devices.split(',').count(); .map_err(|_| env::var("NVIDIA_VISIBLE_DEVICES"))
return Some(n_devices); .ok()?;
} let n_devices = devices.split(',').count();
None Some(n_devices)
} }
#[derive(Deserialize)] #[derive(Deserialize)]
@ -509,9 +509,9 @@ fn find_num_shards(sharded: Option<bool>, num_shard: Option<usize>) -> usize {
let num_shard = match (sharded, num_shard) { let num_shard = match (sharded, num_shard) {
(Some(true), None) => { (Some(true), None) => {
// try to default to the number of available GPUs // try to default to the number of available GPUs
tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES"); tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES");
let n_devices = let n_devices = num_cuda_devices()
num_cuda_devices().expect("--num-shard and CUDA_VISIBLE_DEVICES are not set"); .expect("--num-shard and CUDA_VISIBLE_DEVICES/NVIDIA_VISIBLE_DEVICES are not set");
if n_devices <= 1 { if n_devices <= 1 {
panic!("`sharded` is true but only found {n_devices} CUDA devices"); panic!("`sharded` is true but only found {n_devices} CUDA devices");
} }