drake/chess_dragon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

working workers

Browse Source
This commit is contained in:
drake
2026-08-27 14:43:02 -04:00
parent aa20c44779
commit 3da2554738
6 changed files with 323 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/main.rs
+3
View File
@@ -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
+21 -20
View File
@@ -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<f32> {
}
impl Mcts {
pub fn search(
pub async fn search<F, Fut>(
&mut self,
board_state: &BoardState,
model: &ChessModel,
device: &Device,
model_inference: F,
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,
device: &Device,
) -> f32 {
model_inference: &F,
) -> 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_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<ChessMove> =
MoveGen::new_legal(&arena[node_idx].board_state.board).collect();
let logits = policy_head.into_data().to_vec::<f32>().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::<f32>().unwrap()[0]
value_head
}
}
engine/src/train.rs
+44 -80
View File
@@ -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 {
println!("Episode: {}", episode);
let mut board_state = BoardState::default();
let mut episode_buffer: Vec<MctsResults> = 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<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
+240 -29
View File
@@ -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>,
cmd_rx: Receiver<WorkerCmd>,
data_tx: Sender<Vec<TrainingSample>>,
data_tx: Sender<WorkerData>,
curr_request_id: u64,
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>,
cmd_rx: Receiver<WorkerCmd>,
data_tx: Sender<Vec<TrainingSample>>,
data_tx: Sender<WorkerData>,
device: Device,
mcts_config: MctsConfig,
) -> Worker {
Worker {
id,
games: (0..num_games)
.map(|_| BoardState::default())
.collect::<Vec<BoardState>>(),
// games: (0..num_games)
// .map(|_| BoardState::default())
// .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(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<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() {
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");
}
}
}