start threading

2026-08-26 14:31:45 -04:00
parent c51964a443
commit aa20c44779
4 changed files with 225 additions and 7 deletions
@@ -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");
+            }
+        }
+    }
+}