working workers

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]]

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"

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);

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::encoding::encode_move;
-use crate::net::model::ChessModel;
-use burn::prelude::Device;
 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<f32> {
 }
 
 impl Mcts {
-    pub fn search(
+    pub async fn search<F, Fut>(
         &mut self,
         board_state: &BoardState,
-        model: &ChessModel,
+        model_inference: F,
-        device: &Device,
         training: bool,
-    ) -> MctsResults {
+    ) -> MctsResults
+    where
+        F: Fn(&BoardState) -> Fut,
+        Fut: Future<Output = (Vec<f32>, f32)>,
+    {
         let mut nodes = Vec::<Node>::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<F, Fut>(
         &mut self,
         node_idx: usize,
         arena: &mut Vec<Node>,
-        model: &ChessModel,
+        model_inference: &F,
-        device: &Device,
+    ) -> f32
-    ) -> f32 {
+    where
+        F: Fn(&BoardState) -> Fut,
+        Fut: Future<Output = (Vec<f32>, 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_logits, value_head) = model_inference(&arena[node_idx].board_state).await;
-            let (policy_head, value_head) = model.forward(state.unsqueeze::<4>());
 
             let legal_moves: Vec<ChessMove> =
                 MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
-
+            let probs = softmax(&policy_logits);
-            let logits = policy_head.into_data().to_vec::<f32>().unwrap();
-            let probs = softmax(&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::<f32>().unwrap()[0]
+            value_head
         }
     }
 

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<Device>) {
 
     let mut replay_buffer: VecDeque<TrainingSample> = 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<Device>) {
 
         let mut worker_map: HashMap<usize, Sender<WorkerCmd>> = HashMap::new();
         let (server_cmd_tx, server_cmd_rx) = mpsc::channel::<NNServerCmd>();
-        let (data_tx, data_rx) = mpsc::channel::<Vec<TrainingSample>>();
+        let (data_tx, data_rx) = mpsc::channel::<WorkerData>();
         let (server_request_tx, server_request_rx) = mpsc::channel::<NNRequest>();
 
         let mut nn_channels: VecDeque<Receiver<NNResponse>> = VecDeque::new();
-        let mut nn_tx_map: HashMap<usize, Sender<NNResponse>> = HashMap::new();
 
         for i in 0..training_config.num_workers {
             let (tx, rx) = mpsc::channel::<NNResponse>();
-            nn_tx_map.insert(i, tx);
             nn_channels.push_back(rx);
         }
 
@@ -143,7 +137,6 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
             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<Device>) {
         });
 
         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::<WorkerCmd>();
                 worker_map.insert(i, cmd_tx);
@@ -162,97 +156,67 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
                 );
 
                 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 {
+        while episodes_recieved < training_config.num_episodes {
-            println!("Episode: {}", episode);
+            let data = data_rx.recv().unwrap();
-            let mut board_state = BoardState::default();
-            let mut episode_buffer: Vec<MctsResults> = vec![];
 
-            let mut game_hist = String::new();
+            for sample in data.data {
-
+                replay_buffer.push_back(sample);
-            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,
-                ));
 
                 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.");
+            panic!("Replay buffer is empty, cannot train.");
-            continue;
+            // 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<Device>) {
     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,

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<BoardState>,
+    num_games: usize,
+    // games: Vec<BoardState>,
     request_tx: Sender<NNRequest>,
-    response_rx: Receiver<NNResponse>,
+    data_tx: Sender<WorkerData>,
-    cmd_rx: Receiver<WorkerCmd>,
+    curr_request_id: u64,
-    data_tx: Sender<Vec<TrainingSample>>,
+    device: Device,
+    mcts_config: MctsConfig,
+}
+
+pub struct WorkerData {
+    pub(crate) worker_id: usize,
+    pub(crate) data: Vec<TrainingSample>,
 }
 
 impl Worker {
@@ -19,23 +31,226 @@ impl Worker {
         id: usize,
         num_games: usize,
         sender: Sender<NNRequest>,
-        receiver: Receiver<NNResponse>,
+        data_tx: Sender<WorkerData>,
-        cmd_rx: Receiver<WorkerCmd>,
+        device: Device,
-        data_tx: Sender<Vec<TrainingSample>>,
+        mcts_config: MctsConfig,
     ) -> Worker {
         Worker {
             id,
-            games: (0..num_games)
+            // games: (0..num_games)
-                .map(|_| BoardState::default())
+            //     .map(|_| BoardState::default())
-                .collect::<Vec<BoardState>>(),
+            //     .collect::<Vec<BoardState>>(),
+            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<F, Fut>(
+        worker_id: usize,
+        mcts_config: MctsConfig,
+        model_inference: F,
+        data_tx: Sender<WorkerData>,
+    ) where
+        F: Fn(&BoardState) -> Fut,
+        Fut: Future<Output = (Vec<f32>, 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<MctsResults> = 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<TrainingSample> = 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>, f32) {
+        let state = encode_board_state_perspective(&board_state, &self.device).unsqueeze::<4>();
+        let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
+
+        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::<f32>()
+            .unwrap();
+        let v = response
+            .response_tensors
+            .1
+            .into_data()
+            .to_vec::<f32>()
+            .unwrap()[0];
+
+        (p, v)
+    }
+
+    pub async fn run_worker(self, cmd_rx: Receiver<WorkerCmd>) {
+        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<NNRequest>,
+    device: Device,
+) -> (Vec<f32>, f32) {
+    let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>();
+
+    let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
+
+    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::<f32>()
+        .unwrap();
+
+    let v = response
+        .response_tensors
+        .1
+        .into_data()
+        .to_vec::<f32>()
+        .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<NNServerCmd>,
     request_rx: Receiver<NNRequest>,
-    response_txs: HashMap<usize, Sender<NNResponse>>,
     model: ChessModel,
     device: Device,
     batch_size: usize,
@@ -59,7 +274,6 @@ impl NNServer {
     pub fn new(
         command_rx: Receiver<NNServerCmd>,
         request_rx: Receiver<NNRequest>,
-        response_txs: HashMap<usize, Sender<NNResponse>>,
         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(Debug, Clone, Eq, Hash, PartialEq)]
-#[derive(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<NNResponse>,
 }
 
+#[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() {
+            for (i, job) in jobs_to_send.into_iter().enumerate() {
-                let worker_tx = &self.response_txs[&job.nn_job_data.worker_id];
+                job.response_tx
-
-                worker_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");
             }
         }
     }