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
Cargo.lock
Generated
+26 -14
View File
@@ -149,17 +149,6 @@ dependencies = [
"portable-atomic-util",
]
[[package]]
name = "async-trait"
version = "0.1.92"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667"
dependencies = [
"proc-macro2",
"quote",
"syn 3.0.4",
]
[[package]]
name = "atomic"
version = "0.6.1"
@@ -1219,9 +1208,9 @@ dependencies = [
[[package]]
name = "crossbeam-epoch"
version = "0.9.18"
version = "0.9.20"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e"
checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f"
dependencies = [
"crossbeam-utils",
]
@@ -2177,10 +2166,10 @@ dependencies = [
name = "engine"
version = "0.1.0"
dependencies = [
"async-trait",
"burn",
"chess",
"ctrlc",
"moka",
"rand 0.10.1",
"rand_distr",
"serde",
@@ -3657,6 +3646,23 @@ version = "0.3.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4a0b3262dc837d2513fe2ef31ff8461352ef932dcca31ba0c0abe33547cf6b9b"
[[package]]
name = "moka"
version = "0.12.16"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4293f18e7567a1caf3c584855554377025c65e0aa445344d04171f5ad63d19b9"
dependencies = [
"crossbeam-channel",
"crossbeam-epoch",
"crossbeam-utils",
"equivalent",
"parking_lot",
"portable-atomic",
"smallvec",
"tagptr",
"uuid",
]
[[package]]
name = "moxcms"
version = "0.8.1"
@@ -5609,6 +5615,12 @@ dependencies = [
"windows",
]
[[package]]
name = "tagptr"
version = "0.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
[[package]]
name = "tar"
version = "0.4.46"
engine/Cargo.toml
+1 -1
View File
@@ -12,4 +12,4 @@ ctrlc = "3.5.2"
serde_json = "1.0.150"
serde = { version = "1.0.228", features = ["derive"] }
tokio = { version = "1.52.3", features = ["sync", "rt", "rt-multi-thread"] }
async-trait = "0.1.92"
moka = { version = "0.12.16", features = ["sync"] }
engine/src/main.rs
+3 -2
View File
@@ -24,8 +24,8 @@ fn main() {
max_time_s: None,
num_iters: None,
max_depth: 100, // unused
model_name: String::from("Test1"),
load_model: true,
model_name: String::from("threadedModel"),
load_model: false,
hidden_channels: 48,
num_blocks: 4,
batch_size: 1024, // num positions sampled to a batch (2048)
@@ -38,6 +38,7 @@ fn main() {
num_workers: 4,
nn_batch_size: 64,
episodes_per_worker: 4,
nn_cache_size: 500_000, // 18,736B / item in cache. 100k=1.75GB.
};
train(training_config, device);
engine/src/mcts.rs
+13 -1
View File
@@ -1,13 +1,13 @@
// use crate::networks::PolicyNetwork;
use crate::mcts::BoardStateStatus::{FiftyMove, InsufficientMaterial, Ongoing, Threefold};
use crate::net::encoding::encode_move;
use crate::worker::NNCacheKey;
use chess::BoardStatus::{Checkmate, Stalemate};
use chess::Color::{Black, White};
use chess::Piece::{Bishop, Knight, Pawn, Queen, Rook};
use chess::{Board, ChessMove, Color, MoveGen, Piece, ALL_COLORS, ALL_PIECES};
use rand::SeedableRng;
use std::collections::HashMap;
//TODO retain tree after move selection?
#[derive(Clone)]
@@ -450,6 +450,18 @@ impl BoardState {
}
}
pub fn nn_cache_key(&self) -> NNCacheKey {
let board_hash = self.board.get_hash(); // whatever chess crate API you're using
let repetition_two = self.repetition_table.get(&board_hash).copied().unwrap_or(0) >= 2;
NNCacheKey {
board_hash,
halfmove_clock: self.halfmove_clock,
repetition_two,
}
}
pub fn apply_move(&mut self, mv: ChessMove) {
let is_pawn_move = self.board.piece_on(mv.get_source()) == Some(Pawn);
let is_normal_capture = self.board.piece_on(mv.get_dest()).is_some();
engine/src/net/encoding.rs
+22 -15
View File
@@ -33,14 +33,14 @@ const UNDERPROMO_DIRS: [(i8, i8); 3] = [
const UNDERPROMO_PIECES: [Piece; 3] = [Piece::Knight, Piece::Bishop, Piece::Rook];
pub fn encode_board_state_perspective(state: &BoardState, device: &Device) -> Tensor<3> {
let mut planes = vec![0.0f32; 18 * 64];
let mut planes = vec![0.0f32; 19 * 64];
encode_board_state_to_slice(state, &mut planes);
Tensor::<3>::from_floats(TensorData::new(planes, [18, 8, 8]), device)
Tensor::<3>::from_floats(TensorData::new(planes, [19, 8, 8]), device)
}
/// Write 18x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1152 elements long).
/// Write 19x8x8 encoded planes for `state` directly into `out_slice` (must be at least 1216 elements long).
pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
out_slice[..1152].fill(0.0);
out_slice[..1216].fill(0.0);
let board = &state.board;
let us = board.side_to_move();
@@ -86,23 +86,23 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
let (us_castle, them_castle) = (board.castle_rights(us), board.castle_rights(them));
if us_castle.has_kingside() {
fill_plane(out_slice, 12);
fill_plane_ones(out_slice, 12);
}
if us_castle.has_queenside() {
fill_plane(out_slice, 13);
fill_plane_ones(out_slice, 13);
}
if them_castle.has_kingside() {
fill_plane(out_slice, 14);
fill_plane_ones(out_slice, 14);
}
if them_castle.has_queenside() {
fill_plane(out_slice, 15);
fill_plane_ones(out_slice, 15);
}
// Repetition plane (16)
let current_hash = board.get_hash();
if let Some(count) = state.repetition_table.get(&current_hash) {
if *count >= 2 {
fill_plane(out_slice, 16);
fill_plane_ones(out_slice, 16);
}
}
@@ -116,22 +116,29 @@ pub fn encode_board_state_to_slice(state: &BoardState, out_slice: &mut [f32]) {
}
out_slice[idx(17, rank, file)] = 1.0;
}
fill_plane(out_slice, 18, Some(state.halfmove_clock as f32 / 100.0));
}
pub fn encode_board_states_batch(states: &[&BoardState], device: &Device) -> Tensor<4> {
let b = states.len();
let mut buffer = vec![0.0f32; b * 1152];
let mut buffer = vec![0.0f32; b * 1216];
for (i, state) in states.iter().enumerate() {
let start = i * 1152;
encode_board_state_to_slice(state, &mut buffer[start..start + 1152]);
let start = i * 1216;
encode_board_state_to_slice(state, &mut buffer[start..start + 1216]);
}
Tensor::<4>::from_floats(TensorData::new(buffer, [b, 18, 8, 8]), device)
Tensor::<4>::from_floats(TensorData::new(buffer, [b, 19, 8, 8]), device)
}
#[inline]
fn fill_plane(buffer: &mut [f32], plane: usize) {
fn fill_plane(buffer: &mut [f32], plane: usize, value: Option<f32>) {
let start = plane * 64;
buffer[start..start + 64].fill(1.0);
buffer[start..start + 64].fill(value.unwrap_or(1.0));
}
#[inline]
fn fill_plane_ones(buffer: &mut [f32], plane: usize) {
fill_plane(buffer, plane, Some(1.0));
}
/// Rotate a square into the side-to-move perspective.
engine/src/net/model.rs
+28 -20
View File
@@ -21,6 +21,7 @@ Input planes:
16: opponent queenside
17: rep >= 2
18: En passant square (gap pass over)
19: No-progress plys / 100
Output:
@@ -45,10 +46,10 @@ pub struct ResidualBlock {
activation: Relu,
}
#[derive(Config, Debug)]
pub struct ResidualBlockConfig {
channels: usize,
}
// #[derive(Config, Debug)]
// pub struct ResidualBlockConfig {
// channels: usize,
// }
impl ResidualBlock {
pub fn new(channels: usize, device: &Device) -> Self {
@@ -110,7 +111,7 @@ pub struct ChessModelConfig {
impl ChessModelConfig {
pub fn init(num_blocks: usize, channels: usize, device: &Device) -> ChessModel {
ChessModel {
conv: Conv2dConfig::new([18, channels], [3, 3]) // 18 plane input
conv: Conv2dConfig::new([19, channels], [3, 3]) // 19 plane input
.with_padding(PaddingConfig2d::Explicit(1, 1, 1, 1))
.init(device),
bn: BatchNormConfig::new(channels).init(device),
@@ -205,23 +206,30 @@ pub struct ChessOutput {
pub value_targets: Tensor<2>, // kept 2d for consistency?
pub policy_loss: Tensor<1>,
pub value_loss: Tensor<1>,
pub loss: Tensor<1>, // [sample]
pub loss: Tensor<1>, // [sample]
}
impl ItemLazy for ChessOutput {
fn sync(self) -> Self {
let [policy_logits, value, policy_targets, value_targets, policy_loss, value_loss, loss] =
Transaction::default()
.register(self.policy_logits)
.register(self.value)
.register(self.policy_targets)
.register(self.value_targets)
.register(self.policy_loss)
.register(self.value_loss)
.register(self.loss)
.execute()
.try_into()
.expect("Correct amount of tensor data");
let [
policy_logits,
value,
policy_targets,
value_targets,
policy_loss,
value_loss,
loss,
] = Transaction::default()
.register(self.policy_logits)
.register(self.value)
.register(self.policy_targets)
.register(self.value_targets)
.register(self.policy_loss)
.register(self.value_loss)
.register(self.loss)
.execute()
.try_into()
.expect("Correct amount of tensor data");
let device = &Default::default();
@@ -299,7 +307,7 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
let state_tensors = items
.iter()
.map(|item| {
encode_board_state_perspective(&item.board_state, device).reshape([1, 18, 8, 8])
encode_board_state_perspective(&item.board_state, device).reshape([1, 19, 8, 8])
})
.collect::<Vec<_>>();
@@ -331,7 +339,7 @@ impl Batcher<TrainingSample, ChessBatch> for ChessBatcher {
})
.collect::<Vec<_>>();
let states = Tensor::cat(state_tensors, 0); // [B, 18, 8, 8]
let states = Tensor::cat(state_tensors, 0); // [B, 19, 8, 8]
let policy_targets = Tensor::cat(policy_target_tensors, 0); // [B, 4672]
let value_targets = Tensor::cat(value_target_tensors, 0); // [B, 1]
engine/src/train.rs
+5
View File
@@ -42,6 +42,7 @@ pub struct TrainingConfig {
pub num_workers: usize,
pub nn_batch_size: usize,
pub episodes_per_worker: usize,
pub nn_cache_size: u64,
}
pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
@@ -140,6 +141,7 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
infer_model,
server_device,
training_config.nn_batch_size,
training_config.nn_cache_size,
);
server.run_server()
@@ -211,6 +213,9 @@ pub fn train(training_config: TrainingConfig, device: impl Into<Device>) {
}
}
server_cmd_tx
.send(NNServerCmd::Shutdown)
.expect("Expected shutdown command to be sent");
// train
if replay_buffer.is_empty() {
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,
}