diff --git a/Cargo.lock b/Cargo.lock index 02002f0..f812c8d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -149,17 +149,6 @@ 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" @@ -1219,9 +1208,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -2177,10 +2166,10 @@ dependencies = [ name = "engine" version = "0.1.0" dependencies = [ - "async-trait", "burn", "chess", "ctrlc", + "moka", "rand 0.10.1", "rand_distr", "serde", @@ -3657,6 +3646,23 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b" +[[package]] +name = "moka" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9" +dependencies = [ + "crossbeam-channel", + "crossbeam-epoch", + "crossbeam-utils", + "equivalent", + "parking_lot", + "portable-atomic", + "smallvec", + "tagptr", + "uuid", +] + [[package]] name = "moxcms" version = "0.8.1" @@ -5609,6 +5615,12 @@ dependencies = [ "windows", ] +[[package]] +name = "tagptr" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417" + [[package]] name = "tar" version = "0.4.46" diff --git a/engine/Cargo.toml b/engine/Cargo.toml index 8c53740..b84420a 100644 --- a/engine/Cargo.toml +++ b/engine/Cargo.toml @@ -12,4 +12,4 @@ 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" +moka = { version = "0.12.16", features = ["sync"] } diff --git a/engine/src/main.rs b/engine/src/main.rs index 2117573..98fa028 100644 --- a/engine/src/main.rs +++ b/engine/src/main.rs @@ -24,8 +24,8 @@ fn main() { max_time_s: None, num_iters: None, max_depth: 100, // unused - model_name: String::from("Test1"), - load_model: true, + model_name: String::from("threadedModel"), + load_model: false, hidden_channels: 48, num_blocks: 4, batch_size: 1024, // num positions sampled to a batch (2048) @@ -38,6 +38,7 @@ fn main() { num_workers: 4, nn_batch_size: 64, episodes_per_worker: 4, + nn_cache_size: 500_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 793aeaa..ceb88db 100644 --- a/engine/src/mcts.rs +++ b/engine/src/mcts.rs @@ -1,13 +1,13 @@ // use crate::networks::PolicyNetwork; use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold}; use crate::net::encoding::encode_move; +use crate::worker::NNCacheKey; use chess::BoardStatus::{Checkmate, Stalemate}; use chess::Color::{Black, White}; use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook}; use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES}; use rand::SeedableRng; use std::collections::HashMap; - //TODO retain tree after move selection? #[derive(Clone)] @@ -450,6 +450,18 @@ impl BoardState { } } + pub fn nn_cache_key(&self) -> NNCacheKey { + let board_hash = self.board.get_hash(); // whatever chess crate API you're using + + let repetition_two = self.repetition_table.get(&board_hash).copied().unwrap_or(0) >= 2; + + NNCacheKey { + board_hash, + halfmove_clock: self.halfmove_clock, + repetition_two, + } + } + pub fn apply_move(&mut self, mv: ChessMove) { let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn); let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some(); diff --git a/engine/src/net/encoding.rs b/engine/src/net/encoding.rs index 13e829a..d07c289 100644 --- a/engine/src/net/encoding.rs +++ b/engine/src/net/encoding.rs @@ -33,14 +33,14 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [ const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook]; pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> { - let mut planes = vec![0.0f32; 18 * 64]; + let mut planes = vec![0.0f32; 19 * 64]; encode_board_state_to_slice(state, &mut planes); - Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device) + Tensor::<3>::from_floats(TensorData::new(planes, [19, 8, 8]), device) } -/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long). +/// Write 19x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1216 elements long). pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) { - out_slice[..1152].fill(0.0); + out_slice[..1216].fill(0.0); let board = &state.board; let us = board.side_to_move(); @@ -86,23 +86,23 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) { let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them)); if us_castle.has_kingside() { - fill_plane(out_slice, 12); + fill_plane_ones(out_slice, 12); } if us_castle.has_queenside() { - fill_plane(out_slice, 13); + fill_plane_ones(out_slice, 13); } if them_castle.has_kingside() { - fill_plane(out_slice, 14); + fill_plane_ones(out_slice, 14); } if them_castle.has_queenside() { - fill_plane(out_slice, 15); + fill_plane_ones(out_slice, 15); } // Repetition plane (16) let current_hash = board.get_hash(); if let Some(count) = state.repetition_table.get(¤t_hash) { if *count >= 2 { - fill_plane(out_slice, 16); + fill_plane_ones(out_slice, 16); } } @@ -116,22 +116,29 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) { } out_slice[idx(17, rank, file)] = 1.0; } + + fill_plane(out_slice, 18, Some(state.halfmove_clock as f32 / 100.0)); } pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> { let b = states.len(); - let mut buffer = vec![0.0f32; b * 1152]; + let mut buffer = vec![0.0f32; b * 1216]; for (i, state) in states.iter().enumerate() { - let start = i * 1152; - encode_board_state_to_slice(state, &mut buffer[start..start + 1152]); + let start = i * 1216; + encode_board_state_to_slice(state, &mut buffer[start..start + 1216]); } - Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device) + Tensor::<4>::from_floats(TensorData::new(buffer, [b, 19, 8, 8]), device) } #[inline] -fn fill_plane(buffer: &mut [f32], plane: usize) { +fn fill_plane(buffer: &mut [f32], plane: usize, value: Option) { let start = plane * 64; - buffer[start..start + 64].fill(1.0); + buffer[start..start + 64].fill(value.unwrap_or(1.0)); +} + +#[inline] +fn fill_plane_ones(buffer: &mut [f32], plane: usize) { + fill_plane(buffer, plane, Some(1.0)); } /// Rotate a square into the side-to-move perspective. diff --git a/engine/src/net/model.rs b/engine/src/net/model.rs index 8a6666c..79f9efe 100644 --- a/engine/src/net/model.rs +++ b/engine/src/net/model.rs @@ -21,6 +21,7 @@ Input planes: 16: opponent queenside 17: rep >= 2 18: En passant square (gap pass over) +19: No-progress plys / 100 Output: @@ -45,10 +46,10 @@ pub struct ResidualBlock { activation: Relu, } -#[derive(Config, Debug)] -pub struct ResidualBlockConfig { - channels: usize, -} +// #[derive(Config, Debug)] +// pub struct ResidualBlockConfig { +// channels: usize, +// } impl ResidualBlock { pub fn new(channels: usize, device: &Device) -> Self { @@ -110,7 +111,7 @@ pub struct ChessModelConfig { impl ChessModelConfig { pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel { ChessModel { - conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input + conv: Conv2dConfig::new([19, channels], [3, 3]) // 19 plane input .with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1)) .init(device), bn: BatchNormConfig::new(channels).init(device), @@ -205,23 +206,30 @@ pub struct ChessOutput { pub value_targets: Tensor<2>, // kept 2d for consistency? pub policy_loss: Tensor<1>, pub value_loss: Tensor<1>, - pub loss: Tensor<1>, // [sample] + pub loss: Tensor<1>, // [sample] } impl ItemLazy for ChessOutput { fn sync(self) -> Self { - let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] = - Transaction::default() - .register(self.policy_logits) - .register(self.value) - .register(self.policy_targets) - .register(self.value_targets) - .register(self.policy_loss) - .register(self.value_loss) - .register(self.loss) - .execute() - .try_into() - .expect("Correct amount of tensor data"); + let [ + policy_logits, + value, + policy_targets, + value_targets, + policy_loss, + value_loss, + loss, + ] = Transaction::default() + .register(self.policy_logits) + .register(self.value) + .register(self.policy_targets) + .register(self.value_targets) + .register(self.policy_loss) + .register(self.value_loss) + .register(self.loss) + .execute() + .try_into() + .expect("Correct amount of tensor data"); let device = &Default::default(); @@ -299,7 +307,7 @@ impl Batcher for ChessBatcher { let state_tensors = items .iter() .map(|item| { - encode_board_state_perspective(&item.board_state, device).reshape([1, 18, 8, 8]) + encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 8, 8]) }) .collect::>(); @@ -331,7 +339,7 @@ impl Batcher for ChessBatcher { }) .collect::>(); - let states = Tensor::cat(state_tensors, 0); // [B, 18, 8, 8] + let states = Tensor::cat(state_tensors, 0); // [B, 19, 8, 8] let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672] let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1] diff --git a/engine/src/train.rs b/engine/src/train.rs index 9473dbe..e08ddde 100644 --- a/engine/src/train.rs +++ b/engine/src/train.rs @@ -42,6 +42,7 @@ pub struct TrainingConfig { pub num_workers: usize, pub nn_batch_size: usize, pub episodes_per_worker: usize, + pub nn_cache_size: u64, } pub fn train(training_config: TrainingConfig, device: impl Into) { @@ -140,6 +141,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { infer_model, server_device, training_config.nn_batch_size, + training_config.nn_cache_size, ); server.run_server() @@ -211,6 +213,9 @@ pub fn train(training_config: TrainingConfig, device: impl Into) { } } + server_cmd_tx + .send(NNServerCmd::Shutdown) + .expect("Expected shutdown command to be sent"); // train if replay_buffer.is_empty() { diff --git a/engine/src/worker.rs b/engine/src/worker.rs index 8807890..bd93a76 100644 --- a/engine/src/worker.rs +++ b/engine/src/worker.rs @@ -5,10 +5,13 @@ use crate::train::{apply_temperature, sample_move}; use burn::tensor::Device; use burn::Tensor; use chess::Color; +use moka::sync::Cache; use rand::prelude::SmallRng; use rand::SeedableRng; use std::collections::VecDeque; -use std::sync::mpsc::{Receiver, Sender}; +use std::hash::Hash; +use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender}; +use std::time::Duration; pub struct Worker { id: usize, @@ -134,33 +137,33 @@ impl Worker { 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) - } + // 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"); @@ -223,12 +226,14 @@ async fn model_inference_static( request_tx: Sender, device: Device, ) -> (Vec, f32) { + let cache_key = board_state.nn_cache_key(); let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>(); let (tx, rx) = tokio::sync::oneshot::channel::(); request_tx .send(NNRequest { + cache_key, request_tensor: state, response_tx: tx, }) @@ -236,21 +241,7 @@ async fn model_inference_static( 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) + (response.policy, response.value) } pub enum NNServerCmd { @@ -268,6 +259,7 @@ pub struct NNServer { model: ChessModel, device: Device, batch_size: usize, + cache_size: u64, } impl NNServer { @@ -277,6 +269,7 @@ impl NNServer { model: ChessModel, device: Device, batch_size: usize, + cache_size: u64, ) -> NNServer { NNServer { command_rx, @@ -284,6 +277,7 @@ impl NNServer { model, device, batch_size, + cache_size, } } } @@ -294,31 +288,56 @@ struct NNJobData { // game_id: u32, } pub struct NNRequest { + cache_key: NNCacheKey, request_tensor: Tensor<4>, response_tx: tokio::sync::oneshot::Sender, } -#[derive(Debug)] +#[derive(Debug, Clone)] pub struct NNResponse { - // nn_job_data: NNJobData, - response_tensors: (Tensor<2>, Tensor<2>), + policy: Vec, + value: f32, } impl NNServer { pub fn run_server(&self) { + let mut cache: Cache = Cache::new(self.cache_size); let mut queue: VecDeque = VecDeque::new(); 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"); + break; + } // other commands... + } + } + + // If we have no queued requests, wait for either a request + // or periodically wake up to check for shutdown. if queue.is_empty() { - let first = self - .request_rx - .recv() - .expect("failed to receive NN request"); - queue.push_back(first); + match self.request_rx.recv_timeout(Duration::from_millis(10)) { + Ok(first) => { + if try_use_cache(first, &mut cache, &mut queue) { + continue; + } + } + + Err(RecvTimeoutError::Timeout) => { + continue; + } + + Err(RecvTimeoutError::Disconnected) => { + println!("Request channel disconnected"); + break; + } + } } while let Ok(req) = self.request_rx.try_recv() { - queue.push_back(req); + try_use_cache(req, &mut cache, &mut queue); } let mut jobs_to_send: Vec = Vec::new(); @@ -339,18 +358,49 @@ impl NNServer { 0, ); - println!("Sending {} jobs to gpu", jobs_to_send.len()); + // 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.into_iter().enumerate() { + let p = p_split[i].clone().into_data().to_vec::().unwrap(); + + let v = v_split[i].clone().into_data().to_vec::().unwrap()[0]; + + let res = NNResponse { + policy: p, + value: v, + }; + + cache.insert(job.cache_key.clone(), res.clone()); + job.response_tx - .send(NNResponse { - response_tensors: (p_split[i].clone(), v_split[i].clone()), - }) + .send(res) .expect("worker dropped inference request"); } } } } + +fn try_use_cache( + req: NNRequest, + cache: &mut Cache, + queue: &mut VecDeque, +) -> bool { + if let Some(cached) = cache.get(&req.cache_key) { + // println!("Cache hit"); + req.response_tx.send(cached).unwrap(); + true + } else { + queue.push_back(req); + false + } +} + +#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)] +pub struct NNCacheKey { + pub board_hash: u64, + pub halfmove_clock: u8, + pub repetition_two: bool, +}