fix(xtask): 检测系统 cuda、nccl 支持

Signed-off-by: YdrMaster <[email protected]>

xtask/Cargo.toml

@@ -21,6 +21,9 @@ simple_logger = "4.3"
 colored = "2.1"
 clap = { version = "4.5", features = ["derive"] }
 
+[build-dependencies]
+search-cuda-tools.workspace = true
+
 [features]
 default = ["nvidia"]
 nvidia = ["transformer-nv"]

xtask/build.rs

@@ -0,0 +1,12 @@
+fn main() {
+    use search_cuda_tools::*;
+    if !cfg!(feature = "nvidia") {
+        return;
+    }
+    if find_cuda_root().is_some() {
+        detect_cuda();
+    }
+    if find_nccl_root().is_some() {
+        detect_nccl();
+    }
+}

xtask/src/chat.rs

@@ -19,14 +19,14 @@ impl crate::InferenceArgs {
                 use transformer_cpu::Transformer as M;
                 chat!(M; ());
             }
-            #[cfg(feature = "nvidia")]
+            #[cfg(detected_cuda)]
             &[n] => {
                 use transformer_nv::{cuda, Transformer as M};
                 chat!(M; cuda::Device::new(n));
             }
-            #[cfg(feature = "nvidia")]
+            #[cfg(detected_nccl)]
             _distribute => todo!(),
-            #[cfg(not(feature = "nvidia"))]
+            #[cfg(not(all(detected_cuda, detected_nccl)))]
             _ => panic!("Set \"nvidia\" feature to enablel nvidia support."),
         }
     }

xtask/src/main.rs

@@ -25,14 +25,14 @@ fn main() {
 
 #[inline]
 fn block_on(f: impl Future) {
-    #[cfg(feature = "nvidia")]
+    #[cfg(detected_cuda)]
     {
         transformer_nv::cuda::init();
     }
     let runtime = tokio::runtime::Runtime::new().unwrap();
     runtime.block_on(f);
     runtime.shutdown_background();
-    #[cfg(feature = "nvidia")]
+    #[cfg(detected_cuda)]
     {
         transformer_nv::synchronize();
     }

xtask/src/service.rs

@@ -26,14 +26,14 @@ impl ServiceArgs {
                 use transformer_cpu::Transformer as M;
                 serve!(M; ());
             }
-            #[cfg(feature = "nvidia")]
+            #[cfg(detected_cuda)]
             &[n] => {
                 use transformer_nv::{cuda, Transformer as M};
                 serve!(M; cuda::Device::new(n));
             }
-            #[cfg(feature = "nvidia")]
+            #[cfg(detected_nccl)]
             _distribute => todo!(),
-            #[cfg(not(feature = "nvidia"))]
+            #[cfg(not(all(detected_cuda, detected_nccl)))]
             _ => panic!("Set \"nvidia\" feature to enablel nvidia support."),
         }
     }