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/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");
}
}
}