Test server

reduction_server/src/main.rs

@@ -32,15 +32,7 @@ fn main() {
         .init();
     nccl_net::init();
 
-    let mut args = Args::parse();
-
-    if args.recv_threads == 0 {
-        args.recv_threads = args.nrank
-    }
-
-    if args.send_threads == 0 {
-        args.send_threads = args.nrank
-    }
+    let args = Args::parse();
 
     if args.client {
         client(args);

reduction_server/src/server.rs

@@ -94,7 +94,7 @@ fn reduce_loop<T: Float>(
     info!("reduce thread({}) all ranks get connected!", i);
 
     let mut mems = (0..jobs.len())
-        .map(|_| WorkingMemory::new(args.count, args.recv_threads))
+        .map(|_| WorkingMemory::new(args.count / args.reduce_threads, args.recv_threads))
         .collect::<Vec<_>>();
 
     loop {
@@ -447,6 +447,16 @@ fn recv_loop<T: Float>(
 }
 
 fn do_server<T: Float + 'static>(args: Args) {
+    let mut args = args;
+
+    if args.recv_threads == 0 {
+        args.recv_threads = args.nrank
+    }
+
+    if args.send_threads == 0 {
+        args.send_threads = args.nrank
+    }
+
     let listener =
         TcpListener::bind(format!("{}:{}", args.address, args.port)).expect("failed to bind");
 
@@ -583,3 +593,71 @@ pub(crate) fn server(args: Args) {
         do_server::<bf16>(args);
     }
 }
+
+// test
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::client::client;
+    use crate::utils::tests::initialize;
+    use clap::Parser;
+
+    fn do_test(dt: &str) {
+        initialize();
+        let nrank = 4;
+        let server = {
+            let dt = dt.to_string();
+            std::thread::spawn(move || {
+                let nrank = format!("{}", nrank);
+                let args = Args::parse_from([
+                    "--verbose", // doesn't work without specifying a flag that doesn't take an argument
+                    "--port",
+                    "8080",
+                    "--data-type",
+                    &dt,
+                    "--nrank",
+                    &nrank,
+                    "--nreq",
+                    "1", // when using socket plugin, concurrent recv/send requests doesn't work
+                ]);
+                server(args);
+            })
+        };
+        (0..nrank)
+            .map(|_| {
+                let dt = dt.to_string();
+                std::thread::spawn(move || {
+                    std::thread::sleep(std::time::Duration::from_millis(100));
+                    let args = Args::parse_from([
+                        "--client",
+                        "--address",
+                        "127.0.0.1:8080",
+                        "--data-type",
+                        &dt,
+                        "--nreq",
+                        "1", // when using socket plugin, concurrent recv/send requests doesn't work
+                    ]);
+                    client(args);
+                })
+            })
+            .collect::<Vec<_>>()
+            .into_iter()
+            .for_each(|h| h.join().unwrap());
+        server.join().unwrap();
+    }
+
+    #[test]
+    fn test_server_f32() {
+        do_test("f32");
+    }
+
+    #[test]
+    fn test_server_f16() {
+        do_test("f16");
+    }
+
+    #[test]
+    fn test_server_bf16() {
+        do_test("bf16");
+    }
+}