full threading and fixes

engine/src/worker.rs

@@ -61,7 +61,7 @@ impl Worker {
         F: Fn(&BoardState) -> Fut,
         Fut: Future<Output = (Vec<f32>, f32)>,
     {
-        println!("game started");
+        // println!("game started");
         let now = std::time::SystemTime::now()
             .duration_since(std::time::UNIX_EPOCH)
             .unwrap();
@@ -84,12 +84,11 @@ impl Worker {
             // println!("{}", before.elapsed().as_millis());
             episode_buffer.push(results);
 
-            let temp = if ply_count < 30 { 1.0 } else { 0.0 };
+            let temp = if ply_count < 30 { 1.0 } else { 0.2 };
 
             let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
 
-            let stm = board_state.board.side_to_move();
-            let mv = sample_move(&adjusted, &mut rng, stm).unwrap();
+            let mv = sample_move(&adjusted, &mut rng, &board_state.board).unwrap();
 
             game_hist.push_str(mv.to_string().as_str());
             game_hist.push(' ');
@@ -103,6 +102,8 @@ impl Worker {
             }
         }
 
+        println!("game finished: {}", game_hist);
+
         let mut samples: Vec<TrainingSample> = Vec::new();
         for result in episode_buffer.iter() {
             let stm = result.board_state.board.side_to_move();
@@ -166,14 +167,15 @@ impl Worker {
     // }
 
     pub async fn run_worker(self, cmd_rx: Receiver<WorkerCmd>) {
-        println!("Worker started");
+        let mut tasks = Vec::new();
+        // println!("Worker started");
         for _ in 0..self.num_games {
             let mcts_config = self.mcts_config.clone();
             let data_tx = self.data_tx.clone();
             let request_tx = self.request_tx.clone();
             let device = self.device.clone();
 
-            tokio::spawn(async move {
+            tasks.push(tokio::spawn(async move {
                 Self::run_game(
                     self.id,
                     mcts_config,
@@ -186,19 +188,29 @@ impl Worker {
                     data_tx,
                 )
                 .await;
-            });
+            }));
         }
 
         loop {
-            if let Ok(cmd) = cmd_rx.recv() {
-                match cmd {
+            // Remove completed tasks.
+            tasks.retain(|task| !task.is_finished());
+
+            // If every game has completed, we're done.
+            if tasks.is_empty() {
+                break;
+            }
+
+            match cmd_rx.recv_timeout(Duration::from_millis(10)) {
+                Ok(cmd) => match cmd {
                     WorkerCmd::Shutdown => break,
+
                     WorkerCmd::Dispatch => {
                         let mcts_config = self.mcts_config.clone();
                         let request_tx = self.request_tx.clone();
                         let device = self.device.clone();
                         let data_tx = self.data_tx.clone();
-                        tokio::spawn(async move {
+
+                        tasks.push(tokio::spawn(async move {
                             Self::run_game(
                                 self.id,
                                 mcts_config,
@@ -211,13 +223,26 @@ impl Worker {
                                 data_tx,
                             )
                             .await;
-                        });
+                        }));
                     }
+                },
+
+                Err(RecvTimeoutError::Timeout) => {
+                    continue;
+                }
+
+                Err(RecvTimeoutError::Disconnected) => {
+                    println!("Request channel disconnected");
+                    break;
                 }
-            } else {
-                panic!("worker dropped inference command");
             }
         }
+
+        for task in tasks {
+            task.await.expect("Worker game task panicked");
+        }
+
+        println!("Worker {} shut down", self.id);
     }
 }
 
@@ -303,13 +328,24 @@ impl NNServer {
     pub fn run_server(&self) {
         let mut cache: Cache<NNCacheKey, NNResponse> = Cache::new(self.cache_size);
         let mut queue: VecDeque<NNRequest> = VecDeque::new();
+        let mut request_counter: usize = 0;
+        let mut cache_hits: usize = 0;
+        let mut batches_ran: usize = 0;
 
         loop {
             // Check for shutdown before doing any work.
             if let Ok(cmd) = self.command_rx.try_recv() {
                 match cmd {
                     NNServerCmd::Shutdown => {
-                        println!("NN server shutting down");
+                        println!(
+                            "NN server shutting down. Served {} requests, with {} cache hits ({:2}%). Ran {} batches (avg {}requests/batch).",
+                            request_counter,
+                            cache_hits,
+                            cache_hits as f64 / request_counter as f64 * 100.0,
+                            batches_ran,
+                            (request_counter - cache_hits) as f64 / batches_ran as f64,
+                        );
+
                         break;
                     } // other commands...
                 }
@@ -320,7 +356,9 @@ impl NNServer {
             if queue.is_empty() {
                 match self.request_rx.recv_timeout(Duration::from_millis(10)) {
                     Ok(first) => {
+                        request_counter += 1;
                         if try_use_cache(first, &mut cache, &mut queue) {
+                            cache_hits += 1;
                             continue;
                         }
                     }
@@ -337,7 +375,11 @@ impl NNServer {
             }
 
             while let Ok(req) = self.request_rx.try_recv() {
-                try_use_cache(req, &mut cache, &mut queue);
+                request_counter += 1;
+                if try_use_cache(req, &mut cache, &mut queue) {
+                    cache_hits += 1;
+                    continue;
+                }
             }
 
             let mut jobs_to_send: Vec<NNRequest> = Vec::new();
@@ -359,7 +401,9 @@ impl NNServer {
             );
 
             // println!("Sending {} jobs to gpu", jobs_to_send.len());
+
             let (p, v) = self.model.forward(batched_tensor);
+            batches_ran += 1;
             let p_split = p.chunk(jobs_to_send.len(), 0);
             let v_split = v.chunk(jobs_to_send.len(), 0);