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

cache and fixes

Browse Source
This commit is contained in:
drake
2026-08-27 16:30:43 -04:00
parent 3da2554738
commit 4f854aeca4
8 changed files with 204 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
engine/src/worker.rs
+106 -56
View File
@@ -5,10 +5,13 @@ 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::sync::mpsc::{Receiver, Sender};
use std::hash::Hash;
use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender};
use std::time::Duration;
pub struct Worker {
id: usize,
@@ -134,33 +137,33 @@ impl Worker {
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)
}
// 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");
@@ -223,12 +226,14 @@ async fn model_inference_static(
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,
})
@@ -236,21 +241,7 @@ async fn model_inference_static(
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)
(response.policy, response.value)
}
pub enum NNServerCmd {
@@ -268,6 +259,7 @@ pub struct NNServer {
model: ChessModel,
device: Device,
batch_size: usize,
cache_size: u64,
}
impl NNServer {
@@ -277,6 +269,7 @@ impl NNServer {
model: ChessModel,
device: Device,
batch_size: usize,
cache_size: u64,
) -> NNServer {
NNServer {
command_rx,
@@ -284,6 +277,7 @@ impl NNServer {
model,
device,
batch_size,
cache_size,
}
}
}
@@ -294,31 +288,56 @@ struct NNJobData {
// game_id: u32,
}
pub struct NNRequest {
cache_key: NNCacheKey,
request_tensor: Tensor<4>,
response_tx: tokio::sync::oneshot::Sender<NNResponse>,
}
#[derive(Debug)]
#[derive(Debug, Clone)]
pub struct NNResponse {
// nn_job_data: NNJobData,
response_tensors: (Tensor<2>, Tensor<2>),
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();
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");
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() {
let first = self
.request_rx
.recv()
.expect("failed to receive NN request");
queue.push_back(first);
match self.request_rx.recv_timeout(Duration::from_millis(10)) {
Ok(first) => {
if try_use_cache(first, &mut cache, &mut queue) {
continue;
}
}
Err(RecvTimeoutError::Timeout) => {
continue;
}
Err(RecvTimeoutError::Disconnected) => {
println!("Request channel disconnected");
break;
}
}
}
while let Ok(req) = self.request_rx.try_recv() {
queue.push_back(req);
try_use_cache(req, &mut cache, &mut queue);
}
let mut jobs_to_send: Vec<NNRequest> = Vec::new();
@@ -339,18 +358,49 @@ impl NNServer {
0,
);
println!("Sending {} jobs to gpu", jobs_to_send.len());
// 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.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(NNResponse {
response_tensors: (p_split[i].clone(), v_split[i].clone()),
})
.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,
}