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

full threading and fixes

Browse Source
This commit is contained in:
drake
2026-08-29 10:44:24 -04:00
parent 4f854aeca4
commit d901c9a85d
6 changed files with 248 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/main.rs
+4 -4
View File
@@ -25,7 +25,7 @@ fn main() {
num_iters: None,
max_depth: 100, // unused
model_name: String::from("threadedModel"),
load_model: false,
load_model: true,
hidden_channels: 48,
num_blocks: 4,
batch_size: 1024, // num positions sampled to a batch (2048)
@@ -35,10 +35,10 @@ fn main() {
optimizer: adam_config,
lr: 2e-4,
seed: None,
num_workers: 4,
num_workers: 8,
nn_batch_size: 64,
episodes_per_worker: 4,
nn_cache_size: 500_000, // 18,736B / item in cache. 100k=1.75GB.
episodes_per_worker: 12,
nn_cache_size: 50_000, // 18,736B / item in cache. 100k=1.75GB.
};
train(training_config, device);
engine/src/mcts.rs
+33 -3
View File
@@ -87,7 +87,8 @@ pub struct MctsConfig {
pub c_puct: f32,
pub dirichlet_alpha: f32,
pub dirichlet_epsilon: f32,
pub batch_max: usize,
pub batch_size: usize,
pub num_batches: usize,
}
impl MctsConfig {
@@ -96,14 +97,15 @@ impl MctsConfig {
c_puct: f32,
dirichlet_alpha: f32,
dirichlet_epsilon: f32,
batch_max: usize,
batch_size: usize,
) -> MctsConfig {
MctsConfig {
num_simulations,
c_puct,
dirichlet_alpha,
dirichlet_epsilon,
batch_max,
batch_size,
num_batches: num_simulations / batch_size,
}
}
}
@@ -191,6 +193,34 @@ impl Mcts {
let mut path = vec![root];
let mut current = root;
// for batch in 0..self.config.num_batches {
// let mut pending = Vec::new();
//
// // Select several leaves
// for _ in 0..self.config.batch_size {
// let path = select_leaf_with_virtual_loss(...);
//
// let board = nodes[path.last()].board_state.clone();
//
// pending.push(PendingEval {
// path,
// board,
// });
// }
//
// // Send all boards to NN server
// let evaluations = model_inference_batch(pending.boards()).await;
//
// // Expand and backpropagate
// for (pending, (policy, value)) in
// pending.into_iter().zip(evaluations)
// {
// expand_with_policy(...);
// remove_virtual_loss(...);
// backpropagate(...);
// }
// }
while !nodes[current].children.is_empty() {
current = nodes[current].select_child(&nodes, &self.config.c_puct);
path.push(current);
engine/src/net/encoding.rs
+66 -44
View File
@@ -1,6 +1,6 @@
use crate::mcts::BoardState;
use burn::prelude::*;
use chess::{ChessMove, Color, File, Piece, Rank, Square};
use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
const SLIDING_DIRS: [(i8, i8); 8] = [
(0, 1), // N
@@ -212,7 +212,7 @@ fn encode_non_underpromo(dx: i8, dy: i8) -> Option<usize> {
}
/// Decode an index in [0, 4672) back into a ChessMove.
pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
pub fn decode_move(idx: usize, board: &Board) -> Option<ChessMove> {
if idx >= 4672 {
return None;
}
@@ -250,10 +250,24 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
}
// Default pawn moves to final rank (perspective fr == 6 -> tr == 7) to Queen promotion if not underpromotion
if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
if promo.is_none()
&& fr == 6
&& tr == 7
&& dy == 1
&& dx.abs() <= 1
&& board.piece_on(Square::make_square(
Rank::from_index(fr as usize),
File::from_index(ff as usize),
)) == Some(Piece::Pawn)
{
promo = Some(Piece::Queen);
}
// if fr == 6 && tr == 7 && dy == 1 && dx.abs() <= 1 && promo.is_none() {
// promo = Some(Piece::Queen);
// }
let stm = board.side_to_move();
let from = deorient_square(ff, fr, stm);
let to = deorient_square(tf, tr, stm);
@@ -262,69 +276,77 @@ pub fn decode_move(idx: usize, stm: Color) -> Option<ChessMove> {
#[cfg(test)]
mod tests {
use super::*;
use chess::ALL_COLORS;
use crate::net::encoding::{decode_move, encode_move};
use chess::{Board, ChessMove, Color, File, Piece, Rank, Square};
use std::str::FromStr;
#[test]
fn encoding_roundtrips() {
for color in ALL_COLORS {
for action in 0..4672 {
let decoded = decode_move(action, color);
if decoded.is_none() {
continue;
}
let decoded = decoded.unwrap();
let encoded = encode_move(decoded, color).unwrap();
// eprintln!(
// "orig idx = {}, plane={}, from_idx={}",
// action,
// action / 64,
// action % 64
// );
// eprintln!(
// "move = {}, from={}, to={}, promo={:?}",
// decoded,
// decoded.get_source(),
// decoded.get_dest(),
// decoded.get_promotion()
// );
// eprintln!(
// "decoded = {}, from={}, to={}, promo={:?}",
// decoded,
// decoded.get_source(),
// decoded.get_dest(),
// decoded.get_promotion()
// );
// eprintln!(
// "re-encoded idx = {}, plane={}, from_idx={}",
// encoded,
// encoded / 64,
// encoded % 64
// );
assert_eq!(action, encoded);
let board = Board::default();
for action in 0..4672 {
let decoded = decode_move(action, &board);
if decoded.is_none() {
continue;
}
let decoded = decoded.unwrap();
let encoded = encode_move(decoded, board.side_to_move()).unwrap();
// eprintln!(
// "orig idx = {}, plane={}, from_idx={}",
// action,
// action / 64,
// action % 64
// );
// eprintln!(
// "move = {}, from={}, to={}, promo={:?}",
// decoded,
// decoded.get_source(),
// decoded.get_dest(),
// decoded.get_promotion()
// );
// eprintln!(
// "decoded = {}, from={}, to={}, promo={:?}",
// decoded,
// decoded.get_source(),
// decoded.get_dest(),
// decoded.get_promotion()
// );
// eprintln!(
// "re-encoded idx = {}, plane={}, from_idx={}",
// encoded,
// encoded / 64,
// encoded % 64
// );
assert_eq!(action, encoded);
}
}
#[test]
fn queen_promotion_roundtrip() {
use chess::{File, Rank, Square};
let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
let from = Square::make_square(Rank::Seventh, File::E);
let to = Square::make_square(Rank::Eighth, File::E);
let mv = ChessMove::new(from, to, Some(Piece::Queen));
let enc = encode_move(mv, Color::White).unwrap();
let dec = decode_move(enc, Color::White).unwrap();
let dec = decode_move(enc, &board).unwrap();
assert_eq!(mv, dec);
}
#[test]
fn underpromotion_roundtrip() {
use chess::{File, Rank, Square};
let board = Board::from_str("4k3/4P3/8/8/8/8/8/4K3 w - - 0 1").unwrap();
let from = Square::make_square(Rank::Seventh, File::E);
let to = Square::make_square(Rank::Eighth, File::E);
let mv = ChessMove::new(from, to, Some(Piece::Knight));
let enc = encode_move(mv, Color::White).unwrap();
let dec = decode_move(enc, Color::White).unwrap();
let dec = decode_move(enc, &board).unwrap();
assert_eq!(mv, dec);
}
}
engine/src/train.rs
+71 -42
View File
@@ -9,7 +9,7 @@ use burn::module::AutodiffModule;
use burn::optim::{AdamConfig, GradientsParams};
use burn::prelude::{Device, Module};
use burn::store::ModuleRecord;
use chess::{ChessMove, Color};
use chess::{Board, ChessMove};
use rand::rngs::SmallRng;
use rand::seq::SliceRandom;
use rand::{RngExt, SeedableRng};
@@ -113,6 +113,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
println!("Starting training...");
while train.load(Ordering::Relaxed) {
let iter_start = Instant::now();
println!("Iteration: {}", iter);
let infer_model = model.valid();
// Gen samples
@@ -134,7 +135,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
}
let server_device = device.clone();
thread::spawn(move || {
let server_handle = thread::spawn(move || {
let server = NNServer::new(
server_cmd_rx,
server_request_rx,
@@ -147,8 +148,10 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
server.run_server()
});
let mut worker_handles = Vec::new();
let mut workers_dispatched = 0;
for i in 0..training_config.num_workers {
println!("initializing worker {}", i);
// 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,8 +165,8 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
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...");
worker_handles.push(thread::spawn(move || {
// println!("Starting worker thread...");
let worker = Worker::new(
i,
@@ -177,10 +180,15 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
tokio::runtime::Runtime::new()
.unwrap()
.block_on(worker.run_worker(cmd_rx));
});
}));
episodes_dispatched += episodes_to_dispatch;
workers_dispatched += 1;
}
}
println!(
"Started {} workers, dispatched episodes #1-{}",
workers_dispatched, episodes_dispatched
);
while episodes_recieved < training_config.num_episodes {
let data = data_rx.recv().unwrap();
@@ -203,19 +211,26 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
.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);
}
// else {
// worker_map
// .get(&data.worker_id)
// .unwrap()
// .send(WorkerCmd::Shutdown)
// .expect("Expected shutdown to be sent");
// // println!("Shutdown worker #{}", data.worker_id);
// }
}
server_cmd_tx
.send(NNServerCmd::Shutdown)
.expect("Expected shutdown command to be sent");
server_handle.join().unwrap();
for handle in worker_handles {
handle.join().unwrap()
}
// train
if replay_buffer.is_empty() {
@@ -263,6 +278,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
"Iteration {} loss | Total: {:.4} | Policy: {:.4} | Value: {:.4}",
iter, loss_val, p_loss_val, v_loss_val
);
println!("Iteration took {}s", iter_start.elapsed().as_secs());
let grads = output.loss.backward();
let grads = GradientsParams::from_grads(grads, &model);
@@ -271,9 +287,35 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
iter += 1;
if iter % 100 == 0 {
println!("Completed {} iterations", iter);
}
println!("Saving model...");
let metadata = ModelMetadata {
name: training_config.model_name.clone(),
iterations: iter,
};
std::fs::write(
metadata_path.clone(),
serde_json::to_string_pretty(&metadata)
.expect("Should be able to convert metadata to JSON string"),
)
.expect("Should be able to write metadata");
model
.clone()
.into_record()
.save(&model_path)
.expect("Should be able to save the model");
optim
.clone()
.save(format!("{model_path}_optim"))
.expect("Should be able to save the optim");
println!("Model saved in {:?}.", model_path);
// if iter % 100 == 0 {
// println!("Completed {} iterations", iter);
// }
if training_config.max_time_s.is_some()
&& start_time.elapsed().as_secs() > training_config.max_time_s.unwrap()
@@ -291,29 +333,16 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
}
}
println!("Saving model...");
let metadata = ModelMetadata {
name: training_config.model_name,
iterations: iter,
};
std::fs::write(
metadata_path,
serde_json::to_string_pretty(&metadata)
.expect("Should be able to convert metadata to JSON string"),
)
.expect("Should be able to write metadata");
model
.into_record()
.save(&model_path)
.expect("Should be able to save the model");
optim
.save(format!("{model_path}_optim"))
.expect("Should be able to save the optim");
println!("Model saved in {:?}, exiting training.", model_path);
// model
// .into_record()
// .save(&model_path)
// .expect("Should be able to save the model");
//
// optim
// .save(format!("{model_path}_optim"))
// .expect("Should be able to save the optim");
// println!("Model saved in {:?}, exiting training.", model_path);
println!("Finished training!");
return;
}
@@ -356,18 +385,18 @@ pub(crate) fn apply_temperature(visits: &[(usize, f32)], temperature: f32) -> Ve
pub(crate) fn sample_move(
dist: &[(usize, f32)],
rng: &mut SmallRng,
side_to_move: Color,
board: &Board,
) -> Option<ChessMove> {
let mut r: f32 = rng.random_range(0.0..1.0);
for (idx, p) in dist {
r -= *p;
if r <= 0.0 {
return Some(decode_move(*idx, side_to_move).expect("Invalid move"));
return Some(decode_move(*idx, board).expect("Invalid move"));
}
}
// fallback due to floating point drift
dist.get(0)
.map(|(idx, _)| decode_move(*idx, side_to_move).expect("Invalid move"))
.map(|(idx, _)| decode_move(*idx, board).expect("Invalid move"))
}
engine/src/worker.rs
+59 -15
View File
@@ -61,7 +61,7 @@ impl Worker {
F: Fn(&BoardState) -> Fut,
Fut: Future<Output = (Vec<f32>, f32)>,
{
println!("game started");
// println!("game started");
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap();
@@ -84,12 +84,11 @@ impl Worker {
// println!("{}", before.elapsed().as_millis());
episode_buffer.push(results);
let temp = if ply_count < 30 { 1.0 } else { 0.0 };
let temp = if ply_count < 30 { 1.0 } else { 0.2 };
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();
let mv = sample_move(&adjusted, &mut rng, &board_state.board).unwrap();
game_hist.push_str(mv.to_string().as_str());
game_hist.push(' ');
@@ -103,6 +102,8 @@ impl Worker {
}
}
println!("game finished: {}", game_hist);
let mut samples: Vec<TrainingSample> = Vec::new();
for result in episode_buffer.iter() {
let stm = result.board_state.board.side_to_move();
@@ -166,14 +167,15 @@ impl Worker {
// }
pub async fn run_worker(self, cmd_rx: Receiver<WorkerCmd>) {
println!("Worker started");
let mut tasks = Vec::new();
// 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 {
tasks.push(tokio::spawn(async move {
Self::run_game(
self.id,
mcts_config,
@@ -186,19 +188,29 @@ impl Worker {
data_tx,
)
.await;
});
}));
}
loop {
if let Ok(cmd) = cmd_rx.recv() {
match cmd {
// Remove completed tasks.
tasks.retain(|task| !task.is_finished());
// If every game has completed, we're done.
if tasks.is_empty() {
break;
}
match cmd_rx.recv_timeout(Duration::from_millis(10)) {
Ok(cmd) => 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 {
tasks.push(tokio::spawn(async move {
Self::run_game(
self.id,
mcts_config,
@@ -211,13 +223,26 @@ impl Worker {
data_tx,
)
.await;
});
}));
}
},
Err(RecvTimeoutError::Timeout) => {
continue;
}
Err(RecvTimeoutError::Disconnected) => {
println!("Request channel disconnected");
break;
}
} else {
panic!("worker dropped inference command");
}
}
for task in tasks {
task.await.expect("Worker game task panicked");
}
println!("Worker {} shut down", self.id);
}
}
@@ -303,13 +328,24 @@ impl NNServer {
pub fn run_server(&self) {
let mut cache: Cache<NNCacheKey, NNResponse> = Cache::new(self.cache_size);
let mut queue: VecDeque<NNRequest> = VecDeque::new();
let mut request_counter: usize = 0;
let mut cache_hits: usize = 0;
let mut batches_ran: usize = 0;
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");
println!(
"NN server shutting down. Served {} requests, with {} cache hits ({:2}%). Ran {} batches (avg {}requests/batch).",
request_counter,
cache_hits,
cache_hits as f64 / request_counter as f64 * 100.0,
batches_ran,
(request_counter - cache_hits) as f64 / batches_ran as f64,
);
break;
} // other commands...
}
@@ -320,7 +356,9 @@ impl NNServer {
if queue.is_empty() {
match self.request_rx.recv_timeout(Duration::from_millis(10)) {
Ok(first) => {
request_counter += 1;
if try_use_cache(first, &mut cache, &mut queue) {
cache_hits += 1;
continue;
}
}
@@ -337,7 +375,11 @@ impl NNServer {
}
while let Ok(req) = self.request_rx.try_recv() {
try_use_cache(req, &mut cache, &mut queue);
request_counter += 1;
if try_use_cache(req, &mut cache, &mut queue) {
cache_hits += 1;
continue;
}
}
let mut jobs_to_send: Vec<NNRequest> = Vec::new();
@@ -359,7 +401,9 @@ impl NNServer {
);
// println!("Sending {} jobs to gpu", jobs_to_send.len());
let (p, v) = self.model.forward(batched_tensor);
batches_ran += 1;
let p_split = p.chunk(jobs_to_send.len(), 0);
let v_split = v.chunk(jobs_to_send.len(), 0);