diff --git a/Cargo.lock b/Cargo.lock index 2b78f0a..02002f0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -149,6 +149,17 @@ dependencies = [ "portable-atomic-util", ] +[[package]] +name = "async-trait" +version = "0.1.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.4", +] + [[package]] name = "atomic" version = "0.6.1" @@ -2166,6 +2177,7 @@ dependencies = [ name = "engine" version = "0.1.0" dependencies = [ + "async-trait", "burn", "chess", "ctrlc", @@ -2173,6 +2185,7 @@ dependencies = [ "rand_distr", "serde", "serde_json", + "tokio", ] [[package]] diff --git a/engine/Cargo.toml b/engine/Cargo.toml index 633fb12..8c53740 100644 --- a/engine/Cargo.toml +++ b/engine/Cargo.toml @@ -11,3 +11,5 @@ rand = "0.10.1" ctrlc = "3.5.2" serde_json = "1.0.150" serde = { version = "1.0.228", features = ["derive"] } +tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread"] } +async-trait = "0.1.92" diff --git a/engine/src/main.rs b/engine/src/main.rs index 187f2dd..2117573 100644 --- a/engine/src/main.rs +++ b/engine/src/main.rs @@ -35,6 +35,9 @@ fn main() { optimizer: adam_config, lr: 2e-4, seed: None, + num_workers: 4, + nn_batch_size: 64, + episodes_per_worker: 4, }; train(training_config, device); diff --git a/engine/src/mcts.rs b/engine/src/mcts.rs index 256a9f4..793aeaa 100644 --- a/engine/src/mcts.rs +++ b/engine/src/mcts.rs @@ -1,8 +1,6 @@ // use crate::networks::PolicyNetwork; use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold}; -use crate::net::encoding::{encode_board_state_perspective, encode_move}; -use crate::net::model::ChessModel; -use burn::prelude::Device; +use crate::net::encoding::encode_move; use chess::BoardStatus::{Checkmate, Stalemate}; use chess::Color::{Black, White}; use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook}; @@ -83,7 +81,7 @@ impl MctsResults { } } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct MctsConfig { pub num_simulations: usize, pub c_puct: f32, @@ -165,19 +163,22 @@ fn softmax(logits: &[f32]) -> Vec { } impl Mcts { - pub fn search( + pub async fn search( &mut self, board_state: &BoardState, - model: &ChessModel, - device: &Device, + model_inference: F, training: bool, - ) -> MctsResults { + ) -> MctsResults + where + F: Fn(&BoardState) -> Fut, + Fut: Future, f32)>, + { let mut nodes = Vec::::new(); let root = 0; nodes.push(Node::new(0.0, board_state.clone(), None)); - self.expand(root, &mut nodes, model, device); + self.expand(root, &mut nodes, &model_inference).await; if training { // APPLY DIRICHLET NOISE HERE @@ -195,7 +196,7 @@ impl Mcts { path.push(current); } - let value: f32 = self.expand(current, &mut nodes, model, device); + let value: f32 = self.expand(current, &mut nodes, &model_inference).await; let color = nodes[current].board_state.board.side_to_move(); self.backpropagate(&mut nodes, &path, value, color); @@ -225,27 +226,27 @@ impl Mcts { MctsResults::new(board_state.clone(), move_dist, nodes[root].value()) } - fn expand( + async fn expand( &mut self, node_idx: usize, arena: &mut Vec, - model: &ChessModel, - device: &Device, - ) -> f32 { + model_inference: &F, + ) -> f32 + where + F: Fn(&BoardState) -> Fut, + Fut: Future, f32)>, + { let status = arena[node_idx].board_state.status; let stm = arena[node_idx].board_state.board.side_to_move(); if let Some(term_val) = terminal_value(status, stm) { term_val } else { - let state = encode_board_state_perspective(&arena[node_idx].board_state, device); - let (policy_head, value_head) = model.forward(state.unsqueeze::<4>()); + let (policy_logits, value_head) = model_inference(&arena[node_idx].board_state).await; let legal_moves: Vec = MoveGen::new_legal(&arena[node_idx].board_state.board).collect(); - - let logits = policy_head.into_data().to_vec::().unwrap(); - let probs = softmax(&logits); + let probs = softmax(&policy_logits); let mut move_priors = Vec::with_capacity(legal_moves.len()); let mut prior_sum = 0.0f32; @@ -273,7 +274,7 @@ impl Mcts { arena[node_idx].children.push(child_idx); } - value_head.into_data().to_vec::().unwrap()[0] + value_head } } diff --git a/engine/src/train.rs b/engine/src/train.rs index 7be0f5d..9473dbe 100644 --- a/engine/src/train.rs +++ b/engine/src/train.rs @@ -1,9 +1,9 @@ -use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults}; +use crate::mcts::MctsConfig; use crate::net::encoding::decode_move; use crate::net::model::{ ChessBatcher, ChessModel, ChessModelConfig, ModelMetadata, TrainingSample, }; -use crate::worker::{NNRequest, NNResponse, NNServer, NNServerCmd, Worker, WorkerCmd}; +use crate::worker::{NNRequest, NNResponse, NNServer, NNServerCmd, Worker, WorkerCmd, WorkerData}; use burn::data::dataloader::batcher::Batcher; use burn::module::AutodiffModule; use burn::optim::{AdamConfig, GradientsParams}; @@ -99,10 +99,6 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let mut replay_buffer: VecDeque = VecDeque::new(); - let mut mcts: Mcts = Mcts { - config: training_config.mcts_config, - }; - // Create RNG once and reuse it for sampling and shuffling // Seed from system time (platform default entropy may be unavailable in some contexts) let seed = training_config.seed.unwrap_or( @@ -126,15 +122,13 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let mut worker_map: HashMap> = HashMap::new(); let (server_cmd_tx, server_cmd_rx) = mpsc::channel::(); - let (data_tx, data_rx) = mpsc::channel::>(); + let (data_tx, data_rx) = mpsc::channel::(); let (server_request_tx, server_request_rx) = mpsc::channel::(); let mut nn_channels: VecDeque> = VecDeque::new(); - let mut nn_tx_map: HashMap> = HashMap::new(); for i in 0..training_config.num_workers { let (tx, rx) = mpsc::channel::(); - nn_tx_map.insert(i, tx); nn_channels.push_back(rx); } @@ -143,7 +137,6 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { let server = NNServer::new( server_cmd_rx, server_request_rx, - nn_tx_map, infer_model, server_device, training_config.nn_batch_size, @@ -153,6 +146,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { }); for i in 0..training_config.num_workers { + 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,97 +156,67 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { ); let req_tx = server_request_tx.clone(); - let res_rx = nn_channels.pop_front().unwrap(); + // let res_rx = nn_channels.pop_front().unwrap(); 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..."); + let worker = Worker::new( i, episodes_to_dispatch, req_tx, - res_rx, - cmd_rx, data_tx_thread, + device_thread, + mcts_config_thread, ); - worker.run_worker() + tokio::runtime::Runtime::new() + .unwrap() + .block_on(worker.run_worker(cmd_rx)); }); + episodes_dispatched += episodes_to_dispatch; } } - for episode in 0..training_config.num_episodes { - println!("Episode: {}", episode); - let mut board_state = BoardState::default(); - let mut episode_buffer: Vec = vec![]; + while episodes_recieved < training_config.num_episodes { + let data = data_rx.recv().unwrap(); - let mut game_hist = String::new(); - - let mut ply_count: usize = 0; - - while board_state.status == BoardStateStatus::Ongoing { - // let before = Instant::now(); - let results = mcts.search(&board_state, &infer_model, &device, true); - // println!("{}", before.elapsed().as_millis()); - episode_buffer.push(results); - - let temp = if ply_count < 30 { 1.0 } else { 0.0 }; - - 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(); - // println!("playing move: {}", mv); - if episode == 0 { - game_hist.push_str(mv.to_string().as_str()); - game_hist.push(' '); - } - - board_state.apply_move(mv); - ply_count += 1; - - // truncate long games - if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing { - board_state.status = BoardStateStatus::FiftyMove; - } - } - if episode == 0 { - println!("Game history of first game of iteration: {}", game_hist); - } - - for result in episode_buffer.iter() { - let stm = result.board_state.board.side_to_move(); - let outcome = match board_state.status { - BoardStateStatus::WhiteWinner => { - if stm == Color::White { - 1.0 - } else { - -1.0 - } - } - BoardStateStatus::BlackWinner => { - if stm == Color::Black { - 1.0 - } else { - -1.0 - } - } - _ => 0.0, - }; - - replay_buffer.push_back(TrainingSample::from_mcts_with_outcome( - result.clone(), - outcome, - )); + for sample in data.data { + replay_buffer.push_back(sample); if replay_buffer.len() > training_config.buffer_max_size { replay_buffer.pop_front(); } } + episodes_recieved += 1; + println!("Recieved episode #{}", episodes_recieved); + + if episodes_dispatched < training_config.num_episodes { + worker_map + .get(&data.worker_id) + .unwrap() + .send(WorkerCmd::Dispatch) + .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); + } } + // train if replay_buffer.is_empty() { - println!("Replay buffer is empty, skipping training step."); - continue; + panic!("Replay buffer is empty, cannot train."); + // println!("Replay buffer is empty, skipping training step."); + // continue; } let num_samples = std::cmp::min(training_config.batch_size, replay_buffer.len()); @@ -348,7 +312,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { return; } -fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> { +pub(crate) fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f32)> { if visits.is_empty() { return Vec::new(); } @@ -384,7 +348,7 @@ fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Vec<(usize, f adjusted } -fn sample_move( +pub(crate) fn sample_move( dist: &[(usize, f32)], rng: &mut SmallRng, side_to_move: Color, diff --git a/engine/src/worker.rs b/engine/src/worker.rs index 589ac19..8807890 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -1,17 +1,29 @@ -use crate::mcts::BoardState; +use crate::mcts::{BoardState, BoardStateStatus, Mcts, MctsConfig, MctsResults}; +use crate::net::encoding::encode_board_state_perspective; use crate::net::model::{ChessModel, TrainingSample}; +use crate::train::{apply_temperature, sample_move}; use burn::tensor::Device; use burn::Tensor; -use std::collections::{HashMap, VecDeque}; +use chess::Color; +use rand::prelude::SmallRng; +use rand::SeedableRng; +use std::collections::VecDeque; use std::sync::mpsc::{Receiver, Sender}; pub struct Worker { id: usize, - games: Vec, + num_games: usize, + // games: Vec, request_tx: Sender, - response_rx: Receiver, - cmd_rx: Receiver, - data_tx: Sender>, + data_tx: Sender, + curr_request_id: u64, + device: Device, + mcts_config: MctsConfig, +} + +pub struct WorkerData { + pub(crate) worker_id: usize, + pub(crate) data: Vec, } impl Worker { @@ -19,23 +31,226 @@ impl Worker { id: usize, num_games: usize, sender: Sender, - receiver: Receiver, - cmd_rx: Receiver, - data_tx: Sender>, + data_tx: Sender, + device: Device, + mcts_config: MctsConfig, ) -> Worker { Worker { id, - games: (0..num_games) - .map(|_| BoardState::default()) - .collect::>(), + // games: (0..num_games) + // .map(|_| BoardState::default()) + // .collect::>(), + num_games, request_tx: sender, - response_rx: receiver, - cmd_rx, data_tx, + curr_request_id: 0, + device, + mcts_config, } } - pub fn run_worker() {} + async fn run_game( + worker_id: usize, + mcts_config: MctsConfig, + model_inference: F, + data_tx: Sender, + ) where + F: Fn(&BoardState) -> Fut, + Fut: Future, f32)>, + { + println!("game started"); + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap(); + let seed = now.as_nanos() as u64; + let mut rng = SmallRng::seed_from_u64(seed); + let mut mcts: Mcts = Mcts { + config: mcts_config, + }; + + let mut board_state = BoardState::default(); + let mut episode_buffer: Vec = vec![]; + + let mut game_hist = String::new(); + + let mut ply_count: usize = 0; + + while board_state.status == BoardStateStatus::Ongoing { + // let before = Instant::now(); + let results = mcts.search(&board_state, &model_inference, true).await; + // println!("{}", before.elapsed().as_millis()); + episode_buffer.push(results); + + let temp = if ply_count < 30 { 1.0 } else { 0.0 }; + + 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(); + + game_hist.push_str(mv.to_string().as_str()); + game_hist.push(' '); + + board_state.apply_move(mv); + ply_count += 1; + + // truncate long games + if ply_count >= 400 && board_state.status == BoardStateStatus::Ongoing { + board_state.status = BoardStateStatus::FiftyMove; + } + } + + let mut samples: Vec = Vec::new(); + for result in episode_buffer.iter() { + let stm = result.board_state.board.side_to_move(); + let outcome = match board_state.status { + BoardStateStatus::WhiteWinner => { + if stm == Color::White { + 1.0 + } else { + -1.0 + } + } + BoardStateStatus::BlackWinner => { + if stm == Color::Black { + 1.0 + } else { + -1.0 + } + } + _ => 0.0, + }; + + samples.push(TrainingSample::from_mcts_with_outcome( + result.clone(), + outcome, + )); + } + + let worker_data = WorkerData { + worker_id, + data: samples, + }; + data_tx.send(worker_data).unwrap(); + } + + async fn model_inference(&self, board_state: BoardState) -> (Vec, f32) { + let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>(); + let (tx, rx) = tokio::sync::oneshot::channel::(); + + self.request_tx + .send(NNRequest { + request_tensor: state, + response_tx: tx, + }) + .unwrap(); + + let response = rx.await.expect("worker dropped inference response"); + let p = response + .response_tensors + .0 + .into_data() + .to_vec::() + .unwrap(); + let v = response + .response_tensors + .1 + .into_data() + .to_vec::() + .unwrap()[0]; + + (p, v) + } + + pub async fn run_worker(self, cmd_rx: Receiver) { + 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 { + Self::run_game( + self.id, + mcts_config, + |board_state| { + let request_tx = request_tx.clone(); + let device = device.clone(); + + model_inference_static(board_state.clone(), request_tx, device) + }, + data_tx, + ) + .await; + }); + } + + loop { + if let Ok(cmd) = cmd_rx.recv() { + 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 { + Self::run_game( + self.id, + mcts_config, + |board_state| { + let request_tx = request_tx.clone(); + let device = device.clone(); + + model_inference_static(board_state.clone(), request_tx, device) + }, + data_tx, + ) + .await; + }); + } + } + } else { + panic!("worker dropped inference command"); + } + } + } +} + +async fn model_inference_static( + board_state: BoardState, + request_tx: Sender, + device: Device, +) -> (Vec, f32) { + let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>(); + + let (tx, rx) = tokio::sync::oneshot::channel::(); + + request_tx + .send(NNRequest { + request_tensor: state, + response_tx: tx, + }) + .unwrap(); + + let response = rx.await.expect("worker dropped inference response"); + + let p = response + .response_tensors + .0 + .into_data() + .to_vec::() + .unwrap(); + + let v = response + .response_tensors + .1 + .into_data() + .to_vec::() + .unwrap()[0]; + + (p, v) } pub enum NNServerCmd { @@ -44,12 +259,12 @@ pub enum NNServerCmd { pub enum WorkerCmd { Shutdown, + Dispatch, } pub struct NNServer { command_rx: Receiver, request_rx: Receiver, - response_txs: HashMap>, model: ChessModel, device: Device, batch_size: usize, @@ -59,7 +274,6 @@ impl NNServer { pub fn new( command_rx: Receiver, request_rx: Receiver, - response_txs: HashMap>, model: ChessModel, device: Device, batch_size: usize, @@ -67,27 +281,26 @@ impl NNServer { NNServer { command_rx, request_rx, - response_txs, model, device, batch_size, } } } - -#[derive(Clone, Eq, Hash, PartialEq)] +#[derive(Debug, Clone, Eq, Hash, PartialEq)] struct NNJobData { worker_id: usize, request_id: u64, - game_id: u32, + // game_id: u32, } pub struct NNRequest { - nn_job_data: NNJobData, request_tensor: Tensor<4>, + response_tx: tokio::sync::oneshot::Sender, } +#[derive(Debug)] pub struct NNResponse { - nn_job_data: NNJobData, + // nn_job_data: NNJobData, response_tensors: (Tensor<2>, Tensor<2>), } @@ -126,19 +339,17 @@ impl NNServer { 0, ); + println!("Sending {} jobs to gpu", jobs_to_send.len()); let (p, v) = self.model.forward(batched_tensor); let p_split = p.chunk(jobs_to_send.len(), 0); let v_split = v.chunk(jobs_to_send.len(), 0); - for (i, job) in jobs_to_send.iter().enumerate() { - let worker_tx = &self.response_txs[&job.nn_job_data.worker_id]; - - worker_tx + for (i, job) in jobs_to_send.into_iter().enumerate() { + job.response_tx .send(NNResponse { - nn_job_data: job.nn_job_data.clone(), response_tensors: (p_split[i].clone(), v_split[i].clone()), }) - .expect("Expected response to be sent successfully"); + .expect("worker dropped inference request"); } } }