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

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

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.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/infini/src/test_infer.rs → models/llama/infini/src/infer.rs

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

models/llama/infini/src/lib.rs

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

models/llama/nvidia-gpu/src/test_infer.rs → models/llama/nvidia-gpu/src/infer.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/lib.rs

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

models/llama/opencl/src/test_infer.rs → models/llama/opencl/src/infer.rs

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

models/llama/opencl/src/lib.rs

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

test-utils/Cargo.toml

@@ -6,4 +6,5 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
 gguf.workspace = true
+env_logger.workspace = true
 cli-table = "0.4.9"

test-utils/src/lib.rs

@@ -7,12 +7,14 @@ use std::{
     fmt,
     path::{Path, PathBuf},
     str::FromStr,
+    sync::Once,
     time::{Duration, Instant},
     vec,
 };
 
 pub struct Inference {
     pub model: Box<[Mmap]>,
+    pub devices: Option<String>,
     pub prompt: String,
     pub as_user: bool,
     pub temperature: f32,
@@ -23,6 +25,9 @@ pub struct Inference {
 
 impl Inference {
     pub fn load() -> Option<Self> {
+        static ONCE: Once = Once::new();
+        ONCE.call_once(|| env_logger::init());
+
         let Some(path) = var_os("TEST_MODEL") else {
             println!("TEST_MODEL not set");
             return None;
@@ -42,6 +47,7 @@ impl Inference {
 
         Some(Self {
             model: map_files(path),
+            devices: var("DEVICES").ok(),
             prompt: var("PROMPT").unwrap_or_else(|_| String::from("Once upon a time,")),
             as_user: var("AS_USER").ok().is_some_and(|s| !s.is_empty()),
             temperature: parse("TEMPERATURE", 0.),