start threading

engine/src/lib.rs

@@ -1,6 +1,7 @@
 pub mod mcts;
 mod net;
 pub mod train;
+mod worker;
 
 use chess::Game;
 

engine/src/main.rs

@@ -25,11 +25,11 @@ fn main() {
         num_iters: None,
         max_depth: 100, // unused
         model_name: String::from("Test1"),
-        load_model: false,
+        load_model: true,
         hidden_channels: 48,
         num_blocks: 4,
         batch_size: 1024,        // num positions sampled to a batch (2048)
-        num_episodes: 100,       // num games generated per iteration (5000)
+        num_episodes: 96,        // num games generated per iteration (5000)
         buffer_max_size: 25_000, // max number of samples in the buffer
         mcts_config,
         optimizer: adam_config,

engine/src/train.rs

@@ -3,6 +3,7 @@ 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 burn::data::dataloader::batcher::Batcher;
 use burn::module::AutodiffModule;
 use burn::optim::{AdamConfig, GradientsParams};
@@ -12,11 +13,14 @@ use chess::{ChessMove, Color};
 use rand::rngs::SmallRng;
 use rand::seq::SliceRandom;
 use rand::{RngExt, SeedableRng};
-use std::collections::VecDeque;
+use std::cmp::min;
+use std::collections::{HashMap, VecDeque};
 use std::fs::File;
 use std::io::BufReader;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
+use std::sync::mpsc::{Receiver, Sender};
+use std::sync::{mpsc, Arc};
+use std::thread;
 use std::time::Instant;
 use std::time::{SystemTime, UNIX_EPOCH};
 
@@ -35,6 +39,9 @@ pub struct TrainingConfig {
     pub optimizer: AdamConfig,
     pub lr: f64,
     pub seed: Option<u64>,
+    pub num_workers: usize,
+    pub nn_batch_size: usize,
+    pub episodes_per_worker: usize,
 }
 
 pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
@@ -98,8 +105,12 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
     // Create RNG once and reuse it for sampling and shuffling
     // Seed from system time (platform default entropy may be unavailable in some contexts)
-    let now = SystemTime::now().duration_since(UNIX_EPOCH).unwrap();
+    let seed = training_config.seed.unwrap_or(
-    let seed = training_config.seed.unwrap_or(now.as_nanos() as u64);
+        SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .unwrap()
+            .as_nanos() as u64,
+    );
     println!("using seed: {}", seed);
     let mut rng = SmallRng::seed_from_u64(seed);
 
@@ -109,6 +120,65 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
         let infer_model = model.valid();
         // Gen samples
         println!("Generating {} games...", training_config.num_episodes);
+
+        let mut episodes_dispatched: usize = 0;
+        let mut episodes_recieved: usize = 0;
+
+        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 (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);
+        }
+
+        let server_device = device.clone();
+        thread::spawn(move || {
+            let server = NNServer::new(
+                server_cmd_rx,
+                server_request_rx,
+                nn_tx_map,
+                infer_model,
+                server_device,
+                training_config.nn_batch_size,
+            );
+
+            server.run_server()
+        });
+
+        for i in 0..training_config.num_workers {
+            if training_config.num_episodes > episodes_dispatched {
+                let (cmd_tx, cmd_rx) = mpsc::channel::<WorkerCmd>();
+                worker_map.insert(i, cmd_tx);
+                let episodes_to_dispatch = min(
+                    training_config.episodes_per_worker,
+                    training_config.num_episodes - episodes_dispatched,
+                );
+
+                let req_tx = server_request_tx.clone();
+                let res_rx = nn_channels.pop_front().unwrap();
+                let data_tx_thread = data_tx.clone();
+                thread::spawn(move || {
+                    let worker = Worker::new(
+                        i,
+                        episodes_to_dispatch,
+                        req_tx,
+                        res_rx,
+                        cmd_rx,
+                        data_tx_thread,
+                    );
+
+                    worker.run_worker()
+                });
+            }
+        }
+
         for episode in 0..training_config.num_episodes {
             println!("Episode: {}", episode);
             let mut board_state = BoardState::default();
@@ -120,7 +190,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
             while board_state.status == BoardStateStatus::Ongoing {
                 // let before = Instant::now();
-                let results = mcts.search(&board_state, &infer_model, &device);
+                let results = mcts.search(&board_state, &infer_model, &device, true);
                 // println!("{}", before.elapsed().as_millis());
                 episode_buffer.push(results);
 
@@ -138,6 +208,8 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
 
                 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;
                 }

engine/src/worker.rs

@@ -0,0 +1,145 @@
+use crate::mcts::BoardState;
+use crate::net::model::{ChessModel, TrainingSample};
+use burn::tensor::Device;
+use burn::Tensor;
+use std::collections::{HashMap, VecDeque};
+use std::sync::mpsc::{Receiver, Sender};
+
+pub struct Worker {
+    id: usize,
+    games: Vec<BoardState>,
+    request_tx: Sender<NNRequest>,
+    response_rx: Receiver<NNResponse>,
+    cmd_rx: Receiver<WorkerCmd>,
+    data_tx: Sender<Vec<TrainingSample>>,
+}
+
+impl Worker {
+    pub fn new(
+        id: usize,
+        num_games: usize,
+        sender: Sender<NNRequest>,
+        receiver: Receiver<NNResponse>,
+        cmd_rx: Receiver<WorkerCmd>,
+        data_tx: Sender<Vec<TrainingSample>>,
+    ) -> Worker {
+        Worker {
+            id,
+            games: (0..num_games)
+                .map(|_| BoardState::default())
+                .collect::<Vec<BoardState>>(),
+            request_tx: sender,
+            response_rx: receiver,
+            cmd_rx,
+            data_tx,
+        }
+    }
+
+    pub fn run_worker() {}
+}
+
+pub enum NNServerCmd {
+    Shutdown,
+}
+
+pub enum WorkerCmd {
+    Shutdown,
+}
+
+pub struct NNServer {
+    command_rx: Receiver<NNServerCmd>,
+    request_rx: Receiver<NNRequest>,
+    response_txs: HashMap<usize, Sender<NNResponse>>,
+    model: ChessModel,
+    device: Device,
+    batch_size: usize,
+}
+
+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,
+    ) -> NNServer {
+        NNServer {
+            command_rx,
+            request_rx,
+            response_txs,
+            model,
+            device,
+            batch_size,
+        }
+    }
+}
+
+#[derive(Clone, Eq, Hash, PartialEq)]
+struct NNJobData {
+    worker_id: usize,
+    request_id: u64,
+    game_id: u32,
+}
+pub struct NNRequest {
+    nn_job_data: NNJobData,
+    request_tensor: Tensor<4>,
+}
+
+pub struct NNResponse {
+    nn_job_data: NNJobData,
+    response_tensors: (Tensor<2>, Tensor<2>),
+}
+
+impl NNServer {
+    pub fn run_server(&self) {
+        let mut queue: VecDeque<NNRequest> = VecDeque::new();
+
+        loop {
+            if queue.is_empty() {
+                let first = self
+                    .request_rx
+                    .recv()
+                    .expect("failed to receive NN request");
+                queue.push_back(first);
+            }
+
+            while let Ok(req) = self.request_rx.try_recv() {
+                queue.push_back(req);
+            }
+
+            let mut jobs_to_send: Vec<NNRequest> = Vec::new();
+
+            for i in 0..self.batch_size {
+                if let Some(job) = queue.pop_front() {
+                    jobs_to_send.push(job);
+                } else {
+                    break;
+                }
+            }
+
+            let batched_tensor = Tensor::cat(
+                jobs_to_send
+                    .iter()
+                    .map(|j| j.request_tensor.clone())
+                    .collect(),
+                0,
+            );
+
+            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
+                    .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");
+            }
+        }
+    }
+}