454 lines
13 KiB
Rust
454 lines
13 KiB
Rust
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 chess::Color;
|
|
use moka::sync::Cache;
|
|
use rand::prelude::SmallRng;
|
|
use rand::SeedableRng;
|
|
use std::collections::VecDeque;
|
|
use std::hash::Hash;
|
|
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
|
|
use std::time::Duration;
|
|
|
|
pub struct Worker {
|
|
id: usize,
|
|
num_games: usize,
|
|
// games: Vec<BoardState>,
|
|
request_tx: Sender<NNRequest>,
|
|
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 {
|
|
pub fn new(
|
|
id: usize,
|
|
num_games: usize,
|
|
sender: Sender<NNRequest>,
|
|
data_tx: Sender<WorkerData>,
|
|
device: Device,
|
|
mcts_config: MctsConfig,
|
|
) -> Worker {
|
|
Worker {
|
|
id,
|
|
// games: (0..num_games)
|
|
// .map(|_| BoardState::default())
|
|
// .collect::<Vec<BoardState>>(),
|
|
num_games,
|
|
request_tx: sender,
|
|
data_tx,
|
|
curr_request_id: 0,
|
|
device,
|
|
mcts_config,
|
|
}
|
|
}
|
|
|
|
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 = (1.0 - (ply_count as f32 / 200.0).min(1.0)).max(0.3);
|
|
// let temp = if ply_count < 30 { 1.0 } else { 0.2 };
|
|
|
|
let adjusted = apply_temperature(&episode_buffer.last().unwrap().move_dist, temp);
|
|
|
|
let mv = sample_move(&adjusted, &mut rng, &board_state.board).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;
|
|
}
|
|
}
|
|
|
|
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();
|
|
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,
|
|
};
|
|
|
|
let blended_target = 0.75 * outcome + 0.25 * result.value;
|
|
|
|
samples.push(TrainingSample::from_mcts_with_outcome(
|
|
result.clone(),
|
|
blended_target,
|
|
));
|
|
}
|
|
|
|
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>) {
|
|
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();
|
|
|
|
tasks.push(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 {
|
|
// 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();
|
|
|
|
tasks.push(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;
|
|
}));
|
|
}
|
|
},
|
|
|
|
Err(RecvTimeoutError::Timeout) => {
|
|
continue;
|
|
}
|
|
|
|
Err(RecvTimeoutError::Disconnected) => {
|
|
println!("Request channel disconnected");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
for task in tasks {
|
|
task.await.expect("Worker game task panicked");
|
|
}
|
|
|
|
println!("Worker {} shut down", self.id);
|
|
}
|
|
}
|
|
|
|
async fn model_inference_static(
|
|
board_state: BoardState,
|
|
request_tx: Sender<NNRequest>,
|
|
device: Device,
|
|
) -> (Vec<f32>, f32) {
|
|
let cache_key = board_state.nn_cache_key();
|
|
let state = encode_board_state_perspective(&board_state, &device).unsqueeze::<4>();
|
|
|
|
let (tx, rx) = tokio::sync::oneshot::channel::<NNResponse>();
|
|
|
|
request_tx
|
|
.send(NNRequest {
|
|
cache_key,
|
|
request_tensor: state,
|
|
response_tx: tx,
|
|
})
|
|
.unwrap();
|
|
|
|
let response = rx.await.expect("worker dropped inference response");
|
|
|
|
(response.policy, response.value)
|
|
}
|
|
|
|
pub enum NNServerCmd {
|
|
Shutdown,
|
|
}
|
|
|
|
pub enum WorkerCmd {
|
|
Shutdown,
|
|
Dispatch,
|
|
}
|
|
|
|
pub struct NNServer {
|
|
command_rx: Receiver<NNServerCmd>,
|
|
request_rx: Receiver<NNRequest>,
|
|
model: ChessModel,
|
|
device: Device,
|
|
batch_size: usize,
|
|
cache_size: u64,
|
|
}
|
|
|
|
impl NNServer {
|
|
pub fn new(
|
|
command_rx: Receiver<NNServerCmd>,
|
|
request_rx: Receiver<NNRequest>,
|
|
model: ChessModel,
|
|
device: Device,
|
|
batch_size: usize,
|
|
cache_size: u64,
|
|
) -> NNServer {
|
|
NNServer {
|
|
command_rx,
|
|
request_rx,
|
|
model,
|
|
device,
|
|
batch_size,
|
|
cache_size,
|
|
}
|
|
}
|
|
}
|
|
#[derive(Debug, Clone, Eq, Hash, PartialEq)]
|
|
struct NNJobData {
|
|
worker_id: usize,
|
|
request_id: u64,
|
|
// game_id: u32,
|
|
}
|
|
pub struct NNRequest {
|
|
cache_key: NNCacheKey,
|
|
request_tensor: Tensor<4>,
|
|
response_tx: tokio::sync::oneshot::Sender<NNResponse>,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct NNResponse {
|
|
policy: Vec<f32>,
|
|
value: f32,
|
|
}
|
|
|
|
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. 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...
|
|
}
|
|
}
|
|
|
|
// If we have no queued requests, wait for either a request
|
|
// or periodically wake up to check for shutdown.
|
|
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;
|
|
}
|
|
}
|
|
|
|
Err(RecvTimeoutError::Timeout) => {
|
|
continue;
|
|
}
|
|
|
|
Err(RecvTimeoutError::Disconnected) => {
|
|
println!("Request channel disconnected");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
while let Ok(req) = self.request_rx.try_recv() {
|
|
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();
|
|
|
|
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,
|
|
);
|
|
|
|
// 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);
|
|
|
|
for (i, job) in jobs_to_send.into_iter().enumerate() {
|
|
let p = p_split[i].clone().into_data().to_vec::<f32>().unwrap();
|
|
|
|
let v = v_split[i].clone().into_data().to_vec::<f32>().unwrap()[0];
|
|
|
|
let res = NNResponse {
|
|
policy: p,
|
|
value: v,
|
|
};
|
|
|
|
cache.insert(job.cache_key.clone(), res.clone());
|
|
|
|
job.response_tx
|
|
.send(res)
|
|
.expect("worker dropped inference request");
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn try_use_cache(
|
|
req: NNRequest,
|
|
cache: &mut Cache<NNCacheKey, NNResponse>,
|
|
queue: &mut VecDeque<NNRequest>,
|
|
) -> bool {
|
|
if let Some(cached) = cache.get(&req.cache_key) {
|
|
// println!("Cache hit");
|
|
req.response_tx.send(cached).unwrap();
|
|
true
|
|
} else {
|
|
queue.push_back(req);
|
|
false
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
|
|
pub struct NNCacheKey {
|
|
pub board_hash: u64,
|
|
pub halfmove_clock: u8,
|
|
pub repetition_two: bool,
|
|
}
|