test: 添加环境变量控制硬件选项，合并 CPU 单卡和分布式版本测试

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

Cargo.toml

@@ -25,6 +25,8 @@ test-utils.path = "test-utils"
 
 ggus = "0.3"
 itertools = "0.13"
+regex = "1.11"
+env_logger = "0.11"
 build-script-cfg = "0.0"
 
 ndarray-layout = { git = "https://github.com/YdrMaster/ndarray-layout", rev = "f1fdd24" }

models/llama/common-cpu/Cargo.toml

@@ -13,3 +13,4 @@ operators = { workspace = true, features = ["common-cpu"] }
 [dev-dependencies]
 test-utils.workspace = true
 gguf.workspace = true
+regex.workspace = true

models/llama/common-cpu/src/test_dist.rs models/llama/common-cpu/src/infer.rsmodels/llama/common-cpu/src/test_dist.rs renamed to models/llama/common-cpu/src/infer.rs

@@ -7,6 +7,7 @@ use operators::{
     random_sample::{KVPair, SampleArgs},
     Blob,
 };
+use regex::Regex;
 use std::{
     iter::zip,
     ptr::copy_nonoverlapping,
@@ -22,9 +23,10 @@ use test_utils::{Inference, TokenizerAndPrompt};
 type Worker<'w> = LlamaWorker<Operators<InprocNode<usize>, AllReduce>, Weights<'w>>;
 
 #[test]
-fn test_dist() {
+fn test_infer() {
     let Some(Inference {
         model,
+        devices,
         prompt,
         as_user,
         temperature,
@@ -49,7 +51,17 @@ fn test_dist() {
     let sample_args = SampleArgs::new(temperature, top_p, top_k).expect("invalid sample args");
     println!("{sample_args:?}");
 
-    let lens = [1; 4];
+    let lens = match devices {
+        Some(devices) => {
+            let regex = Regex::new(r"\d+").unwrap();
+            regex
+                .find_iter(&devices)
+                .map(|c| c.as_str().parse::<usize>().unwrap())
+                .collect::<Vec<_>>()
+        }
+        None => vec![1],
+    };
+    println!("distribution: {lens:?}");
     let count = lens.iter().sum();
     let (seeds, senders) = WorkerSeed::new(lens.len());
     thread::scope(|s| {

models/llama/common-cpu/src/lib.rs

@@ -287,7 +287,4 @@ impl WeightLoader for Weights<'_> {
 }
 
 #[cfg(test)]
-mod test_infer;
-
-#[cfg(test)]
-mod test_dist;
+mod infer;

models/llama/common-cpu/src/test_infer.rs

@@ -22,6 +22,7 @@ fn test_infer() {
         top_p,
         top_k,
         max_steps,
+        ..
     }) = Inference::load()
     else {
         return;

models/llama/infini/src/test_infer.rs models/llama/infini/src/infer.rsmodels/llama/infini/src/test_infer.rs renamed to models/llama/infini/src/infer.rs

@@ -169,4 +169,4 @@ impl WeightLoader for Weights {
 }
 
 #[cfg(test)]
-mod test_infer;
+mod infer;

models/llama/infini/src/lib.rs

@@ -23,13 +23,14 @@ fn test_infer() {
         top_p,
         top_k,
         max_steps,
+        ..
     }) = Inference::load()
     else {
         return;
     };
 
     let roll_cache_size = load_roll_cache_size();
-    println!("roll_cache_size: {}", roll_cache_size);
+    println!("roll_cache_size: {roll_cache_size}");
     let gguf = GGufModel::read(model.iter().map(|s| &**s));
 
     let TokenizerAndPrompt {

models/llama/nvidia-gpu/src/test_infer.rs models/llama/nvidia-gpu/src/infer.rsmodels/llama/nvidia-gpu/src/test_infer.rs renamed to models/llama/nvidia-gpu/src/infer.rs

@@ -316,4 +316,4 @@ impl<'ctx> WeightLoader for Weights<'ctx> {
 }
 
 #[cfg(test)]
-mod test_infer;
+mod infer;

models/llama/nvidia-gpu/src/lib.rs

@@ -22,6 +22,7 @@ fn test_infer() {
         top_p,
         top_k,
         max_steps,
+        ..
     }) = Inference::load()
     else {
         return;