From d901c9a85db92cc601866e48bea6c09daa4a1d79 Mon Sep 17 00:00:00 2001 From: DragonDuck24 Date: Sat, 29 Aug 2026 10:44:24 -0400 Subject: [PATCH] full threading and fixes --- Cargo.lock | 30 +++++----- engine/src/main.rs | 8 +-- engine/src/mcts.rs | 36 +++++++++++- engine/src/net/encoding.rs | 110 +++++++++++++++++++++--------------- engine/src/train.rs | 113 +++++++++++++++++++++++-------------- engine/src/worker.rs | 74 +++++++++++++++++++----- 6 files changed, 248 insertions(+), 123 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f812c8d..6e49b24 100644 --- a/Cargo.lock +++ b/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", diff --git a/engine/src/main.rs b/engine/src/main.rs index 98fa028..f11def3 100644 --- a/engine/src/main.rs +++ b/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, - nn_cache_size: 500_000, // 18,736B / item in cache. 100k=1.75GB. + episodes_per_worker: 12, + nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB. }; train(training_config, device); diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index ceb88db..4da924a 100644 --- a/engine/src/mcts.rs +++ b/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); diff --git a/engine/src/net/encoding.rs b/engine/src/net/encoding.rs index d07c289..f95fb9e 100644 --- a/engine/src/net/encoding.rs +++ b/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 { } /// Decode an index in [0, 4672) back into a ChessMove. -pub fn decode_move(idx: usize, stm: Color) -> Option { +pub fn decode_move(idx: usize, board: &Board) -> Option { if idx >= 4672 { return None; } @@ -250,10 +250,24 @@ pub fn decode_move(idx: usize, stm: Color) -> Option { } // 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,69 +276,77 @@ pub fn decode_move(idx: usize, stm: Color) -> Option { #[cfg(test)] mod tests { - use super::*; - use chess::ALL_COLORS; + use crate::net::encoding::{decode_move, encode_move}; + use chess::{Board, ChessMove, Color, File, Piece, Rank, Square}; + use std::str::FromStr; #[test] fn encoding_roundtrips() { - for color in ALL_COLORS { - for action in 0..4672 { - let decoded = decode_move(action, color); - if decoded.is_none() { - continue; - } - let decoded = decoded.unwrap(); - let encoded = encode_move(decoded, color).unwrap(); - // eprintln!( - // "orig idx = {}, plane={}, from_idx={}", - // action, - // action / 64, - // action % 64 - // ); - // eprintln!( - // "move = {}, from={}, to={}, promo={:?}", - // decoded, - // decoded.get_source(), - // decoded.get_dest(), - // decoded.get_promotion() - // ); - // eprintln!( - // "decoded = {}, from={}, to={}, promo={:?}", - // decoded, - // decoded.get_source(), - // decoded.get_dest(), - // decoded.get_promotion() - // ); - // eprintln!( - // "re-encoded idx = {}, plane={}, from_idx={}", - // encoded, - // encoded / 64, - // encoded % 64 - // ); - assert_eq!(action, encoded); + let board = Board::default(); + for action in 0..4672 { + let decoded = decode_move(action, &board); + if decoded.is_none() { + continue; } + let decoded = decoded.unwrap(); + let encoded = encode_move(decoded, board.side_to_move()).unwrap(); + // eprintln!( + // "orig idx = {}, plane={}, from_idx={}", + // action, + // action / 64, + // action % 64 + // ); + // eprintln!( + // "move = {}, from={}, to={}, promo={:?}", + // decoded, + // decoded.get_source(), + // decoded.get_dest(), + // decoded.get_promotion() + // ); + // eprintln!( + // "decoded = {}, from={}, to={}, promo={:?}", + // decoded, + // decoded.get_source(), + // decoded.get_dest(), + // decoded.get_promotion() + // ); + // eprintln!( + // "re-encoded idx = {}, plane={}, from_idx={}", + // encoded, + // encoded / 64, + // encoded % 64 + // ); + 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); } } diff --git a/engine/src/train.rs b/engine/src/train.rs index e08ddde..7dd19c8 100644 --- a/engine/src/train.rs +++ b/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) { 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) { } 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) { 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::(); worker_map.insert(i, cmd_tx); @@ -162,8 +165,8 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let data_tx_thread = data_tx.clone(); let mcts_config_thread = training_config.mcts_config.clone(); let device_thread = device.clone(); - thread::spawn(move || { - println!("Starting worker thread..."); + worker_handles.push(thread::spawn(move || { + // println!("Starting worker thread..."); let worker = Worker::new( i, @@ -177,10 +180,15 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { 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) { .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) { "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) { iter += 1; - if iter % 100 == 0 { - println!("Completed {} iterations", iter); - } + println!("Saving model..."); + + 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) { } } - println!("Saving model..."); - - let metadata = ModelMetadata { - name: training_config.model_name, - iterations: iter, - }; - - std::fs::write( - metadata_path, - serde_json::to_string_pretty(&metadata) - .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); + // 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); + println!("Finished training!"); 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 { 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")) } diff --git a/engine/src/worker.rs b/engine/src/worker.rs index bd93a76..da43794 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -61,7 +61,7 @@ impl Worker { F: Fn(&BoardState) -> Fut, Fut: Future, 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 = 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) { - 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 = Cache::new(self.cache_size); let mut queue: VecDeque = 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 = 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);