full threading and fixes

Cargo.lock

@@ -2480,9 +2480,9 @@ dependencies = [
 
 [[package]]
 name = "futures-channel"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d"
+checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4"
 dependencies = [
  "futures-core",
  "futures-sink",
@@ -2490,15 +2490,15 @@ dependencies = [
 
 [[package]]
 name = "futures-core"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d"
+checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
 
 [[package]]
 name = "futures-io"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718"
+checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
 
 [[package]]
 name = "futures-lite"
@@ -2515,26 +2515,26 @@ dependencies = [
 
 [[package]]
 name = "futures-macro"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b"
+checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44"
 dependencies = [
  "proc-macro2",
  "quote",
- "syn 2.0.117",
+ "syn 3.0.4",
 ]
 
 [[package]]
 name = "futures-sink"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893"
+checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d"
 
 [[package]]
 name = "futures-task"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393"
+checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd"
 
 [[package]]
 name = "futures-timer"
@@ -2544,9 +2544,9 @@ checksum = "af43fadb8a98512d547e37b4e92e0ced13e205c061b87b4623eff01d918d6968"
 
 [[package]]
 name = "futures-util"
-version = "0.3.32"
+version = "0.3.34"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6"
+checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
 dependencies = [
  "futures-core",
  "futures-io",

engine/src/main.rs

@@ -25,7 +25,7 @@ fn main() {
         num_iters: None,
         max_depth: 100, // unused
         model_name: String::from("threadedModel"),
-        load_model: false,
+        load_model: true,
         hidden_channels: 48,
         num_blocks: 4,
         batch_size: 1024,        // num positions sampled to a batch (2048)
@@ -35,10 +35,10 @@ fn main() {
         optimizer: adam_config,
         lr: 2e-4,
         seed: None,
-        num_workers: 4,
+        num_workers: 8,
         nn_batch_size: 64,
-        episodes_per_worker: 4,
+        episodes_per_worker: 12,
-        nn_cache_size: 500_000, // 18,736B / item in cache. 100k=1.75GB.
+        nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB.
     };
 
     train(training_config, device);

engine/src/mcts.rs

@@ -87,7 +87,8 @@ pub struct MctsConfig {
     pub c_puct: f32,
     pub dirichlet_alpha: f32,
     pub dirichlet_epsilon: f32,
-    pub batch_max: usize,
+    pub batch_size: usize,
+    pub num_batches: usize,
 }
 
 impl MctsConfig {
@@ -96,14 +97,15 @@ impl MctsConfig {
         c_puct: f32,
         dirichlet_alpha: f32,
         dirichlet_epsilon: f32,
-        batch_max: usize,
+        batch_size: usize,
     ) -> MctsConfig {
         MctsConfig {
             num_simulations,
             c_puct,
             dirichlet_alpha,
             dirichlet_epsilon,
-            batch_max,
+            batch_size,
+            num_batches: num_simulations / batch_size,
         }
     }
 }
@@ -191,6 +193,34 @@ impl Mcts {
             let mut path = vec![root];
             let mut current = root;
 
+            // for batch in 0..self.config.num_batches {
+            //     let mut pending = Vec::new();
+            //
+            //     // Select several leaves
+            //     for _ in 0..self.config.batch_size {
+            //         let path = select_leaf_with_virtual_loss(...);
+            //
+            //         let board = nodes[path.last()].board_state.clone();
+            //
+            //         pending.push(PendingEval {
+            //             path,
+            //             board,
+            //         });
+            //     }
+            //
+            //     // Send all boards to NN server
+            //     let evaluations = model_inference_batch(pending.boards()).await;
+            //
+            //     // Expand and backpropagate
+            //     for (pending, (policy, value)) in
+            //         pending.into_iter().zip(evaluations)
+            //     {
+            //         expand_with_policy(...);
+            //         remove_virtual_loss(...);
+            //         backpropagate(...);
+            //     }
+            // }
+
             while !nodes[current].children.is_empty() {
                 current = nodes[current].select_child(&nodes, &self.config.c_puct);
                 path.push(current);

engine/src/net/encoding.rs

@@ -1,6 +1,6 @@
 use crate::mcts::BoardState;
 use burn::prelude::*;
-use chess::{ChessMove, Color, File, Piece, Rank, Square};
+use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
 
 const SLIDING_DIRS: [(i8, i8); 8] = [
     (0, 1),   // N
@@ -212,7 +212,7 @@ fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
 }
 
 /// Decode an index in [0, 4672) back into a ChessMove.
-pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
+pub fn decode_move(idx: usize, board: &Board) -> Option<ChessMove> {
     if idx >= 4672 {
         return None;
     }
@@ -250,10 +250,24 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
     }
 
     // Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
-    if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
+    if promo.is_none()
+        && fr == 6
+        && tr == 7
+        && dy == 1
+        && dx.abs() <= 1
+        && board.piece_on(Square::make_square(
+            Rank::from_index(fr as usize),
+            File::from_index(ff as usize),
+        )) == Some(Piece::Pawn)
+    {
         promo = Some(Piece::Queen);
     }
 
+    // if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
+    //     promo = Some(Piece::Queen);
+    // }
+
+    let stm = board.side_to_move();
     let from = deorient_square(ff, fr, stm);
     let to = deorient_square(tf, tr, stm);
 
@@ -262,19 +276,20 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
 
 #[cfg(test)]
 mod tests {
-    use super::*;
+    use crate::net::encoding::{decode_move, encode_move};
-    use chess::ALL_COLORS;
+    use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
+    use std::str::FromStr;
 
     #[test]
     fn encoding_roundtrips() {
-        for color in ALL_COLORS {
+        let board = Board::default();
         for action in 0..4672 {
-                let decoded = decode_move(action, color);
+            let decoded = decode_move(action, &board);
             if decoded.is_none() {
                 continue;
             }
             let decoded = decoded.unwrap();
-                let encoded = encode_move(decoded, color).unwrap();
+            let encoded = encode_move(decoded, board.side_to_move()).unwrap();
             // eprintln!(
             //     "orig idx = {}, plane={}, from_idx={}",
             //     action,
@@ -304,27 +319,34 @@ mod tests {
             assert_eq!(action, encoded);
         }
     }
-    }
 
     #[test]
     fn queen_promotion_roundtrip() {
-        use chess::{File, Rank, Square};
+        let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
+
         let from = Square::make_square(Rank::Seventh, File::E);
         let to = Square::make_square(Rank::Eighth, File::E);
+
         let mv = ChessMove::new(from, to, Some(Piece::Queen));
+
         let enc = encode_move(mv, Color::White).unwrap();
-        let dec = decode_move(enc, Color::White).unwrap();
+        let dec = decode_move(enc, &board).unwrap();
+
         assert_eq!(mv, dec);
     }
 
     #[test]
     fn underpromotion_roundtrip() {
-        use chess::{File, Rank, Square};
+        let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
+
         let from = Square::make_square(Rank::Seventh, File::E);
         let to = Square::make_square(Rank::Eighth, File::E);
+
         let mv = ChessMove::new(from, to, Some(Piece::Knight));
+
         let enc = encode_move(mv, Color::White).unwrap();
-        let dec = decode_move(enc, Color::White).unwrap();
+        let dec = decode_move(enc, &board).unwrap();
+
         assert_eq!(mv, dec);
     }
 }

engine/src/train.rs

@@ -9,7 +9,7 @@ use burn::module::AutodiffModule;
 use burn::optim::{AdamConfig, GradientsParams};
 use burn::prelude::{Device, Module};
 use burn::store::ModuleRecord;
-use chess::{ChessMove, Color};
+use chess::{Board, ChessMove};
 use rand::rngs::SmallRng;
 use rand::seq::SliceRandom;
 use rand::{RngExt, SeedableRng};
@@ -113,6 +113,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
     println!("Starting training...");
     while train.load(Ordering::Relaxed) {
+        let iter_start = Instant::now();
         println!("Iteration: {}", iter);
         let infer_model = model.valid();
         // Gen samples
@@ -134,7 +135,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
         }
 
         let server_device = device.clone();
-        thread::spawn(move || {
+        let server_handle = thread::spawn(move || {
             let server = NNServer::new(
                 server_cmd_rx,
                 server_request_rx,
@@ -147,8 +148,10 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
             server.run_server()
         });
 
+        let mut worker_handles = Vec::new();
+        let mut workers_dispatched = 0;
         for i in 0..training_config.num_workers {
-            println!("initializing worker {}", i);
+            // println!("initializing worker {}", i);
             if training_config.num_episodes > episodes_dispatched {
                 let (cmd_tx, cmd_rx) = mpsc::channel::<WorkerCmd>();
                 worker_map.insert(i, cmd_tx);
@@ -162,8 +165,8 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                 let data_tx_thread = data_tx.clone();
                 let mcts_config_thread = training_config.mcts_config.clone();
                 let device_thread = device.clone();
-                thread::spawn(move || {
+                worker_handles.push(thread::spawn(move || {
-                    println!("Starting worker thread...");
+                    // println!("Starting worker thread...");
 
                     let worker = Worker::new(
                         i,
@@ -177,10 +180,15 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                     tokio::runtime::Runtime::new()
                         .unwrap()
                         .block_on(worker.run_worker(cmd_rx));
-                });
+                }));
                 episodes_dispatched += episodes_to_dispatch;
+                workers_dispatched += 1;
             }
         }
+        println!(
+            "Started {} workers, dispatched episodes #1-{}",
+            workers_dispatched, episodes_dispatched
+        );
 
         while episodes_recieved < training_config.num_episodes {
             let data = data_rx.recv().unwrap();
@@ -203,19 +211,26 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                     .expect("Expected dispatch to be sent");
                 episodes_dispatched += 1;
                 println!("Dispatched episode #{}", episodes_dispatched);
-            } else if episodes_recieved == training_config.num_episodes {
-                worker_map
-                    .get(&data.worker_id)
-                    .unwrap()
-                    .send(WorkerCmd::Shutdown)
-                    .expect("Expected shutdown to be sent");
-                println!("Shutdown worker #{}", data.worker_id);
             }
+            // else {
+            //     worker_map
+            //         .get(&data.worker_id)
+            //         .unwrap()
+            //         .send(WorkerCmd::Shutdown)
+            //         .expect("Expected shutdown to be sent");
+            //     // println!("Shutdown worker #{}", data.worker_id);
+            // }
         }
 
         server_cmd_tx
             .send(NNServerCmd::Shutdown)
             .expect("Expected shutdown command to be sent");
+
+        server_handle.join().unwrap();
+        for handle in worker_handles {
+            handle.join().unwrap()
+        }
+
         // train
 
         if replay_buffer.is_empty() {
@@ -263,6 +278,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
             "Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
             iter, loss_val, p_loss_val, v_loss_val
         );
+        println!("Iteration took {}s", iter_start.elapsed().as_secs());
 
         let grads = output.loss.backward();
         let grads = GradientsParams::from_grads(grads, &model);
@@ -271,9 +287,35 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
         iter += 1;
 
-        if iter % 100 == 0 {
+        println!("Saving model...");
-            println!("Completed {} iterations", iter);
+
-        }
+        let metadata = ModelMetadata {
+            name: training_config.model_name.clone(),
+            iterations: iter,
+        };
+
+        std::fs::write(
+            metadata_path.clone(),
+            serde_json::to_string_pretty(&metadata)
+                .expect("Should be able to convert metadata to JSON string"),
+        )
+        .expect("Should be able to write metadata");
+
+        model
+            .clone()
+            .into_record()
+            .save(&model_path)
+            .expect("Should be able to save the model");
+
+        optim
+            .clone()
+            .save(format!("{model_path}_optim"))
+            .expect("Should be able to save the optim");
+        println!("Model saved in {:?}.", model_path);
+
+        // if iter % 100 == 0 {
+        //     println!("Completed {} iterations", iter);
+        // }
 
         if training_config.max_time_s.is_some()
             && start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
@@ -291,29 +333,16 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
         }
     }
 
-    println!("Saving model...");
+    // model
-
+    //     .into_record()
-    let metadata = ModelMetadata {
+    //     .save(&model_path)
-        name: training_config.model_name,
+    //     .expect("Should be able to save the model");
-        iterations: iter,
+    //
-    };
+    // optim
-
+    //     .save(format!("{model_path}_optim"))
-    std::fs::write(
+    //     .expect("Should be able to save the optim");
-        metadata_path,
+    // println!("Model saved in {:?}, exiting training.", model_path);
-        serde_json::to_string_pretty(&metadata)
+    println!("Finished training!");
-            .expect("Should be able to convert metadata to JSON string"),
-    )
-    .expect("Should be able to write metadata");
-
-    model
-        .into_record()
-        .save(&model_path)
-        .expect("Should be able to save the model");
-
-    optim
-        .save(format!("{model_path}_optim"))
-        .expect("Should be able to save the optim");
-    println!("Model saved in {:?}, exiting training.", model_path);
     return;
 }
 
@@ -356,18 +385,18 @@ pub(crate) fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Ve
 pub(crate) fn sample_move(
     dist: &[(usize, f32)],
     rng: &mut SmallRng,
-    side_to_move: Color,
+    board: &Board,
 ) -> Option<ChessMove> {
     let mut r: f32 = rng.random_range(0.0..1.0);
 
     for (idx, p) in dist {
         r -= *p;
         if r <= 0.0 {
-            return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
+            return Some(decode_move(*idx, board).expect("Invalid move"));
         }
     }
 
     // fallback due to floating point drift
     dist.get(0)
-        .map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
+        .map(|(idx, _)| decode_move(*idx, board).expect("Invalid move"))
 }

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, &board_state.board).unwrap();
-            let mv = sample_move(&adjusted, &mut rng, stm).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() {
+            // Remove completed tasks.
-                match cmd {
+            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);