diff --git a/launcher/src/main.rs b/launcher/src/main.rs index e170db2b..1317868c 100644 --- a/launcher/src/main.rs +++ b/launcher/src/main.rs @@ -115,7 +115,8 @@ fn main() -> ExitCode { None => { // try to default to the number of available GPUs tracing::info!("Parsing num_shard from CUDA_VISIBLE_DEVICES"); - let n_devices = num_cuda_devices().expect("--num-shard and CUDA_VISIBLE_DEVICES are not set"); + let n_devices = num_cuda_devices() + .expect("--num-shard and CUDA_VISIBLE_DEVICES are not set"); if n_devices <= 1 { panic!("`sharded` is true but only found {n_devices} CUDA devices"); } @@ -144,7 +145,7 @@ fn main() -> ExitCode { match num_shard { // get num_shard from CUDA_VISIBLE_DEVICES or default to a single shard None => num_cuda_devices().unwrap_or(1), - Some(num_shard) => num_shard + Some(num_shard) => num_shard, } }; tracing::info!("Sharding model on {num_shard} processes"); @@ -155,7 +156,7 @@ fn main() -> ExitCode { ctrlc::set_handler(move || { r.store(false, Ordering::SeqCst); }) - .expect("Error setting Ctrl-C handler"); + .expect("Error setting Ctrl-C handler"); // Check if model_id is a local model let local_path = Path::new(&model_id); @@ -671,11 +672,10 @@ fn shutdown_shards(shutdown: Arc>, shutdown_receiver: &mpsc::Receive let _ = shutdown_receiver.recv(); } - fn num_cuda_devices() -> Option { if let Ok(cuda_visible_devices) = env::var("CUDA_VISIBLE_DEVICES") { - let n_devices = cuda_visible_devices.split(",").count(); + let n_devices = cuda_visible_devices.split(',').count(); return Some(n_devices); } None -} \ No newline at end of file +}