diff --git a/.gitignore b/.gitignore index 19ef6e0..4160fe2 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ build/ .venv/ .idea/ +runs/ diff --git a/AGENTS.md b/AGENTS.md index 98229c1..d364e10 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -9,4 +9,10 @@ our search impl will also be in rust. check out /mnt/github/TheConverseEngineer/AlphaSnake/ for last years impl -`just test` runs tests +`just test` runs tests. u should use this rather than doign it yourself bc this auto installs the package + +- this just tests the rust board impl. do not use it for eg training code changes + +`nix develop` GETS YOU CUDA +`maturin develop --release` to bring python bindings in. you can `--manifest-path training/Cargo.toml` etc to choose the package to reinstall +`PYTHONPATH=python` is uesful sometimes diff --git a/Cargo.lock b/Cargo.lock index b523419..8c81c59 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,26 +2,113 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "alpha_paint" version = "0.1.0" dependencies = [ + "criterion", "pyo3", - "rand", + "rand 0.10.0", + "safetensors", + "serde", + "serde_json", ] +[[package]] +name = "alphapaint_training" +version = "0.1.0" +dependencies = [ + "alpha_paint", + "cudarc", + "event-listener", + "ndarray", + "numpy", + "pyo3", + "rand 0.10.0", + "rand_chacha 0.10.0", + "rand_distr", + "rayon", + "rstest", + "serde", + "serde_json", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "anyhow" version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + [[package]] name = "bitflags" version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + +[[package]] +name = "cc" +version = "1.2.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1e928d4b69e3077709075a938a05ffbedfa53a84c8f766efbf8220bb1ff60e1" +dependencies = [ + "find-msvc-tools", + "shlex", +] + [[package]] name = "cfg-if" version = "1.0.4" @@ -36,7 +123,68 @@ checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" dependencies = [ "cfg-if", "cpufeatures", - "rand_core", + "rand_core 0.10.0", +] + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", ] [[package]] @@ -48,18 +196,176 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "cudarc" +version = "0.18.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3aa12038120eb13347a6ae2ffab1d34efe78150125108627fd85044dd4d6ff1e" +dependencies = [ + "libloading", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "equivalent" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + [[package]] name = "foldhash" version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-timer" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f288b0a4f20f9a56b5d1da57e2227c661b7b16168e2f72365f57b63326e29b24" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + [[package]] name = "getrandom" version = "0.4.2" @@ -68,19 +374,36 @@ checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" dependencies = [ "cfg-if", "libc", - "r-efi", - "rand_core", + "r-efi 6.0.0", + "rand_core 0.10.0", "wasip2", "wasip3", ] +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ - "foldhash", + "foldhash 0.1.5", ] [[package]] @@ -88,6 +411,13 @@ name = "hashbrown" version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", + "serde", + "serde_core", +] [[package]] name = "heck" @@ -113,12 +443,31 @@ dependencies = [ "serde_core", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" +[[package]] +name = "js-sys" +version = "0.3.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4c90f45aa2e6eacbe8645f77fdea542ac97a494bcd117a67df9ff4d611f995" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + [[package]] name = "leb128fmt" version = "0.1.0" @@ -131,30 +480,189 @@ version = "0.2.182" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6800badb6cb2082ffd7b6a67e6125bb39f18782f793520caee8cb8846be06112" +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "log" version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "numpy" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778da78c64ddc928ebf5ad9df5edf0789410ff3bdbf3619aed51cd789a6af1e2" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + [[package]] name = "once_cell" version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" +[[package]] +name = "portable-atomic-util" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "091397be61a01d4be58e7841595bd4bfedb15f1cd54977d79b8271e94ed799a3" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + [[package]] name = "prettyplease" version = "0.2.37" @@ -165,6 +673,15 @@ dependencies = [ "syn", ] +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -241,12 +758,28 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "r-efi" version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + [[package]] name = "rand" version = "0.10.0" @@ -254,8 +787,37 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" dependencies = [ "chacha20", - "getrandom", - "rand_core", + "getrandom 0.4.2", + "rand_core 0.10.0", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e6af7f3e25ded52c41df4e0b1af2d047e45896c2f3281792ed68a1c243daedb" +dependencies = [ + "ppv-lite86", + "rand_core 0.10.0", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", ] [[package]] @@ -264,6 +826,147 @@ version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" +[[package]] +name = "rand_distr" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" +dependencies = [ + "num-traits", + "rand 0.9.2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + +[[package]] +name = "relative-path" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba39f3699c378cd8970968dcbff9c43159ea4cfbd88d43c00b22f2ef10a435d2" + +[[package]] +name = "rstest" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5a3193c063baaa2a95a33f03035c8a72b83d97a54916055ba22d35ed3839d49" +dependencies = [ + "futures-timer", + "futures-util", + "rstest_macros", +] + +[[package]] +name = "rstest_macros" +version = "0.26.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c845311f0ff7951c5506121a9ad75aec44d083c31583b2ea5a30bcb0b0abba0" +dependencies = [ + "cfg-if", + "glob", + "proc-macro-crate", + "proc-macro2", + "quote", + "regex", + "relative-path", + "rustc_version", + "syn", + "unicode-ident", +] + +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "safetensors" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675656c1eabb620b921efea4f9199f97fc86e36dd6ffd1fbbe48d0f59a4987f5" +dependencies = [ + "hashbrown 0.16.1", + "serde", + "serde_json", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "semver" version = "1.0.27" @@ -277,6 +980,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -312,6 +1016,18 @@ dependencies = [ "zmij", ] +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "syn" version = "2.0.117" @@ -329,6 +1045,46 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "toml_datetime" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97251a7c317e03ad83774a8752a7e81fb6067740609f75ea2b585b569a59198f" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.25.8+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16bff38f1d86c47f9ff0647e6838d7bb362522bdf44006c7068c2b1e606f1f3c" +dependencies = [ + "indexmap", + "toml_datetime", + "toml_parser", + "winnow", +] + +[[package]] +name = "toml_parser" +version = "1.1.0+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2334f11ee363607eb04df9b8fc8a13ca1715a72ba8662a26ac285c98aabb4011" +dependencies = [ + "winnow", +] + [[package]] name = "tooling" version = "0.1.0" @@ -348,6 +1104,16 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "wasip2" version = "1.0.2+wasi-0.2.9" @@ -366,6 +1132,51 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasm-bindgen" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6523d69017b7633e396a89c5efab138161ed5aafcbc8d3e5c5a42ae38f50495a" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e3a6c758eb2f701ed3d052ff5737f5bfe6614326ea7f3bbac7156192dc32e67" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "921de2737904886b52bcbb237301552d05969a6f9c40d261eb0533c8b055fedf" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.115" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a93e946af942b58934c604527337bad9ae33ba1d5c6900bbb41c2c07c2364a93" +dependencies = [ + "unicode-ident", +] + [[package]] name = "wasm-encoder" version = "0.244.0" @@ -400,6 +1211,71 @@ dependencies = [ "semver", ] +[[package]] +name = "web-sys" +version = "0.3.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84cde8507f4d7cfcb1185b8cb5890c494ffea65edbe1ba82cfd63661c805ed94" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winnow" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a90e88e4667264a994d34e6d1ab2d26d398dcdca8b7f52bec8668957517fc7d8" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" @@ -488,6 +1364,26 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "zerocopy" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efbb2a062be311f2ba113ce66f697a4dc589f85e78a4aea276200804cea0ed87" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e8bc7269b54418e7aeeef514aa68f8690b8c0489a06b0136e5f57c4c5ccab89" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "zmij" version = "1.0.21" diff --git a/Cargo.toml b/Cargo.toml index 72a8140..82d1b8d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,5 @@ [workspace] -members = ["alpha_paint", "tooling"] +members = ["alpha_paint", "tooling", "training"] resolver = "3" [profile.dev] @@ -7,4 +7,4 @@ opt-level = 3 debug_assertions = true [profile.release] -debug = true \ No newline at end of file +debug = true diff --git a/alpha_paint/Cargo.toml b/alpha_paint/Cargo.toml index 19f5a00..f4da249 100644 --- a/alpha_paint/Cargo.toml +++ b/alpha_paint/Cargo.toml @@ -7,9 +7,19 @@ edition = "2024" name = "alpha_paint" crate-type = ["cdylib", "rlib"] +[[bench]] +name = "nn_latency" +harness = false + [dependencies] pyo3 = { version = "0.28.2", features = ["extension-module"] } rand = "0.10.0" +safetensors = "0.7.0" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.149" [features] debug=[] + +[dev-dependencies] +criterion = "0.8.2" diff --git a/alpha_paint/assets/value_net.safetensors b/alpha_paint/assets/value_net.safetensors new file mode 100644 index 0000000..79184e6 Binary files /dev/null and b/alpha_paint/assets/value_net.safetensors differ diff --git a/alpha_paint/benches/nn_latency.rs b/alpha_paint/benches/nn_latency.rs new file mode 100644 index 0000000..2ecd24f --- /dev/null +++ b/alpha_paint/benches/nn_latency.rs @@ -0,0 +1,63 @@ +use std::hint::black_box; +use std::time::Duration; + +use alpha_paint::board::Board; +use alpha_paint::evaluation::Evaluator; +use alpha_paint::nn::features; +use alpha_paint::nn::model::ValueModel; +use alpha_paint::search::GameSearchTree; +use criterion::{Criterion, criterion_group, criterion_main}; + +fn bench_feature_extraction(c: &mut Criterion) { + let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[0]).unwrap(); + c.bench_function("board_to_features", |b| { + b.iter(|| { + let features = features::extract(black_box(&board)); + black_box(features); + }); + }); +} + +fn bench_single_eval(c: &mut Criterion) { + let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[1]).unwrap(); + let evaluator = Evaluator::benchmark(); + c.bench_function("single_position_eval", |b| { + b.iter(|| { + let value = evaluator.evaluate(black_box(&board)); + black_box(value); + }); + }); +} + +fn bench_forward_pass(c: &mut Criterion) { + let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[1]).unwrap(); + let features = features::extract(&board); + let model = ValueModel::benchmark_model(); + c.bench_function("value_model_forward", |b| { + b.iter(|| { + let value = black_box(&model).evaluate(black_box(&features)); + black_box(value); + }); + }); +} + +fn bench_tree_step(c: &mut Criterion) { + let board = Board::from_fen(alpha_paint::TRAINING_START_FENS[2]).unwrap(); + let evaluator = Evaluator::benchmark(); + c.bench_function("ubfm_single_iteration", |b| { + b.iter(|| { + let mut tree = GameSearchTree::new(black_box(&board), &evaluator); + tree.run_descent_for_iter(1, Duration::from_millis(5)); + black_box(tree.get_best_action_index()); + }); + }); +} + +criterion_group!( + benches, + bench_feature_extraction, + bench_single_eval, + bench_forward_pass, + bench_tree_step +); +criterion_main!(benches); diff --git a/alpha_paint/build.rs b/alpha_paint/build.rs new file mode 100644 index 0000000..1b734eb --- /dev/null +++ b/alpha_paint/build.rs @@ -0,0 +1,8 @@ +fn main() { + // pyo3 with `extension-module` doesn't link libpython on macOS. + // maturin handles this automatically, but plain cargo bench/build does not. + if std::env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("macos") { + println!("cargo:rustc-cdylib-link-arg=-undefined"); + println!("cargo:rustc-cdylib-link-arg=dynamic_lookup"); + } +} diff --git a/alpha_paint/src/bindings.rs b/alpha_paint/src/bindings.rs index 4f1c078..7d2f8f8 100644 --- a/alpha_paint/src/bindings.rs +++ b/alpha_paint/src/bindings.rs @@ -1,7 +1,7 @@ -use pyo3::{pyclass, pymethods}; +use pyo3::{exceptions::PyRuntimeError, pyclass, pymethods, PyResult}; use rand::{rngs::StdRng, seq::SliceRandom, RngExt, SeedableRng}; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use crate::board::{board_structs::*, structs::*, *}; use crate::evaluation::Evaluator; @@ -67,7 +67,7 @@ impl PyBoard { 0: [[u16::MAX; 32]; 32], }), hills: Arc::new(vec![Vec::new(); hill_owners.len()]), - dist: Arc::new(DoubleArray32x32(vec![])), // computed after tiles are set + dist: Arc::new(DoubleArray32x32::splat(2000)), // computed after tiles are set tiles: TileMap::default(), powerups: Array32x32 { @@ -139,7 +139,11 @@ impl PyBoard { }) .collect(), ); - board.dist = Arc::new(compute_distances(&board.tiles)); + let dist = Arc::new(compute_distances(&board.tiles)); + board + .tiles + .attach_dist(dist.clone(), board.rows, board.cols); + board.dist = dist; PyBoard(board) } @@ -147,9 +151,14 @@ impl PyBoard { fn perft(&self, max_depth: usize) -> Vec { let mut all_turns = Vec::new(); let mut move_stack = Vec::new(); - enumerate_single_turns(&mut self.0.clone(), max_depth, &mut move_stack, &mut |turn| { - all_turns.push(turn); - }); + enumerate_single_turns( + &mut self.0.clone(), + max_depth, + &mut move_stack, + &mut |turn| { + all_turns.push(turn); + }, + ); all_turns } @@ -161,9 +170,14 @@ impl PyBoard { let mut move_stack = Vec::new(); let mut sampler = ReservoirSampler::new(max_samples, seed); - enumerate_single_turns(&mut self.0.clone(), max_depth, &mut move_stack, &mut |turn| { - sampler.record(turn); - }); + enumerate_single_turns( + &mut self.0.clone(), + max_depth, + &mut move_stack, + &mut |turn| { + sampler.record(turn); + }, + ); sampler.finish() } @@ -171,41 +185,30 @@ impl PyBoard { self.0.to_string() } - fn search(&mut self, time_left: f32) -> Vec { + fn search(&mut self, time_left: f32) -> PyResult> { let mut local_board = self.0.clone(); let mut py_actions: Vec = vec![]; - let evaluator = Evaluator::new(&local_board); - - // Time control - // For reference, on my laptop we can hit around 20k iterations per second (with incremental logic) + let evaluator = Evaluator::new(&local_board).map_err(PyRuntimeError::new_err)?; let max_moves_left = ((2000 - local_board.turn_count) / 2) as u32; - let (max_duration, iterations) = if time_left < 5.0 { - // At 100 iterations per move, we could finish the entire game in ~5 seconds - (time_left / (max_moves_left as f32 * 5.0), 100) + let turn_budget_secs = if time_left < 5.0 { + time_left / (max_moves_left.max(1) as f32 * 5.0) } else if time_left < 30.0 { - // Speed mode (try to leave 2.5s buffer while still finishing all moves) - let iter_budget = ((time_left - 2.5) * 18000.0) as u32; - ( - time_left / ((max_moves_left as f32 - 2.5) * 5.0), - iter_budget / (max_moves_left * 5), - ) + (time_left - 2.5).max(0.0) / (max_moves_left.max(1) as f32 * 5.0) } else if time_left < 90.0 { - // Late game mode (try to leave 20s buffer) - let iter_budget = ((time_left - 20.0) * 18000.0) as u32; - ( - time_left / ((max_moves_left as f32 - 20.0) * 5.0), - iter_budget / (max_moves_left * 5), - ) + (time_left - 20.0).max(0.0) / (max_moves_left.max(1) as f32 * 5.0) } else { - // For the first half of our time, let's use it more liberally. We can - // finish most games before we get under 90s anyway - (0.150, 15_000) + 0.150 }; + let turn_budget = Duration::from_secs_f32(turn_budget_secs.max(0.005)); + let turn_started_at = Instant::now(); let mut tree = GameSearchTree::new(&local_board, &evaluator); loop { - tree.run_descent_for_iter(15_000, Duration::from_secs_f32(max_duration)); + let remaining_budget = turn_budget.saturating_sub(turn_started_at.elapsed()); + if !remaining_budget.is_zero() { + tree.run_descent_for_iter(u32::MAX, remaining_budget); + } let (action_id, action) = tree.get_best_action_and_index(); let player_coord = local_board.current_player_coord(); @@ -215,33 +218,31 @@ impl PyBoard { ApplyActionOutcome::Ongoing => { py_actions.push(action.to_python_primitives(&player_coord)); if action.is_final() { - return py_actions; + return Ok(py_actions); } } ApplyActionOutcome::Terminal { .. } => { py_actions.push(action.to_python_primitives(&player_coord)); - return py_actions; + return Ok(py_actions); } ApplyActionOutcome::PlayInstead { play_instead, .. } => { py_actions.push(play_instead.to_python_primitives(&player_coord)); - return py_actions; + return Ok(py_actions); } - ApplyActionOutcome::Killshot { terminal: _, moves } => { + ApplyActionOutcome::Killshot { + terminal: _, + actions: ks_actions, + } => { py_actions.push(action.to_python_primitives(&player_coord)); if !action.is_final() { // Killshot for us - add all killshot moves let mut coord = local_board.current_player_coord(); - for (i, mv) in moves.iter().enumerate() { - let a = if i == moves.len() - 1 { - Action::FinalMove(*mv) - } else { - Action::Move(*mv) - }; - py_actions.push(a.to_python_primitives(&coord)); - coord = mv.target; + for ks_action in ks_actions.iter() { + py_actions.push(ks_action.to_python_primitives(&coord)); + coord = ks_action.target(); } } - return py_actions; + return Ok(py_actions); } }; @@ -250,6 +251,20 @@ impl PyBoard { } } + fn bench_eval(&self, reps: usize) -> PyResult { + let reps = reps.max(1); + let evaluator = Evaluator::new(&self.0).map_err(PyRuntimeError::new_err)?; + let started_at = std::time::Instant::now(); + let mut sink = 0.0f32; + for _ in 0..reps { + sink += evaluator.evaluate(&self.0); + } + let elapsed = started_at.elapsed().as_secs_f64(); + let avg_micros = elapsed * 1_000_000.0 / reps as f64; + let _ = sink; + Ok(avg_micros) + } + fn apply_turn(&mut self, turn: &Turn) { let len = turn.actions.len(); for (i, (action, _)) in turn.actions.iter().enumerate() { @@ -269,8 +284,6 @@ impl PyBoard { } } } - - } fn enumerate_single_turns( @@ -330,28 +343,31 @@ fn enumerate_single_turns( }); move_stack.pop(); } + // if we played a final action, this killshot is for the OPPONENT, so we dont include + // it in the single turn perft. + ApplyActionOutcome::Killshot { .. } if action.is_final() => { + move_stack.push((action, player_coord)); + save_turn(Turn { + actions: move_stack.clone(), + is_terminal: None, + }); + move_stack.pop(); + } ApplyActionOutcome::Killshot { - terminal, moves, .. + terminal, + actions: ks_actions, } => { - debug_assert!(!action.is_final()); move_stack.push((action, player_coord)); - // Add all killshot intermediate moves let mut coord = board.current_player_coord(); - for (i, mv) in moves.iter().enumerate() { - let a = if i == moves.len() - 1 { - Action::FinalMove(*mv) - } else { - Action::Move(*mv) - }; - move_stack.push((a, coord)); - coord = mv.target; + for &ks_action in ks_actions.iter() { + move_stack.push((ks_action, coord)); + coord = ks_action.target(); } save_turn(Turn { actions: move_stack.clone(), is_terminal: Some(terminal), }); - // Pop all the killshot moves + the original action - for _ in 0..moves.len() + 1 { + for _ in 0..ks_actions.len() + 1 { move_stack.pop(); } } diff --git a/alpha_paint/src/board/action_generation.rs b/alpha_paint/src/board/action_generation.rs index 738f945..f3e4e1d 100644 --- a/alpha_paint/src/board/action_generation.rs +++ b/alpha_paint/src/board/action_generation.rs @@ -106,7 +106,7 @@ impl Board { self.black_stamina }; - for &beacon_target in self.tiles.get_beacon_iterator::() { + for &beacon_target in self.tiles.get_beacon_iterator::() { if stamina >= move_cost + EXTRA_MOVE_COST { actions.push(Action::Move(Move::new( beacon_target, diff --git a/alpha_paint/src/board/actions.rs b/alpha_paint/src/board/actions.rs index 619cf9c..9e19544 100644 --- a/alpha_paint/src/board/actions.rs +++ b/alpha_paint/src/board/actions.rs @@ -135,12 +135,14 @@ pub struct ActionList { impl ActionList { pub fn new() -> Self { Self { - local_buf: [Action::Paint(Paint{ target: Coordinate::new(0, 0) }); 20], + local_buf: [Action::Paint(Paint { + target: Coordinate::new(0, 0), + }); 20], overflow: None, cnt: 0, } } - + pub fn len(&self) -> usize { self.cnt } @@ -186,4 +188,4 @@ impl<'a> IntoIterator for &'a ActionList { index: 0, } } -} \ No newline at end of file +} diff --git a/alpha_paint/src/board/bitboard.rs b/alpha_paint/src/board/bitboard.rs new file mode 100644 index 0000000..f8af802 --- /dev/null +++ b/alpha_paint/src/board/bitboard.rs @@ -0,0 +1,127 @@ +use crate::board::structs::{Array32x32, Coordinate}; + +/// 32×32 bitboard: row `y` holds bits for `x = 0..31`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Bitboard { + pub rows: [u32; 32], +} + +impl Default for Bitboard { + fn default() -> Self { + Self::empty() + } +} + +impl Bitboard { + pub const fn empty() -> Self { + Self { rows: [0; 32] } + } + + pub fn from_array32x32(arr: &Array32x32, conv: impl Fn(&T) -> bool) -> Self { + let mut rows = [0u32; 32]; + // TODO: we can simd this if compiler doesnt do it alr + for y in 0..32 { + for x in 0..32 { + if conv(&arr[Coordinate::new(x, y)]) { + rows[y as usize] |= 1u32 << x; + } + } + } + Self { rows } + } + + #[inline] + pub const fn get(self, c: Coordinate) -> bool { + (self.rows[c.y as usize] >> c.x) & 1 != 0 + } + + #[inline] + pub fn set_bit(&mut self, c: Coordinate) { + self.rows[c.y as usize] |= 1u32 << c.x; + } + + #[inline] + pub fn clear_bit(&mut self, c: Coordinate) { + self.rows[c.y as usize] &= !(1u32 << c.x); + } + + pub fn and(self, other: Self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = self.rows[i] & other.rows[i]; + } + Self { rows } + } + + pub fn or(self, other: Self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = self.rows[i] | other.rows[i]; + } + Self { rows } + } + + pub fn not(self) -> Self { + let mut rows = [0u32; 32]; + for i in 0..32 { + rows[i] = !self.rows[i]; + } + Self { rows } + } + + pub fn any(self) -> bool { + self.rows.iter().any(|&r| r != 0) + } + + pub fn count_ones(self) -> u32 { + self.rows.iter().map(|r| r.count_ones()).sum() + } + + /// First set bit in row-major order (y then x), or `None` if empty. + pub fn first_set(self) -> Option { + for y in 0u8..32 { + let r = self.rows[y as usize]; + if r != 0 { + let x = r.trailing_zeros() as u8; + return Some(Coordinate::new(x, y)); + } + } + None + } + + pub fn iter_set(self) -> IterSet { + IterSet { + bb: self, + y: 0, + cur_row: self.rows[0], + } + } +} + +pub struct IterSet { + bb: Bitboard, + y: u8, + cur_row: u32, +} + +impl Iterator for IterSet { + type Item = Coordinate; + + fn next(&mut self) -> Option { + loop { + if self.y >= 32 { + return None; + } + if self.cur_row != 0 { + let x = self.cur_row.trailing_zeros() as u8; + let c = Coordinate::new(x, self.y); + self.cur_row &= self.cur_row - 1; + return Some(c); + } + self.y += 1; + if self.y < 32 { + self.cur_row = self.bb.rows[self.y as usize]; + } + } + } +} diff --git a/alpha_paint/src/board/board_impl.rs b/alpha_paint/src/board/board_impl.rs index 04df5c8..2ed9594 100644 --- a/alpha_paint/src/board/board_impl.rs +++ b/alpha_paint/src/board/board_impl.rs @@ -2,11 +2,11 @@ use crate::board::actions::{Action, Move, MoveKind, Paint}; use crate::board::board_structs::{Player, Powerup, TerminalState}; use crate::board::consts::*; use crate::board::structs::{ - Array32x32, Coordinate, DoubleArray32x32, killshot_stamina_cost, min_path_to, + cost_for_movement, min_path_to, Array32x32, Coordinate, DoubleArray32x32, }; use crate::board::tile::Tile; use crate::board::tile_map::TileMap; -use std::cmp::{Ordering, min}; +use std::cmp::{min, Ordering}; use std::sync::Arc; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -30,12 +30,18 @@ impl Rollback { original_coord: Coordinate, original_hill_owner: Option, white_stamina: u16, - black_stamina: u16 + black_stamina: u16, ) -> Self { Self { - original_tile, original_move_count, event_pointer, - consumed_p1: false, consumed_p2: false, - original_hill_owner, original_coord, white_stamina, black_stamina + original_tile, + original_move_count, + event_pointer, + consumed_p1: false, + consumed_p2: false, + original_hill_owner, + original_coord, + white_stamina, + black_stamina, } } } @@ -54,12 +60,13 @@ pub enum ApplyActionOutcome { terminal: TerminalState, play_instead: Action, }, - /// The current player can win by collision via a sequence of moves. - /// The moves lead the current player to the opponent's position for a kill. - /// Only the last move is terminal; the rest are intermediate steps. + /// The player whose turn it is to move (which is evaluated AFTER a potential `Final` move + /// variant) can force a collision win using `actions`. NOTE: this does not necessarily mean + /// that the player who played the most recent action is WINNING and it is important to check + /// the `terminal` value. Killshot { terminal: TerminalState, - moves: Vec, + actions: Vec, }, } @@ -91,7 +98,7 @@ pub struct Board { pub beacon_tile_rollback_stack: Vec<(Tile, Option)>, /// Internally used for rollbacks. - /// Stores a list of (usize, coordinates) for all powerups that were placed at the start + /// Stores a list of (usize, coordinates) for all powerups that were placed at the start /// of the n-th turn (where n is the first number in the tuple) pub powerup_rollback_stack: Vec<(usize, Coordinate)>, @@ -128,52 +135,71 @@ impl Board { self.apply_action_generic::(action) } } - + pub fn rollback(&mut self, action: Action, rollback: Rollback) { // Undo the effects of end-turn (placing powerups) if action.is_final() { - while self.powerup_rollback_stack.last().is_some_and(|&(i, _)| {i == self.turn_count}) { + while self + .powerup_rollback_stack + .last() + .is_some_and(|&(i, _)| i == self.turn_count) + { self.powerups[self.powerup_rollback_stack.pop().unwrap().1] = false; } - + self.turn_count -= 1; } match action { - Action::Move(Move { target, kind: _, place_beacon }) | - Action::FinalMove(Move { target, kind: _, place_beacon }) => { - if rollback.consumed_p2 { self.powerups[target] = true; } - if rollback.consumed_p1 { self.powerups[rollback.original_coord] = true; } + Action::Move(Move { + target, + kind: _, + place_beacon, + }) + | Action::FinalMove(Move { + target, + kind: _, + place_beacon, + }) => { + if rollback.consumed_p2 { + self.powerups[target] = true; + } + if rollback.consumed_p1 { + self.powerups[rollback.original_coord] = true; + } if place_beacon { for &coord in target.region_3x3_with_self().iter().rev() { if !self.in_bounds(coord) || self.tiles[coord].is_wall() { continue; } - let (original_tile, hill_owner) = self.beacon_tile_rollback_stack.pop().unwrap(); + let (original_tile, hill_owner) = + self.beacon_tile_rollback_stack.pop().unwrap(); self.tiles.set(coord, original_tile); self.tiles.override_hill_owner(coord, hill_owner); } } - + self.tiles.set(target, rollback.original_tile); - self.tiles.override_hill_owner(target, rollback.original_hill_owner); + self.tiles + .override_hill_owner(target, rollback.original_hill_owner); } - - Action::Paint(Paint{ target }) | Action::FinalPaint(Paint{ target }) => { + + Action::Paint(Paint { target }) | Action::FinalPaint(Paint { target }) => { self.tiles.set(target, rollback.original_tile); - self.tiles.override_hill_owner(target, rollback.original_hill_owner); + self.tiles + .override_hill_owner(target, rollback.original_hill_owner); } } - + // Now do all the normal rollback things that apply to all moves self.white_stamina = rollback.white_stamina as usize; self.black_stamina = rollback.black_stamina as usize; self.consecutives_moves_so_far = rollback.original_move_count as usize; self.event_pointer = rollback.event_pointer as usize; - - if self.is_white_turn() { - self.white_coord = rollback.original_coord; + + if self.is_white_turn() { + self.white_coord = rollback.original_coord; } else { self.black_coord = rollback.original_coord; } @@ -228,16 +254,10 @@ impl Board { (None, Some(_)) => unreachable!("we never have a finalizer for a non terminal"), }; - // If ongoing, check if current player can win by collision (killshot) - if matches!(outcome, ApplyActionOutcome::Ongoing) && !action.is_final() { - if let Some((ks_terminal, ks_moves)) = self.check_killshot() { - return ( - ApplyActionOutcome::Killshot { - terminal: ks_terminal, - moves: ks_moves, - }, - rollback, - ); + // NOTE: this can be the NEXT player's killshot, because we run this AFTER `end_turn` + if matches!(outcome, ApplyActionOutcome::Ongoing) { + if let Some((terminal, actions)) = self.check_killshot() { + return (ApplyActionOutcome::Killshot { terminal, actions }, rollback); } } @@ -252,11 +272,13 @@ impl Board { let player_coord = self.player_coord::(); let tile = self.tiles[target]; let rollback = Rollback::new( - tile, self.consecutives_moves_so_far as u8, + tile, + self.consecutives_moves_so_far as u8, self.event_pointer as u16, player_coord, self.tiles.get_hill_owner(target), - self.white_stamina as u16, self.black_stamina as u16 + self.white_stamina as u16, + self.black_stamina as u16, ); let stamina = if IS_WHITE { @@ -303,10 +325,13 @@ impl Board { } = mv; let is_collision = target == opponent_coord; let mut rollback = Rollback::new( - self.tiles[target], self.consecutives_moves_so_far as u8, + self.tiles[target], + self.consecutives_moves_so_far as u8, self.event_pointer as u16, - *player_coord, self.tiles.get_hill_owner(target), - self.white_stamina as u16, self.black_stamina as u16 + *player_coord, + self.tiles.get_hill_owner(target), + self.white_stamina as u16, + self.black_stamina as u16, ); let stamina = if IS_WHITE { &mut self.white_stamina @@ -314,7 +339,6 @@ impl Board { &mut self.black_stamina }; - // make sure we can reach it debug_assert!( target.manhattan_dist(*player_coord) == 1 @@ -410,7 +434,8 @@ impl Board { if !self.in_bounds(coord) || self.tiles[coord].is_wall() { continue; } - self.beacon_tile_rollback_stack.push((self.tiles[coord], self.tiles.get_hill_owner(coord))); + self.beacon_tile_rollback_stack + .push((self.tiles[coord], self.tiles.get_hill_owner(coord))); self.tiles.apply_beacon_placement_effect::(coord); } @@ -445,8 +470,9 @@ impl Board { } if !self.powerups[powerup.coord] { self.powerups[powerup.coord] = true; - self.powerup_rollback_stack.push((self.turn_count, powerup.coord)); - } + self.powerup_rollback_stack + .push((self.turn_count, powerup.coord)); + } self.event_pointer += 1; } if IS_WHITE { @@ -499,8 +525,6 @@ impl Board { let white_hills = self.tiles.controlled_hill_count::(); let black_hills = self.tiles.controlled_hill_count::(); let total_hills = self.hills.len(); - // TODO: is this fine to assert lol - assert!(total_hills > 0); let domination_threshold = DOMINATION_WIN_THRESHOLD * (total_hills as f64); let domination_win = if (white_hills as f64) >= domination_threshold { @@ -610,54 +634,245 @@ impl Board { regen } - /// Check if the current player can win by collision (killshot). - /// Returns the terminal state and the sequence of moves to reach the opponent. - fn check_killshot(&self) -> Option<(TerminalState, Vec)> { - let is_white_turn = self.is_white_turn(); - let (player_coord, player_stamina, opponent_coord) = if is_white_turn { - (self.white_coord, self.white_stamina, self.black_coord) + + /// Returns if the player to move (e.g. based on is_white_turn) can win via collision + fn check_killshot(&self) -> Option<(TerminalState, Vec)> { + if self.is_white_turn() { + self.check_killshot_generic::() + } else { + self.check_killshot_generic::() + } + } + + /// Three ways this player (per `IS_WHITE`) can execute a killshot: + /// 1. walk straight into the opponent + /// 2. walk to a beacon, teleport, then walk into the opponent + /// 3. walk, place a beacon, teleport through it, then walk into the opponent + fn check_killshot_generic(&self) -> Option<(TerminalState, Vec)> { + use crate::board::bitboard::Bitboard; + + let moves_so_far = self.consecutives_moves_so_far; + let (player_coord, stamina, opponent_coord, beacons) = if IS_WHITE { + ( + self.white_coord, + self.white_stamina, + self.black_coord, + &self.tiles.white_beacons, + ) } else { - (self.black_coord, self.black_stamina, self.white_coord) + ( + self.black_coord, + self.black_stamina, + self.white_coord, + &self.tiles.black_beacons, + ) }; let pdist = self.dist[(player_coord, opponent_coord)] as usize; + // we r already colliding, handled by terminal() if pdist == 0 { - return None; // already colliding, handled by terminal() + return None; } - // Check if the opponent's tile is favorable for collision - // (if opponent stands on their own paint, they win the collision) - if is_white_turn { + // make sure we'd acc win the collision before doing more work + if IS_WHITE { if self.tiles[opponent_coord].is_owned_by::() { return None; } - } else { - if self.tiles[opponent_coord].is_owned_by::() { - return None; + } else if self.tiles[opponent_coord].is_owned_by::() { + return None; + } + + let win = TerminalState::win_for(IS_WHITE); + + // walk only + let cost_walk = cost_for_movement(pdist, moves_so_far); + if cost_walk <= stamina { + let path = min_path_to(&self.dist, player_coord, opponent_coord); + if path.is_empty() { + return None; // Unreachable but defensive. } + + let moves: Vec = path + .iter() + .map(|&target| Move { + target, + kind: MoveKind::Regular, + place_beacon: false, + }) + .collect(); + + return Some((win, moves_to_killshot_actions(moves))); + } + + // both case 2 and 3 require already placed beacons + if beacons.is_empty() { + return None; + } + + // both beacon_dist lookups use OUR beacons — we want "how far are + // we from our nearest beacon" and "how far is our nearest beacon from + // the opponent" + let (us_to_beacon, opp_to_beacon) = if IS_WHITE { + ( + self.tiles.white_beacon_dist[player_coord], + self.tiles.white_beacon_dist[opponent_coord], + ) + } else { + ( + self.tiles.black_beacon_dist[player_coord], + self.tiles.black_beacon_dist[opponent_coord], + ) + }; + debug_assert!(us_to_beacon != u16::MAX); + debug_assert!(opp_to_beacon != u16::MAX); + + // case 2: walk to beacon + teleport + walk + let consecutive_at_beacon = moves_so_far + us_to_beacon as usize; + let stamina_required = cost_for_movement(us_to_beacon as usize, moves_so_far) + + EXTRA_MOVE_COST * consecutive_at_beacon // teleport cost + + cost_for_movement(opp_to_beacon as usize, 1); + if stamina_required <= stamina { + let (b1, b2) = ( + self.find_beacon_by_dist::(player_coord, us_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"), + self.find_beacon_by_dist::(opponent_coord, opp_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"), + ); + // walk to b1 (may be empty if already on it) + let to_b1 = min_path_to(&self.dist, player_coord, b1); + let mut moves: Vec<_> = to_b1 + .into_iter() + .map(|coord| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: coord, + }) + .collect(); + // teleport from b1 to b2 (appended, not replacing) + moves.push(Move { + kind: MoveKind::BeaconTravel, + place_beacon: false, + target: b2, + }); + // walk from b2 to opponent + let to_opp = min_path_to(&self.dist, b2, opponent_coord); + moves.extend(to_opp.into_iter().map(|coord| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: coord, + })); + + return Some((win, moves_to_killshot_actions(moves))); } - let stamina_cost = killshot_stamina_cost(pdist, self.consecutives_moves_so_far); - if stamina_cost > player_stamina { + // case 3: walk to X, place beacon, teleport to existing beacon near opp, walk to opp + // budget = stamina minus (teleport from X + walk from target_beacon to opp) + // the teleport cost depends on how far we walk to X, so we use a conservative + // lower bound: teleport costs at least EXTRA_MOVE_COST * (moves_so_far + 1) + let min_teleport_cost = EXTRA_MOVE_COST * (moves_so_far + 1); + let stamina_budget = stamina + .saturating_sub(cost_for_movement(opp_to_beacon as usize, 1)) + .saturating_sub(min_teleport_cost); + if stamina_budget < EXTRA_MOVE_COST * moves_so_far { + return None; + } + let moves_we_can_afford = max_possible_move_amt(stamina_budget, moves_so_far) as u16; + if moves_we_can_afford == 0 { return None; } + let reach = Bitboard::from_array32x32(self.dist.get(player_coord), |&d| { + 0 < d && d <= moves_we_can_afford + }); - // We can reach the opponent with enough stamina - compute the path - let path = min_path_to(&self.dist, player_coord, opponent_coord); - if path.is_empty() { - return None; // unreachable but defensive + let placement = if IS_WHITE { + self.tiles.white_placement_valid + } else { + self.tiles.black_placement_valid + }; + let candidates = reach.and(placement); + if !candidates.any() { + return None; } + let target_beacon = self + .find_beacon_by_dist::(opponent_coord, opp_to_beacon) + .expect("couldnt find a beacon that our dist array says exists"); + let target_to_opp = min_path_to(&self.dist, target_beacon, opponent_coord); - let terminal = TerminalState::win_for(is_white_turn); - let moves: Vec = path - .iter() - .map(|&target| Move { - target, + for x in candidates.iter_set() { + if x == player_coord || self.tiles[x].is_beacon() { + continue; + } + let dx = self.dist[(player_coord, x)] as usize; + let consecutive_at_x = moves_so_far + dx; + let total_cost = cost_for_movement(dx, moves_so_far) + + EXTRA_MOVE_COST * consecutive_at_x + + cost_for_movement(opp_to_beacon as usize, 1); + if total_cost > stamina { + continue; + } + + let path_to_x = min_path_to(&self.dist, player_coord, x); + let mut moves: Vec<_> = path_to_x + .into_iter() + .map(|t| Move { + kind: MoveKind::Regular, + place_beacon: false, + target: t, + }) + .collect(); + // place beacon on the last step (landing on X) + moves + .last_mut() + .expect("path_to_x should not be empty") + .place_beacon = true; + // teleport from X (our new beacon) to target_beacon (existing beacon near opp) + moves.push(Move { + kind: MoveKind::BeaconTravel, + place_beacon: false, + target: target_beacon, + }); + // walk from target_beacon to opponent + moves.extend(target_to_opp.iter().map(|&t| Move { kind: MoveKind::Regular, place_beacon: false, - }) - .collect(); + target: t, + })); + return Some((win, moves_to_killshot_actions(moves))); + } - Some((terminal, moves)) + None } + + fn find_beacon_by_dist( + &self, + coord: Coordinate, + dist: u16, + ) -> Option { + self.tiles + .get_beacon_iterator::() + .find(|&&b| self.dist[(coord, b)] == dist) + .copied() + } +} + +/// Walking steps as [`Action::Move`], collision as [`Action::FinalMove`]. +fn moves_to_killshot_actions(mut moves: Vec) -> Vec { + if moves.is_empty() { + return Vec::new(); + } + let last = moves.pop().expect("non-empty"); + let mut out: Vec = moves.into_iter().map(Action::Move).collect(); + out.push(Action::FinalMove(last)); + out +} + +/// Cost is 5n(2c+n-1) <= B. we use quadratic equation to get the following equation we use +/// where n=moves we can afford, B=stamina, c=moves so far +fn max_possible_move_amt(stamina: usize, moves_so_far: usize) -> usize { + let c = moves_so_far as isize; + let b = stamina as isize; + + let n = (5 - 10 * c + ((10 * c - 5).pow(2) + 20 * b).isqrt()) / 10; + n as usize } diff --git a/alpha_paint/src/board/fen.rs b/alpha_paint/src/board/fen.rs index 1bdbd32..512e7e7 100644 --- a/alpha_paint/src/board/fen.rs +++ b/alpha_paint/src/board/fen.rs @@ -1,6 +1,6 @@ use crate::board::board_impl::Board; use crate::board::board_structs::{Player, Powerup}; -use crate::board::structs::{Array32x32, Coordinate, compute_distances}; +use crate::board::structs::{compute_distances, Array32x32, Coordinate}; use crate::board::tile::Tile; use crate::board::tile_map::TileMap; use std::fmt::{self, Display, Write}; @@ -237,6 +237,9 @@ impl Board { let hills = Arc::new(hills_vec); tiles.attach_hills(hill_id.clone(), hills.clone(), hill_owners); + let dist = Arc::new(compute_distances(&tiles)); + tiles.attach_dist(dist.clone(), rows, cols); + Ok(Board { rows, cols, @@ -244,7 +247,7 @@ impl Board { event_pointer, hill_id, hills, - dist: Arc::new(compute_distances(&tiles)), + dist, tiles, powerups, white_coord, diff --git a/alpha_paint/src/board/mod.rs b/alpha_paint/src/board/mod.rs index 0b5a107..9566d30 100644 --- a/alpha_paint/src/board/mod.rs +++ b/alpha_paint/src/board/mod.rs @@ -1,6 +1,7 @@ mod action_generation; #[allow(dead_code)] pub mod actions; +pub mod bitboard; pub mod board_impl; pub mod board_structs; pub mod consts; diff --git a/alpha_paint/src/board/structs.rs b/alpha_paint/src/board/structs.rs index 71f27a2..935bbe0 100644 --- a/alpha_paint/src/board/structs.rs +++ b/alpha_paint/src/board/structs.rs @@ -1,3 +1,4 @@ +use std::array; use std::collections::VecDeque; use std::ops::{Index, IndexMut}; @@ -111,6 +112,12 @@ impl Array32x32 { return default; } } + pub fn splat(value: T) -> Self + where + T: Copy, + { + Array32x32([[value; 32]; 32]) + } } impl Index for Array32x32 { @@ -128,27 +135,44 @@ impl IndexMut for Array32x32 { } #[derive(Debug, Clone)] -pub struct DoubleArray32x32(pub Vec); +pub struct DoubleArray32x32(pub Box>>); impl DoubleArray32x32 { + #[allow(dead_code)] fn flat_index(index: (Coordinate, Coordinate)) -> usize { (((index.0.x as usize) * 32 + index.0.y as usize) * 32 + index.1.x as usize) * 32 + index.1.y as usize } + pub fn get(&self, index: Coordinate) -> &Array32x32 { + &self.0[index] + } + pub fn splat(value: T) -> Self + where + T: Copy, + { + // Allocate on heap to avoid stack overflow (~2MB array) + let vec: Vec<[Array32x32; 32]> = (0..32) + .map(|_| array::from_fn(|_| Array32x32::splat(value))) + .collect(); + let boxed_slice: Box<[[Array32x32; 32]]> = vec.into_boxed_slice(); + assert_eq!(boxed_slice.len(), 32); + let ptr = Box::into_raw(boxed_slice); + let boxed_arr: Box<[[Array32x32; 32]; 32]> = unsafe { Box::from_raw(ptr as *mut _) }; + DoubleArray32x32(Box::new(Array32x32(*boxed_arr))) + } } impl Index<(Coordinate, Coordinate)> for DoubleArray32x32 { type Output = T; fn index(&self, index: (Coordinate, Coordinate)) -> &Self::Output { - &self.0[Self::flat_index(index)] + &self.0[index.0][index.1] } } impl IndexMut<(Coordinate, Coordinate)> for DoubleArray32x32 { fn index_mut(&mut self, index: (Coordinate, Coordinate)) -> &mut Self::Output { - let idx = Self::flat_index(index); - &mut self.0[idx] + &mut self.0[index.0][index.1] } } @@ -188,7 +212,7 @@ pub fn min_path_to( } /// Compute the stamina cost for `pdist` consecutive moves starting from `consecutive_so_far`. -pub fn killshot_stamina_cost(pdist: usize, consecutive_so_far: usize) -> usize { +pub fn cost_for_movement(pdist: usize, consecutive_so_far: usize) -> usize { // Each move i (0-indexed) costs EXTRA_MOVE_COST * (consecutive_so_far + i) // Total = EXTRA_MOVE_COST * sum(consecutive_so_far + i for i in 0..pdist) // = EXTRA_MOVE_COST * (pdist * consecutive_so_far + pdist*(pdist-1)/2) @@ -201,7 +225,7 @@ pub fn compute_distances(tiles: &T) -> DoubleArray32x32 where T: Index, { - let mut dist = DoubleArray32x32(vec![2000; 32 * 32 * 32 * 32]); + let mut dist = DoubleArray32x32::splat(2000); for i in 0..32u8 { for j in 0..32u8 { diff --git a/alpha_paint/src/board/tile_map.rs b/alpha_paint/src/board/tile_map.rs index 246a001..e2c482d 100644 --- a/alpha_paint/src/board/tile_map.rs +++ b/alpha_paint/src/board/tile_map.rs @@ -2,9 +2,10 @@ use std::ops::Index; use std::slice::Iter; use std::sync::Arc; +use crate::board::bitboard::Bitboard; use crate::board::board_structs::{HillData, Player}; use crate::board::consts::HILL_CONTROL_THRESHOLD; -use crate::board::structs::{Array32x32, Coordinate}; +use crate::board::structs::{Array32x32, Coordinate, DoubleArray32x32}; use crate::board::tile::Tile; /// The idea of the TileMap is to provide a useful abstraction layer between the game @@ -21,11 +22,22 @@ pub struct TileMap { coverage_eval: i32, white_tiles: usize, black_tiles: usize, + wall_tiles: usize, hill_id: Arc>, hills: Arc>>, hill_metadata: Vec, - white_beacons: Vec, - black_beacons: Vec, + pub white_beacons: Vec, + pub black_beacons: Vec, + + /// All-pairs walking distances (walls only). Set via [`Self::attach_dist`]. + pub dist: Option>>, + pub rows: u8, + pub cols: u8, + /// Min walking distance from any friendly beacon to each tile (`u16::MAX` if no beacons). + pub white_beacon_dist: Array32x32, + pub black_beacon_dist: Array32x32, + pub white_placement_valid: Bitboard, + pub black_placement_valid: Bitboard, } impl Default for TileMap { @@ -37,6 +49,7 @@ impl Default for TileMap { coverage_eval: 0, white_tiles: 0, black_tiles: 0, + wall_tiles: 0, hill_id: Arc::new(Array32x32 { 0: [[u16::MAX; 32]; 32], }), @@ -44,6 +57,17 @@ impl Default for TileMap { hill_metadata: Vec::new(), white_beacons: Vec::new(), black_beacons: Vec::new(), + dist: None, + rows: 32, + cols: 32, + white_beacon_dist: Array32x32 { + 0: [[u16::MAX; 32]; 32], + }, + black_beacon_dist: Array32x32 { + 0: [[u16::MAX; 32]; 32], + }, + white_placement_valid: Bitboard::empty(), + black_placement_valid: Bitboard::empty(), } } } @@ -77,7 +101,9 @@ impl TileMap { /// Returns None if the hill is unowned or if no hill is present here pub fn get_hill_owner(&self, coord: Coordinate) -> Option { let hill_id = self.hill_id[coord]; - if hill_id == u16::MAX { return None; } + if hill_id == u16::MAX { + return None; + } self.hill_metadata[hill_id as usize].owner } @@ -90,7 +116,9 @@ impl TileMap { /// should own this hill. It is ___highly recommended___ that the set command be used instead. pub fn override_hill_owner(&mut self, coord: Coordinate, hill_owner: Option) { let hill_id = self.hill_id[coord]; - if hill_id == u16::MAX { return; } + if hill_id == u16::MAX { + return; + } self.hill_metadata[hill_id as usize].owner = hill_owner; } @@ -102,6 +130,145 @@ impl TileMap { } } + pub fn on_edge(&self, coord: Coordinate) -> bool { + coord.x == 0 || coord.x + 1 == self.cols || coord.y == 0 || coord.y + 1 == self.rows + } + + /// Must be called after the tile map is fully built; wires walking distances and + /// initializes beacon distance and beacon placement bitboards. + pub fn attach_dist(&mut self, dist: Arc>, rows: u8, cols: u8) { + self.dist = Some(dist); + self.rows = rows; + self.cols = cols; + let d = self.dist.as_ref().unwrap().clone(); + self.recompute_beacon_dist::(d.as_ref()); + self.recompute_beacon_dist::(d.as_ref()); + self.recompute_all_placement_valid(); + } + + fn recompute_beacon_dist(&mut self, dist: &DoubleArray32x32) { + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + let mut m = u16::MAX; + for &b in self.get_beacon_iterator::() { + let dbt = dist[(b, c)]; + if dbt < m { + m = dbt; + } + } + if IS_WHITE { + self.white_beacon_dist[c] = m; + } else { + self.black_beacon_dist[c] = m; + } + } + } + } + + fn apply_beacon_add( + &mut self, + dist: &DoubleArray32x32, + b: Coordinate, + ) { + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + let d = dist[(b, c)]; + if IS_WHITE { + if d < self.white_beacon_dist[c] { + self.white_beacon_dist[c] = d; + } + } else if d < self.black_beacon_dist[c] { + self.black_beacon_dist[c] = d; + } + } + } + } + + fn update_beacon_dist_after_tile_change(&mut self, old: Tile, new: Tile, coord: Coordinate) { + let Some(dist) = self.dist.as_ref().cloned() else { + return; + }; + let d = dist.as_ref(); + if old.is_beacon_of::() && !new.is_beacon_of::() { + self.recompute_beacon_dist::(d); + } else if !old.is_beacon_of::() && new.is_beacon_of::() { + self.apply_beacon_add::(d, coord); + } + if old.is_beacon_of::() && !new.is_beacon_of::() { + self.recompute_beacon_dist::(d); + } else if !old.is_beacon_of::() && new.is_beacon_of::() { + self.apply_beacon_add::(d, coord); + } + } + + fn recompute_all_placement_valid(&mut self) { + self.white_placement_valid = Bitboard::empty(); + self.black_placement_valid = Bitboard::empty(); + for i in 0..32u8 { + for j in 0..32u8 { + let c = Coordinate::new(i, j); + if self.placement_valid_tile::(c) { + self.white_placement_valid.set_bit(c); + } + if self.placement_valid_tile::(c) { + self.black_placement_valid.set_bit(c); + } + } + } + } + + fn placement_valid_tile(&self, target: Coordinate) -> bool { + if !target.in_bounds() { + return false; + } + let tile = self[target]; + if tile.is_wall() { + return false; + } + if self.on_edge(target) { + return false; + } + if tile.is_beacon_of::() { + return false; + } + let (valid_cells, controlled_cells) = self.beacon_window_counts::(target); + if controlled_cells * 3 < valid_cells * 2 { + return false; + } + let mut post = tile; + if tile.is_erasable::() { + post.erase::(); + } else { + post.maybe_erase1::(); + } + let opponent = if IS_WHITE { + Player::Black + } else { + Player::White + }; + post.paint_owner() != Some(opponent) + } + + fn update_placement_neighborhood(&mut self, coord: Coordinate) { + for center in coord.region_3x3_with_self() { + if !center.in_bounds() { + continue; + } + if self.placement_valid_tile::(center) { + self.white_placement_valid.set_bit(center); + } else { + self.white_placement_valid.clear_bit(center); + } + if self.placement_valid_tile::(center) { + self.black_placement_valid.set_bit(center); + } else { + self.black_placement_valid.clear_bit(center); + } + } + } + /// Set the value of a specified tile, and update hill ownership accordingly /// Returns the original tile at that location pub fn set(&mut self, coord: Coordinate, tile: Tile) -> Tile { @@ -122,6 +289,9 @@ impl TileMap { if old.is_owned_by::() { self.black_tiles -= 1; } + if old.is_wall() { + self.wall_tiles -= 1; + } if tile.is_owned_by::() { self.white_tiles += 1; @@ -129,6 +299,9 @@ impl TileMap { if tile.is_owned_by::() { self.black_tiles += 1; } + if tile.is_wall() { + self.wall_tiles += 1; + } if old.is_beacon() || tile.is_beacon() { if old.is_beacon_of::() && !tile.is_beacon_of::() { @@ -149,6 +322,10 @@ impl TileMap { self.tiles[coord] = tile; self.apply_owner_transition(coord, before_owner, after_owner); + if self.dist.is_some() { + self.update_beacon_dist_after_tile_change(old, tile, coord); + self.update_placement_neighborhood(coord); + } old } @@ -194,6 +371,10 @@ impl TileMap { self.coverage_eval } + pub fn wall_count(&self) -> usize { + self.wall_tiles + } + pub fn hill_metadata(&self) -> &[HillData] { &self.hill_metadata } diff --git a/alpha_paint/src/evaluation.rs b/alpha_paint/src/evaluation.rs index 9900b32..297b703 100644 --- a/alpha_paint/src/evaluation.rs +++ b/alpha_paint/src/evaluation.rs @@ -1,107 +1,43 @@ -use crate::board::board_structs::Player; +use std::sync::{Arc, OnceLock}; + use crate::board::Board; -use std::cmp::min; +use crate::nn::model::ValueModel; + +static SHARED_MODEL: OnceLock, String>> = OnceLock::new(); -pub struct Evaluator; +#[derive(Clone)] +pub struct Evaluator { + model: Arc, +} impl Evaluator { - pub fn new(_board: &Board) -> Evaluator { - Evaluator + pub fn new(_board: &Board) -> Result { + let model = load_shared_model()?; + Ok(Evaluator { model }) } - pub fn evaluate(&self, board: &Board) -> i32 { - let mut white_hills = 0; - let mut black_hills = 0; - for md in board.tiles.hill_metadata().iter() { - match md.owner { - Some(Player::White) => white_hills += 1, - Some(Player::Black) => black_hills += 1, - None => (), - }; + pub fn from_model(model: ValueModel) -> Evaluator { + Evaluator { + model: Arc::new(model), } - if white_hills * 4 >= board.hills.len() * 3 { - return 2_000_000_000 - 5 * (board.turn_count as i32); - } else if black_hills * 4 >= board.hills.len() * 3 { - return -2_000_000_000 + 5 * (board.turn_count as i32); - } - - // Deci-tile scale (10 eval = 1 painted tile) - let mut evaluation = 0; - - evaluation += Self::hill_paint_eval(board); // Hill paint gets a large point bonus - evaluation += Self::paint_coverage_eval(board); - - evaluation += board.white_stamina as i32 * 5; - evaluation -= board.black_stamina as i32 * 5; - - evaluation += Self::find_target_hill_eval(board); - - evaluation } - fn find_target_hill_eval(board: &Board) -> i32 { - let mut white_best_score = usize::MAX; - let mut black_best_score = usize::MAX; - - let mut total_eval: i32 = 0; - - for (hill_tiles, hill_data) in board.hills.iter().zip(board.tiles.hill_metadata().iter()) { - let mut white_dist = usize::MAX; - let mut black_dist = usize::MAX; - - for &tile in hill_tiles.iter() { - if !board.tiles[tile].is_owned_by::() { - white_dist = min(white_dist, board.dist[(board.white_coord, tile)] as usize); - } - if !board.tiles[tile].is_owned_by::() { - black_dist = min(black_dist, board.dist[(board.black_coord, tile)] as usize); - } - } - - match hill_data.owner { - Some(Player::White) => { - total_eval += 200_000; - black_best_score = min(black_dist, black_best_score); - } - Some(Player::Black) => { - total_eval -= 200_000; - white_best_score = min(white_dist, white_best_score); - } - None => { - white_best_score = min(white_dist, white_best_score); - black_best_score = min(black_dist, black_best_score); - } - } - } - - if white_best_score != usize::MAX { - total_eval += -20 * (white_best_score as i32); - } - - if black_best_score != usize::MAX { - total_eval -= -20 * (black_best_score as i32); - } - - total_eval + pub fn benchmark() -> Evaluator { + Self::from_model(ValueModel::benchmark_model()) } - fn hill_paint_eval(board: &Board) -> i32 { - let mut evaluation = 0; - - for hill in board.hills.iter() { - for tile in hill.iter() { - if board.tiles[*tile].is_owned_by::() { - evaluation += 25_000; - } else if board.tiles[*tile].is_owned_by::() { - evaluation -= 25_000; - } - } - } - - evaluation + pub fn evaluate(&self, board: &Board) -> f32 { + let value = self.model.evaluate(&crate::nn::features::extract(board)); + if board.is_white_turn() { value } else { -value } } +} - fn paint_coverage_eval(board: &Board) -> i32 { - board.tiles.coverage_eval() * 5 - } +fn load_shared_model() -> Result, String> { + SHARED_MODEL + .get_or_init(|| { + ValueModel::from_embedded() + .map(Arc::new) + .map_err(|err| format!("failed to load embedded model weights: {err}")) + }) + .clone() } diff --git a/alpha_paint/src/lib.rs b/alpha_paint/src/lib.rs index c04d31b..40dd0a6 100644 --- a/alpha_paint/src/lib.rs +++ b/alpha_paint/src/lib.rs @@ -1,16 +1,15 @@ -use std::time::Instant; -use pyo3::prelude::*; use crate::board::{ApplyActionOutcome, Board}; +use pyo3::prelude::*; mod bindings; pub mod board; -mod evaluation; -mod search; - +pub mod evaluation; +pub mod nn; +pub mod search; fn perft(board: &mut Board, depth: usize) -> usize { if depth == 1 { - return board.get_valid_actions().len() + return board.get_valid_actions().len(); } let mut sum = 0; @@ -48,17 +47,24 @@ const PERFT_STR: [&str; 8] = [ "ap2|31x31|tc:0|cm:0|ep:0|w:15,0,99|b:15,30,100|h:n@14,14+16,14+15,15+14,16+16,16;n@0,13+0,14+1,14+2,14+2,15+0,16+1,16+2,16+0,17;n@30,13+28,14+29,14+30,14+28,15+28,16+29,16+30,16+30,17;n@13,6+14,6+15,6+16,6+17,6;n@13,24+14,24+15,24+16,24+17,24;n@7,1+7,8+7,22+7,29;n@23,1+23,8+23,22+23,29|pu:-|ps:2,21,11+2,21,19+2,19,6+2,19,24+2,29,8+2,29,22+2,4,10+2,4,20+2,13,4+2,13,26+2,23,10+2,23,20+2,20,0+2,20,30+2,18,3+2,18,27+2,19,9+2,19,21+2,22,5+2,22,25+52,16,14+52,16,16+52,26,10+52,26,20+52,29,10+52,29,20+52,26,12+52,26,18+52,12,6+52,12,24+102,4,6+102,4,24+102,29,14+102,29,16+102,2,14+102,2,16+102,23,12+102,23,18+102,21,4+102,21,26+152,16,10+152,16,20+152,27,8+152,27,22+152,5,10+152,5,20+152,9,8+152,9,22+152,4,12+152,4,18+202,12,11+202,12,19+202,30,14+202,30,16+202,7,6+202,7,24+202,23,2+202,23,28+202,4,11+202,4,19+252,8,3+252,8,27+252,25,0+252,25,30+252,5,2+252,5,28+252,26,4+252,26,26+252,26,9+252,26,21+302,1,12+302,1,18+302,10,1+302,10,29+302,30,13+302,30,17+302,28,7+302,28,23+302,1,0+302,1,30+352,9,11+352,9,19+352,20,0+352,20,30+352,23,13+352,23,17+352,9,6+352,9,24+352,23,8+352,23,22+402,22,11+402,22,19+402,14,13+402,14,17+402,30,10+402,30,20+402,28,11+402,28,19+402,20,7+402,20,23+452,5,12+452,5,18+452,1,8+452,1,22+452,13,2+452,13,28+452,1,7+452,1,23+452,23,10+452,23,20+502,21,11+502,21,19+502,9,0+502,9,30+502,30,1+502,30,29+502,26,1+502,26,29+502,19,5+502,19,25+552,15,2+552,15,28+552,28,3+552,28,27+552,15,8+552,15,22+552,14,7+552,14,23+552,30,1+552,30,29+602,7,6+602,7,24+602,20,5+602,20,25+602,11,5+602,11,25+602,15,6+602,15,24+602,2,0+602,2,30+652,1,12+652,1,18+652,17,1+652,17,29+652,15,8+652,15,22+652,15,6+652,15,24+652,27,13+652,27,17+702,6,11+702,6,19+702,24,8+702,24,22+702,14,13+702,14,17+702,11,5+702,11,25+702,30,2+702,30,28+752,30,3+752,30,27+752,3,15+752,28,1+752,28,29+752,28,3+752,28,27+752,18,12+752,18,18+802,8,3+802,8,27+802,15,2+802,15,28+802,12,12+802,12,18+802,17,1+802,17,29+802,19,3+802,19,27+852,19,11+852,19,19+852,22,0+852,22,30+852,10,11+852,10,19+852,9,5+852,9,25+852,22,13+852,22,17+902,24,14+902,24,16+902,1,13+902,1,17+902,3,5+902,3,25+902,28,7+902,28,23+902,19,12+902,19,18+952,26,15+952,21,10+952,21,20+952,21,8+952,21,22+952,11,3+952,11,27+952,13,11+952,13,19+1002,24,10+1002,24,20+1002,7,13+1002,7,17+1002,8,14+1002,8,16+1002,22,10+1002,22,20+1002,11,11+1002,11,19+1052,19,14+1052,19,16+1052,11,3+1052,11,27+1052,15,8+1052,15,22+1052,5,12+1052,5,18+1052,17,8+1052,17,22+1102,22,9+1102,22,21+1102,27,12+1102,27,18+1102,14,15+1102,26,1+1102,26,29+1102,23,4+1102,23,26+1152,3,8+1152,3,22+1152,30,3+1152,30,27+1152,17,5+1152,17,25+1152,4,0+1152,4,30+1152,6,11+1152,6,19+1202,17,6+1202,17,24+1202,19,9+1202,19,21+1202,1,0+1202,1,30+1202,14,6+1202,14,24+1202,27,13+1202,27,17+1252,19,5+1252,19,25+1252,18,0+1252,18,30+1252,24,10+1252,24,20+1252,4,2+1252,4,28+1252,21,10+1252,21,20+1302,21,15+1302,8,11+1302,8,19+1302,24,8+1302,24,22+1302,9,9+1302,9,21+1302,8,9+1302,8,21+1352,16,7+1352,16,23+1352,10,11+1352,10,19+1352,2,11+1352,2,19+1352,27,1+1352,27,29+1352,2,3+1352,2,27+1402,3,5+1402,3,25+1402,13,7+1402,13,23+1402,12,13+1402,12,17+1402,10,8+1402,10,22+1402,27,0+1402,27,30+1452,17,6+1452,17,24+1452,5,12+1452,5,18+1452,1,10+1452,1,20+1452,4,6+1452,4,24+1452,8,7+1452,8,23+1502,20,0+1502,20,30+1502,13,1+1502,13,29+1502,20,9+1502,20,21+1502,26,10+1502,26,20+1502,20,15+1552,27,2+1552,27,28+1552,17,15+1552,11,15+1552,21,6+1552,21,24+1552,4,14+1552,4,16+1552,20,1+1552,20,29+1602,23,6+1602,23,24+1602,19,15+1602,17,1+1602,17,29+1602,7,0+1602,7,30+1602,20,1+1602,20,29+1652,19,9+1652,19,21+1652,22,3+1652,22,27+1652,21,9+1652,21,21+1652,3,7+1652,3,23+1652,0,4+1652,0,26+1702,10,9+1702,10,21+1702,11,15+1702,15,6+1702,15,24+1702,30,14+1702,30,16+1702,27,4+1702,27,26+1752,23,8+1752,23,22+1752,21,7+1752,21,23+1752,20,12+1752,20,18+1752,4,2+1752,4,28+1752,10,11+1752,10,19+1802,21,8+1802,21,22+1802,19,7+1802,19,23+1802,11,4+1802,11,26+1802,19,9+1802,19,21+1802,1,14+1802,1,16+1852,20,8+1852,20,22+1852,1,3+1852,1,27+1852,19,14+1852,19,16+1852,12,3+1852,12,27+1852,17,10+1852,17,20+1902,6,3+1902,6,27+1902,30,11+1902,30,19+1902,15,1+1902,15,29+1902,10,8+1902,10,22+1902,30,3+1902,30,27+1952,5,7+1952,5,23+1952,9,5+1952,9,25+1952,30,9+1952,30,21+1952,27,1+1952,27,29+1952,28,4+1952,28,26|bd:#10#7#10#/1#3#5#7#5#3#1/2#9#5#9#2/3#1#7#3#7#1#3/10#1#5#1#10/5#19#5/3#1#4#9#4#1#3/12#5#12/5#6#5#6#5/5#19#5/3#6#9#6#3/5#1#15#1#5/13#####13/5#3#3#3#3#3#5/3#1#19#1#3/##27##/3#1#19#1#3/5#3#3#3#3#3#5/13#####13/5#1#15#1#5/3#6#9#6#3/5#19#5/5#6#5#6#5/12#5#12/3#1#4#9#4#1#3/5#19#5/10#1#5#1#10/3#1#7#3#7#1#3/2#9#5#9#2/1#3#5#7#5#3#1/#10#7#10#", ]; +pub const TRAINING_START_FENS: &[&str; 8] = &PERFT_STR; + pub fn perft_test() { - use std::time::Instant; use crate::board::Board; - use crate::{perft, PERFT_STR}; + use crate::{PERFT_STR, perft}; + use std::time::Instant; for &perft_str in PERFT_STR.iter() { let mut board = Board::from_fen(perft_str).unwrap(); let start = Instant::now(); let res = perft(&mut board, 8); let elapsed = start.elapsed(); - println!("{} - {} ({} nodes/sec)", res, elapsed.as_secs_f32(), res as f32 / elapsed.as_secs_f32()); + println!( + "{} - {} ({} nodes/sec)", + res, + elapsed.as_secs_f32(), + res as f32 / elapsed.as_secs_f32() + ); } } @@ -67,4 +73,3 @@ mod alpha_paint { #[pymodule_export] use crate::bindings::PyBoard; } - diff --git a/alpha_paint/src/nn/features.rs b/alpha_paint/src/nn/features.rs new file mode 100644 index 0000000..1a09a65 --- /dev/null +++ b/alpha_paint/src/nn/features.rs @@ -0,0 +1,260 @@ +use crate::board::Board; +use crate::board::board_structs::Player; +use crate::board::structs::Coordinate; + +pub const BOARD_SIDE: usize = 32; +pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; +pub const BOARD_PLANES: usize = 17; +pub const INTRINSIC_COUNT: usize = 10; +pub const INTRINSIC_SCALE: [f32; INTRINSIC_COUNT] = [ + 420.0, 420.0, 8.0, 8.0, 2000.0, 1024.0, 1024.0, 1024.0, 1024.0, 64.0, +]; + +pub const CURRENT_PAINT_L1: usize = 0; +pub const CURRENT_PAINT_L2: usize = 1; +pub const CURRENT_PAINT_L3: usize = 2; +pub const CURRENT_PAINT_L4: usize = 3; +pub const OPPONENT_PAINT_L1: usize = 4; +pub const OPPONENT_PAINT_L2: usize = 5; +pub const OPPONENT_PAINT_L3: usize = 6; +pub const OPPONENT_PAINT_L4: usize = 7; +pub const WALL: usize = 8; +pub const POWERUP: usize = 9; +pub const BEACON_CURRENT: usize = 10; +pub const BEACON_OPPONENT: usize = 11; +pub const HILL_NEUTRAL: usize = 12; +pub const HILL_CURRENT: usize = 13; +pub const HILL_OPPONENT: usize = 14; +pub const CURRENT_PLAYER: usize = 15; +pub const OPPONENT_PLAYER: usize = 16; + +pub const INTRINSIC_CURRENT_STAMINA: usize = 0; +pub const INTRINSIC_OPPONENT_STAMINA: usize = 1; +pub const INTRINSIC_CURRENT_HILLS: usize = 2; +pub const INTRINSIC_OPPONENT_HILLS: usize = 3; +pub const INTRINSIC_TURN_COUNT: usize = 4; +pub const INTRINSIC_CURRENT_TERRITORY: usize = 5; +pub const INTRINSIC_OPPONENT_TERRITORY: usize = 6; +pub const INTRINSIC_CURRENT_BEACONS: usize = 7; +pub const INTRINSIC_OPPONENT_BEACONS: usize = 8; +pub const INTRINSIC_CONSECUTIVE_MOVES: usize = 9; + +#[derive(Clone, Debug)] +pub struct Features { + pub board: Vec, + pub intrinsics: [f32; INTRINSIC_COUNT], +} + +impl Features { + pub fn zeros() -> Self { + let mut board = vec![0.0; BOARD_PLANES * BOARD_CELLS]; + for y in 0..BOARD_SIDE { + for x in 0..BOARD_SIDE { + board[plane_index(WALL, x as u8, y as u8)] = 1.0; + } + } + Self { + board, + intrinsics: [0.0; INTRINSIC_COUNT], + } + } +} + +pub fn extract(board: &Board) -> Features { + let mut features = Features::zeros(); + fill_board_planes(board, &mut features.board); + features.intrinsics = extract_intrinsics(board); + features +} + +pub fn fill_board_planes(board: &Board, out: &mut [f32]) { + assert_eq!( + out.len(), + BOARD_PLANES * BOARD_CELLS, + "feature board buffer size mismatch" + ); + + out.fill(0.0); + for y in 0..BOARD_SIDE { + for x in 0..BOARD_SIDE { + out[plane_index(WALL, x as u8, y as u8)] = 1.0; + } + } + + let current_is_white = board.is_white_turn(); + let current_coord = board.current_player_coord(); + let opponent_coord = if current_is_white { + board.black_coord + } else { + board.white_coord + }; + + for y in 0..board.rows { + for x in 0..board.cols { + let coord = Coordinate::new(x, y); + let tile = board.tiles[coord]; + let paint = tile.paint_value(); + let strength = paint.unsigned_abs() as usize; + let is_enemy_paint = if current_is_white { + paint < 0 + } else { + paint > 0 + }; + + out[plane_index(WALL, x, y)] = if tile.is_wall() { 1.0 } else { 0.0 }; + out[plane_index(POWERUP, x, y)] = if board.powerups[coord] { 1.0 } else { 0.0 }; + + if strength > 0 { + let offset = if is_enemy_paint { + OPPONENT_PAINT_L1 + } else { + CURRENT_PAINT_L1 + }; + for level in 0..strength.min(4) { + out[plane_index(offset + level, x, y)] = 1.0; + } + } + + match tile.beacon_owner() { + Some(Player::White) if current_is_white => { + out[plane_index(BEACON_CURRENT, x, y)] = 1.0; + } + Some(Player::White) => { + out[plane_index(BEACON_OPPONENT, x, y)] = 1.0; + } + Some(Player::Black) if current_is_white => { + out[plane_index(BEACON_OPPONENT, x, y)] = 1.0; + } + Some(Player::Black) => { + out[plane_index(BEACON_CURRENT, x, y)] = 1.0; + } + None => {} + } + + match board.hill_id[coord] { + u16::MAX => {} + hill_id => match board.tiles.hill_metadata()[hill_id as usize].owner { + None => out[plane_index(HILL_NEUTRAL, x, y)] = 1.0, + Some(Player::White) if current_is_white => { + out[plane_index(HILL_CURRENT, x, y)] = 1.0 + } + Some(Player::White) => out[plane_index(HILL_OPPONENT, x, y)] = 1.0, + Some(Player::Black) if current_is_white => { + out[plane_index(HILL_OPPONENT, x, y)] = 1.0 + } + Some(Player::Black) => out[plane_index(HILL_CURRENT, x, y)] = 1.0, + }, + } + + if coord == current_coord { + out[plane_index(CURRENT_PLAYER, x, y)] = 1.0; + } + if coord == opponent_coord { + out[plane_index(OPPONENT_PLAYER, x, y)] = 1.0; + } + } + } +} + +pub fn extract_intrinsics(board: &Board) -> [f32; INTRINSIC_COUNT] { + let current_is_white = board.is_white_turn(); + let mut out = [0.0; INTRINSIC_COUNT]; + + if current_is_white { + out[INTRINSIC_CURRENT_STAMINA] = board.white_stamina as f32; + out[INTRINSIC_OPPONENT_STAMINA] = board.black_stamina as f32; + out[INTRINSIC_CURRENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + } else { + out[INTRINSIC_CURRENT_STAMINA] = board.black_stamina as f32; + out[INTRINSIC_OPPONENT_STAMINA] = board.white_stamina as f32; + out[INTRINSIC_CURRENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_OPPONENT_HILLS] = board.tiles.controlled_hill_count::() as f32; + out[INTRINSIC_CURRENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_OPPONENT_TERRITORY] = board.tiles.territory_count::() as f32; + out[INTRINSIC_CURRENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + out[INTRINSIC_OPPONENT_BEACONS] = board.tiles.get_beacon_iterator::().count() as f32; + } + + out[INTRINSIC_TURN_COUNT] = board.turn_count as f32; + out[INTRINSIC_CONSECUTIVE_MOVES] = board.consecutives_moves_so_far as f32; + + for (value, scale) in out.iter_mut().zip(INTRINSIC_SCALE) { + *value /= scale; + } + + out +} + +#[inline] +pub const fn plane_index(plane: usize, x: u8, y: u8) -> usize { + plane * BOARD_CELLS + y as usize * BOARD_SIDE + x as usize +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_features_for_white_turn() { + let board = Board::from_fen( + "ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:1B1/d1O/3", + ) + .unwrap(); + + let features = extract(&board); + + assert_eq!(features.board[plane_index(CURRENT_PAINT_L1, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PAINT_L2, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L4, 0, 1)], 1.0); + assert_eq!(features.board[plane_index(POWERUP, 0, 1)], 1.0); + assert_eq!(features.board[plane_index(HILL_NEUTRAL, 1, 1)], 1.0); + assert_eq!(features.board[plane_index(BEACON_CURRENT, 2, 1)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PLAYER, 0, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 2, 2)], 1.0); + assert_eq!(features.board[plane_index(WALL, 31, 31)], 1.0); + + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 99.0 / 420.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 88.0 / 420.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_CURRENT_TERRITORY], + 1.0 / 1024.0 + ); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_TERRITORY], + 1.0 / 1024.0 + ); + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_BEACONS], 1.0 / 1024.0); + assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 2.0 / 64.0); + } + + #[test] + fn extracts_features_for_black_turn() { + let board = Board::from_fen( + "ap2|3x3|tc:1|cm:5|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/2o", + ) + .unwrap(); + + let features = extract(&board); + + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L1, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PAINT_L2, 1, 0)], 1.0); + assert_eq!(features.board[plane_index(HILL_CURRENT, 1, 1)], 1.0); + assert_eq!(features.board[plane_index(CURRENT_PLAYER, 2, 2)], 1.0); + assert_eq!(features.board[plane_index(OPPONENT_PLAYER, 0, 0)], 1.0); + + assert_eq!(features.intrinsics[INTRINSIC_CURRENT_STAMINA], 88.0 / 420.0); + assert_eq!( + features.intrinsics[INTRINSIC_OPPONENT_STAMINA], + 99.0 / 420.0 + ); + assert_eq!(features.intrinsics[INTRINSIC_CONSECUTIVE_MOVES], 5.0 / 64.0); + } +} diff --git a/alpha_paint/src/nn/mod.rs b/alpha_paint/src/nn/mod.rs new file mode 100644 index 0000000..c572841 --- /dev/null +++ b/alpha_paint/src/nn/mod.rs @@ -0,0 +1,2 @@ +pub mod features; +pub mod model; diff --git a/alpha_paint/src/nn/model.rs b/alpha_paint/src/nn/model.rs new file mode 100644 index 0000000..db52097 --- /dev/null +++ b/alpha_paint/src/nn/model.rs @@ -0,0 +1,575 @@ +use std::cell::UnsafeCell; + +use safetensors::tensor::Metadata; +use safetensors::{Dtype, SafeTensors}; + +use crate::nn::features::{BOARD_CELLS, BOARD_PLANES, BOARD_SIDE, Features, INTRINSIC_COUNT}; + +const DEFAULT_EPS: f32 = 1.0e-5; +const MODEL_WIDTH: usize = 8; +const MODEL_BLOCKS: usize = 1; +const MODEL_HIDDEN_DIM: usize = 32; +const KERNEL_SIZE: usize = 9; +const ACTIVATION_LEN: usize = MODEL_WIDTH * BOARD_CELLS; +const STEM_WEIGHT_LEN: usize = MODEL_WIDTH * BOARD_PLANES * KERNEL_SIZE; +const BLOCK_WEIGHT_LEN: usize = MODEL_WIDTH * MODEL_WIDTH * KERNEL_SIZE; +const HEAD0_IN_DIM: usize = MODEL_WIDTH + INTRINSIC_COUNT; +const HEAD0_WEIGHT_LEN: usize = MODEL_HIDDEN_DIM * HEAD0_IN_DIM; + +thread_local! { + static SCRATCH: UnsafeCell = const { UnsafeCell::new(Scratch::new()) }; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ModelConfig { + pub width: usize, + pub num_blocks: usize, + pub hidden_dim: usize, +} + +impl Default for ModelConfig { + fn default() -> Self { + Self { + width: MODEL_WIDTH, + num_blocks: MODEL_BLOCKS, + hidden_dim: MODEL_HIDDEN_DIM, + } + } +} + +#[derive(Clone, Debug)] +pub struct ValueModel { + pub config: ModelConfig, + stem_weight: [f32; STEM_WEIGHT_LEN], + stem_bias: [f32; MODEL_WIDTH], + block0_norm1_scale: [f32; MODEL_WIDTH], + block0_norm1_bias: [f32; MODEL_WIDTH], + block0_conv1_weight: [f32; BLOCK_WEIGHT_LEN], + block0_conv1_bias: [f32; MODEL_WIDTH], + block0_conv2_weight: [f32; BLOCK_WEIGHT_LEN], + head0_weight: [f32; HEAD0_WEIGHT_LEN], + head0_bias: [f32; MODEL_HIDDEN_DIM], + head1_weight: [f32; MODEL_HIDDEN_DIM], + head1_bias: f32, +} + +struct Scratch { + a: [f32; ACTIVATION_LEN], + b: [f32; ACTIVATION_LEN], + residual: [f32; ACTIVATION_LEN], + pooled: [f32; MODEL_WIDTH], + head_input: [f32; HEAD0_IN_DIM], + hidden: [f32; MODEL_HIDDEN_DIM], +} + +impl Scratch { + const fn new() -> Self { + Self { + a: [0.0; ACTIVATION_LEN], + b: [0.0; ACTIVATION_LEN], + residual: [0.0; ACTIVATION_LEN], + pooled: [0.0; MODEL_WIDTH], + head_input: [0.0; HEAD0_IN_DIM], + hidden: [0.0; MODEL_HIDDEN_DIM], + } + } +} + +struct BatchNormStats { + weight: [f32; C], + bias: [f32; C], + running_mean: [f32; C], + running_var: [f32; C], +} + +impl ValueModel { + pub fn zeroed(config: ModelConfig) -> Self { + validate_config(config).unwrap(); + Self { + config, + stem_weight: [0.0; STEM_WEIGHT_LEN], + stem_bias: [0.0; MODEL_WIDTH], + block0_norm1_scale: [1.0; MODEL_WIDTH], + block0_norm1_bias: [0.0; MODEL_WIDTH], + block0_conv1_weight: [0.0; BLOCK_WEIGHT_LEN], + block0_conv1_bias: [0.0; MODEL_WIDTH], + block0_conv2_weight: [0.0; BLOCK_WEIGHT_LEN], + head0_weight: [0.0; HEAD0_WEIGHT_LEN], + head0_bias: [0.0; MODEL_HIDDEN_DIM], + head1_weight: [0.0; MODEL_HIDDEN_DIM], + head1_bias: 0.0, + } + } + + pub fn benchmark_model() -> Self { + Self::zeroed(ModelConfig::default()) + } + + pub fn from_embedded() -> Result { + Self::from_bytes(include_bytes!(concat!( + env!("CARGO_MANIFEST_DIR"), + "/assets/value_net.safetensors" + ))) + } + + pub fn from_bytes(bytes: &[u8]) -> Result { + if bytes == b"placeholder\n" || bytes == b"placeholder" { + return Err( + "embedded weights are still the placeholder file; run python/scripts/export_value_net.py to overwrite alpha_paint/assets/value_net.safetensors".to_string(), + ); + } + let (_, metadata) = SafeTensors::read_metadata(bytes) + .map_err(|err| format!("failed to read safetensors metadata: {err}"))?; + let config = load_config(&metadata)?; + let tensors = SafeTensors::deserialize(bytes) + .map_err(|err| format!("failed to deserialize safetensors: {err}"))?; + Self::from_safetensors(&tensors, config) + } + + pub fn from_safetensors( + tensors: &SafeTensors<'_>, + config: ModelConfig, + ) -> Result { + validate_config(config)?; + + let stem_weight = tensor_array::( + tensors, + "stem.conv.weight", + &[MODEL_WIDTH, BOARD_PLANES, 3, 3], + )?; + let stem_bn = batch_norm_from_prefix::(tensors, "stem.bn")?; + let (stem_weight, stem_bias) = fold_conv_bn(stem_weight, stem_bn, BOARD_PLANES); + + let block0_norm1 = batch_norm_from_prefix::(tensors, "blocks.0.norm1")?; + let (block0_norm1_scale, block0_norm1_bias) = bn_to_affine(block0_norm1); + + let block0_conv1_weight = tensor_array::( + tensors, + "blocks.0.conv1.weight", + &[MODEL_WIDTH, MODEL_WIDTH, 3, 3], + )?; + let block0_norm2 = batch_norm_from_prefix::(tensors, "blocks.0.norm2")?; + let (block0_conv1_weight, block0_conv1_bias) = + fold_conv_bn(block0_conv1_weight, block0_norm2, MODEL_WIDTH); + + let block0_conv2_weight = tensor_array::( + tensors, + "blocks.0.conv2.weight", + &[MODEL_WIDTH, MODEL_WIDTH, 3, 3], + )?; + + let head0_weight = tensor_array::( + tensors, + "head.fc1.weight", + &[MODEL_HIDDEN_DIM, HEAD0_IN_DIM], + )?; + let head0_bias = + tensor_array::(tensors, "head.fc1.bias", &[MODEL_HIDDEN_DIM])?; + let head1_weight = + tensor_array::(tensors, "head.fc2.weight", &[1, MODEL_HIDDEN_DIM])?; + let head1_bias = tensor_array::<1>(tensors, "head.fc2.bias", &[1])?[0]; + + Ok(Self { + config, + stem_weight, + stem_bias, + block0_norm1_scale, + block0_norm1_bias, + block0_conv1_weight, + block0_conv1_bias, + block0_conv2_weight, + head0_weight, + head0_bias, + head1_weight, + head1_bias, + }) + } + + #[inline(always)] + pub fn evaluate(&self, features: &Features) -> f32 { + SCRATCH.with(|scratch| { + // SAFETY: each thread has its own scratch buffer, and evaluate does not retain + // references beyond this call. + let scratch = unsafe { &mut *scratch.get() }; + self.forward(&features.board, &features.intrinsics, scratch) + }) + } + + #[inline(always)] + fn forward( + &self, + board: &[f32], + intrinsics: &[f32; INTRINSIC_COUNT], + scratch: &mut Scratch, + ) -> f32 { + debug_assert_eq!(board.len(), BOARD_PLANES * BOARD_CELLS); + + conv3x3_bias::( + &self.stem_weight, + &self.stem_bias, + board, + &mut scratch.a, + ); + relu_inplace(&mut scratch.a); + + scratch.residual.copy_from_slice(&scratch.a); + affine_relu_inplace::( + &mut scratch.a, + &self.block0_norm1_scale, + &self.block0_norm1_bias, + ); + conv3x3_bias::( + &self.block0_conv1_weight, + &self.block0_conv1_bias, + &scratch.a, + &mut scratch.b, + ); + relu_inplace(&mut scratch.b); + conv3x3_no_bias::( + &self.block0_conv2_weight, + &scratch.b, + &mut scratch.a, + ); + for idx in 0..ACTIVATION_LEN { + scratch.a[idx] += scratch.residual[idx]; + } + + global_avg_pool::(&scratch.a, &mut scratch.pooled); + scratch.head_input[..MODEL_WIDTH].copy_from_slice(&scratch.pooled); + scratch.head_input[MODEL_WIDTH..].copy_from_slice(intrinsics); + linear_relu::( + &self.head0_weight, + &self.head0_bias, + &scratch.head_input, + &mut scratch.hidden, + ); + linear_scalar::(&self.head1_weight, self.head1_bias, &scratch.hidden) + } +} + +#[inline(always)] +fn batch_norm_from_prefix( + tensors: &SafeTensors<'_>, + prefix: &str, +) -> Result, String> { + Ok(BatchNormStats { + weight: tensor_array::(tensors, &format!("{prefix}.weight"), &[C])?, + bias: tensor_array::(tensors, &format!("{prefix}.bias"), &[C])?, + running_mean: tensor_array::(tensors, &format!("{prefix}.running_mean"), &[C])?, + running_var: tensor_array::(tensors, &format!("{prefix}.running_var"), &[C])?, + }) +} + +#[inline(always)] +fn bn_to_affine(bn: BatchNormStats) -> ([f32; C], [f32; C]) { + let mut scale = [0.0; C]; + let mut bias = [0.0; C]; + for channel in 0..C { + let channel_scale = bn.weight[channel] / (bn.running_var[channel] + DEFAULT_EPS).sqrt(); + scale[channel] = channel_scale; + bias[channel] = bn.bias[channel] - bn.running_mean[channel] * channel_scale; + } + (scale, bias) +} + +#[inline(always)] +fn fold_conv_bn( + mut conv_weight: [f32; W], + bn: BatchNormStats, + in_channels: usize, +) -> ([f32; W], [f32; C]) { + let (scale, bias) = bn_to_affine(bn); + let channel_stride = in_channels * KERNEL_SIZE; + for channel in 0..C { + let start = channel * channel_stride; + let end = start + channel_stride; + for value in &mut conv_weight[start..end] { + *value *= scale[channel]; + } + } + (conv_weight, bias) +} + +#[inline(always)] +fn conv3x3_bias( + weight: &[f32], + bias: &[f32; OUT_C], + input: &[f32], + output: &mut [f32; ACTIVATION_LEN], +) { + for out_channel in 0..OUT_C { + let out_plane = &mut output[out_channel * BOARD_CELLS..(out_channel + 1) * BOARD_CELLS]; + out_plane.fill(bias[out_channel]); + accumulate_conv3x3::(weight, out_channel, input, out_plane); + } +} + +#[inline(always)] +fn conv3x3_no_bias( + weight: &[f32], + input: &[f32; ACTIVATION_LEN], + output: &mut [f32; ACTIVATION_LEN], +) { + for out_channel in 0..OUT_C { + let out_plane = &mut output[out_channel * BOARD_CELLS..(out_channel + 1) * BOARD_CELLS]; + out_plane.fill(0.0); + accumulate_conv3x3::(weight, out_channel, input, out_plane); + } +} + +#[inline(always)] +fn accumulate_conv3x3( + weight: &[f32], + out_channel: usize, + input: &[f32], + out_plane: &mut [f32], +) { + let out_ptr = out_plane.as_mut_ptr(); + for in_channel in 0..IN_C { + let in_plane = &input[in_channel * BOARD_CELLS..(in_channel + 1) * BOARD_CELLS]; + let kernel = &weight[(out_channel * IN_C + in_channel) * KERNEL_SIZE..][..KERNEL_SIZE]; + let k00 = kernel[0]; + let k01 = kernel[1]; + let k02 = kernel[2]; + let k10 = kernel[3]; + let k11 = kernel[4]; + let k12 = kernel[5]; + let k20 = kernel[6]; + let k21 = kernel[7]; + let k22 = kernel[8]; + let in_ptr = in_plane.as_ptr(); + + unsafe { + for y in 1..BOARD_SIDE - 1 { + let row = y * BOARD_SIDE; + let row_up = row - BOARD_SIDE; + let row_down = row + BOARD_SIDE; + for x in 1..BOARD_SIDE - 1 { + let idx = row + x; + *out_ptr.add(idx) += *in_ptr.add(row_up + x - 1) * k00 + + *in_ptr.add(row_up + x) * k01 + + *in_ptr.add(row_up + x + 1) * k02 + + *in_ptr.add(row + x - 1) * k10 + + *in_ptr.add(row + x) * k11 + + *in_ptr.add(row + x + 1) * k12 + + *in_ptr.add(row_down + x - 1) * k20 + + *in_ptr.add(row_down + x) * k21 + + *in_ptr.add(row_down + x + 1) * k22; + } + } + } + + for x in 0..BOARD_SIDE { + out_plane[x] += conv3x3_border(in_plane, kernel, x, 0); + out_plane[(BOARD_SIDE - 1) * BOARD_SIDE + x] += + conv3x3_border(in_plane, kernel, x, BOARD_SIDE - 1); + } + for y in 1..BOARD_SIDE - 1 { + let row = y * BOARD_SIDE; + out_plane[row] += conv3x3_border(in_plane, kernel, 0, y); + out_plane[row + BOARD_SIDE - 1] += conv3x3_border(in_plane, kernel, BOARD_SIDE - 1, y); + } + } +} + +#[inline(always)] +fn conv3x3_border(in_plane: &[f32], kernel: &[f32], x: usize, y: usize) -> f32 { + let mut sum = 0.0; + for ky in 0..3 { + let input_y = y as isize + ky as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&input_y) { + continue; + } + let row = input_y as usize * BOARD_SIDE; + for kx in 0..3 { + let input_x = x as isize + kx as isize - 1; + if !(0..BOARD_SIDE as isize).contains(&input_x) { + continue; + } + sum += in_plane[row + input_x as usize] * kernel[ky * 3 + kx]; + } + } + sum +} + +#[inline(always)] +fn affine_relu_inplace( + values: &mut [f32; ACTIVATION_LEN], + scale: &[f32; C], + bias: &[f32; C], +) { + for channel in 0..C { + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + let channel_scale = scale[channel]; + let channel_bias = bias[channel]; + for value in &mut values[start..end] { + *value = (*value * channel_scale + channel_bias).max(0.0); + } + } +} + +#[inline(always)] +fn relu_inplace(values: &mut [f32]) { + for value in values { + *value = value.max(0.0); + } +} + +#[inline(always)] +fn global_avg_pool(input: &[f32; ACTIVATION_LEN], output: &mut [f32; C]) { + let scale = 1.0 / BOARD_CELLS as f32; + for channel in 0..C { + let start = channel * BOARD_CELLS; + let end = start + BOARD_CELLS; + let mut sum = 0.0; + for value in &input[start..end] { + sum += *value; + } + output[channel] = sum * scale; + } +} + +#[inline(always)] +fn linear_relu( + weight: &[f32], + bias: &[f32; OUT], + input: &[f32; IN], + output: &mut [f32; OUT], +) { + for out_idx in 0..OUT { + let mut sum = bias[out_idx]; + let row = &weight[out_idx * IN..(out_idx + 1) * IN]; + for in_idx in 0..IN { + sum += row[in_idx] * input[in_idx]; + } + output[out_idx] = sum.max(0.0); + } +} + +#[inline(always)] +fn linear_scalar(weight: &[f32; IN], bias: f32, input: &[f32; IN]) -> f32 { + let mut sum = bias; + for idx in 0..IN { + sum += weight[idx] * input[idx]; + } + sum +} + +fn validate_config(config: ModelConfig) -> Result<(), String> { + if config == ModelConfig::default() { + Ok(()) + } else { + Err(format!( + "only the default value model is supported (got width={}, num_blocks={}, hidden_dim={})", + config.width, config.num_blocks, config.hidden_dim + )) + } +} + +fn load_config(metadata: &Metadata) -> Result { + let metadata = metadata + .metadata() + .as_ref() + .ok_or_else(|| "missing safetensors metadata".to_string())?; + let parse_usize = |key: &str| -> Result { + metadata + .get(key) + .ok_or_else(|| format!("missing metadata key {key}"))? + .parse::() + .map_err(|err| format!("invalid metadata {key}: {err}")) + }; + Ok(ModelConfig { + width: parse_usize("width")?, + num_blocks: parse_usize("num_blocks")?, + hidden_dim: parse_usize("hidden_dim")?, + }) +} + +fn tensor_array( + tensors: &SafeTensors<'_>, + name: &str, + expected_shape: &[usize], +) -> Result<[f32; N], String> { + let values = tensor_f32(tensors, name, expected_shape)?; + values + .try_into() + .map_err(|got: Vec| format!("tensor {name} has {} values, expected {N}", got.len())) +} + +fn tensor_f32( + tensors: &SafeTensors<'_>, + name: &str, + expected_shape: &[usize], +) -> Result, String> { + let tensor = tensors + .tensor(name) + .map_err(|err| format!("missing tensor {name}: {err}"))?; + if tensor.dtype() != Dtype::F32 { + return Err(format!( + "tensor {name} has dtype {:?}, expected F32", + tensor.dtype() + )); + } + if tensor.shape() != expected_shape { + return Err(format!( + "tensor {name} has shape {:?}, expected {:?}", + tensor.shape(), + expected_shape + )); + } + let data = tensor.data(); + if data.len() % 4 != 0 { + return Err(format!( + "tensor {name} has invalid byte length {}", + data.len() + )); + } + let mut out = Vec::with_capacity(data.len() / 4); + for chunk in data.chunks_exact(4) { + out.push(f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]])); + } + Ok(out) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::nn::features::{Features, plane_index}; + + #[test] + fn forward_matches_simple_hand_computation() { + let mut model = ValueModel::zeroed(ModelConfig::default()); + model.stem_weight[4] = 1.0; + model.head0_weight[0] = 1.0; + model.head1_weight[0] = 0.5; + + let mut features = Features::zeros(); + for y in 0..BOARD_SIDE as u8 { + for x in 0..BOARD_SIDE as u8 { + features.board[plane_index(0, x, y)] = 1.0; + } + } + + let expected = 0.5f32; + assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); + } + + #[test] + fn residual_block_preserves_input_when_convs_are_zero() { + let mut model = ValueModel::zeroed(ModelConfig::default()); + model.stem_weight[4] = 1.0; + model.head0_weight[0] = 1.0; + model.head1_weight[0] = 1.0; + + let mut features = Features::zeros(); + for y in 0..BOARD_SIDE as u8 { + for x in 0..BOARD_SIDE as u8 { + features.board[plane_index(0, x, y)] = 1.0; + } + } + + let expected = 1.0f32; + assert!((model.evaluate(&features) - expected).abs() < 1.0e-5); + } +} diff --git a/alpha_paint/src/search.rs b/alpha_paint/src/search.rs index 8e317f5..ec06da5 100644 --- a/alpha_paint/src/search.rs +++ b/alpha_paint/src/search.rs @@ -1,14 +1,18 @@ -use crate::board::actions::Move; -use crate::board::{Action, ApplyActionOutcome, Board, TerminalState}; -use crate::evaluation::Evaluator; +use std::cmp::Ordering; +use std::time::Duration; + use rand::rngs::SmallRng; use rand::{Rng, SeedableRng}; -use std::time::Duration; + +use crate::board::{Action, ApplyActionOutcome, Board, TerminalState}; +use crate::evaluation::Evaluator; + +const AVG_GAME_LENGTH: f32 = 500.0; #[derive(Debug)] pub struct ChildData { action: Action, - child_value: i32, + child_value: f32, entrance_count: usize, pub node: Option>, } @@ -16,7 +20,7 @@ pub struct ChildData { impl ChildData { fn completion_value(&self) -> i32 { match &self.node { - Some(c) => c.completion_value, + Some(child) => child.completion_value, None => 0, } } @@ -24,14 +28,14 @@ impl ChildData { #[derive(Debug)] pub struct SearchNode { - pub value: i32, + pub value: f32, pub completion_value: i32, resolved: bool, pub children: Vec, } impl SearchNode { - fn new(value: i32, completion_value: i32, is_resolved: bool) -> SearchNode { + fn new(value: f32, completion_value: i32, is_resolved: bool) -> SearchNode { SearchNode { value, completion_value, @@ -40,74 +44,86 @@ impl SearchNode { } } - /// Prefers moves with low entrance counts fn completed_best_action_dual( &self, is_max_player: bool, rng: &mut SmallRng, ) -> (usize, Action) { - let res = if is_max_player { + let result = if is_max_player { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { + .max_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action_dual without any valid actions!") + .expect("called completed_best_action_dual without any valid actions") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .min_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action_dual without any valid actions!") + .expect("called completed_best_action_dual without any valid actions") }; - - (res.0, res.1.action) + (result.0, result.1.action) } - /// Prefers moves with high entrance counts fn completed_best_action(&self, is_max_player: bool, rng: &mut SmallRng) -> (usize, Action) { - let res = if is_max_player { + let result = if is_max_player { self.children .iter() .enumerate() - .max_by_key(|&(_, child)| { - ( - child.completion_value(), - child.child_value, - child.entrance_count, - rng.next_u32(), - ) + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action without any valid actions!") + .expect("called completed_best_action without any valid actions") } else { self.children .iter() .enumerate() - .min_by_key(|&(_, child)| { + .min_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - child.child_value, - -(child.entrance_count as i32), - rng.next_u32(), + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) }) - .expect("Called completed_best_action without any valid actions!") + .expect("called completed_best_action without any valid actions") }; - - (res.0, res.1.action) + (result.0, result.1.action) } fn backup_resolution(&self) -> bool { @@ -118,42 +134,34 @@ impl SearchNode { child .node .as_ref() - .and_then(|c| Some(c.resolved)) + .map(|node| node.resolved) .unwrap_or(false) }) } } - /// Build a chain of resolved SearchNodes for a killshot move sequence. - /// All nodes are resolved. Only the final node is terminal (no children). - fn build_killshot_chain(board: &Board, terminal: TerminalState, moves: &[Move]) -> SearchNode { + fn build_killshot_chain( + board: &Board, + terminal: TerminalState, + actions: &[Action], + ) -> SearchNode { let term_value = Self::value_from_term(board, terminal); - let comp_value = terminal.value(); - - // Start with the terminal leaf (the collision result) - let mut node = SearchNode::new(term_value, comp_value, true); - - // Build chain from last move to first - for i in (0..moves.len()).rev() { - let action = if i == moves.len() - 1 { - Action::FinalMove(moves[i]) - } else { - Action::Move(moves[i]) - }; + let completion_value = terminal.value(); + let mut node = SearchNode::new(term_value, completion_value, true); - let parent = SearchNode { + for index in (0..actions.len()).rev() { + let action = actions[index]; + node = SearchNode { value: term_value, - completion_value: comp_value, + completion_value, resolved: true, children: vec![ChildData { action, - child_value: comp_value, + child_value: term_value, entrance_count: 0, node: Some(Box::new(node)), }], }; - - node = parent; } node @@ -167,7 +175,7 @@ impl SearchNode { ) -> SearchNode { match outcome { ApplyActionOutcome::Ongoing => { - let mut new_node = SearchNode::new(0, 0, false); + let mut new_node = SearchNode::new(0.0, 0, false); let actions = board.get_valid_actions(); new_node.children.reserve(actions.len()); @@ -178,8 +186,8 @@ impl SearchNode { for action in actions.into_iter().copied() { let mut local_board = board.clone(); - let (outcome, _) = local_board.apply_action(action); - match outcome { + let (child_outcome, _) = local_board.apply_action(action); + match child_outcome { ApplyActionOutcome::Ongoing => { new_node.children.push(ChildData { action, @@ -188,29 +196,29 @@ impl SearchNode { node: None, }); } - // we can only guarantee that the play instead action is a valid action if - // its a Move (e.g. collision) or a Paint after a move - // Therefore, we just do the same thing as we did above, where we just - // treat this move as a terminal move too, and handle it when we consume - // the SearchTree (e.g. in bindings.rs) ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => { + let child_value = Self::value_from_term(&local_board, terminal); new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value, entrance_count: 0, node: Some(Box::new(SearchNode::new( - Self::value_from_term(&local_board, terminal), + child_value, terminal.value(), true, ))), }); } - ApplyActionOutcome::Killshot { terminal, moves } => { - let chain = Self::build_killshot_chain(&local_board, terminal, &moves); + ApplyActionOutcome::Killshot { + terminal, + actions: ks_actions, + } => { + let chain = + Self::build_killshot_chain(&local_board, terminal, &ks_actions); new_node.children.push(ChildData { action, - child_value: terminal.value(), + child_value: chain.value, entrance_count: 0, node: Some(Box::new(chain)), }); @@ -225,24 +233,41 @@ impl SearchNode { new_node.resolved = new_node.backup_resolution(); new_node } - // in search, we don't really CARE about the distinction between these two. - // e.g. the search tree doesn't _really_ care that we need to play a `Final` variant of - // the action instead. It just cares that this is a terminal state. - // We can just fix this in the consumers of the search tree, e.g. in bindings.rs ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } => SearchNode::new( Self::value_from_term(board, terminal), terminal.value(), true, ), - ApplyActionOutcome::Killshot { terminal, moves } => { - Self::build_killshot_chain(board, terminal, &moves) + ApplyActionOutcome::Killshot { terminal, actions } => { + Self::build_killshot_chain(board, terminal, &actions) } } } - fn value_from_term(board: &Board, term: TerminalState) -> i32 { - term.value() * (2_000_000_000 - 5 * (board.turn_count as i32)) + fn value_from_term(board: &Board, term: TerminalState) -> f32 { + let sign = term.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + + let progress = board.turn_count.max(1) as f32; + let depth = (AVG_GAME_LENGTH / progress).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) } fn create_child( @@ -251,20 +276,17 @@ impl SearchNode { action: Action, evaluator: &Evaluator, rng: &mut SmallRng, - ) -> i32 { + ) -> f32 { let (outcome, _) = state.apply_action(action); - let node = Box::new(SearchNode::build_self(&state, outcome, evaluator, rng)); let value = node.value; - - if let Some(id) = self + if let Some(index) = self .children .iter() .position(|child| child.action == action) { - self.children[id].node = Some(node); + self.children[index].node = Some(node); } - value } @@ -274,12 +296,11 @@ impl SearchNode { outcome: ApplyActionOutcome, evaluator: &Evaluator, rng: &mut SmallRng, - ) -> i32 { + ) -> f32 { let white_turn = state.is_white_turn(); - match outcome { ApplyActionOutcome::Ongoing => { - if self.children.len() == 0 { + if self.children.is_empty() { let loss = TerminalState::loss_for(white_turn); self.resolved = true; self.completion_value = loss.value(); @@ -290,12 +311,11 @@ impl SearchNode { if !self.resolved { let (best_action_id, best_action) = self.completed_best_action_dual(white_turn, rng); - self.children[best_action_id].entrance_count += 1; - if let Some(child_val) = self.children[best_action_id].node.as_mut() { - let (outcome, _) = state.apply_action(best_action); - child_val.ubfms_iteration(state, outcome, evaluator, rng); + if let Some(child) = self.children[best_action_id].node.as_mut() { + let (next_outcome, _) = state.apply_action(best_action); + child.ubfms_iteration(state, next_outcome, evaluator, rng); } else { self.children[best_action_id].child_value = self.create_child(state, best_action, evaluator, rng); @@ -307,7 +327,6 @@ impl SearchNode { self.resolved = self.backup_resolution(); } } - ApplyActionOutcome::Terminal { terminal } | ApplyActionOutcome::PlayInstead { terminal, .. } | ApplyActionOutcome::Killshot { terminal, .. } => { @@ -316,7 +335,6 @@ impl SearchNode { self.value = Self::value_from_term(&state, terminal); } } - self.value } } @@ -330,46 +348,51 @@ pub struct GameSearchTree<'a> { impl GameSearchTree<'_> { fn safest_action(&mut self) -> (usize, Action) { - let val = if self.root_state.is_white_turn() { + let result = if self.root_state.is_white_turn() { self.root_node .children .iter() .enumerate() - .max_by_key(|(_, child)| { - ( - child.completion_value(), - child.entrance_count, - child.child_value, - self.rng.next_u32(), - ) + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.entrance_count as i32, a.child_value) + .partial_cmp(&( + b.completion_value(), + b.entrance_count as i32, + b.child_value, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) - .expect("No valid action at root state!") + .expect("no valid action at root state") } else { self.root_node .children .iter() .enumerate() - .min_by_key(|(_, child)| { + .min_by(|&(_, a), &(_, b)| { ( - child.completion_value(), - -(child.entrance_count as isize), - child.child_value, - self.rng.next_u32(), + a.completion_value(), + -(a.entrance_count as i32), + a.child_value, ) + .partial_cmp(&( + b.completion_value(), + -(b.entrance_count as i32), + b.child_value, + )) + .unwrap_or(Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) }) - .expect("No valid action at root state!") + .expect("no valid action at root state") }; - - (val.0, val.1.action) + (result.0, result.1.action) } - /// This has the invariant that the board is NOT in a terminal state. pub fn new<'a>(board: &Board, evaluator: &'a Evaluator) -> GameSearchTree<'a> { - let mut cpy = board.clone(); let mut rng = SmallRng::seed_from_u64(123312); GameSearchTree { root_node: Box::new(SearchNode::build_self( - &mut cpy, + board, ApplyActionOutcome::Ongoing, evaluator, &mut rng, @@ -382,47 +405,49 @@ impl GameSearchTree<'_> { pub fn step_tree(&mut self, new_board: &Board, action_id: usize, outcome: ApplyActionOutcome) { self.root_state = new_board.clone(); - if self.root_node.children[action_id].node.is_some() { self.root_node = self.root_node.children[action_id].node.take().unwrap(); } else { self.root_node = Box::new(SearchNode::build_self( - &mut self.root_state.clone(), + &self.root_state, outcome, self.evaluator, &mut self.rng, - )) + )); } } pub fn run_descent_for_iter(&mut self, iterations: u32, max_duration: Duration) { - let dcv = if self.root_state.is_white_turn() { + let desired_completion = if self.root_state.is_white_turn() { 1 } else { -1 }; - - let start_time = std::time::Instant::now(); + if max_duration.is_zero() { + return; + } + let started_at = std::time::Instant::now(); if self.root_node.children.len() <= 1 { - return; // No need to compute with only one option. + return; } + for epoch in 0..iterations { - if start_time.elapsed() > max_duration && epoch >= 50 { + if epoch > 0 && started_at.elapsed() >= max_duration { break; } - let ba = self.get_best_action_index(); - if self.root_node.children[ba].entrance_count >= 6000 - && self.root_node.children[ba].completion_value() != -dcv + let best_action = self.get_best_action_index(); + if self.root_node.children[best_action].entrance_count >= 6000 + && self.root_node.children[best_action].completion_value() != -desired_completion { break; } - if self.root_node.children[ba].completion_value() == dcv { + if self.root_node.children[best_action].completion_value() == desired_completion { break; } self.root_node.ubfms_iteration( self.root_state.clone(), ApplyActionOutcome::Ongoing, - &self.evaluator, + self.evaluator, &mut self.rng, ); } @@ -440,3 +465,49 @@ impl GameSearchTree<'_> { self.safest_action() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::board::board_structs::Player; + + fn sample_board() -> Board { + Board::from_fen(crate::TRAINING_START_FENS[0]).unwrap() + } + + #[test] + fn value_from_term_compresses_terminal_scale() { + let mut board = sample_board(); + board.turn_count = 10; + let win = SearchNode::value_from_term(&board, TerminalState::Win(Player::White)); + let loss = SearchNode::value_from_term(&board, TerminalState::Win(Player::Black)); + assert!(win > 0.0); + assert!(loss < 0.0); + assert!(win.abs() <= 1.0); + assert!(loss.abs() <= 1.0); + } + + #[test] + fn faster_wins_score_higher() { + let mut fast = sample_board(); + let mut slow = sample_board(); + fast.turn_count = 20; + slow.turn_count = 120; + assert!( + SearchNode::value_from_term(&fast, TerminalState::Win(Player::White)) + > SearchNode::value_from_term(&slow, TerminalState::Win(Player::White)) + ); + } + + #[test] + fn faster_losses_score_lower() { + let mut fast = sample_board(); + let mut slow = sample_board(); + fast.turn_count = 20; + slow.turn_count = 120; + assert!( + SearchNode::value_from_term(&fast, TerminalState::Win(Player::Black)) + < SearchNode::value_from_term(&slow, TerminalState::Win(Player::Black)) + ); + } +} diff --git a/flake.lock b/flake.lock new file mode 100644 index 0000000..f194334 --- /dev/null +++ b/flake.lock @@ -0,0 +1,61 @@ +{ + "nodes": { + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1774386573, + "narHash": "sha256-4hAV26quOxdC6iyG7kYaZcM3VOskcPUrdCQd/nx8obc=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "46db2e09e1d3f113a13c0d7b81e2f221c63b8ce9", + "type": "github" + }, + "original": { + "owner": "NixOS", + "ref": "nixos-unstable", + "repo": "nixpkgs", + "type": "github" + } + }, + "root": { + "inputs": { + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/flake.nix b/flake.nix new file mode 100644 index 0000000..5694219 --- /dev/null +++ b/flake.nix @@ -0,0 +1,50 @@ +{ + description = "Development shell for AlphaPaint"; + + inputs = { + nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable"; + flake-utils.url = "github:numtide/flake-utils"; + }; + + outputs = { + self, + nixpkgs, + flake-utils, + }: + flake-utils.lib.eachDefaultSystem (system: let + pkgs = import nixpkgs { + inherit system; + config = { + allowUnfree = true; + cudaSupport = true; + }; + }; + + cudaPkgs = pkgs.cudaPackages_12_9; + in { + devShells.default = pkgs.mkShell { + buildInputs = with pkgs; [ + rustc + cargo + maturin + cmake + cudaPkgs.cudatoolkit + cudaPkgs.cudnn + cudaPkgs.libcublas + stdenv.cc.cc.lib + ]; + + shellHook = '' + export CUDA_PATH=${cudaPkgs.cudatoolkit} + export CUDA_HOME=$CUDA_PATH + export CUDA_ROOT=$CUDA_PATH + export CUDNN_PATH=${cudaPkgs.cudnn.lib} + export RUST_MIN_STACK=67108864 + export LD_LIBRARY_PATH=$CUDA_PATH/lib:$CUDA_PATH/lib64:$CUDNN_PATH/lib:/run/opengl-driver/lib:${pkgs.stdenv.cc.cc.lib}/lib:$LD_LIBRARY_PATH + export TRITON_LIBCUDA_PATH=/run/opengl-driver/lib + export PATH=$CUDA_PATH/bin:$PATH + export PYTHONPATH=python + ''; + }; + }); +} diff --git a/pyproject.toml b/pyproject.toml index a92efbf..94f9b9d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,9 @@ dependencies = [ "psutil==5.9.0", "cython==3.0.11", "py-cpuinfo", + "torch==2.10.0", + "triton>=3.6.0; sys_platform == 'linux'", + "wandb>=0.25.1", ] classifiers = [ "Programming Language :: Rust", diff --git a/python/alphapaint_training/__init__.py b/python/alphapaint_training/__init__.py new file mode 100644 index 0000000..d9e640d --- /dev/null +++ b/python/alphapaint_training/__init__.py @@ -0,0 +1,44 @@ +"""AlphaPaint training infrastructure - Python bindings.""" + +from pkgutil import extend_path + +__path__ = extend_path(__path__, __name__) + +from .alphapaint_training import ( + EphemeralReplayBuffer, + SelfPlay, + sample_random_terminal_batch, +) +from .logger import TrainingLogger +from .model import NnueValueNet, PackedValueModel +from .packed_obs import ( + BOARD_CELLS, + BOARD_SIDE, + GLOBAL_FEATURES, + GLOBAL_SCALES, + LOCAL_FEATURES, + LOCAL_WINDOW_TILES, + OBS_WORDS, + TILE_PLANES, + TOTAL_TILE_FEATURES, + decode_nnue_obs, +) + +__all__ = [ + "BOARD_CELLS", + "BOARD_SIDE", + "EphemeralReplayBuffer", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", + "NnueValueNet", + "OBS_WORDS", + "PackedValueModel", + "SelfPlay", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", + "TrainingLogger", + "decode_nnue_obs", + "sample_random_terminal_batch", +] diff --git a/python/alphapaint_training/alphapaint_training.pyi b/python/alphapaint_training/alphapaint_training.pyi new file mode 100644 index 0000000..0a8cd21 --- /dev/null +++ b/python/alphapaint_training/alphapaint_training.pyi @@ -0,0 +1,58 @@ +from __future__ import annotations + +import numpy as np +import numpy.typing as npt + +class EphemeralReplayBuffer: + def __init__(self, capacity: int) -> None: ... + def __len__(self) -> int: ... + @property + def capacity(self) -> int: ... + def sample( + self, n: int, seed: int + ) -> tuple[npt.NDArray[np.uint16], npt.NDArray[np.float32]]: ... + +class SelfPlay: + def __init__( + self, + replay_buffer: EphemeralReplayBuffer, + num_threads: int, + workers_per_thread: int, + seed: int, + *, + max_gpu_evals_per_move: int = 4096, + model: object, + selfplay_precision: str = "bf16", + num_gpus: int = 1, + ) -> None: ... + def start(self) -> None: ... + def wait_for(self, target_samples: int) -> int: ... + def samples(self) -> int: ... + def games(self) -> int: ... + def action_steps(self) -> int: ... + def final_actions(self) -> int: ... + def nonfinal_actions(self) -> int: ... + def completed_turns(self) -> int: ... + def action_turn_count_total(self) -> int: ... + def max_turn_count_seen(self) -> int: ... + def completed_game_actions_total(self) -> int: ... + def max_actions_in_completed_game(self) -> int: ... + def completed_game_turn_count_total(self) -> int: ... + def max_turn_count_in_completed_game(self) -> int: ... + def take_completed_game_turn_counts(self) -> list[int]: ... + def tree_build_nanos(self) -> int: ... + def descent_nanos(self) -> int: ... + def sample_collect_nanos(self) -> int: ... + def replay_push_nanos(self) -> int: ... + def descent_expand_cpu_nanos(self) -> int: ... + def descent_apply_action_nanos(self) -> int: ... + def descent_eval_submit_nanos(self) -> int: ... + def descent_eval_await_nanos(self) -> int: ... + def descent_backup_nanos(self) -> int: ... + def gpu_batches(self) -> int: ... + def gpu_evals(self) -> int: ... + def drop(self) -> None: ... + +def sample_random_terminal_batch( + n: int, seed: int +) -> tuple[npt.NDArray[np.uint16], npt.NDArray[np.float32]]: ... diff --git a/python/alphapaint_training/cudagraph_backend.py b/python/alphapaint_training/cudagraph_backend.py new file mode 100644 index 0000000..c6e3e19 --- /dev/null +++ b/python/alphapaint_training/cudagraph_backend.py @@ -0,0 +1,147 @@ +"""CUDA graph capture for AlphaPaint value-only inference. + +Captures a CUDA graph that runs: H2D copy -> model forward -> D2H copy. +Observations are packed uint16 words; the model is expected to decode them. +""" + +from typing import Optional + +import torch +import torch.utils.dlpack as dlpack + + +def _validate_lane_tensors( + obs_host: torch.Tensor, + obs_device: torch.Tensor, + value_host: torch.Tensor, + value_device: torch.Tensor, + gpu_id: int, +) -> None: + if obs_host.device.type != "cpu": + raise ValueError("obs_host must be a CPU tensor") + if value_host.device.type != "cpu": + raise ValueError("value_host must be a CPU tensor") + if obs_device.device.type != "cuda": + raise ValueError("obs_device must be a CUDA tensor") + if value_device.device.type != "cuda": + raise ValueError("value_device must be a CUDA tensor") + + if obs_host.dtype != torch.uint16: + raise ValueError(f"obs_host must be uint16, got {obs_host.dtype}") + if obs_device.dtype != torch.uint16: + raise ValueError(f"obs_device must be uint16, got {obs_device.dtype}") + if value_host.dtype != torch.float32: + raise ValueError(f"value_host must be float32, got {value_host.dtype}") + if value_device.dtype != torch.float32: + raise ValueError(f"value_device must be float32, got {value_device.dtype}") + + if obs_host.shape != obs_device.shape: + raise ValueError( + f"obs_host and obs_device shape mismatch: {obs_host.shape} vs {obs_device.shape}" + ) + batch = obs_host.shape[0] + if value_host.shape != (batch,): + raise ValueError(f"value_host must be shape ({batch},), got {value_host.shape}") + if value_device.shape != (batch,): + raise ValueError( + f"value_device must be shape ({batch},), got {value_device.shape}" + ) + + for name, tensor in (("obs_device", obs_device), ("value_device", value_device)): + if tensor.device.index != gpu_id: + raise ValueError(f"{name} must be on cuda:{gpu_id}, got {tensor.device}") + + +def _autocast_dtype(precision: str) -> Optional[torch.dtype]: + if precision == "fp32": + return None + if precision == "fp16": + return torch.float16 + if precision == "bf16": + if torch.cuda.is_bf16_supported(): + return torch.bfloat16 + return torch.float16 + raise ValueError(f"Unsupported precision: {precision}") + + +def capture_lane_graph( + model, + obs_host_dlpack, + obs_device_dlpack, + value_host_dlpack, + value_device_dlpack, + stream_handle: int, + precision: str = "bf16", + gpu_id: int = 0, +) -> tuple[int, object]: + """Capture a CUDA graph for value-only inference. + + The model should accept observations and return a scalar value tensor. + No policy head is used (Athenan-style training). + + Args: + model: PyTorch model that takes packed obs tensor and returns value tensor. + obs_host_dlpack: DLPack capsule for pinned host observation buffer. + obs_device_dlpack: DLPack capsule for device observation buffer. + value_host_dlpack: DLPack capsule for pinned host value buffer. + value_device_dlpack: DLPack capsule for device value buffer. + stream_handle: Raw CUDA stream handle. + precision: "bf16", "fp16", or "fp32". + gpu_id: CUDA device index to capture the graph on. + + Returns: + Tuple of (cudaGraphExec_t handle as int, owner object keeping things alive). + """ + if gpu_id < 0: + raise ValueError(f"gpu_id must be >= 0, got {gpu_id}") + + obs_host = dlpack.from_dlpack(obs_host_dlpack) + obs_device = dlpack.from_dlpack(obs_device_dlpack) + value_host = dlpack.from_dlpack(value_host_dlpack) + value_device = dlpack.from_dlpack(value_device_dlpack) + + _validate_lane_tensors(obs_host, obs_device, value_host, value_device, gpu_id) + + with torch.cuda.device(gpu_id): + model = model.to(f"cuda:{gpu_id}") + model = model.to(memory_format=torch.channels_last) + model.eval() + torch.backends.cudnn.benchmark = True + stream = torch.cuda.ExternalStream(stream_handle) + graph = torch.cuda.CUDAGraph(keep_graph=True) + dtype = _autocast_dtype(precision) + + def run_step() -> None: + obs_device.copy_(obs_host, non_blocking=True) + if dtype is None: + value = model(obs_device) + else: + with torch.autocast(device_type="cuda", dtype=dtype): + value = model(obs_device) + if value.ndim == 2: + value = value.squeeze(-1) + value_device.copy_(value, non_blocking=True) + value_host.copy_(value_device, non_blocking=True) + + with torch.inference_mode(): + with torch.cuda.stream(stream): + for _ in range(3): + run_step() + torch.cuda.synchronize(device=gpu_id) + + with torch.cuda.graph( + graph, stream=stream, capture_error_mode="thread_local" + ): + run_step() + + graph.instantiate() + owner = ( + graph, + model, + obs_host, + obs_device, + value_host, + value_device, + stream, + ) + return int(graph.raw_cuda_graph_exec()), owner diff --git a/python/alphapaint_training/logger.py b/python/alphapaint_training/logger.py new file mode 100644 index 0000000..2ae46b6 --- /dev/null +++ b/python/alphapaint_training/logger.py @@ -0,0 +1,494 @@ +"""Training logger for AlphaPaint - handles wandb and console output.""" + +from __future__ import annotations + +import time +from dataclasses import asdict +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch +import wandb + +if TYPE_CHECKING: + from alphapaint_training.train import TrainConfig + from alphapaint_training.model import PackedValueModel + + +class TrainingLogger: + """Handles all logging for training runs including wandb and console output.""" + + def __init__( + self, + config: TrainConfig, + run_dir: Path, + checkpoint_dir: Path, + model: PackedValueModel, + param_count: int, + enable_wandb: bool = True, + ) -> None: + self.config = config + self.run_dir = run_dir + self.checkpoint_dir = checkpoint_dir + self.enable_wandb = enable_wandb + self.run = None + + if enable_wandb: + self.run = wandb.init( + project="alphapaint", + name=run_dir.name, + dir=run_dir, + config=asdict(config), + ) + self.run.summary["model_params"] = param_count + + self._pretrain_loss_window: list[float] = [] + self._pretrain_started_at: float | None = None + self._previous_samples = 0 + self._previous_terminal_mix_samples = 0 + self._previous_games = 0 + self._previous_gpu_batches = 0 + self._previous_gpu_evals = 0 + self._previous_action_steps = 0 + self._previous_final_actions = 0 + self._previous_nonfinal_actions = 0 + self._previous_completed_turns = 0 + self._previous_action_turn_count_total = 0 + self._previous_completed_game_actions_total = 0 + self._previous_completed_game_turn_count_total = 0 + self._previous_tree_build_nanos = 0 + self._previous_descent_nanos = 0 + self._previous_sample_collect_nanos = 0 + self._previous_replay_push_nanos = 0 + self._previous_descent_expand_cpu_nanos = 0 + self._previous_descent_apply_action_nanos = 0 + self._previous_descent_eval_submit_nanos = 0 + self._previous_descent_eval_await_nanos = 0 + self._previous_descent_backup_nanos = 0 + + def start_pretrain(self) -> None: + """Call at the start of pretraining.""" + self._pretrain_started_at = time.perf_counter() + self._pretrain_loss_window.clear() + print( + f"pretrain samples={self.config.pretrain_terminal_samples} " + f"batch={self.config.pretrain_batch_size}" + ) + + def log_pretrain_step( + self, + step_idx: int, + step_batch_size: int, + loss_value: float, + pretrain_steps: int, + device: torch.device, + ) -> None: + """Log a single pretrain step.""" + self._pretrain_loss_window.append(loss_value) + pretrain_log_interval = max(1, self.config.pretrain_log_interval) + if len(self._pretrain_loss_window) > pretrain_log_interval: + self._pretrain_loss_window.pop(0) + + step_number = step_idx + 1 + if step_number % pretrain_log_interval == 0 or step_number == pretrain_steps: + if device.type == "cuda": + torch.cuda.synchronize(device) + samples_done = step_idx * self.config.pretrain_batch_size + step_batch_size + elapsed = ( + time.perf_counter() - self._pretrain_started_at + if self._pretrain_started_at + else 0.0 + ) + window_mean_loss = float(np.mean(self._pretrain_loss_window)) + print( + f"pretrain step={step_number}/{pretrain_steps} " + f"samples={samples_done}/{self.config.pretrain_terminal_samples} " + f"loss={window_mean_loss:.6f} elapsed={elapsed:.2f}s" + ) + if self.run is not None: + wandb.log( + { + "pretrain_step": step_number, + "pretrain_samples_total": samples_done, + "pretrain_loss_mean": window_mean_loss, + "pretrain_seconds": elapsed, + "learning_rate": self._get_current_lr(), + } + ) + + def finish_pretrain(self, pretrain_steps: int, device: torch.device) -> None: + """Call at the end of pretraining.""" + if device.type == "cuda": + torch.cuda.synchronize(device) + elapsed = ( + time.perf_counter() - self._pretrain_started_at + if self._pretrain_started_at + else 0.0 + ) + print( + f"pretrain_complete samples={self.config.pretrain_terminal_samples} " + f"steps={pretrain_steps} elapsed={elapsed:.2f}s" + ) + + def _get_current_lr(self) -> float: + """Get current learning rate from optimizer (stored externally).""" + return 0.0 + + def set_lr_getter(self, getter: Any) -> None: + """Set a function to get the current learning rate.""" + self._get_current_lr = getter + + def update( + self, + selfplay: Any, + round_number: int, + collected: int, + replay_size: int, + mean_loss: float, + collect_seconds: float, + train_seconds: float, + train_sample_seconds: float, + train_h2d_seconds: float, + train_forward_seconds: float, + train_backward_seconds: float, + train_optimizer_seconds: float, + diag_stats: dict[str, float], + diag_histograms: dict[str, wandb.Histogram], + ) -> dict[str, object]: + """Update logs for a completed round. Returns the record dict for checkpointing.""" + # Extract all metrics from selfplay + terminal_mix_samples = selfplay.terminal_mix_samples_added() + games = selfplay.games() + gpu_batches = selfplay.gpu_batches() + gpu_evals = selfplay.gpu_evals() + action_steps = selfplay.action_steps() + final_actions = selfplay.final_actions() + nonfinal_actions = selfplay.nonfinal_actions() + completed_turns = selfplay.completed_turns() + action_turn_count_total = selfplay.action_turn_count_total() + max_turn_count_seen = selfplay.max_turn_count_seen() + completed_game_actions_total = selfplay.completed_game_actions_total() + max_turn_count_in_completed_game = selfplay.max_turn_count_in_completed_game() + completed_game_turn_counts = selfplay.take_completed_game_turn_counts() + completed_game_turn_count_total = selfplay.completed_game_turn_count_total() + max_actions_in_completed_game = selfplay.max_actions_in_completed_game() + tree_build_nanos = selfplay.tree_build_nanos() + descent_nanos = selfplay.descent_nanos() + sample_collect_nanos = selfplay.sample_collect_nanos() + replay_push_nanos = selfplay.replay_push_nanos() + descent_expand_cpu_nanos = selfplay.descent_expand_cpu_nanos() + descent_apply_action_nanos = selfplay.descent_apply_action_nanos() + descent_eval_submit_nanos = selfplay.descent_eval_submit_nanos() + descent_eval_await_nanos = selfplay.descent_eval_await_nanos() + descent_backup_nanos = selfplay.descent_backup_nanos() + + # Calculate deltas + samples_added = collected - self._previous_samples + terminal_mix_samples_added = ( + terminal_mix_samples - self._previous_terminal_mix_samples + ) + games_added = games - self._previous_games + gpu_batches_added = gpu_batches - self._previous_gpu_batches + gpu_evals_added = gpu_evals - self._previous_gpu_evals + action_steps_added = action_steps - self._previous_action_steps + final_actions_added = final_actions - self._previous_final_actions + nonfinal_actions_added = nonfinal_actions - self._previous_nonfinal_actions + completed_turns_added = completed_turns - self._previous_completed_turns + action_turn_count_added = ( + action_turn_count_total - self._previous_action_turn_count_total + ) + completed_game_actions_added = ( + completed_game_actions_total - self._previous_completed_game_actions_total + ) + completed_game_turn_count_added = ( + completed_game_turn_count_total + - self._previous_completed_game_turn_count_total + ) + + tree_build_seconds = (tree_build_nanos - self._previous_tree_build_nanos) / 1e9 + descent_seconds = (descent_nanos - self._previous_descent_nanos) / 1e9 + sample_collect_seconds = ( + sample_collect_nanos - self._previous_sample_collect_nanos + ) / 1e9 + replay_push_seconds = ( + replay_push_nanos - self._previous_replay_push_nanos + ) / 1e9 + descent_expand_cpu_seconds = ( + descent_expand_cpu_nanos - self._previous_descent_expand_cpu_nanos + ) / 1e9 + descent_apply_action_seconds = ( + descent_apply_action_nanos - self._previous_descent_apply_action_nanos + ) / 1e9 + descent_eval_submit_seconds = ( + descent_eval_submit_nanos - self._previous_descent_eval_submit_nanos + ) / 1e9 + descent_eval_await_seconds = ( + descent_eval_await_nanos - self._previous_descent_eval_await_nanos + ) / 1e9 + descent_backup_seconds = ( + descent_backup_nanos - self._previous_descent_backup_nanos + ) / 1e9 + descent_other_seconds = max( + 0.0, + descent_seconds + - descent_expand_cpu_seconds + - descent_apply_action_seconds + - descent_eval_submit_seconds + - descent_eval_await_seconds + - descent_backup_seconds, + ) + + avg_turn_count = ( + action_turn_count_added / action_steps_added if action_steps_added else 0.0 + ) + actions_per_turn = ( + action_steps_added / completed_turns_added if completed_turns_added else 0.0 + ) + actions_per_game = ( + completed_game_actions_added / games_added if games_added else 0.0 + ) + + if completed_game_turn_counts: + completed_game_turn_counts_arr = np.asarray( + completed_game_turn_counts, dtype=np.float32 + ) + completed_game_turn_count_mean = float( + completed_game_turn_counts_arr.mean() + ) + completed_game_turn_count_p50 = float( + np.percentile(completed_game_turn_counts_arr, 50) + ) + completed_game_turn_count_p90 = float( + np.percentile(completed_game_turn_counts_arr, 90) + ) + else: + completed_game_turn_count_mean = float("nan") + completed_game_turn_count_p50 = float("nan") + completed_game_turn_count_p90 = float("nan") + + build_ms_per_action = ( + tree_build_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_ms_per_action = ( + descent_seconds * 1000.0 / action_steps_added if action_steps_added else 0.0 + ) + collect_ms_per_action = ( + sample_collect_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + push_ms_per_action = ( + replay_push_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_expand_ms_per_action = ( + descent_expand_cpu_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_apply_ms_per_action = ( + descent_apply_action_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_ms_per_action = ( + descent_eval_submit_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_wait_ms_per_action = ( + descent_eval_await_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_backup_ms_per_action = ( + descent_backup_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_other_ms_per_action = ( + descent_other_seconds * 1000.0 / action_steps_added + if action_steps_added + else 0.0 + ) + descent_submit_us_per_eval = ( + descent_eval_submit_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) + descent_wait_us_per_eval = ( + descent_eval_await_seconds * 1e6 / gpu_evals_added + if gpu_evals_added + else 0.0 + ) + + record: dict[str, Any] = { + "round": round_number, + "samples_total": collected, + "samples_added": samples_added, + "terminal_mix_samples_total": terminal_mix_samples, + "terminal_mix_samples_added": terminal_mix_samples_added, + "games_total": games, + "games_added": games_added, + "gpu_batches_total": gpu_batches, + "gpu_batches_added": gpu_batches_added, + "gpu_batches_per_second": gpu_batches_added / max(collect_seconds, 1e-9), + "gpu_evals_total": gpu_evals, + "gpu_evals_added": gpu_evals_added, + "gpu_evals_per_second": gpu_evals_added / max(collect_seconds, 1e-9), + "action_steps_total": action_steps, + "action_steps_added": action_steps_added, + "final_actions_total": final_actions, + "final_actions_added": final_actions_added, + "nonfinal_actions_total": nonfinal_actions, + "nonfinal_actions_added": nonfinal_actions_added, + "completed_turns_total": completed_turns, + "completed_turns_added": completed_turns_added, + "action_turn_count_total": action_turn_count_total, + "action_turn_count_added": action_turn_count_added, + "avg_turn_count": avg_turn_count, + "max_turn_count_seen": max_turn_count_seen, + "actions_per_turn": actions_per_turn, + "completed_game_actions_total": completed_game_actions_total, + "completed_game_actions_added": completed_game_actions_added, + "actions_per_completed_game": actions_per_game, + "max_actions_in_completed_game": max_actions_in_completed_game, + "completed_game_turn_count_total": completed_game_turn_count_total, + "completed_game_turn_count_added": completed_game_turn_count_added, + "completed_game_turn_count_mean": completed_game_turn_count_mean, + "completed_game_turn_count_p50": completed_game_turn_count_p50, + "completed_game_turn_count_p90": completed_game_turn_count_p90, + "max_turn_count_in_completed_game": max_turn_count_in_completed_game, + "replay_size": replay_size, + "collection_seconds": collect_seconds, + "training_seconds": train_seconds, + "samples_per_second": samples_added / max(collect_seconds, 1e-9), + "train_steps_per_second": self.config.train_steps_per_round + / max(train_seconds, 1e-9), + "selfplay_tree_build_seconds": tree_build_seconds, + "selfplay_descent_seconds": descent_seconds, + "selfplay_sample_collect_seconds": sample_collect_seconds, + "selfplay_replay_push_seconds": replay_push_seconds, + "selfplay_tree_build_ms_per_action": build_ms_per_action, + "selfplay_descent_ms_per_action": descent_ms_per_action, + "selfplay_sample_collect_ms_per_action": collect_ms_per_action, + "selfplay_replay_push_ms_per_action": push_ms_per_action, + "descent_expand_cpu_seconds": descent_expand_cpu_seconds, + "descent_apply_action_seconds": descent_apply_action_seconds, + "descent_eval_submit_seconds": descent_eval_submit_seconds, + "descent_eval_await_seconds": descent_eval_await_seconds, + "descent_backup_seconds": descent_backup_seconds, + "descent_other_seconds": descent_other_seconds, + "descent_expand_cpu_ms_per_action": descent_expand_ms_per_action, + "descent_apply_action_ms_per_action": descent_apply_ms_per_action, + "descent_eval_submit_ms_per_action": descent_submit_ms_per_action, + "descent_eval_await_ms_per_action": descent_wait_ms_per_action, + "descent_backup_ms_per_action": descent_backup_ms_per_action, + "descent_other_ms_per_action": descent_other_ms_per_action, + "descent_eval_submit_us_per_eval": descent_submit_us_per_eval, + "descent_eval_await_us_per_eval": descent_wait_us_per_eval, + "train_sample_seconds": train_sample_seconds, + "train_h2d_seconds": train_h2d_seconds, + "train_forward_seconds": train_forward_seconds, + "train_backward_seconds": train_backward_seconds, + "train_optimizer_seconds": train_optimizer_seconds, + "loss_mean": mean_loss, + "learning_rate": self._get_current_lr(), + "timestamp": time.time(), + **diag_stats, + } + + print( + f"round={round_number} samples={collected} games={games} " + f"gpu_evals/s={gpu_evals_added / max(collect_seconds, 1e-9):.1f} " + f"replay={replay_size} loss={mean_loss:.6f}" + ) + + if self.run is not None: + wandb.log({**record, **diag_histograms}) + + self._update_previous_values( + collected, + terminal_mix_samples, + games, + gpu_batches, + gpu_evals, + action_steps, + final_actions, + nonfinal_actions, + completed_turns, + action_turn_count_total, + completed_game_actions_total, + completed_game_turn_count_total, + tree_build_nanos, + descent_nanos, + sample_collect_nanos, + replay_push_nanos, + descent_expand_cpu_nanos, + descent_apply_action_nanos, + descent_eval_submit_nanos, + descent_eval_await_nanos, + descent_backup_nanos, + ) + + return record + + def _update_previous_values( + self, + collected: int, + terminal_mix_samples: int, + games: int, + gpu_batches: int, + gpu_evals: int, + action_steps: int, + final_actions: int, + nonfinal_actions: int, + completed_turns: int, + action_turn_count_total: int, + completed_game_actions_total: int, + completed_game_turn_count_total: int, + tree_build_nanos: int, + descent_nanos: int, + sample_collect_nanos: int, + replay_push_nanos: int, + descent_expand_cpu_nanos: int, + descent_apply_action_nanos: int, + descent_eval_submit_nanos: int, + descent_eval_await_nanos: int, + descent_backup_nanos: int, + ) -> None: + """Update previous values for delta calculations.""" + self._previous_samples = collected + self._previous_terminal_mix_samples = terminal_mix_samples + self._previous_games = games + self._previous_gpu_batches = gpu_batches + self._previous_gpu_evals = gpu_evals + self._previous_action_steps = action_steps + self._previous_final_actions = final_actions + self._previous_nonfinal_actions = nonfinal_actions + self._previous_completed_turns = completed_turns + self._previous_action_turn_count_total = action_turn_count_total + self._previous_completed_game_actions_total = completed_game_actions_total + self._previous_completed_game_turn_count_total = completed_game_turn_count_total + self._previous_tree_build_nanos = tree_build_nanos + self._previous_descent_nanos = descent_nanos + self._previous_sample_collect_nanos = sample_collect_nanos + self._previous_replay_push_nanos = replay_push_nanos + self._previous_descent_expand_cpu_nanos = descent_expand_cpu_nanos + self._previous_descent_apply_action_nanos = descent_apply_action_nanos + self._previous_descent_eval_submit_nanos = descent_eval_submit_nanos + self._previous_descent_eval_await_nanos = descent_eval_await_nanos + self._previous_descent_backup_nanos = descent_backup_nanos + + def log_checkpoint(self, checkpoint_path: Path) -> None: + """Log a checkpoint save.""" + print(f"checkpoint={checkpoint_path}") + + def finish(self) -> None: + """Clean up logging.""" + if self.run is not None: + self.run.finish() diff --git a/python/alphapaint_training/model.py b/python/alphapaint_training/model.py new file mode 100644 index 0000000..cf9f015 --- /dev/null +++ b/python/alphapaint_training/model.py @@ -0,0 +1,184 @@ +"""NNUE value model for GPU training. + +Architecture mirrors the CPU NNUE evaluator: + - Accumulator: shared linear 17408 -> acc_dim for both perspectives + - Local MLP: shared linear 375 -> local_hidden_dim (ReLU) for both players + - Global features: 20 normalized scalars + - Head: concat all -> fc1 (ReLU) -> fc2 (ReLU) -> fc3 -> scalar + +The same weight matrix is applied to both the current-player and +opponent perspectives, matching the CPU NNUE design where one set of +feature weights serves both the white and black accumulators. +""" + +from __future__ import annotations + +from typing import cast + +import torch +from torch import nn + +from .packed_obs import ( + GLOBAL_FEATURES, + GLOBAL_SCALES, + LOCAL_FEATURES, + OPP_PLANE_PERM, + TILE_BITMASK_PLANES, + TOTAL_TILE_FEATURES, + decode_nnue_obs, +) + + +class NnueValueNet(nn.Module): + """NNUE-style value network for GPU training. + + Input components: + acc_mine: [B, acc_dim] from accumulator (my perspective) + acc_opp: [B, acc_dim] from accumulator (opp perspective) + local_my: [B, local_hidden] from local MLP (my position) + local_opp: [B, local_hidden] from local MLP (opp position) + globals: [B, 20] normalized global features + + Architecture: + concat -> fc1 (ReLU) -> fc2 (ReLU) -> fc3 -> scalar + """ + + def __init__( + self, + *, + acc_dim: int = 128, + local_hidden_dim: int = 64, + fc1_dim: int = 128, + fc2_dim: int = 32, + ): + super().__init__() + self.acc_dim = acc_dim + self.local_hidden_dim = local_hidden_dim + + # Accumulator: shared for both perspectives + self.acc_linear = nn.Linear(TOTAL_TILE_FEATURES, acc_dim) + + # Local MLP: shared for both player windows + self.local_mlp = nn.Sequential( + nn.Linear(LOCAL_FEATURES, local_hidden_dim), + nn.ReLU(inplace=True), + ) + + # Final head + fc1_in = acc_dim * 2 + local_hidden_dim * 2 + GLOBAL_FEATURES + self.fc1 = nn.Linear(fc1_in, fc1_dim) + self.fc2 = nn.Linear(fc1_dim, fc2_dim) + self.fc3 = nn.Linear(fc2_dim, 1) + + def forward( + self, + my_features: torch.Tensor, + opp_features: torch.Tensor, + local_my: torch.Tensor, + local_opp: torch.Tensor, + globals_norm: torch.Tensor, + ) -> torch.Tensor: + # Shared accumulator + acc_mine = self.acc_linear(my_features) + acc_opp = self.acc_linear(opp_features) + + # Shared local MLP + local_mine_out = self.local_mlp(local_my) + local_opp_out = self.local_mlp(local_opp) + + # Concat and head + x = torch.cat( + [acc_mine, acc_opp, local_mine_out, local_opp_out, globals_norm], dim=1 + ) + x = torch.relu(self.fc1(x)) + x = torch.relu(self.fc2(x)) + return self.fc3(x) + + +class PackedValueModel(nn.Module): + """Wraps NnueValueNet with packed observation decoding. + + Takes [B, OBS_WORDS] uint16 packed observations, decodes them on-device, + runs the NNUE value net, and returns [B] scalar values from White's + perspective. + """ + + def __init__( + self, + *, + acc_dim: int = 128, + local_hidden_dim: int = 64, + fc1_dim: int = 128, + fc2_dim: int = 32, + board_dtype: torch.dtype = torch.bfloat16, + ): + super().__init__() + self.board_dtype = board_dtype + self.value_net = NnueValueNet( + acc_dim=acc_dim, + local_hidden_dim=local_hidden_dim, + fc1_dim=fc1_dim, + fc2_dim=fc2_dim, + ) + + # Decode buffers — registered so they follow .to(device) and are + # captured in CUDA graphs. + self.register_buffer( + "_plane_shift", + torch.arange(TILE_BITMASK_PLANES, dtype=torch.int32), + persistent=False, + ) + self.register_buffer( + "_opp_perm", + torch.tensor(OPP_PLANE_PERM, dtype=torch.long), + persistent=False, + ) + self.register_buffer( + "_global_scales", + torch.tensor(GLOBAL_SCALES, dtype=torch.float32), + persistent=False, + ) + + def decode( + self, packed_obs: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ]: + return decode_nnue_obs( + packed_obs, + plane_shift=cast(torch.Tensor, self._plane_shift), + opp_perm=cast(torch.Tensor, self._opp_perm), + global_scales=cast(torch.Tensor, self._global_scales), + dtype=self.board_dtype, + ) + + def forward(self, packed_obs: torch.Tensor) -> torch.Tensor: + my_features, opp_features, local_my, local_opp, globals_norm, turn_count = ( + self.decode(packed_obs) + ) + + if not torch.is_autocast_enabled(): + param_dtype = next(self.value_net.parameters()).dtype + my_features = my_features.to(dtype=param_dtype) + opp_features = opp_features.to(dtype=param_dtype) + local_my = local_my.to(dtype=param_dtype) + local_opp = local_opp.to(dtype=param_dtype) + globals_norm = globals_norm.to(dtype=param_dtype) + + value = self.value_net( + my_features, opp_features, local_my, local_opp, globals_norm + ) + value = value.squeeze(-1) + + # Convert from current-player perspective to White perspective. + white_to_move = (turn_count & 1) == 0 + sign = torch.where(white_to_move, 1.0, -1.0).to(dtype=value.dtype) + return (value * sign).tanh() + + +__all__ = ["NnueValueNet", "PackedValueModel"] diff --git a/python/alphapaint_training/packed_obs.py b/python/alphapaint_training/packed_obs.py new file mode 100644 index 0000000..6089baf --- /dev/null +++ b/python/alphapaint_training/packed_obs.py @@ -0,0 +1,204 @@ +"""Packed observation decoding for NNUE training. + +The Rust encoder packs board state into a flat u16 array with this layout: + [0, 1024) — per-tile 15-bit bitmask (32×32 board) + [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) + [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) + [1078, 1098) — 20 global scalar features + +Tile bitmask planes (from current player's perspective): + bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) + bits 4-7: opponent paint thermometer + bit 8: wall + bit 9: powerup + bit 10: current player beacon + bit 11: opponent beacon + bit 12: hill neutral + bit 13: hill current player + bit 14: hill opponent + +Additional planes added during decoding: + plane 15: one-hot my player location (derived from positions) + plane 16: one-hot opponent player location (derived from positions) +""" + +from __future__ import annotations + +import torch + +# Layout constants ---------------------------------------------------------- + +BOARD_SIDE = 32 +BOARD_CELLS = BOARD_SIDE * BOARD_SIDE # 1024 +# Bits 0–14 from each u16 tile word (Rust encoder); +2 derived position planes → accumulator +TILE_BITMASK_PLANES = 15 +TILE_PLANES = TILE_BITMASK_PLANES + 2 # 17 +LOCAL_WINDOW_TILES = 25 +GLOBAL_FEATURES = 20 +TOTAL_TILE_FEATURES = BOARD_CELLS * TILE_PLANES # 17408 +LOCAL_TILE_PLANES = TILE_BITMASK_PLANES # Local windows don't have player position planes +LOCAL_FEATURES = LOCAL_WINDOW_TILES * LOCAL_TILE_PLANES # 375 + +OFFSET_LOCALS = BOARD_CELLS # 1024 +OFFSET_POSITIONS = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2 # 1074 +OFFSET_GLOBALS = OFFSET_POSITIONS + 4 # 1078 +OBS_WORDS = OFFSET_GLOBALS + GLOBAL_FEATURES # 1098 + +# Index of turn_count within the global section +GLOBAL_TURN_COUNT_IDX = 17 + +# Plane permutation for opponent perspective: +# swap my_paint[0:4] <-> opp_paint[4:8] +# swap my_beacon[10] <-> opp_beacon[11] +# swap my_hill[13] <-> opp_hill[14] +# swap my_pos[15] <-> opp_pos[16] +# wall[8], powerup[9], hill_neutral[12] stay +OPP_PLANE_PERM = [4, 5, 6, 7, 0, 1, 2, 3, 8, 9, 11, 10, 12, 14, 13, 16, 15] + +# Normalization divisors for the 20 global features. +# Brings raw u16 values into roughly [0, 1] range. +GLOBAL_SCALES: tuple[float, ...] = ( + 420.0, # 0 my_stamina + 420.0, # 1 opp_stamina + 420.0, # 2 my_max_stamina + 420.0, # 3 opp_max_stamina + 1024.0, # 4 my_territory + 1024.0, # 5 opp_territory + 8.0, # 6 my_hills + 8.0, # 7 opp_hills + 1024.0, # 8 my_hill_tiles + 1024.0, # 9 opp_hill_tiles + 8.0, # 10 contested_hills + 1024.0, # 11 my_beacons + 1024.0, # 12 opp_beacons + 64.0, # 13 player_dist + 64.0, # 14 my_hill_dist + 64.0, # 15 opp_hill_dist + 32.0, # 16 consecutive_moves + 2000.0, # 17 turn_count + 32.0, # 18 rows + 32.0, # 19 cols +) + + +def _check_packed_obs(packed_obs: torch.Tensor) -> None: + if packed_obs.dtype != torch.uint16: + raise ValueError(f"packed_obs must be uint16, got {packed_obs.dtype}") + if packed_obs.ndim != 2: + raise ValueError(f"packed_obs must be rank-2, got {packed_obs.ndim}") + if packed_obs.shape[1] != OBS_WORDS: + raise ValueError( + f"packed_obs must have shape (B, {OBS_WORDS}), got {tuple(packed_obs.shape)}" + ) + + +def _unpack_bits(words: torch.Tensor, shift: torch.Tensor) -> torch.Tensor: + """Extract binary planes from u16 words (one channel per shift index). + + Args: + words: [...] int32 tensor of packed bitmasks + shift: [K] int32 tensor of bit indices (e.g. arange(15) for tile bitmasks) + + Returns: + [..., K] float tensor of 0/1 values + """ + return ((words.unsqueeze(-1) >> shift) & 1).float() + + +def decode_nnue_obs( + packed_obs: torch.Tensor, + *, + plane_shift: torch.Tensor, + opp_perm: torch.Tensor, + global_scales: torch.Tensor, + dtype: torch.dtype = torch.bfloat16, +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + """Decode packed NNUE observation into model inputs. + + All buffers (plane_shift, opp_perm, global_scales) should be registered + as model buffers so they live on the right device and are captured by + CUDA graphs. + + Returns: + my_features: [B, 17408] dtype — accumulator input (my perspective) + opp_features: [B, 17408] dtype — accumulator input (opponent perspective) + local_my: [B, 375] dtype — local window around my position + local_opp: [B, 375] dtype — local window around opponent position + globals_norm: [B, 20] dtype — normalized global scalars + turn_count: [B] int32 — for white-to-move sign + """ + B = packed_obs.shape[0] + + # --- Tile features: [B, 1024] u16 -> [B, 17408] ---------------------- + tiles = packed_obs[:, :BOARD_CELLS].to(torch.int32) + # Only bits 0–14 are defined on tile words; planes 15–16 come from positions below. + bits = _unpack_bits(tiles, plane_shift[:TILE_BITMASK_PLANES]) # [B, 1024, 15] + + # --- Player positions -> one-hot planes [B, 1024, 2] ----------------- + my_x = packed_obs[:, OFFSET_POSITIONS].long() # [B] + my_y = packed_obs[:, OFFSET_POSITIONS + 1].long() # [B] + opp_x = packed_obs[:, OFFSET_POSITIONS + 2].long() # [B] + opp_y = packed_obs[:, OFFSET_POSITIONS + 3].long() # [B] + + my_pos_flat = my_x * BOARD_SIDE + my_y # [B] + opp_pos_flat = opp_x * BOARD_SIDE + opp_y # [B] + + my_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + opp_pos_plane = torch.zeros( + B, BOARD_CELLS, dtype=torch.float32, device=packed_obs.device + ) + my_pos_plane.scatter_(1, my_pos_flat.unsqueeze(1), 1.0) + opp_pos_plane.scatter_(1, opp_pos_flat.unsqueeze(1), 1.0) + + # Concatenate: [B, 1024, 15] + [B, 1024, 1] + [B, 1024, 1] = [B, 1024, 17] + bits_with_pos = torch.cat( + [bits, my_pos_plane.unsqueeze(-1), opp_pos_plane.unsqueeze(-1)], dim=-1 + ) # [B, 1024, 17] + + my_features = bits_with_pos.reshape(B, -1).to(dtype) + + # Opponent perspective: permute planes + opp_features = bits_with_pos[:, :, opp_perm].reshape(B, -1).to(dtype) + + # --- Local windows: [B, 50] u16 -> [B, 375] each --------------------- + # Local windows only have 15 planes (no player positions) + locals_raw = packed_obs[:, OFFSET_LOCALS:OFFSET_POSITIONS].to(torch.int32) + local_shift = plane_shift[:TILE_BITMASK_PLANES] + local_bits = _unpack_bits(locals_raw, local_shift) # [B, 50, 15] + local_my = local_bits[:, :LOCAL_WINDOW_TILES, :].reshape(B, -1).to(dtype) + local_opp = local_bits[:, LOCAL_WINDOW_TILES:, :].reshape(B, -1).to(dtype) + + # --- Global features: [B, 20] u16 -> [B, 20] float normalized -------- + globals_raw = packed_obs[:, OFFSET_GLOBALS:].to(torch.float32) + globals_norm = (globals_raw / global_scales).clamp(0.0, 1.0).to(dtype) + + # --- Turn count for white-to-move sign -------------------------------- + turn_count = packed_obs[:, OFFSET_GLOBALS + GLOBAL_TURN_COUNT_IDX].to(torch.int32) + + return my_features, opp_features, local_my, local_opp, globals_norm, turn_count + + +__all__ = [ + "BOARD_CELLS", + "BOARD_SIDE", + "GLOBAL_FEATURES", + "GLOBAL_SCALES", + "GLOBAL_TURN_COUNT_IDX", + "LOCAL_FEATURES", + "LOCAL_WINDOW_TILES", + "OBS_WORDS", + "OPP_PLANE_PERM", + "TILE_BITMASK_PLANES", + "TILE_PLANES", + "TOTAL_TILE_FEATURES", + "decode_nnue_obs", +] diff --git a/python/alphapaint_training/train.py b/python/alphapaint_training/train.py new file mode 100644 index 0000000..0d55710 --- /dev/null +++ b/python/alphapaint_training/train.py @@ -0,0 +1,602 @@ +from __future__ import annotations + +import argparse +import json +import os +import time +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import cast + +import numpy as np +import torch +import torch.nn.functional as F +import wandb + +from alphapaint_training import ( + EphemeralReplayBuffer, + SelfPlay, + sample_random_terminal_batch, +) +from alphapaint_training.logger import TrainingLogger +from alphapaint_training.model import PackedValueModel + + +def _int_env(name: str) -> int | None: + value = os.environ.get(name) + if not value: + return None + try: + return int(value) + except ValueError: + return None + + +def _slurm_gpu_count() -> int | None: + for name in ("SLURM_GPUS_ON_NODE", "SLURM_GPUS"): + value = _int_env(name) + if value is not None and value > 0: + return value + + job_gpus = os.environ.get("SLURM_JOB_GPUS") + if not job_gpus: + return None + + gpu_ids = [gpu_id.strip() for gpu_id in job_gpus.split(",") if gpu_id.strip()] + if not gpu_ids: + return None + return len(gpu_ids) + + +def _default_num_threads() -> int: + cpus_per_gpu = _int_env("SLURM_CPUS_PER_GPU") + slurm_gpu_count = _slurm_gpu_count() + if cpus_per_gpu is not None and cpus_per_gpu > 0 and slurm_gpu_count is not None: + return cpus_per_gpu * slurm_gpu_count + return 32 + + +def _default_num_gpus() -> int: + slurm_gpu_count = _slurm_gpu_count() + if slurm_gpu_count is None: + return 0 + return slurm_gpu_count + + +@dataclass(slots=True) +class TrainConfig: + rounds: int = 200 + samples_per_round: int = 1024 * 1024 * 2 + train_steps_per_round: int = 256 + batch_size: int = 8_192 + replay_capacity: int = 16_000_000 + num_threads: int = field(default_factory=_default_num_threads) + workers_per_thread: int = 8 + max_gpu_evals_per_move: int = 16 * 4096 + lr: float = 3e-4 + pretrain_lr: float | None = None + selfplay_lr: float | None = 5e-5 + weight_decay: float = 1e-4 + pretrain_terminal_samples: int = 4_000_000 + pretrain_batch_size: int = 8_192 + pretrain_log_interval: int = 10 + seed: int = 42 + selfplay_precision: str = "bf16" + num_gpus: int = field(default_factory=_default_num_gpus) + checkpoint_interval: int = 2 + run_dir: str = "runs/latest" + wandb: bool = True + + +@dataclass(slots=True) +class CudaSectionTiming: + start: torch.cuda.Event + end: torch.cuda.Event + + def seconds(self) -> float: + return float(self.start.elapsed_time(self.end) / 1000.0) + + +@dataclass(slots=True) +class TrainStepResult: + loss: torch.Tensor + sample_seconds: float + h2d_seconds: float + forward_timing: CudaSectionTiming + backward_timing: CudaSectionTiming + optimizer_timing: CudaSectionTiming + + +def _numpy_batch_to_device( + obs_np: np.ndarray, + values_np: np.ndarray, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float]: + h2d_started_at = time.perf_counter() + obs = torch.from_numpy(np.asarray(obs_np, dtype=np.uint16)).to( + device=device, + dtype=torch.uint16, + non_blocking=True, + ) + values = torch.from_numpy(np.asarray(values_np, dtype=np.float32)).to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + h2d_seconds = time.perf_counter() - h2d_started_at + return obs, values, h2d_seconds + + +def _sample_replay_batch( + replay_buffer: EphemeralReplayBuffer, + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float, float]: + actual_batch_size = min(batch_size, len(replay_buffer)) + if actual_batch_size == 0: + raise RuntimeError("replay buffer is empty") + + sample_started_at = time.perf_counter() + obs_np, values_np = replay_buffer.sample(actual_batch_size, seed) + sample_seconds = time.perf_counter() - sample_started_at + + obs, values, h2d_seconds = _numpy_batch_to_device(obs_np, values_np, device) + return obs, values, sample_seconds, h2d_seconds + + +def _sample_terminal_batch( + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, float, float]: + sample_started_at = time.perf_counter() + obs_np, values_np = sample_random_terminal_batch(batch_size, seed) + sample_seconds = time.perf_counter() - sample_started_at + + obs, values, h2d_seconds = _numpy_batch_to_device(obs_np, values_np, device) + return obs, values, sample_seconds, h2d_seconds + + +def _train_step_from_batch( + model: PackedValueModel, + optimizer: torch.optim.Optimizer, + obs: torch.Tensor, + target: torch.Tensor, + sample_seconds: float, + h2d_seconds: float, +) -> TrainStepResult: + model.train() + optimizer.zero_grad(set_to_none=True) + + forward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + backward_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + optimizer_timing = CudaSectionTiming( + start=torch.cuda.Event(enable_timing=True), + end=torch.cuda.Event(enable_timing=True), + ) + + forward_timing.start.record() + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + pred = model(obs) + loss = F.huber_loss(pred.float(), target) + forward_timing.end.record() + + backward_timing.start.record() + loss.backward() + backward_timing.end.record() + + optimizer_timing.start.record() + optimizer.step() + optimizer_timing.end.record() + + return TrainStepResult( + loss=loss.detach(), + sample_seconds=sample_seconds, + h2d_seconds=h2d_seconds, + forward_timing=forward_timing, + backward_timing=backward_timing, + optimizer_timing=optimizer_timing, + ) + + +def train_step( + model: PackedValueModel, + replay_buffer: EphemeralReplayBuffer, + optimizer: torch.optim.Optimizer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> TrainStepResult: + obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( + replay_buffer, batch_size, seed, device + ) + return _train_step_from_batch( + model, optimizer, obs, target, sample_seconds, h2d_seconds + ) + + +def pretrain_step( + model: PackedValueModel, + optimizer: torch.optim.Optimizer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> TrainStepResult: + obs, target, sample_seconds, h2d_seconds = _sample_terminal_batch( + batch_size, seed, device + ) + return _train_step_from_batch( + model, optimizer, obs, target, sample_seconds, h2d_seconds + ) + + +def _prepare_run_dir(config: TrainConfig) -> tuple[Path, Path]: + run_dir = Path(config.run_dir) + checkpoint_dir = run_dir / "checkpoints" + run_dir.mkdir(parents=True, exist_ok=True) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + (run_dir / "config.json").write_text(json.dumps(asdict(config), indent=2) + "\n") + return run_dir, checkpoint_dir + + +def _set_optimizer_lr(optimizer: torch.optim.Optimizer, lr: float) -> None: + for group in optimizer.param_groups: + group["lr"] = lr + + +def _save_checkpoint( + *, + checkpoint_dir: Path, + model: PackedValueModel, + optimizer: torch.optim.Optimizer, + config: TrainConfig, + round_idx: int, + samples: int, + games: int, + replay_size: int, + record: dict[str, object], +) -> Path: + state = { + "model": model.state_dict(), + "optimizer": optimizer.state_dict(), + "config": asdict(config), + "round": round_idx, + "samples": samples, + "games": games, + "replay_size": replay_size, + "metrics": record, + } + checkpoint_path = checkpoint_dir / f"round_{round_idx:05d}.pt" + torch.save(state, checkpoint_path) + torch.save(state, checkpoint_dir / "latest.pt") + return checkpoint_path + + +def _tensor_summary(prefix: str, values: torch.Tensor) -> dict[str, float]: + values = values.detach().float().cpu() + return { + f"{prefix}_mean": float(values.mean().item()), + f"{prefix}_std": float(values.std(unbiased=False).item()), + f"{prefix}_min": float(values.min().item()), + f"{prefix}_max": float(values.max().item()), + } + + +def _collect_round_diagnostics( + model: PackedValueModel, + replay_buffer: EphemeralReplayBuffer, + *, + batch_size: int, + seed: int, + device: torch.device, +) -> tuple[dict[str, float], dict[str, wandb.Histogram]]: + obs, target, sample_seconds, h2d_seconds = _sample_replay_batch( + replay_buffer, batch_size, seed, device + ) + model.eval() + with torch.no_grad(): + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + pred = model(obs) + torch.cuda.synchronize(device) + + pred = pred.float() + residual = pred - target + stats = { + "diag_sample_seconds": sample_seconds, + "diag_h2d_seconds": h2d_seconds, + **_tensor_summary("diag_target", target), + **_tensor_summary("diag_pred", pred), + **_tensor_summary("diag_residual", residual), + } + histograms = { + "diag_target_hist": wandb.Histogram( + target.detach().float().cpu().numpy().tolist() + ), + "diag_pred_hist": wandb.Histogram(pred.detach().float().cpu().numpy().tolist()), + "diag_residual_hist": wandb.Histogram( + residual.detach().float().cpu().numpy().tolist() + ), + } + return stats, histograms + + +def run_training(config: TrainConfig) -> tuple[PackedValueModel, list[float]]: + device = torch.device("cuda:0") + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required for the training runner") + torch.manual_seed(config.seed) + pretrain_lr = config.lr if config.pretrain_lr is None else config.pretrain_lr + selfplay_lr = config.lr if config.selfplay_lr is None else config.selfplay_lr + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = True + + run_dir, checkpoint_dir = _prepare_run_dir(config) + + model = cast( + PackedValueModel, + PackedValueModel().to(device), + ) + model.eval() + param_count = sum(param.numel() for param in model.parameters()) + print(f"model (nnue) params={param_count}") + + initial_lr = pretrain_lr if config.pretrain_terminal_samples > 0 else selfplay_lr + optimizer = torch.optim.AdamW( + model.parameters(), + lr=initial_lr, + weight_decay=config.weight_decay, + ) + + logger = TrainingLogger( + config=config, + run_dir=run_dir, + checkpoint_dir=checkpoint_dir, + model=model, + param_count=param_count, + enable_wandb=config.wandb, + ) + logger.set_lr_getter(lambda: optimizer.param_groups[0]["lr"]) + + if config.pretrain_terminal_samples > 0: + pretrain_steps = ( + config.pretrain_terminal_samples + config.pretrain_batch_size - 1 + ) // config.pretrain_batch_size + logger.start_pretrain() + for step_idx in range(pretrain_steps): + samples_done = step_idx * config.pretrain_batch_size + step_batch_size = min( + config.pretrain_batch_size, + config.pretrain_terminal_samples - samples_done, + ) + result = pretrain_step( + model, + optimizer, + batch_size=step_batch_size, + seed=config.seed + step_idx, + device=device, + ) + loss_value = float(result.loss.float().cpu().item()) + logger.log_pretrain_step( + step_idx=step_idx, + step_batch_size=step_batch_size, + loss_value=loss_value, + pretrain_steps=pretrain_steps, + device=device, + ) + logger.finish_pretrain(pretrain_steps=pretrain_steps, device=device) + if selfplay_lr != pretrain_lr: + _set_optimizer_lr(optimizer, selfplay_lr) + print(f"selfplay_lr={selfplay_lr:.6g}") + + replay_buffer = EphemeralReplayBuffer(config.replay_capacity) + available_gpus = torch.cuda.device_count() + if available_gpus < 1: + raise RuntimeError("AlphaPaint self-play requires at least one CUDA GPU") + num_gpus = config.num_gpus if config.num_gpus > 0 else available_gpus + if num_gpus > available_gpus: + raise ValueError( + f"Requested --num-gpus={num_gpus}, but only {available_gpus} GPUs are visible" + ) + print( + "selfplay " + f"threads={config.num_threads} workers_per_thread={config.workers_per_thread} " + f"gpus={num_gpus} max_gpu_evals_per_move={config.max_gpu_evals_per_move}" + ) + selfplay = SelfPlay( + replay_buffer, + config.num_threads, + config.workers_per_thread, + config.seed, + max_gpu_evals_per_move=config.max_gpu_evals_per_move, + model=model, + selfplay_precision=config.selfplay_precision, + num_gpus=num_gpus, + ) + # Graph capture via cudaSetDevice may leave the CUDA default on another GPU; + # pin it back to the training device so H2D transfers / diagnostics land on cuda:0. + torch.cuda.set_device(device) + + losses: list[float] = [] + target_samples = 0 + try: + for round_idx in range(config.rounds): + round_number = round_idx + 1 + collect_started_at = time.perf_counter() + target_samples += config.samples_per_round + collected = selfplay.wait_for(target_samples) + collect_seconds = time.perf_counter() - collect_started_at + + round_results: list[TrainStepResult] = [] + train_started_at = time.perf_counter() + if collected > config.replay_capacity // 2: + for step_idx in range(config.train_steps_per_round): + result = train_step( + model, + replay_buffer, + optimizer, + batch_size=config.batch_size, + seed=config.seed + round_idx * 10_000 + step_idx, + device=device, + ) + round_results.append(result) + if round_results: + torch.cuda.synchronize(device) + model.eval() + train_seconds = time.perf_counter() - train_started_at + + train_sample_seconds = sum( + result.sample_seconds for result in round_results + ) + train_h2d_seconds = sum(result.h2d_seconds for result in round_results) + train_forward_seconds = sum( + result.forward_timing.seconds() for result in round_results + ) + train_backward_seconds = sum( + result.backward_timing.seconds() for result in round_results + ) + train_optimizer_seconds = sum( + result.optimizer_timing.seconds() for result in round_results + ) + + if round_results: + round_loss_values = ( + torch.stack([result.loss for result in round_results]).float().cpu() + ) + loss_values = round_loss_values.tolist() + losses.extend(loss_values) + mean_loss = float(round_loss_values.mean().item()) + else: + mean_loss = 0.0 + + diag_stats, diag_histograms = _collect_round_diagnostics( + model, + replay_buffer, + batch_size=config.batch_size, + seed=config.seed + round_idx * 10_000 + config.train_steps_per_round, + device=device, + ) + + record = logger.update( + selfplay=selfplay, + round_number=round_number, + collected=collected, + replay_size=len(replay_buffer), + mean_loss=mean_loss, + collect_seconds=collect_seconds, + train_seconds=train_seconds, + train_sample_seconds=train_sample_seconds, + train_h2d_seconds=train_h2d_seconds, + train_forward_seconds=train_forward_seconds, + train_backward_seconds=train_backward_seconds, + train_optimizer_seconds=train_optimizer_seconds, + diag_stats=diag_stats, + diag_histograms=diag_histograms, + ) + + if config.checkpoint_interval > 0 and ( + round_number % config.checkpoint_interval == 0 + or round_number == config.rounds + ): + checkpoint_path = _save_checkpoint( + checkpoint_dir=checkpoint_dir, + model=model, + optimizer=optimizer, + config=config, + round_idx=round_number, + samples=collected, + games=selfplay.games(), + replay_size=len(replay_buffer), + record=record, + ) + logger.log_checkpoint(checkpoint_path) + finally: + selfplay.drop() + logger.finish() + + return model, losses + + +def _parse_args() -> TrainConfig: + defaults = TrainConfig() + parser = argparse.ArgumentParser(description="Run AlphaPaint value training") + parser.add_argument("--rounds", type=int, default=defaults.rounds) + parser.add_argument( + "--samples-per-round", type=int, default=defaults.samples_per_round + ) + parser.add_argument( + "--train-steps-per-round", type=int, default=defaults.train_steps_per_round + ) + parser.add_argument("--batch-size", type=int, default=defaults.batch_size) + parser.add_argument("--replay-capacity", type=int, default=defaults.replay_capacity) + parser.add_argument( + "--num-threads", + type=int, + default=defaults.num_threads, + help="Self-play executor threads (defaults to Slurm CPU allocation when available)", + ) + parser.add_argument( + "--workers-per-thread", type=int, default=defaults.workers_per_thread + ) + parser.add_argument( + "--max-gpu-evals-per-move", + type=int, + default=defaults.max_gpu_evals_per_move, + ) + parser.add_argument("--lr", type=float, default=defaults.lr) + parser.add_argument("--pretrain-lr", type=float, default=defaults.pretrain_lr) + parser.add_argument("--selfplay-lr", type=float, default=defaults.selfplay_lr) + parser.add_argument("--weight-decay", type=float, default=defaults.weight_decay) + parser.add_argument( + "--pretrain-terminal-samples", + type=int, + default=defaults.pretrain_terminal_samples, + ) + parser.add_argument( + "--pretrain-batch-size", + type=int, + default=defaults.pretrain_batch_size, + ) + parser.add_argument( + "--pretrain-log-interval", + type=int, + default=defaults.pretrain_log_interval, + ) + parser.add_argument("--seed", type=int, default=defaults.seed) + parser.add_argument( + "--selfplay-precision", + default=defaults.selfplay_precision, + choices=["bf16", "fp16", "fp32"], + ) + parser.add_argument( + "--num-gpus", + type=int, + default=defaults.num_gpus, + help="Number of GPUs for self-play (0 uses all visible CUDA GPUs; defaults to Slurm allocation when available)", + ) + parser.add_argument( + "--checkpoint-interval", type=int, default=defaults.checkpoint_interval + ) + parser.add_argument("--run-dir", default=defaults.run_dir) + parser.add_argument("--wandb", dest="wandb", action="store_true") + parser.add_argument("--no-wandb", dest="wandb", action="store_false") + parser.set_defaults(wandb=defaults.wandb) + args = parser.parse_args() + return TrainConfig(**vars(args)) + + +def main() -> None: + run_training(_parse_args()) + + +if __name__ == "__main__": + main() diff --git a/python/scripts/export_value_net.py b/python/scripts/export_value_net.py new file mode 100644 index 0000000..21ae238 --- /dev/null +++ b/python/scripts/export_value_net.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch +from safetensors.torch import save_file + + +def _infer_model_config(state_dict: dict[str, torch.Tensor]) -> dict[str, int]: + width = int(state_dict["value_net.stem.0.weight"].shape[0]) + hidden_dim = int(state_dict["value_net.head.0.weight"].shape[0]) + block_ids = { + int(key.split(".")[2]) + for key in state_dict + if key.startswith("value_net.blocks.") and key.endswith("conv1.weight") + } + return { + "width": width, + "hidden_dim": hidden_dim, + "num_blocks": len(block_ids), + } + + +def _rename_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + renamed: dict[str, torch.Tensor] = {} + + def copy(src: str, dst: str) -> None: + tensor = ( + state_dict[src].detach().to(dtype=torch.float32, device="cpu").contiguous() + ) + renamed[dst] = tensor + + copy("value_net.stem.0.weight", "stem.conv.weight") + copy("value_net.stem.1.weight", "stem.bn.weight") + copy("value_net.stem.1.bias", "stem.bn.bias") + copy("value_net.stem.1.running_mean", "stem.bn.running_mean") + copy("value_net.stem.1.running_var", "stem.bn.running_var") + + model_config = _infer_model_config(state_dict) + block_ids = sorted( + { + int(key.split(".")[2]) + for key in state_dict + if key.startswith("value_net.blocks.") and key.endswith("conv1.weight") + } + ) + + for block_id in block_ids: + prefix = f"value_net.blocks.{block_id}" + copy(f"{prefix}.norm1.weight", f"blocks.{block_id}.norm1.weight") + copy(f"{prefix}.norm1.bias", f"blocks.{block_id}.norm1.bias") + copy(f"{prefix}.norm1.running_mean", f"blocks.{block_id}.norm1.running_mean") + copy(f"{prefix}.norm1.running_var", f"blocks.{block_id}.norm1.running_var") + copy(f"{prefix}.conv1.weight", f"blocks.{block_id}.conv1.weight") + copy(f"{prefix}.norm2.weight", f"blocks.{block_id}.norm2.weight") + copy(f"{prefix}.norm2.bias", f"blocks.{block_id}.norm2.bias") + copy(f"{prefix}.norm2.running_mean", f"blocks.{block_id}.norm2.running_mean") + copy(f"{prefix}.norm2.running_var", f"blocks.{block_id}.norm2.running_var") + copy(f"{prefix}.conv2.weight", f"blocks.{block_id}.conv2.weight") + + copy("value_net.head.0.weight", "head.fc1.weight") + copy("value_net.head.0.bias", "head.fc1.bias") + copy("value_net.head.2.weight", "head.fc2.weight") + copy("value_net.head.2.bias", "head.fc2.bias") + + renamed["__meta.width"] = torch.tensor([model_config["width"]], dtype=torch.float32) + renamed["__meta.hidden_dim"] = torch.tensor( + [model_config["hidden_dim"]], dtype=torch.float32 + ) + renamed["__meta.num_blocks"] = torch.tensor( + [model_config["num_blocks"]], dtype=torch.float32 + ) + return renamed + + +def _extract_state_dict(checkpoint: dict[str, Any]) -> dict[str, torch.Tensor]: + maybe_model = checkpoint.get("model") + if isinstance(maybe_model, dict): + return maybe_model + return checkpoint + + +def _build_fixtures( + checkpoint_path: Path, + fixture_path: Path, + fixture_count: int, + seed: int, +) -> None: + if fixture_count <= 0: + return + + from alphapaint_training import ( + decode_packed_observation, + sample_random_terminal_batch, + ) + + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + state_dict = _extract_state_dict(checkpoint) + from alphapaint_training.model import PackedValueModel + + model = PackedValueModel( + board_dtype=torch.float32, + **_infer_model_config(state_dict), + ) + model.load_state_dict(state_dict) + model.eval() + + packed_obs_np, _ = sample_random_terminal_batch(fixture_count, seed) + packed_obs = torch.from_numpy(packed_obs_np).to(dtype=torch.uint16) + board, intrinsics = decode_packed_observation( + packed_obs, + board_dtype=torch.float32, + intrinsic_dtype=torch.float32, + ) + with torch.no_grad(): + outputs = model(packed_obs.to(dtype=torch.uint16)).to(dtype=torch.float32) + + fixture_payload = { + "seed": seed, + "samples": [ + { + "board": board[idx].reshape(-1).tolist(), + "intrinsics": intrinsics[idx].tolist(), + "output": float(outputs[idx].item()), + } + for idx in range(fixture_count) + ], + } + fixture_path.write_text(json.dumps(fixture_payload, indent=2) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Export AlphaPaint value network checkpoints for Rust CPU inference." + ) + parser.add_argument( + "--checkpoint", type=Path, required=True, help="Input .pt checkpoint" + ) + parser.add_argument( + "--output", + type=Path, + default=Path("alpha_paint/assets/value_net.safetensors"), + help="Output safetensors path", + ) + parser.add_argument( + "--fixtures", + type=Path, + default=None, + help="Optional JSON fixture output with unpacked features and model outputs", + ) + parser.add_argument( + "--fixture-count", + type=int, + default=0, + help="Number of random terminal samples to export as fixtures", + ) + parser.add_argument("--seed", type=int, default=42, help="Fixture sampling seed") + args = parser.parse_args() + + checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + state_dict = _extract_state_dict(checkpoint) + renamed = _rename_state_dict(state_dict) + + metadata = { + "width": str(int(renamed.pop("__meta.width").item())), + "num_blocks": str(int(renamed.pop("__meta.num_blocks").item())), + "hidden_dim": str(int(renamed.pop("__meta.hidden_dim").item())), + "source_checkpoint": str(args.checkpoint), + } + + args.output.parent.mkdir(parents=True, exist_ok=True) + save_file(renamed, str(args.output), metadata=metadata) + + if args.fixtures is not None: + args.fixtures.parent.mkdir(parents=True, exist_ok=True) + _build_fixtures(args.checkpoint, args.fixtures, args.fixture_count, args.seed) + + print(f"exported={args.output}") + if args.fixtures is not None and args.fixture_count > 0: + print(f"fixtures={args.fixtures}") + + +if __name__ == "__main__": + main() diff --git a/tooling/src/main.rs b/tooling/src/main.rs index 2f01c09..62aa592 100644 --- a/tooling/src/main.rs +++ b/tooling/src/main.rs @@ -2,4 +2,4 @@ use alpha_paint::perft_test; fn main() { perft_test() -} \ No newline at end of file +} diff --git a/training/Cargo.toml b/training/Cargo.toml new file mode 100644 index 0000000..b8b586f --- /dev/null +++ b/training/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "alphapaint_training" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html +[lib] +name = "alphapaint_training" +crate-type = ["cdylib"] + +[dependencies] +alpha_paint = { path = "../alpha_paint" } +ndarray = "0.17.1" +pyo3 = { version = "0.28.2", features = ["extension-module"] } +rand = "0.10.0" +rand_chacha = "0.10.0" +rand_distr = "0.5" +event-listener = "5" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +rayon = "1.11.0" +cudarc = { version = "0.18.2", features = ["cuda-version-from-build-system"] } +numpy = "0.28.0" + +[dev-dependencies] +rstest = "0.26.1" diff --git a/training/src/cudagraph.rs b/training/src/cudagraph.rs new file mode 100644 index 0000000..edbc4f5 --- /dev/null +++ b/training/src/cudagraph.rs @@ -0,0 +1,510 @@ +//! Rust-launched CUDA graph backend for AlphaPaint value-only inference. +//! +//! No policy head - the model outputs only a scalar value per position. + +use std::ffi::{c_void, CStr}; +use std::mem::size_of; +use std::slice; +use std::sync::atomic::{AtomicU64, Ordering}; + +use cudarc::runtime::sys as cuda; +use ndarray::{ArrayView, Ix2}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::ffi; +use pyo3::prelude::*; + +use crate::queue::BatchCompletion; + +#[allow(non_camel_case_types)] +type cudaStream_t = cuda::cudaStream_t; + +#[allow(non_camel_case_types)] +type cudaGraphExec_t = cuda::cudaGraphExec_t; + +type CudaError = cuda::cudaError_t; + +const DL_TENSOR_NAME: &[u8] = b"dltensor\0"; +const DL_DEVICE_CPU: i32 = 1; +const DL_DEVICE_CUDA: i32 = 2; +const DL_DTYPE_FLOAT: u8 = 2; +const DL_DTYPE_UINT: u8 = 1; + +pub const OBS_WORDS: usize = crate::observation::OBS_WORDS; + +#[repr(C)] +struct DLDevice { + device_type: i32, + device_id: i32, +} + +#[repr(C)] +struct DLDataType { + code: u8, + bits: u8, + lanes: u16, +} + +#[repr(C)] +struct DLTensor { + data: *mut c_void, + device: DLDevice, + ndim: i32, + dtype: DLDataType, + shape: *mut i64, + strides: *mut i64, + byte_offset: usize, +} + +#[repr(C)] +struct DLManagedTensor { + dl_tensor: DLTensor, + manager_ctx: *mut c_void, + deleter: Option, +} + +struct DLPackContext { + shape: Box<[i64]>, +} + +unsafe extern "C" fn dlpack_capsule_destructor(capsule: *mut ffi::PyObject) { + if capsule.is_null() { + return; + } + + let name = ffi::PyCapsule_GetName(capsule); + if name.is_null() { + return; + } + + let c_name = CStr::from_ptr(name); + if c_name.to_bytes() != b"dltensor" { + return; + } + + let ptr = ffi::PyCapsule_GetPointer(capsule, DL_TENSOR_NAME.as_ptr() as *const i8); + if ptr.is_null() { + return; + } + + let managed = ptr as *mut DLManagedTensor; + if let Some(deleter) = unsafe { (*managed).deleter } { + deleter(managed); + } +} + +extern "C" fn dlpack_deleter(ptr: *mut DLManagedTensor) { + if ptr.is_null() { + return; + } + + unsafe { + let ctx_ptr = (*ptr).manager_ctx as *mut DLPackContext; + if !ctx_ptr.is_null() { + drop(Box::from_raw(ctx_ptr)); + } + drop(Box::from_raw(ptr)); + } +} + +fn dlpack_capsule( + py: Python<'_>, + data: *mut c_void, + shape: &[i64], + device_type: i32, + device_id: i32, + dtype_code: u8, + dtype_bits: u8, +) -> PyResult> { + let ctx = Box::new(DLPackContext { + shape: shape.to_vec().into_boxed_slice(), + }); + let shape_ptr = ctx.shape.as_ptr() as *mut i64; + let ctx_ptr = Box::into_raw(ctx); + + let managed = Box::new(DLManagedTensor { + dl_tensor: DLTensor { + data, + device: DLDevice { + device_type, + device_id, + }, + ndim: shape.len() as i32, + dtype: DLDataType { + code: dtype_code, + bits: dtype_bits, + lanes: 1, + }, + shape: shape_ptr, + strides: std::ptr::null_mut(), + byte_offset: 0, + }, + manager_ctx: ctx_ptr as *mut c_void, + deleter: Some(dlpack_deleter), + }); + + let managed_ptr = Box::into_raw(managed); + let capsule = unsafe { + ffi::PyCapsule_New( + managed_ptr as *mut c_void, + DL_TENSOR_NAME.as_ptr() as *const i8, + Some(dlpack_capsule_destructor), + ) + }; + + if capsule.is_null() { + dlpack_deleter(managed_ptr); + return Err(PyErr::new::( + "failed to create DLPack capsule", + )); + } + + Ok(unsafe { Py::from_owned_ptr(py, capsule) }) +} + +fn check_cuda(code: CudaError, context: &str) -> PyResult<()> { + if code == cuda::cudaError::cudaSuccess { + Ok(()) + } else { + Err(PyErr::new::(format!( + "{} failed with CUDA error {:?}", + context, code + ))) + } +} + +fn check_cuda_or_panic(code: CudaError, context: &str) { + if code != cuda::cudaError::cudaSuccess { + panic!("{} failed with CUDA error {:?}", context, code); + } +} + +fn validate_cuda_device(gpu_id: usize) -> PyResult { + let gpu_device_id = i32::try_from(gpu_id) + .map_err(|_| PyErr::new::("gpu_id must fit in i32"))?; + + let mut device_count = 0i32; + unsafe { + check_cuda( + cuda::cudaGetDeviceCount(&mut device_count as *mut i32), + "cudaGetDeviceCount", + )?; + } + + if device_count <= 0 { + return Err(PyErr::new::( + "no CUDA devices available for AlphaPaint CUDA graph runner", + )); + } + + if gpu_device_id >= device_count { + return Err(PyErr::new::(format!( + "gpu_id {} out of range for {} CUDA devices", + gpu_id, device_count + ))); + } + + Ok(gpu_device_id) +} + +fn cuda_malloc_host_f32(count: usize, context: &str) -> PyResult<*mut f32> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("host allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMallocHost(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr.cast::()) +} + +fn cuda_malloc_host_u16(count: usize, context: &str) -> PyResult<*mut u16> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("host allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMallocHost(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr.cast::()) +} + +fn cuda_malloc_device_f32(count: usize, context: &str) -> PyResult<*mut c_void> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("device allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMalloc(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr) +} + +fn cuda_malloc_device_u16(count: usize, context: &str) -> PyResult<*mut c_void> { + let mut ptr: *mut c_void = std::ptr::null_mut(); + let bytes = count + .checked_mul(size_of::()) + .ok_or_else(|| PyErr::new::("device allocation size overflow"))?; + unsafe { + check_cuda( + cuda::cudaMalloc(&mut ptr as *mut *mut c_void, bytes), + context, + )?; + } + Ok(ptr) +} + +struct CudaGraphLane { + gpu_device_id: i32, + stream: cudaStream_t, + graph_exec: cudaGraphExec_t, + /// Owns Python-side graph/tensor objects for this lane. + _py_owner: Py, + obs_host: *mut u16, + obs_dev: *mut c_void, + value_host: *mut f32, + value_dev: *mut c_void, +} + +struct LaneCompletionContext { + value_host: *const f32, + batch_size: usize, + completion: Option>, +} + +unsafe impl Send for LaneCompletionContext {} + +unsafe extern "C" fn lane_completion_callback(user_data: *mut c_void) { + if user_data.is_null() { + return; + } + + let mut ctx = unsafe { Box::from_raw(user_data.cast::()) }; + let value_src = unsafe { slice::from_raw_parts(ctx.value_host, ctx.batch_size) }; + + let outputs: Vec = value_src.to_vec(); + + if let Some(completion) = ctx.completion.take() { + completion.complete(&outputs); + } +} + +impl Drop for CudaGraphLane { + fn drop(&mut self) { + unsafe { + let _ = cuda::cudaSetDevice(self.gpu_device_id); + let _ = cuda::cudaFree(self.obs_dev); + let _ = cuda::cudaFree(self.value_dev); + let _ = cuda::cudaFreeHost(self.obs_host.cast::()); + let _ = cuda::cudaFreeHost(self.value_host.cast::()); + let _ = cuda::cudaStreamDestroy(self.stream); + } + } +} + +/// Per-lane CUDA graph executor for AlphaPaint value-only inference. +pub struct CudaGraphRunner { + gpu_device_id: i32, + batch_size: usize, + lanes: Vec, + dispatched_batches: AtomicU64, + dispatched_evals: AtomicU64, +} + +// SAFETY: Lane buffers/streams are independent per batch_idx and queue dispatch +// ensures a lane is not reused before dispatch returns for that lane. +unsafe impl Send for CudaGraphRunner {} +unsafe impl Sync for CudaGraphRunner {} + +impl CudaGraphRunner { + pub fn new( + py: Python<'_>, + model: Py, + gpu_id: usize, + num_lanes: usize, + batch_size: usize, + precision: &str, + ) -> PyResult { + if num_lanes == 0 { + return Err(PyErr::new::("num_lanes must be > 0")); + } + if batch_size == 0 { + return Err(PyErr::new::("batch_size must be > 0")); + } + + let gpu_device_id = validate_cuda_device(gpu_id)?; + unsafe { + check_cuda( + cuda::cudaSetDevice(gpu_device_id), + "cudaSetDevice in CudaGraphRunner::new", + )?; + } + + let module = PyModule::import(py, "alphapaint_training.cudagraph_backend")?; + let capture_fn = module.getattr("capture_lane_graph")?; + + let obs_count = batch_size * OBS_WORDS; + let value_count = batch_size; + + let obs_shape = [batch_size as i64, OBS_WORDS as i64]; + let value_shape = [batch_size as i64]; + + let mut lanes = Vec::with_capacity(num_lanes); + + for lane_idx in 0..num_lanes { + let mut stream: cudaStream_t = std::ptr::null_mut(); + unsafe { + check_cuda( + cuda::cudaStreamCreate(&mut stream as *mut cudaStream_t), + "cudaStreamCreate", + )?; + } + + let obs_host = + cuda_malloc_host_u16(obs_count, &format!("cudaMallocHost obs lane {}", lane_idx))?; + let value_host = cuda_malloc_host_f32( + value_count, + &format!("cudaMallocHost value lane {}", lane_idx), + )?; + + let obs_dev = + cuda_malloc_device_u16(obs_count, &format!("cudaMalloc obs lane {}", lane_idx))?; + let value_dev = cuda_malloc_device_f32( + value_count, + &format!("cudaMalloc value lane {}", lane_idx), + )?; + + let obs_host_capsule = dlpack_capsule( + py, + obs_host.cast::(), + &obs_shape, + DL_DEVICE_CPU, + 0, + DL_DTYPE_UINT, + 16, + )?; + let obs_dev_capsule = dlpack_capsule( + py, + obs_dev, + &obs_shape, + DL_DEVICE_CUDA, + gpu_device_id, + DL_DTYPE_UINT, + 16, + )?; + let value_host_capsule = dlpack_capsule( + py, + value_host.cast::(), + &value_shape, + DL_DEVICE_CPU, + 0, + DL_DTYPE_FLOAT, + 32, + )?; + let value_dev_capsule = dlpack_capsule( + py, + value_dev, + &value_shape, + DL_DEVICE_CUDA, + gpu_device_id, + DL_DTYPE_FLOAT, + 32, + )?; + + let (exec_handle, py_owner): (u64, Py) = capture_fn + .call1(( + model.clone_ref(py), + obs_host_capsule, + obs_dev_capsule, + value_host_capsule, + value_dev_capsule, + stream as u64, + precision, + gpu_id, + ))? + .extract()?; + + let lane = CudaGraphLane { + gpu_device_id, + stream, + graph_exec: exec_handle as cudaGraphExec_t, + _py_owner: py_owner, + obs_host, + obs_dev, + value_host, + value_dev, + }; + lanes.push(lane); + } + + Ok(Self { + gpu_device_id, + batch_size, + lanes, + dispatched_batches: AtomicU64::new(0), + dispatched_evals: AtomicU64::new(0), + }) + } + + pub fn dispatched_batches(&self) -> u64 { + self.dispatched_batches.load(Ordering::Relaxed) + } + + pub fn dispatched_evals(&self) -> u64 { + self.dispatched_evals.load(Ordering::Relaxed) + } + + pub fn dispatch_async( + &self, + batch_idx: usize, + obs_view: ArrayView, + completion: BatchCompletion, + ) { + debug_assert_eq!(obs_view.shape(), &[self.batch_size, OBS_WORDS]); + + let lane = &self.lanes[batch_idx % self.lanes.len()]; + + let obs_src = obs_view + .as_slice() + .expect("observation batch must be contiguous"); + let obs_dst = + unsafe { slice::from_raw_parts_mut(lane.obs_host, self.batch_size * OBS_WORDS) }; + obs_dst.copy_from_slice(obs_src); + self.dispatched_batches.fetch_add(1, Ordering::Relaxed); + self.dispatched_evals + .fetch_add(self.batch_size as u64, Ordering::Relaxed); + + unsafe { + check_cuda_or_panic( + cuda::cudaSetDevice(self.gpu_device_id), + "cudaSetDevice before cudaGraphLaunch", + ); + check_cuda_or_panic( + cuda::cudaGraphLaunch(lane.graph_exec, lane.stream), + "cudaGraphLaunch", + ); + + let ctx = Box::new(LaneCompletionContext { + value_host: lane.value_host, + batch_size: self.batch_size, + completion: Some(completion), + }); + check_cuda_or_panic( + cuda::cudaLaunchHostFunc( + lane.stream, + Some(lane_completion_callback), + Box::into_raw(ctx).cast::(), + ), + "cudaLaunchHostFunc", + ); + } + } +} diff --git a/training/src/descent.rs b/training/src/descent.rs new file mode 100644 index 0000000..21fe59b --- /dev/null +++ b/training/src/descent.rs @@ -0,0 +1,1153 @@ +//! Async Descent/UBFM search for Athénan-style training. +//! +//! Uses GPU neural network evaluation for leaf expansion. Values are f32. +//! Implements tree learning (collect training samples from internal nodes) +//! and ordinal distribution for action selection. + +use alpha_paint::board::{Action, ApplyActionOutcome, Board, Rollback, TerminalState}; +use rand::rngs::SmallRng; +#[cfg(test)] +use rand::SeedableRng; +use rand::{Rng, RngExt}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Instant; + +use crate::eval::{EvalCountTracker, Evaluator}; + +const AVG_GAME_LENGTH: f32 = 500.0; + +pub struct TreeLearningSample { + pub board: Board, + pub value: f32, + pub terminal_mix: bool, +} + +#[derive(Clone)] +pub struct SearchTimingMetrics { + pub expand_cpu_nanos: Arc, + pub apply_action_nanos: Arc, + pub eval_submit_nanos: Arc, + pub eval_await_nanos: Arc, + pub backup_nanos: Arc, +} + +impl SearchTimingMetrics { + fn record(counter: &AtomicU64, started_at: Instant) { + counter.fetch_add( + started_at.elapsed().as_nanos().min(u64::MAX as u128) as u64, + Ordering::AcqRel, + ); + } + + fn record_expand_cpu(&self, started_at: Instant) { + Self::record(&self.expand_cpu_nanos, started_at); + } + + fn record_apply_action(&self, started_at: Instant) { + Self::record(&self.apply_action_nanos, started_at); + } + + fn record_eval_submit(&self, started_at: Instant) { + Self::record(&self.eval_submit_nanos, started_at); + } + + fn record_eval_await(&self, started_at: Instant) { + Self::record(&self.eval_await_nanos, started_at); + } + + fn record_backup(&self, started_at: Instant) { + Self::record(&self.backup_nanos, started_at); + } +} + +#[derive(Debug)] +pub struct ChildData { + pub action: Action, + pub child_value: f32, + entrance_count: usize, + pub node: Option>, +} + +impl ChildData { + fn completion_value(&self) -> i32 { + match &self.node { + Some(c) => c.completion_value, + None => 0, + } + } +} + +#[derive(Debug)] +pub struct SearchNode { + pub value: f32, + pub completion_value: i32, + resolved: bool, + pub children: Vec, +} + +impl SearchNode { + fn new(value: f32, completion_value: i32, is_resolved: bool) -> SearchNode { + SearchNode { + value, + completion_value, + resolved: is_resolved, + children: vec![], + } + } + + /// Prefers moves with low entrance counts (exploration). + fn completed_best_action_dual( + &self, + is_max_player: bool, + rng: &mut SmallRng, + ) -> (usize, Action) { + let res = if is_max_player { + self.children + .iter() + .enumerate() + .max_by(|&(_, a), &(_, b)| { + ( + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), + ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) + }) + .expect("Called completed_best_action_dual without any valid actions!") + } else { + self.children + .iter() + .enumerate() + .min_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) + }) + .expect("Called completed_best_action_dual without any valid actions!") + }; + + (res.0, res.1.action) + } + + /// Prefers moves with high entrance counts (exploitation). + fn completed_best_action(&self, is_max_player: bool, rng: &mut SmallRng) -> (usize, Action) { + let res = if is_max_player { + self.children + .iter() + .enumerate() + .max_by(|&(_, a), &(_, b)| { + (a.completion_value(), a.child_value, a.entrance_count as i32) + .partial_cmp(&( + b.completion_value(), + b.child_value, + b.entrance_count as i32, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) + }) + .expect("Called completed_best_action without any valid actions!") + } else { + self.children + .iter() + .enumerate() + .min_by(|&(_, a), &(_, b)| { + ( + a.completion_value(), + a.child_value, + -(a.entrance_count as i32), + ) + .partial_cmp(&( + b.completion_value(), + b.child_value, + -(b.entrance_count as i32), + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| rng.next_u32().cmp(&rng.next_u32())) + }) + .expect("Called completed_best_action without any valid actions!") + }; + + (res.0, res.1.action) + } + + fn backup_resolution(&self) -> bool { + if self.completion_value.abs() == 1 { + true + } else { + self.children + .iter() + .all(|child| child.node.as_ref().map(|c| c.resolved).unwrap_or(false)) + } + } + + /// Build a chain of resolved SearchNodes for a killshot action sequence. + fn build_killshot_chain(board: &Board, terminal: TerminalState, actions: &[Action]) -> SearchNode { + let term_value = Self::value_from_term(board, terminal); + let comp_value = terminal.value(); + + let mut node = SearchNode::new(term_value, comp_value, true); + + for i in (0..actions.len()).rev() { + let action = actions[i]; + + let parent = SearchNode { + value: term_value, + completion_value: comp_value, + resolved: true, + children: vec![ChildData { + action, + child_value: term_value, + entrance_count: 0, + node: Some(Box::new(node)), + }], + }; + + node = parent; + } + + node + } + + fn apply_play_instead( + mut board: Board, + action: Action, + rollback: Rollback, + play_instead: Action, + ) -> Board { + board.rollback(action, rollback); + let (outcome, _) = board.apply_action(play_instead); + debug_assert!(matches!(outcome, ApplyActionOutcome::Terminal { .. })); + board + } + + fn apply_killshot_actions(mut board: Board, actions: &[Action]) -> Board { + let last_idx = actions.len().saturating_sub(1); + for (idx, action) in actions.iter().copied().enumerate() { + let (outcome, _) = board.apply_action(action); + if idx == last_idx { + debug_assert!(matches!(outcome, ApplyActionOutcome::Terminal { .. })); + } else { + debug_assert!(matches!(outcome, ApplyActionOutcome::Ongoing)); + } + } + board + } + + fn terminal_leaf_sample_board( + state: &Board, + action: Action, + rollback: Rollback, + outcome: ApplyActionOutcome, + ) -> Board { + match outcome { + ApplyActionOutcome::Terminal { .. } => state.clone(), + ApplyActionOutcome::PlayInstead { play_instead, .. } => { + Self::apply_play_instead(state.clone(), action, rollback, play_instead) + } + ApplyActionOutcome::Killshot { actions, .. } => { + Self::apply_killshot_actions(state.clone(), &actions) + } + ApplyActionOutcome::Ongoing => { + unreachable!("terminal leaf sampling only applies to resolved terminal outcomes") + } + } + } + + /// Expand a node: evaluate all children with the neural network. + async fn build_self( + board: &Board, + outcome: ApplyActionOutcome, + evaluator: &E, + rng: &mut SmallRng, + timing: &SearchTimingMetrics, + ) -> SearchNode { + match outcome { + ApplyActionOutcome::Ongoing => { + let mut new_node = SearchNode::new(0.0, 0, false); + let actions_started_at = Instant::now(); + let actions = board.get_valid_actions(); + timing.record_expand_cpu(actions_started_at); + new_node.children.reserve(actions.len()); + + if actions.len() == 0 { + let loss = TerminalState::loss_for(board.is_white_turn()); + return SearchNode::new(Self::value_from_term(board, loss), loss.value(), true); + } + + // Evaluate all children with GPU neural net. + // We create all futures first (submitting observations eagerly), + // then await them. This enables batching across children. + struct ChildEvalResult { + action: Action, + value: f32, + } + + let mut eval_results = Vec::new(); + let batch: Vec<_> = actions + .into_iter() + .map(|action| { + let prep_started_at = Instant::now(); + let mut local_board = board.clone(); + timing.record_expand_cpu(prep_started_at); + + let apply_started_at = Instant::now(); + let (child_outcome, rollback) = local_board.apply_action(*action); + timing.record_apply_action(apply_started_at); + // NOTE: our 'futures' from the GPU Queue do not follow normal rust future semantics + // rust futures are normally lazily evaluated when you .await them for the + // first time. Ours are eager, which is why this code actually works as we expect. + let submit_started_at = Instant::now(); + let value = evaluator.evaluate(local_board.clone()); + timing.record_eval_submit(submit_started_at); + (action, local_board, child_outcome, rollback, value) + }) + .collect(); + + for (&action, local_board, child_outcome, rollback, value) in batch { + match child_outcome { + ApplyActionOutcome::Ongoing => { + // Evaluate this child with neural net + let await_started_at = Instant::now(); + let value = value.await; + timing.record_eval_await(await_started_at); + eval_results.push(ChildEvalResult { action, value }); + } + ApplyActionOutcome::Terminal { terminal } => { + let term_value = Self::value_from_term(&local_board, terminal); + new_node.children.push(ChildData { + action, + child_value: term_value, + entrance_count: 0, + node: Some(Box::new(SearchNode::new( + term_value, + terminal.value(), + true, + ))), + }); + } + ApplyActionOutcome::PlayInstead { + terminal, + play_instead, + } => { + let terminal_board = Self::apply_play_instead( + local_board, + action, + rollback, + play_instead, + ); + let term_value = Self::value_from_term(&terminal_board, terminal); + new_node.children.push(ChildData { + action, + child_value: term_value, + entrance_count: 0, + node: Some(Box::new(SearchNode::new( + term_value, + terminal.value(), + true, + ))), + }); + } + ApplyActionOutcome::Killshot { terminal, actions } => { + let terminal_board = Self::apply_killshot_actions(local_board, &actions); + let term_value = Self::value_from_term(&terminal_board, terminal); + let chain = + Self::build_killshot_chain(&terminal_board, terminal, &actions); + new_node.children.push(ChildData { + action, + child_value: term_value, + entrance_count: 0, + node: Some(Box::new(chain)), + }); + } + } + } + + // Now store the neural net evaluation results + for result in eval_results { + new_node.children.push(ChildData { + action: result.action, + child_value: result.value, + entrance_count: 0, + node: None, + }); + } + + if new_node.children.is_empty() { + let loss = TerminalState::loss_for(board.is_white_turn()); + return SearchNode::new(Self::value_from_term(board, loss), loss.value(), true); + } + + let (best_action_id, _) = + new_node.completed_best_action(board.is_white_turn(), rng); + new_node.completion_value = new_node.children[best_action_id].completion_value(); + new_node.value = new_node.children[best_action_id].child_value; + new_node.resolved = new_node.backup_resolution(); + new_node + } + ApplyActionOutcome::Terminal { terminal } + | ApplyActionOutcome::PlayInstead { terminal, .. } => SearchNode::new( + Self::value_from_term(board, terminal), + terminal.value(), + true, + ), + ApplyActionOutcome::Killshot { terminal, actions } => { + Self::build_killshot_chain(board, terminal, &actions) + } + } + } + + /// Reinforcement heuristic: preserve faster-win ordering without exploding. + fn value_from_term(board: &Board, term: TerminalState) -> f32 { + let sign = term.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + + let p = board.turn_count.max(1) as f32; + let depth = (AVG_GAME_LENGTH / p).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + (sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth)).clamp(-1.0, 1.0) + } + + async fn ubfms_iteration( + &mut self, + mut state: Board, + outcome: ApplyActionOutcome, + evaluator: &E, + rng: &mut SmallRng, + timing: &SearchTimingMetrics, + ) -> f32 { + let white_turn = state.is_white_turn(); + + match outcome { + ApplyActionOutcome::Ongoing => { + if self.children.is_empty() { + let loss = TerminalState::loss_for(white_turn); + self.resolved = true; + self.completion_value = loss.value(); + self.value = Self::value_from_term(&state, loss); + return self.value; + } + + if !self.resolved { + let (best_action_id, best_action) = + self.completed_best_action_dual(white_turn, rng); + + self.children[best_action_id].entrance_count += 1; + + if self.children[best_action_id].node.is_some() { + let apply_started_at = Instant::now(); + let (outcome, _) = state.apply_action(best_action); + timing.record_apply_action(apply_started_at); + let child = self.children[best_action_id].node.as_mut().unwrap(); + self.children[best_action_id].child_value = + Box::pin(child.ubfms_iteration(state, outcome, evaluator, rng, timing)) + .await; + } else { + let apply_started_at = Instant::now(); + let (child_outcome, _) = state.apply_action(best_action); + timing.record_apply_action(apply_started_at); + let should_descend = matches!(child_outcome, ApplyActionOutcome::Ongoing); + let child_node = Box::new( + Box::pin(SearchNode::build_self( + &state, + child_outcome, + evaluator, + rng, + timing, + )) + .await, + ); + + self.children[best_action_id].child_value = child_node.value; + self.children[best_action_id].node = Some(child_node); + + if should_descend { + let child = self.children[best_action_id].node.as_mut().unwrap(); + self.children[best_action_id].child_value = + Box::pin(child.ubfms_iteration( + state, + ApplyActionOutcome::Ongoing, + evaluator, + rng, + timing, + )) + .await; + } + } + + let backup_started_at = Instant::now(); + let (best_action_id, _) = self.completed_best_action(white_turn, rng); + self.completion_value = self.children[best_action_id].completion_value(); + self.value = self.children[best_action_id].child_value; + self.resolved = self.backup_resolution(); + timing.record_backup(backup_started_at); + } + } + + ApplyActionOutcome::Terminal { terminal } + | ApplyActionOutcome::PlayInstead { terminal, .. } + | ApplyActionOutcome::Killshot { terminal, .. } => { + self.resolved = true; + self.completion_value = terminal.value(); + self.value = Self::value_from_term(&state, terminal); + } + } + + self.value + } + + /// Compute the minimax value considering only expanded children. + /// This is used for training targets to avoid including unexpanded + /// children's raw network predictions in the value signal. + fn compute_backed_value(&self, is_white_turn: bool) -> f32 { + if is_white_turn { + self.children + .iter() + .filter(|c| c.node.is_some()) + .map(|c| c.child_value) + .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .expect("compute_backed_value called without expanded children") + } else { + self.children + .iter() + .filter(|c| c.node.is_some()) + .map(|c| c.child_value) + .min_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) + .expect("compute_backed_value called without expanded children") + } + } + + /// Collect tree learning samples from all internal nodes. + /// + /// An internal node is one that has children and at least one expanded child. + /// Non-terminal leaf nodes (where the network estimate was used without + /// minimax backing) are excluded per Athénan's tree learning rules. + fn collect_samples_with_terminal_mix( + &self, + state: &mut Board, + out: &mut Vec, + rng: &mut R, + terminal_mix_prob: f32, + ) { + let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); + if has_expanded_child { + // Use backed value (only expanded children) for training targets + let backed_value = self.compute_backed_value(state.is_white_turn()); + out.push(TreeLearningSample { + board: state.clone(), + value: backed_value, + terminal_mix: false, + }); + } else if self.resolved + && self.children.is_empty() + && state.get_valid_actions().len() == 0 + && rng.random::() <= terminal_mix_prob + { + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + terminal_mix: true, + }); + } + + for child in &self.children { + if let Some(node) = child.node.as_ref() { + let (outcome, rollback) = state.apply_action(child.action); + if node.children.is_empty() + && matches!( + outcome, + ApplyActionOutcome::Terminal { .. } + | ApplyActionOutcome::PlayInstead { .. } + ) + { + if rng.random::() <= terminal_mix_prob { + out.push(TreeLearningSample { + board: Self::terminal_leaf_sample_board( + state, + child.action, + rollback, + outcome, + ), + value: node.value, + terminal_mix: true, + }); + } + } else { + node.collect_samples_with_terminal_mix(state, out, rng, terminal_mix_prob); + } + state.rollback(child.action, rollback); + } + } + } + + #[cfg(test)] + fn collect_samples(&self, state: &mut Board, out: &mut Vec) { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_samples_with_terminal_mix(state, out, &mut rng, 0.0); + } + + fn collect_samples_excluding_child_with_terminal_mix( + &self, + state: &mut Board, + excluded_child: usize, + out: &mut Vec, + rng: &mut R, + terminal_mix_prob: f32, + ) { + let has_expanded_child = self.children.iter().any(|c| c.node.is_some()); + if has_expanded_child { + // Use backed value (only expanded children) for training targets + let backed_value = self.compute_backed_value(state.is_white_turn()); + out.push(TreeLearningSample { + board: state.clone(), + value: backed_value, + terminal_mix: false, + }); + } else if self.resolved + && self.children.is_empty() + && state.get_valid_actions().len() == 0 + && rng.random::() <= terminal_mix_prob + { + out.push(TreeLearningSample { + board: state.clone(), + value: self.value, + terminal_mix: true, + }); + } + + for (child_idx, child) in self.children.iter().enumerate() { + if child_idx == excluded_child { + continue; + } + if let Some(node) = child.node.as_ref() { + let (outcome, rollback) = state.apply_action(child.action); + if node.children.is_empty() + && matches!( + outcome, + ApplyActionOutcome::Terminal { .. } + | ApplyActionOutcome::PlayInstead { .. } + ) + { + if rng.random::() <= terminal_mix_prob { + out.push(TreeLearningSample { + board: Self::terminal_leaf_sample_board( + state, + child.action, + rollback, + outcome, + ), + value: node.value, + terminal_mix: true, + }); + } + } else { + node.collect_samples_with_terminal_mix(state, out, rng, terminal_mix_prob); + } + state.rollback(child.action, rollback); + } + } + } + + #[cfg(test)] + fn collect_samples_excluding_child( + &self, + state: &mut Board, + excluded_child: usize, + out: &mut Vec, + ) { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_samples_excluding_child_with_terminal_mix( + state, + excluded_child, + out, + &mut rng, + 0.0, + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alpha_paint::board::board_structs::Player; + use rand::rngs::SmallRng; + use rand::SeedableRng; + + struct NeverEval; + + impl Evaluator for NeverEval { + fn evaluate(&self, _board: Board) -> impl std::future::Future { + std::future::ready(0.0) + } + } + + impl EvalCountTracker for NeverEval { + fn reset_eval_count(&self) {} + + fn eval_count(&self) -> u64 { + 0 + } + } + + fn zero_timing() -> SearchTimingMetrics { + SearchTimingMetrics { + expand_cpu_nanos: Arc::new(AtomicU64::new(0)), + apply_action_nanos: Arc::new(AtomicU64::new(0)), + eval_submit_nanos: Arc::new(AtomicU64::new(0)), + eval_await_nanos: Arc::new(AtomicU64::new(0)), + backup_nanos: Arc::new(AtomicU64::new(0)), + } + } + + fn board_with_turn_count(turn_count: usize) -> Board { + Board::from_fen(&format!( + "ap2|3x3|tc:{turn_count}|cm:0|ep:0|w:0,0,99|b:2,2,99|h:-|pu:-|ps:-|bd:3/3/3" + )) + .expect("board FEN must parse") + } + + fn first_nonfinal_action(board: &Board) -> Action { + *board + .get_valid_actions() + .into_iter() + .find(|action| !action.is_final()) + .expect("test board should have a non-final action") + } + + #[test] + fn value_from_term_compresses_terminal_scale() { + let board = board_with_turn_count(1); + let value = SearchNode::value_from_term(&board, TerminalState::Win(Player::White)); + + assert!(value > 0.3); + assert!(value <= 1.0); + } + + #[test] + fn value_from_term_preserves_faster_win_ordering() { + let early = board_with_turn_count(1); + let late = board_with_turn_count(500); + + let early_value = SearchNode::value_from_term(&early, TerminalState::Win(Player::White)); + let late_value = SearchNode::value_from_term(&late, TerminalState::Win(Player::White)); + + assert!(early_value > late_value); + } + + #[test] + fn value_from_term_preserves_faster_loss_ordering() { + let early = board_with_turn_count(1); + let late = board_with_turn_count(500); + + let early_value = SearchNode::value_from_term(&early, TerminalState::Win(Player::Black)); + let late_value = SearchNode::value_from_term(&late, TerminalState::Win(Player::Black)); + + assert!(early_value < late_value); + } + + #[test] + fn collect_samples_excluding_child_keeps_root_and_skips_reused_subtree() { + let mut board = board_with_turn_count(0); + let root_actions: Vec = board + .get_valid_actions() + .into_iter() + .filter(|action| !action.is_final()) + .take(2) + .copied() + .collect(); + assert_eq!(root_actions.len(), 2); + + let mut reused_board = board.clone(); + let _ = reused_board.apply_action(root_actions[0]); + let reused_action = first_nonfinal_action(&reused_board); + + let mut dropped_board = board.clone(); + let _ = dropped_board.apply_action(root_actions[1]); + let dropped_action = first_nonfinal_action(&dropped_board); + + let leaf = SearchNode::new(0.25, 0, false); + + let mut reused_child = SearchNode::new(0.5, 0, false); + reused_child.children.push(ChildData { + action: reused_action, + child_value: 0.25, + entrance_count: 0, + node: Some(Box::new(SearchNode::new(0.25, 0, false))), + }); + + let mut dropped_child = SearchNode::new(-0.5, 0, false); + dropped_child.children.push(ChildData { + action: dropped_action, + child_value: -0.25, + entrance_count: 0, + node: Some(Box::new(leaf)), + }); + + let root = SearchNode::new(1.0, 0, false); + let mut root = root; + root.children.push(ChildData { + action: root_actions[0], + child_value: 0.5, + entrance_count: 0, + node: Some(Box::new(reused_child)), + }); + root.children.push(ChildData { + action: root_actions[1], + child_value: -0.5, + entrance_count: 0, + node: Some(Box::new(dropped_child)), + }); + + let mut full_samples = Vec::new(); + root.collect_samples(&mut board, &mut full_samples); + assert_eq!(full_samples.len(), 3); + + let mut dropped_samples = Vec::new(); + root.collect_samples_excluding_child(&mut board, 0, &mut dropped_samples); + assert_eq!(dropped_samples.len(), 2); + } + + #[test] + fn resolved_forced_loss_stops_descent() { + let board = board_with_turn_count(0); + let root_actions: Vec = board + .get_valid_actions() + .into_iter() + .filter(|action| !action.is_final()) + .take(2) + .copied() + .collect(); + assert_eq!(root_actions.len(), 2); + + let mut root = SearchNode::new(-1.0, -1, true); + for &action in &root_actions { + root.children.push(ChildData { + action, + child_value: -1.0, + entrance_count: 0, + node: Some(Box::new(SearchNode::new(-1.0, -1, true))), + }); + } + + let evaluator = NeverEval; + let mut tree = GameSearchTree { + root_node: Box::new(root), + root_state: board, + rng: SmallRng::seed_from_u64(7), + evaluator: &evaluator, + timing: zero_timing(), + }; + + let dcv = if tree.root_state.is_white_turn() { + 1 + } else { + -1 + }; + let best_action = tree.get_best_action_index(); + assert!(tree.should_stop_descent(dcv, best_action)); + } +} + +pub struct GameSearchTree<'a, E: Evaluator> { + pub root_node: Box, + pub root_state: Board, + rng: SmallRng, + evaluator: &'a E, + timing: SearchTimingMetrics, +} + +impl<'a, E: Evaluator> GameSearchTree<'a, E> { + fn should_stop_descent(&self, dcv: i32, best_action_id: usize) -> bool { + if self.root_node.resolved { + return true; + } + + let best_child = &self.root_node.children[best_action_id]; + if best_child.entrance_count >= 6000 && best_child.completion_value() != -dcv { + return true; + } + + best_child.completion_value() == dcv + } + + fn safest_action(&mut self) -> (usize, Action) { + let val = if self.root_state.is_white_turn() { + self.root_node + .children + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| { + (a.completion_value(), a.entrance_count, a.child_value) + .partial_cmp(&(b.completion_value(), b.entrance_count, b.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) + }) + .expect("No valid action at root state!") + } else { + self.root_node + .children + .iter() + .enumerate() + .min_by(|(_, a), (_, b)| { + ( + a.completion_value(), + -(a.entrance_count as isize), + a.child_value, + ) + .partial_cmp(&( + b.completion_value(), + -(b.entrance_count as isize), + b.child_value, + )) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| self.rng.next_u32().cmp(&self.rng.next_u32())) + }) + .expect("No valid action at root state!") + }; + + (val.0, val.1.action) + } + + /// Create a new search tree. The board must NOT be in a terminal state. + pub async fn new( + board: &Board, + evaluator: &'a E, + rng: SmallRng, + timing: SearchTimingMetrics, + ) -> GameSearchTree<'a, E> { + let mut local_rng = rng; + GameSearchTree { + root_node: Box::new( + SearchNode::build_self( + board, + ApplyActionOutcome::Ongoing, + evaluator, + &mut local_rng, + &timing, + ) + .await, + ), + root_state: board.clone(), + rng: local_rng, + evaluator, + timing, + } + } + + pub fn collect_tree_learning_samples_with_terminal_mix( + &self, + rng: &mut R, + terminal_mix_prob: f32, + ) -> Vec { + let mut state = self.root_state.clone(); + let mut samples = Vec::new(); + self.root_node.collect_samples_with_terminal_mix( + &mut state, + &mut samples, + rng, + terminal_mix_prob, + ); + samples + } + + #[cfg(test)] + pub fn collect_tree_learning_samples(&self) -> Vec { + let mut rng = SmallRng::seed_from_u64(0); + self.collect_tree_learning_samples_with_terminal_mix(&mut rng, 0.0) + } + + pub fn step_tree_and_collect_dropped_samples_with_terminal_mix( + &mut self, + new_board: &Board, + action_id: usize, + rng: &mut R, + terminal_mix_prob: f32, + ) -> Vec { + let mut state = self.root_state.clone(); + let mut samples = Vec::new(); + self.root_node + .collect_samples_excluding_child_with_terminal_mix( + &mut state, + action_id, + &mut samples, + rng, + terminal_mix_prob, + ); + + self.root_state = new_board.clone(); + + if self.root_node.children[action_id].node.is_some() { + self.root_node = self.root_node.children[action_id].node.take().unwrap(); + } else { + // Can't reuse tree - would need async expansion. + // This shouldn't happen in normal play since we always expand before selecting. + panic!("step_tree called on unexpanded child"); + } + + samples + } + + #[cfg(test)] + pub fn step_tree_and_collect_dropped_samples( + &mut self, + new_board: &Board, + action_id: usize, + ) -> Vec { + let mut rng = SmallRng::seed_from_u64(0); + self.step_tree_and_collect_dropped_samples_with_terminal_mix( + new_board, action_id, &mut rng, 0.0, + ) + } + + pub async fn run_descent_for_iter(&mut self, iterations: u32) { + let dcv: i32 = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; + + if self.root_node.children.len() <= 1 { + return; // No need to compute with only one option. + } + for _epoch in 0..iterations { + let ba = self.get_best_action_index(); + if self.should_stop_descent(dcv, ba) { + break; + } + self.root_node + .ubfms_iteration( + self.root_state.clone(), + ApplyActionOutcome::Ongoing, + self.evaluator, + &mut self.rng, + &self.timing, + ) + .await; + } + } + + pub fn get_best_action(&mut self) -> Action { + self.safest_action().1 + } + + pub fn get_best_action_index(&mut self) -> usize { + self.safest_action().0 + } + + pub fn get_best_action_and_index(&mut self) -> (usize, Action) { + self.safest_action() + } + + /// Select an action using Athénan's ordinal distribution. + /// + /// The ordinal distribution depends only on rank ordering of moves. + /// The exploitation rate is drawn uniformly from [0, 1] each time. + pub fn ordinal_select(&mut self) -> usize { + let n = self.root_node.children.len(); + if n <= 1 { + return 0; + } + + let is_white = self.root_state.is_white_turn(); + + // Sort children indices by value (best first for current player) + let mut indices: Vec = (0..n).collect(); + indices.sort_by(|&a, &b| { + let ca = &self.root_node.children[a]; + let cb = &self.root_node.children[b]; + if is_white { + (cb.completion_value(), cb.child_value) + .partial_cmp(&(ca.completion_value(), ca.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + } else { + (ca.completion_value(), ca.child_value) + .partial_cmp(&(cb.completion_value(), cb.child_value)) + .unwrap_or(std::cmp::Ordering::Equal) + } + }); + + // Draw exploitation rate uniformly from [0, 1] + let eps: f32 = self.rng.random(); + + // Ordinal distribution: P(c_i) = (eps + (1-eps)/(n-i)) * (1 - sum of previous) + let mut remaining = 1.0f32; + let r: f32 = self.rng.random(); + let mut cumulative = 0.0f32; + + for (rank, &idx) in indices.iter().enumerate() { + let slots_left = (n - rank) as f32; + let p = remaining * (eps + (1.0 - eps) / slots_left); + cumulative += p; + remaining -= p; + + if r < cumulative { + return idx; + } + } + + // Fallback: return the last (worst) action + *indices.last().unwrap() + } +} + +impl<'a, E: EvalCountTracker> GameSearchTree<'a, E> { + pub async fn run_descent_to_eval_limit(&mut self, max_gpu_evals: u64) { + let dcv: i32 = if self.root_state.is_white_turn() { + 1 + } else { + -1 + }; + + if self.root_node.children.len() <= 1 { + return; + } + + self.evaluator.reset_eval_count(); + + loop { + if self.evaluator.eval_count() >= max_gpu_evals { + break; + } + + let ba = self.get_best_action_index(); + if self.should_stop_descent(dcv, ba) { + break; + } + self.root_node + .ubfms_iteration( + self.root_state.clone(), + ApplyActionOutcome::Ongoing, + self.evaluator, + &mut self.rng, + &self.timing, + ) + .await; + } + } +} diff --git a/training/src/eval.rs b/training/src/eval.rs new file mode 100644 index 0000000..f3e124e --- /dev/null +++ b/training/src/eval.rs @@ -0,0 +1,96 @@ +//! Async evaluator for GPU inference. Value-only (no policy head). + +use std::future::Future; +use std::sync::atomic::{AtomicU64, Ordering}; + +use alpha_paint::board::Board; +use ndarray::Ix1; + +use crate::observation; +use crate::queue::GpuJobQueue; + +/// Async evaluator trait for neural network inference. +/// +/// Returns a scalar value estimate for the given board position. +/// Value is from White's perspective. +pub trait Evaluator { + /// Evaluate the board and return a value estimate. + fn evaluate(&self, board: Board) -> impl Future; +} + +/// Evaluators that can report how many GPU eval submissions they have issued. +pub trait EvalCountTracker: Evaluator { + fn reset_eval_count(&self); + fn eval_count(&self) -> u64; +} + +/// GPU-backed evaluator that batches inference requests. +/// +/// Wraps a GpuJobQueue and serializes Board state into observations +/// for GPU inference. Returns scalar value. +pub struct GpuEvaluator<'a> { + queue: &'a GpuJobQueue, +} + +impl<'a> GpuEvaluator<'a> { + pub fn new(queue: &'a GpuJobQueue) -> Self { + Self { queue } + } +} + +impl Evaluator for GpuEvaluator<'_> { + fn evaluate(&self, board: Board) -> impl Future { + let future = self.queue.eval(|mut out| { + observation::encode_into_slice( + &board, + out.as_slice_mut() + .expect("packed observation slot must be contiguous"), + ); + }); + + async move { future.await } + } +} + +/// Wrapper that counts evaluator submissions for a single worker. +pub struct CountingEvaluator<'a, E> { + inner: &'a E, + eval_count: AtomicU64, +} + +impl<'a, E> CountingEvaluator<'a, E> { + pub fn new(inner: &'a E) -> Self { + Self { + inner, + eval_count: AtomicU64::new(0), + } + } +} + +impl Evaluator for CountingEvaluator<'_, E> { + fn evaluate(&self, board: Board) -> impl Future { + self.eval_count.fetch_add(1, Ordering::AcqRel); + self.inner.evaluate(board) + } +} + +impl EvalCountTracker for CountingEvaluator<'_, E> { + fn reset_eval_count(&self) { + self.eval_count.store(0, Ordering::Release); + } + + fn eval_count(&self) -> u64 { + self.eval_count.load(Ordering::Acquire) + } +} + +/// Synchronous CPU evaluator for testing. +/// +/// Returns zero value for all positions. +pub struct UniformEvaluator; + +impl Evaluator for UniformEvaluator { + fn evaluate(&self, _: Board) -> impl Future { + std::future::ready(0.0) + } +} diff --git a/training/src/executor.rs b/training/src/executor.rs new file mode 100644 index 0000000..5f29537 --- /dev/null +++ b/training/src/executor.rs @@ -0,0 +1,431 @@ +//! Single-threaded async executor for GPU inference workers. +//! +//! This executor is designed for polling many workers that submit GPU inference +//! requests. It doesn't use wakers - instead, it polls all futures in a tight +//! loop and parks when no progress is made. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker}; + +use event_listener::{EventListener, Listener}; + +use crate::future::{signal_progress, take_progress}; + +/// Create a dummy waker that does nothing. +/// We don't use wakers for signaling - we use event_listener + progress tracking. +fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + + // SAFETY: The vtable functions are valid and the data pointer is null (unused) + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } +} + +/// A single-threaded executor for GPU inference workers. +/// +/// Polls all futures in round-robin until no progress is made, then parks +/// waiting for GPU batch completion. +pub struct Executor<'a> { + /// Function to get an event listener for parking. + listen_fn: Box EventListener + 'a>, +} + +impl<'a> Executor<'a> { + /// Create a new executor that will park using the given listen function. + /// + /// The listen function should return an EventListener from the GPU queue's + /// completion_event. + pub fn new(listen_fn: F) -> Self + where + F: Fn() -> EventListener + 'a, + { + Self { + listen_fn: Box::new(listen_fn), + } + } + + /// Run futures until completion or cancellation. + /// + /// Polls all futures in round-robin. When no future makes progress, + /// parks until the GPU signals batch completion. + /// + /// Completed futures are removed (swap_remove). Pending futures remain + /// alive in the vec. Returns when `cancel()` returns true or all futures + /// complete. This allows preserving in-progress game state across + /// pause/resume cycles. + pub fn run(&self, futures: &mut Vec>>, cancel: &mut C) + where + F: Future + ?Sized, + C: FnMut() -> bool, + { + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + loop { + if cancel() || futures.is_empty() { + return; + } + + // Poll all futures until no progress + loop { + if cancel() || futures.is_empty() { + return; + } + + // Clear progress flag before polling round + take_progress(); + + // Poll all pending futures + let mut i = 0; + while i < futures.len() { + if cancel() { + return; + } + + let poll_result = futures[i].as_mut().poll(&mut cx); + match poll_result { + Poll::Ready(()) => { + futures.swap_remove(i); + signal_progress(); + } + Poll::Pending => { + i += 1; + } + } + } + + if cancel() || futures.is_empty() { + return; + } + + // If no progress was made, break to park + if !take_progress() { + break; + } + } + + if futures.is_empty() { + return; + } + + // Set up listener BEFORE re-checking (avoid race). + let listener = (self.listen_fn)(); + + // Double-check before parking and re-evaluate cancellation. + take_progress(); + let mut i = 0; + while i < futures.len() { + let poll_result = futures[i].as_mut().poll(&mut cx); + match poll_result { + Poll::Ready(()) => { + futures.swap_remove(i); + signal_progress(); + } + Poll::Pending => { + i += 1; + } + } + } + + if cancel() || futures.is_empty() { + return; + } + + if !take_progress() { + listener.wait(); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::cell::Cell; + use std::rc::Rc; + + /// A simple counter future that completes after N polls + struct CountdownFuture { + remaining: Cell, + } + + impl CountdownFuture { + fn new(count: usize) -> Self { + Self { + remaining: Cell::new(count), + } + } + } + + impl Future for CountdownFuture { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + let remaining = self.remaining.get(); + if remaining == 0 { + Poll::Ready(()) + } else { + self.remaining.set(remaining - 1); + signal_progress(); + Poll::Pending + } + } + } + + #[test] + fn test_executor_runs_single_future() { + let completed = Rc::new(Cell::new(false)); + let completed_clone = completed.clone(); + + let fut = async move { + completed_clone.set(true); + }; + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + assert!(completed.get()); + } + + #[test] + fn test_executor_runs_multiple_futures() { + let count = Rc::new(Cell::new(0)); + + let mut futures: Vec>>> = (0..10) + .map(|_| { + let count = count.clone(); + let fut = async move { + count.set(count.get() + 1); + }; + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut futures, &mut || false); + + assert_eq!(count.get(), 10); + } + + #[test] + fn test_executor_cancels_pending_futures() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let event = Arc::new(event_listener::Event::new()); + + // A future that never completes. + let fut = async move { + std::future::pending::<()>().await; + }; + + let event_for_exec = event.clone(); + let executor = Executor::new(move || event_for_exec.listen()); + + // Cancel shortly after starting and notify to wake the executor. + std::thread::spawn({ + let event = event.clone(); + let cancelled = cancelled.clone(); + move || { + std::thread::sleep(std::time::Duration::from_millis(5)); + cancelled.store(true, Ordering::Relaxed); + event.notify(usize::MAX); + } + }); + + executor.run(&mut vec![Box::pin(fut)], &mut || { + cancelled_clone.load(Ordering::Relaxed) + }); + } + + #[test] + fn test_executor_handles_multi_poll_futures() { + let completed = Rc::new(Cell::new(0)); + + let mut futures: Vec>>> = (0..5) + .map(|i| { + let completed = completed.clone(); + let countdown = CountdownFuture::new(i + 1); + let fut = async move { + // Wrap countdown in a custom future that signals progress + std::future::poll_fn(|_cx| { + // Need to poll countdown, but poll requires Pin<&mut Self> + // so we'll inline the countdown logic + let remaining = countdown.remaining.get(); + if remaining == 0 { + Poll::Ready(()) + } else { + countdown.remaining.set(remaining - 1); + signal_progress(); + Poll::Pending + } + }) + .await; + completed.set(completed.get() + 1); + }; + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + executor.run(&mut futures, &mut || false); + + assert_eq!(completed.get(), 5); + } + + #[test] + fn test_executor_with_event_notification() { + use std::sync::Arc; + + let event = Arc::new(event_listener::Event::new()); + let event_clone = event.clone(); + + // Future that waits for event then completes + let completed = Rc::new(Cell::new(false)); + let completed_clone = completed.clone(); + + // Spawn a thread that will notify after a short delay + let notify_thread = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(10)); + event_clone.notify(usize::MAX); + }); + + // Future that needs external notification to complete + let polls = Rc::new(Cell::new(0)); + let polls_clone = polls.clone(); + + let fut = std::future::poll_fn(move |_cx| { + let p = polls_clone.get(); + polls_clone.set(p + 1); + + // Complete after being woken up (second poll after notification) + if p > 0 { + completed_clone.set(true); + signal_progress(); + Poll::Ready(()) + } else { + Poll::Pending + } + }); + + let executor = Executor::new(move || event.listen()); + executor.run(&mut vec![Box::pin(fut)], &mut || false); + + notify_thread.join().unwrap(); + assert!(completed.get()); + } + + #[test] + fn test_run_preserves_futures() { + // Test that run does NOT drop pending futures. + // We use futures that always signal progress and complete after enough polls. + // Cancel immediately on first check to guarantee futures are still pending. + + let completed = Rc::new(Cell::new(0usize)); + + let mut futures: Vec>>> = (0..3) + .map(|_| { + let completed = completed.clone(); + let polls = Cell::new(0usize); + let fut = std::future::poll_fn(move |_cx| { + let p = polls.get(); + polls.set(p + 1); + signal_progress(); + if p >= 10 { + completed.set(completed.get() + 1); + Poll::Ready(()) + } else { + Poll::Pending + } + }); + Box::pin(fut) as Pin>> + }) + .collect(); + + let executor = Executor::new(|| event_listener::Event::new().listen()); + + // Cancel immediately - futures should not be polled at all. + executor.run(&mut futures, &mut || true); + + // Futures should still be alive (cancel was true from the start). + assert_eq!( + futures.len(), + 3, + "all futures must be preserved when cancel is immediate" + ); + assert_eq!(completed.get(), 0, "no futures should have completed"); + + // Now let them finish. + executor.run(&mut futures, &mut || false); + + assert_eq!(completed.get(), 3, "all futures should have completed"); + assert!(futures.is_empty(), "completed futures should be removed"); + } + + #[test] + fn test_run_does_not_drop_pending() { + // Verify a future's drop impl is NOT called between run calls. + + let was_dropped = Rc::new(Cell::new(false)); + let was_dropped_clone = was_dropped.clone(); + + struct DropDetector { + flag: Rc>, + polls: Cell, + } + impl Drop for DropDetector { + fn drop(&mut self) { + self.flag.set(true); + } + } + impl Future for DropDetector { + type Output = (); + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + let this = unsafe { self.get_unchecked_mut() }; + let p = this.polls.get(); + this.polls.set(p + 1); + signal_progress(); + if p >= 5 { + Poll::Ready(()) + } else { + Poll::Pending + } + } + } + + let mut futures: Vec>>> = vec![Box::pin(DropDetector { + flag: was_dropped_clone, + polls: Cell::new(0), + })]; + + let executor = Executor::new(|| event_listener::Event::new().listen()); + + // Cancel immediately + executor.run(&mut futures, &mut || true); + + // The future should NOT have been dropped + assert!( + !was_dropped.get(), + "future must not be dropped between run calls" + ); + assert_eq!(futures.len(), 1, "future should still be in the vec"); + + // Now let it finish + executor.run(&mut futures, &mut || false); + + assert!(futures.is_empty()); + assert!( + was_dropped.get(), + "future should be dropped after completion" + ); + } +} diff --git a/training/src/future.rs b/training/src/future.rs new file mode 100644 index 0000000..b9e3bee --- /dev/null +++ b/training/src/future.rs @@ -0,0 +1,233 @@ +//! Future implementation for GPU evaluation requests. +//! +//! GpuEvalFuture represents a pending GPU inference job. The observation +//! is submitted immediately when the future is created, so the future +//! just polls for completion. + +use std::future::Future; +use std::pin::Pin; +use std::task::{Context, Poll}; + +use ndarray::Dimension; + +use crate::queue::GpuJobQueue; +use crate::BatchDim; + +// Thread-local flag for tracking whether any future made progress. +// Used by the executor to decide whether to park. +std::thread_local! { + static MADE_PROGRESS: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Signal that progress was made (a future submitted or completed work). +pub fn signal_progress() { + MADE_PROGRESS.with(|p| p.set(true)); +} + +/// Check if progress was made and reset the flag. +pub fn take_progress() -> bool { + MADE_PROGRESS.with(|p| p.replace(false)) +} + +/// A future representing a GPU evaluation request. +/// +/// The observation is submitted when this future is created (not on first poll). +/// Polling checks if the batch containing this job is complete. +pub struct GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + queue: &'a GpuJobQueue, + ticket: u64, + completed: bool, +} + +impl<'a, A, D, O> GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + /// Create a new future with an already-submitted ticket. + pub fn new(queue: &'a GpuJobQueue, ticket: u64) -> Self { + // Signal progress on creation since we just submitted + signal_progress(); + Self { + queue, + ticket, + completed: false, + } + } +} + +impl<'a, A, D, O> Future for GpuEvalFuture<'a, A, D, O> +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + type Output = O; + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll { + // SAFETY: We don't move out of self, just update completed flag + let this = unsafe { self.get_unchecked_mut() }; + + if this.completed { + panic!("GpuEvalFuture polled after completion"); + } + + if let Some(&output) = this.queue.poll(this.ticket) { + this.completed = true; + signal_progress(); + Poll::Ready(output) + } else { + Poll::Pending + } + } +} + +impl GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + D::Larger: Dimension, + O: Copy + Default + Send + Sync, +{ + /// Submit an observation and create a future that will resolve to the output. + /// + /// The observation is written immediately via the callback. + /// The returned future polls for the batch to complete. + pub fn eval(&self, write_obs: F) -> GpuEvalFuture<'_, A, D, O> + where + F: FnOnce(ndarray::ArrayViewMut), + { + let ticket = self.submit(write_obs); + GpuEvalFuture::new(self, ticket) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::queue::BATCH_SIZE; + use ndarray::Ix0; + use std::sync::Arc; + use std::task::{RawWaker, RawWakerVTable, Waker}; + + // Create a dummy waker that does nothing + fn dummy_waker() -> Waker { + fn clone(_: *const ()) -> RawWaker { + RawWaker::new(std::ptr::null(), &VTABLE) + } + fn wake(_: *const ()) {} + fn wake_by_ref(_: *const ()) {} + fn drop(_: *const ()) {} + + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, wake, wake_by_ref, drop); + + unsafe { Waker::from_raw(RawWaker::new(std::ptr::null(), &VTABLE)) } + } + + #[test] + fn test_future_submits_immediately() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let submit_count = Arc::new(AtomicUsize::new(0)); + let submit_count_clone = submit_count.clone(); + + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + move |_batch_idx, inputs, completion| { + submit_count_clone.fetch_add(1, Ordering::SeqCst); + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + assert_eq!(submit_count.load(Ordering::SeqCst), 0); + + // Create futures - they submit immediately + let mut futures: Vec<_> = (0..BATCH_SIZE as u64) + .map(|i| queue.eval(|mut out| out[()] = i)) + .collect(); + + // Batch should have been dispatched (last eval triggered it) + assert_eq!(submit_count.load(Ordering::SeqCst), 1); + + // Poll all futures - they should all be ready + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + for fut in &mut futures { + match Pin::new(fut).poll(&mut cx) { + Poll::Ready(_) => {} + Poll::Pending => panic!("future should be ready"), + } + } + } + + #[test] + fn test_future_returns_correct_result() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input + 100; + } + completion.complete(&outputs); + }, + )); + + let mut futures: Vec<_> = (0..BATCH_SIZE as u64) + .map(|i| queue.eval(|mut out| out[()] = i)) + .collect(); + + let waker = dummy_waker(); + let mut cx = Context::from_waker(&waker); + + // All should be ready immediately (batch was triggered) + for (i, fut) in futures.iter_mut().enumerate() { + match Pin::new(fut).poll(&mut cx) { + Poll::Ready(result) => { + assert_eq!(result, (i as u64) + 100); + } + Poll::Pending => panic!("future {} should be ready", i), + } + } + } + + #[test] + fn test_progress_tracking() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input; + } + completion.complete(&outputs); + }, + )); + + // Clear any previous progress + take_progress(); + + // Create one future (partial batch) - should signal progress on creation + let _fut = queue.eval(|mut out| out[()] = 0); + + // Should have made progress (submitted) + assert!(take_progress()); + } +} diff --git a/training/src/lib.rs b/training/src/lib.rs new file mode 100644 index 0000000..0e720a2 --- /dev/null +++ b/training/src/lib.rs @@ -0,0 +1,616 @@ +use alpha_paint::board::{Action, ApplyActionOutcome, Board, TerminalState}; +use alpha_paint::TRAINING_START_FENS; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::{Mutex, OnceLock}; + +use cudagraph::CudaGraphRunner; +use queue::{queue_shape_for_workers, BATCH_SIZE}; +use replay_buffer::ReplayBuffer; +use training::{SelfPlaySession, SessionConfig}; +use worker::WorkerConfig; + +use ndarray::{ArrayView, Dimension, Ix0, Ix1, Ix2, Ix3, Ix4, Ix5, Ix6, RemoveAxis}; +use numpy::{PyArray, PyArrayMethods}; +use pyo3::prelude::*; +use rand::{Rng, RngExt, SeedableRng}; +use rand_chacha::ChaCha8Rng; +use rayon::prelude::*; + +pub mod cudagraph; +pub mod descent; +pub mod eval; +pub mod executor; +pub mod future; +mod observation; +pub mod queue; +pub mod replay_buffer; +pub mod training; +pub mod worker; + +/// Extension trait for prepending a batch dimension to a shape. +pub trait BatchDim: Dimension + Clone { + type BatchedDim: Dimension + RemoveAxis; + + fn with_batch(batch_size: usize, obs_shape: Self) -> Self::BatchedDim; +} + +impl BatchDim for Ix0 { + type BatchedDim = Ix1; + + fn with_batch(batch_size: usize, _obs: Self) -> Ix1 { + Ix1(batch_size) + } +} + +impl BatchDim for Ix1 { + type BatchedDim = Ix2; + + fn with_batch(batch_size: usize, obs: Self) -> Ix2 { + Ix2(batch_size, obs[0]) + } +} + +impl BatchDim for Ix2 { + type BatchedDim = Ix3; + + fn with_batch(batch_size: usize, obs: Self) -> Ix3 { + Ix3(batch_size, obs[0], obs[1]) + } +} + +impl BatchDim for Ix3 { + type BatchedDim = Ix4; + + fn with_batch(batch_size: usize, obs: Self) -> Ix4 { + Ix4(batch_size, obs[0], obs[1], obs[2]) + } +} + +impl BatchDim for Ix4 { + type BatchedDim = Ix5; + + fn with_batch(batch_size: usize, obs: Self) -> Ix5 { + Ix5(batch_size, obs[0], obs[1], obs[2], obs[3]) + } +} + +impl BatchDim for Ix5 { + type BatchedDim = Ix6; + + fn with_batch(batch_size: usize, obs: Self) -> Ix6 { + Ix6(batch_size, obs[0], obs[1], obs[2], obs[3], obs[4]) + } +} + +struct GraphCacheEntry { + model_ptr: usize, + num_batches: usize, + precision: String, + num_gpus: usize, + runners: HashMap>, + models: HashMap>, +} + +static GRAPH_CACHE: OnceLock>> = OnceLock::new(); + +fn graph_cache() -> &'static Mutex> { + GRAPH_CACHE.get_or_init(|| Mutex::new(None)) +} + +const AVG_GAME_LENGTH: f32 = 500.0; +const PRETRAIN_SAMPLE_CHUNK: usize = 256; + +fn terminal_value_from_term(board: &Board, terminal: TerminalState) -> f32 { + let sign = terminal.value() as f32; + if sign == 0.0 { + return 0.0; + } + + let total_hills = board.hills.len().max(1) as f32; + let white_hills = board.tiles.controlled_hill_count::() as f32; + let black_hills = board.tiles.controlled_hill_count::() as f32; + + let total_cells = ((board.rows as usize * board.cols as usize) + .saturating_sub(board.tiles.wall_count())) + .max(1) as f32; + let white_terr = board.tiles.territory_count::() as f32; + let black_terr = board.tiles.territory_count::() as f32; + + let hill_margin = (sign * (white_hills - black_hills) / total_hills).max(0.0); + let terr_margin = (sign * (white_terr - black_terr) / total_cells).max(0.0); + + let p = board.turn_count.max(1) as f32; + let depth = (AVG_GAME_LENGTH / p).ln_1p() / AVG_GAME_LENGTH.ln_1p(); + + sign * (0.3 + 0.4 * hill_margin + 0.2 * terr_margin + 0.1 * depth) +} + +fn sample_random_action(board: &Board, rng: &mut R) -> Option { + let actions = board.get_valid_actions(); + let action_idx = (actions.len() > 0).then(|| rng.random_range(0..actions.len()))?; + actions.into_iter().nth(action_idx).copied() +} + +fn random_terminal_board(rng: &mut R) -> (Board, TerminalState) { + let mut board = + Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) + .expect("training start FEN must parse"); + + loop { + let Some(action) = sample_random_action(&board, rng) else { + let is_white_turn = board.is_white_turn(); + return (board, TerminalState::loss_for(is_white_turn)); + }; + + let (outcome, rollback) = board.apply_action(action); + match outcome { + ApplyActionOutcome::Ongoing => {} + ApplyActionOutcome::Terminal { terminal } => return (board, terminal), + ApplyActionOutcome::PlayInstead { + terminal, + play_instead, + } => { + board.rollback(action, rollback); + let (play_outcome, _) = board.apply_action(play_instead); + debug_assert!(matches!(play_outcome, ApplyActionOutcome::Terminal { .. })); + return (board, terminal); + } + ApplyActionOutcome::Killshot { terminal, actions } => { + for (idx, action) in actions.iter().copied().enumerate() { + let (killshot_outcome, _) = board.apply_action(action); + if idx == actions.len() - 1 { + debug_assert!(matches!( + killshot_outcome, + ApplyActionOutcome::Terminal { .. } + )); + } else { + debug_assert!(matches!(killshot_outcome, ApplyActionOutcome::Ongoing)); + } + } + return (board, terminal); + } + } + } +} + +fn sample_random_terminal_examples(n: usize, seed: u64) -> (Vec, Vec) { + let mut obs_data = vec![0u16; n * cudagraph::OBS_WORDS]; + let mut values = vec![0.0; n]; + + obs_data + .par_chunks_mut(PRETRAIN_SAMPLE_CHUNK * cudagraph::OBS_WORDS) + .zip(values.par_chunks_mut(PRETRAIN_SAMPLE_CHUNK)) + .enumerate() + .for_each(|(chunk_idx, (obs_chunk, values_chunk))| { + let mut rng = ChaCha8Rng::seed_from_u64(seed.wrapping_add(chunk_idx as u64)); + for (obs_sample, value_out) in obs_chunk + .chunks_mut(cudagraph::OBS_WORDS) + .zip(values_chunk.iter_mut()) + { + let (board, terminal) = random_terminal_board(&mut rng); + observation::encode_into_slice(&board, obs_sample); + *value_out = terminal_value_from_term(&board, terminal); + } + }); + + (obs_data, values) +} + +#[pyfunction] +fn sample_random_terminal_batch<'py>( + py: Python<'py>, + n: usize, + seed: u64, +) -> PyResult<(Bound<'py, PyArray>, Bound<'py, PyArray>)> { + let (obs_data, values) = py.detach(|| sample_random_terminal_examples(n, seed)); + let obs = PyArray::from_vec(py, obs_data).reshape(Ix2(n, cudagraph::OBS_WORDS))?; + let values = PyArray::from_vec(py, values); + Ok((obs, values)) +} + +/// Replay buffer storing (packed observation, value) pairs. +/// +/// Observations are flat u16 tensors, values are f32 scalars (no policy). +#[pyclass] +struct EphemeralReplayBuffer { + inner: Arc>, +} + +#[pymethods] +impl EphemeralReplayBuffer { + #[new] + fn new(capacity: usize) -> Self { + let obs_shape = Ix1(cudagraph::OBS_WORDS); + Self { + inner: Arc::new(ReplayBuffer::new(capacity, obs_shape)), + } + } + + fn __len__(&self) -> usize { + self.inner.len() + } + + #[getter] + fn capacity(&self) -> usize { + self.inner.capacity() + } + + /// Sample `n` items and return (observations, values) as numpy arrays. + fn sample<'py>( + &self, + py: Python<'py>, + n: usize, + seed: u64, + ) -> PyResult<(Bound<'py, PyArray>, Bound<'py, PyArray>)> { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let batch = self.inner.sample(n, &mut rng); + let num_samples = batch.values.len(); + + let obs_data = batch.observations.into_raw_vec_and_offset().0; + + let obs = + PyArray::from_vec(py, obs_data).reshape(Ix2(num_samples, cudagraph::OBS_WORDS))?; + let values = PyArray::from_vec(py, batch.values); + + Ok((obs, values)) + } +} + +impl EphemeralReplayBuffer { + pub fn inner(&self) -> &Arc> { + &self.inner + } +} + +/// Persistent self-play session using Athénan (Descent + tree learning). +/// +/// Uses the CUDA graph runner for GPU dispatch (value-only, no policy). +#[pyclass] +struct SelfPlay { + session: Option, + runners: HashMap>, + source_model: Py, + source_model_ptr: usize, +} + +impl SelfPlay { + fn session(&self) -> PyResult<&SelfPlaySession> { + self.session.as_ref().ok_or_else(|| { + PyErr::new::("session already dropped") + }) + } +} + +#[pymethods] +impl SelfPlay { + #[new] + #[pyo3(signature = ( + replay_buffer, + num_threads, + workers_per_thread, + seed, + *, + max_gpu_evals_per_move = 4096, + model, + selfplay_precision = "bf16", + num_gpus = 1 + ))] + fn new( + py: Python<'_>, + replay_buffer: &EphemeralReplayBuffer, + num_threads: usize, + workers_per_thread: usize, + seed: u64, + max_gpu_evals_per_move: u64, + model: Py, + selfplay_precision: &str, + num_gpus: usize, + ) -> PyResult { + if num_gpus == 0 { + return Err(PyErr::new::( + "num_gpus must be >= 1", + )); + } + + let config = SessionConfig { + num_threads, + workers_per_thread, + seed, + thread_stack_size_bytes: SessionConfig::default().thread_stack_size_bytes, + worker: WorkerConfig { + max_gpu_evals_per_move, + }, + }; + + let total_workers = num_threads.checked_mul(workers_per_thread).ok_or_else(|| { + PyErr::new::( + "num_threads * workers_per_thread overflow", + ) + })?; + let (num_batches, _total_slots) = queue_shape_for_workers(total_workers); + let model_ptr = model.bind(py).as_ptr() as usize; + + // Build or reuse the CUDA graph runners. + let runners = { + let cache = graph_cache(); + let mut guard = cache.lock().expect("graph cache mutex poisoned"); + + let needs_rebuild = match guard.as_ref() { + Some(entry) => { + entry.model_ptr != model_ptr + || entry.num_batches != num_batches + || entry.precision != selfplay_precision + || entry.num_gpus != num_gpus + } + None => true, + }; + + if needs_rebuild { + let copy = PyModule::import(py, "copy")?; + let deepcopy = copy.getattr("deepcopy")?; + let mut runners = HashMap::new(); + let mut models = HashMap::new(); + for gpu_id in 0..num_gpus { + let model_copy: Py = deepcopy.call1((model.clone_ref(py),))?.into(); + let runner = Arc::new(CudaGraphRunner::new( + py, + model_copy.clone_ref(py), + gpu_id, + num_batches, + BATCH_SIZE, + selfplay_precision, + )?); + runners.insert(gpu_id, runner); + models.insert(gpu_id, model_copy); + } + *guard = Some(GraphCacheEntry { + model_ptr, + num_batches, + precision: selfplay_precision.to_string(), + num_gpus, + runners: runners.clone(), + models, + }); + runners + } else { + guard + .as_ref() + .expect("cached runners should exist") + .runners + .clone() + } + }; + + let runners_for_dispatch = runners.clone(); + let dispatch = move |batch_idx: usize, + obs_view: ArrayView, + completion: queue::BatchCompletion| { + let gpu_id = batch_idx % num_gpus; + runners_for_dispatch[&gpu_id].dispatch_async(batch_idx, obs_view, completion); + }; + + let session = SelfPlaySession::new(config, replay_buffer.inner().clone(), dispatch); + + Ok(Self { + session: Some(session), + runners, + source_model: model, + source_model_ptr: model_ptr, + }) + } + + /// Start self-play with no sample limit. + fn start(&self, py: Python<'_>) -> PyResult<()> { + self.sync_model_replicas(py)?; + self.session()?.start(); + Ok(()) + } + + /// Block until absolute target_samples is reached, then pause and quiesce. + fn wait_for(&self, py: Python<'_>, target_samples: usize) -> PyResult { + self.sync_model_replicas(py)?; + let session = self.session()?; + let result = py.detach(|| session.wait_for(target_samples)); + Ok(result) + } + + /// Return the current absolute sample count. + fn samples(&self) -> PyResult { + Ok(self.session()?.samples()) + } + + /// Return the current absolute game count. + fn games(&self) -> PyResult { + Ok(self.session()?.games()) + } + + /// Return the total number of terminal leaf samples injected into replay. + fn terminal_mix_samples_added(&self) -> PyResult { + Ok(self.session()?.terminal_mix_samples_added()) + } + + /// Return the total number of selected actions. + fn action_steps(&self) -> PyResult { + Ok(self.session()?.action_steps()) + } + + /// Return the total number of selected final actions. + fn final_actions(&self) -> PyResult { + Ok(self.session()?.final_actions()) + } + + /// Return the total number of selected non-final actions. + fn nonfinal_actions(&self) -> PyResult { + Ok(self.session()?.nonfinal_actions()) + } + + /// Return the total number of completed turns. + fn completed_turns(&self) -> PyResult { + Ok(self.session()?.completed_turns()) + } + + /// Return the sum of `board.turn_count` seen before each selected action. + fn action_turn_count_total(&self) -> PyResult { + Ok(self.session()?.action_turn_count_total()) + } + + /// Return the largest `board.turn_count` seen in any in-progress game. + fn max_turn_count_seen(&self) -> PyResult { + Ok(self.session()?.max_turn_count_seen()) + } + + /// Return the total number of selected actions in completed games. + fn completed_game_actions_total(&self) -> PyResult { + Ok(self.session()?.completed_game_actions_total()) + } + + /// Return the largest number of selected actions in a completed game. + fn max_actions_in_completed_game(&self) -> PyResult { + Ok(self.session()?.max_actions_in_completed_game()) + } + + /// Return the total completed turn count accumulated across finished games. + fn completed_game_turn_count_total(&self) -> PyResult { + Ok(self.session()?.completed_game_turn_count_total()) + } + + /// Return the largest completed turn count seen in a finished game. + fn max_turn_count_in_completed_game(&self) -> PyResult { + Ok(self.session()?.max_turn_count_in_completed_game()) + } + + /// Drain completed game turn counts collected since the last call. + fn take_completed_game_turn_counts(&self) -> PyResult> { + Ok(self.session()?.take_completed_game_turn_counts()) + } + + /// Return total time spent building fresh search trees, in nanoseconds. + fn tree_build_nanos(&self) -> PyResult { + Ok(self.session()?.tree_build_nanos()) + } + + /// Return total time spent running descent iterations, in nanoseconds. + fn descent_nanos(&self) -> PyResult { + Ok(self.session()?.descent_nanos()) + } + + /// Return total time spent collecting replay samples, in nanoseconds. + fn sample_collect_nanos(&self) -> PyResult { + Ok(self.session()?.sample_collect_nanos()) + } + + /// Return total time spent encoding and pushing replay samples, in nanoseconds. + fn replay_push_nanos(&self) -> PyResult { + Ok(self.session()?.replay_push_nanos()) + } + + /// Return total CPU expansion prep time inside Descent, in nanoseconds. + fn descent_expand_cpu_nanos(&self) -> PyResult { + Ok(self.session()?.descent_expand_cpu_nanos()) + } + + /// Return total time spent applying actions during Descent, in nanoseconds. + fn descent_apply_action_nanos(&self) -> PyResult { + Ok(self.session()?.descent_apply_action_nanos()) + } + + /// Return total time spent submitting evals to the GPU queue, in nanoseconds. + fn descent_eval_submit_nanos(&self) -> PyResult { + Ok(self.session()?.descent_eval_submit_nanos()) + } + + /// Return total time spent awaiting queued evals, in nanoseconds. + fn descent_eval_await_nanos(&self) -> PyResult { + Ok(self.session()?.descent_eval_await_nanos()) + } + + /// Return total time spent backing up/searching best children, in nanoseconds. + fn descent_backup_nanos(&self) -> PyResult { + Ok(self.session()?.descent_backup_nanos()) + } + + /// Return the total number of CUDA graph launches completed so far. + fn gpu_batches(&self) -> u64 { + self.runners + .values() + .map(|runner| runner.dispatched_batches()) + .sum() + } + + /// Return the total number of packed observations sent to GPU so far. + fn gpu_evals(&self) -> u64 { + self.runners + .values() + .map(|runner| runner.dispatched_evals()) + .sum() + } + + /// Shut down the session. Idempotent. + #[pyo3(name = "drop")] + fn py_drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } +} + +impl Drop for SelfPlay { + fn drop(&mut self) { + if let Some(mut session) = self.session.take() { + session.shutdown(); + } + } +} + +impl SelfPlay { + fn sync_model_replicas(&self, py: Python<'_>) -> PyResult<()> { + let replicas = { + let cache = graph_cache(); + let guard = cache.lock().expect("graph cache mutex poisoned"); + let entry = guard.as_ref().ok_or_else(|| { + PyErr::new::( + "graph cache missing while syncing model replicas", + ) + })?; + + if entry.model_ptr != self.source_model_ptr { + return Err(PyErr::new::( + "graph cache model mismatch while syncing model replicas", + )); + } + + let mut models: Vec<(usize, Py)> = entry + .models + .iter() + .map(|(gpu_id, model)| (*gpu_id, model.clone_ref(py))) + .collect(); + models.sort_by_key(|(gpu_id, _)| *gpu_id); + models + .into_iter() + .map(|(_, model)| model) + .collect::>() + }; + + let state_dict: Py = self + .source_model + .bind(py) + .call_method0("state_dict")? + .into(); + for replica in replicas { + let _ = replica + .bind(py) + .call_method1("load_state_dict", (state_dict.clone_ref(py),))?; + } + + Ok(()) + } +} + +#[pymodule] +fn alphapaint_training(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + m.add_function(pyo3::wrap_pyfunction!(sample_random_terminal_batch, m)?)?; + Ok(()) +} diff --git a/training/src/observation.rs b/training/src/observation.rs new file mode 100644 index 0000000..22ec701 --- /dev/null +++ b/training/src/observation.rs @@ -0,0 +1,381 @@ +use alpha_paint::board::board_structs::{HillData, Player}; +use alpha_paint::board::consts::{BASE_MAX_STAMINA, HILL_MAX_STAMINA_BONUS}; +use alpha_paint::board::structs::Coordinate; +use alpha_paint::board::Board; + +pub const BOARD_SIDE: usize = 32; +pub const BOARD_CELLS: usize = BOARD_SIDE * BOARD_SIDE; +pub const TILE_PLANES: usize = 15; +pub const LOCAL_WINDOW_TILES: usize = 25; +pub const GLOBAL_FEATURES: usize = 20; + +/// 32×32 tiles + 2×25 local windows + 4 positions + 20 globals +pub const OBS_WORDS: usize = BOARD_CELLS + (LOCAL_WINDOW_TILES * 2) + 4 + GLOBAL_FEATURES; + +const OFFSET_LOCALS: usize = BOARD_CELLS; +const OFFSET_POSITIONS: usize = OFFSET_LOCALS + LOCAL_WINDOW_TILES * 2; +const OFFSET_GLOBALS: usize = OFFSET_POSITIONS + 4; + +/// Index of turn_count within the global features section. +pub const GLOBAL_TURN_COUNT_OFFSET: usize = 17; + +/// Encodes the observation to ship off to the GPU for training. +/// +/// Layout (all from current player's perspective): +/// [0, 1024) — per-tile 15-bit bitmask (32×32 board) +/// [1024, 1074) — 2×25 local window tile bitmasks (my pos, opp pos) +/// [1074, 1078) — player positions (my_x, my_y, opp_x, opp_y) +/// [1078, 1098) — 20 global scalar features +/// +/// Tile bitmask planes: +/// bits 0-3: current player paint thermometer (≥1, ≥2, ≥3, ≥4) +/// bits 4-7: opponent paint thermometer +/// bit 8: wall +/// bit 9: powerup +/// bit 10: current player beacon +/// bit 11: opponent beacon +/// bit 12: hill neutral +/// bit 13: hill current player +/// bit 14: hill opponent +pub fn encode_into_slice(board: &Board, out: &mut [u16]) { + debug_assert!(out.len() >= OBS_WORDS); + + let is_white = board.is_white_turn(); + let (my_coord, opp_coord) = if is_white { + (board.white_coord, board.black_coord) + } else { + (board.black_coord, board.white_coord) + }; + let (my_stamina, opp_stamina) = if is_white { + (board.white_stamina, board.black_stamina) + } else { + (board.black_stamina, board.white_stamina) + }; + + let hill_meta = board.tiles.hill_metadata(); + + // 1. Tile bitmasks [0, 1024) + for x in 0..32u8 { + for y in 0..32u8 { + let coord = Coordinate::new(x, y); + out[(x as usize) * 32 + (y as usize)] = encode_tile(board, is_white, hill_meta, coord); + } + } + + // 2. Local windows [1024, 1074) + encode_local_window( + board, + is_white, + hill_meta, + my_coord, + &mut out[OFFSET_LOCALS..], + ); + encode_local_window( + board, + is_white, + hill_meta, + opp_coord, + &mut out[OFFSET_LOCALS + 25..], + ); + + // 3. Player positions [1074, 1078) + out[OFFSET_POSITIONS] = my_coord.x as u16; + out[OFFSET_POSITIONS + 1] = my_coord.y as u16; + out[OFFSET_POSITIONS + 2] = opp_coord.x as u16; + out[OFFSET_POSITIONS + 3] = opp_coord.y as u16; + + // 4. Global features [1078, 1098) + let my_hills = if is_white { + board.tiles.controlled_hill_count::() + } else { + board.tiles.controlled_hill_count::() + }; + let opp_hills = if is_white { + board.tiles.controlled_hill_count::() + } else { + board.tiles.controlled_hill_count::() + }; + + let my_territory = if is_white { + board.tiles.territory_count::() + } else { + board.tiles.territory_count::() + }; + let opp_territory = if is_white { + board.tiles.territory_count::() + } else { + board.tiles.territory_count::() + }; + + let mut my_hill_tiles: usize = 0; + let mut opp_hill_tiles: usize = 0; + let mut contested_hills: usize = 0; + for hd in hill_meta { + let (mine, theirs) = if is_white { + (hd.white_count, hd.black_count) + } else { + (hd.black_count, hd.white_count) + }; + my_hill_tiles += mine; + opp_hill_tiles += theirs; + if mine > 0 && theirs > 0 && hd.owner.is_none() { + contested_hills += 1; + } + } + + let my_beacons = if is_white { + board.tiles.get_beacon_iterator::().count() + } else { + board.tiles.get_beacon_iterator::().count() + }; + let opp_beacons = if is_white { + board.tiles.get_beacon_iterator::().count() + } else { + board.tiles.get_beacon_iterator::().count() + }; + + let player_dist = board.dist[(my_coord, opp_coord)]; + let my_hill_dist = nearest_non_controlled_hill_dist(board, my_coord, is_white, hill_meta); + let opp_hill_dist = nearest_non_controlled_hill_dist(board, opp_coord, !is_white, hill_meta); + + let my_max_stamina = BASE_MAX_STAMINA + my_hills * HILL_MAX_STAMINA_BONUS; + let opp_max_stamina = BASE_MAX_STAMINA + opp_hills * HILL_MAX_STAMINA_BONUS; + + let g = &mut out[OFFSET_GLOBALS..]; + g[0] = my_stamina as u16; + g[1] = opp_stamina as u16; + g[2] = my_max_stamina as u16; + g[3] = opp_max_stamina as u16; + g[4] = my_territory as u16; + g[5] = opp_territory as u16; + g[6] = my_hills as u16; + g[7] = opp_hills as u16; + g[8] = my_hill_tiles as u16; + g[9] = opp_hill_tiles as u16; + g[10] = contested_hills as u16; + g[11] = my_beacons as u16; + g[12] = opp_beacons as u16; + g[13] = player_dist; + g[14] = my_hill_dist; + g[15] = opp_hill_dist; + g[16] = board.consecutives_moves_so_far as u16; + g[17] = board.turn_count as u16; + g[18] = board.rows as u16; + g[19] = board.cols as u16; +} + +#[inline] +fn encode_tile(board: &Board, is_white: bool, hill_meta: &[HillData], coord: Coordinate) -> u16 { + let tile = board.tiles[coord]; + let mut bits: u16 = 0; + + let paint = tile.paint_value(); + let (my_paint, opp_paint) = if is_white { + (paint.max(0) as u16, (-paint).max(0) as u16) + } else { + ((-paint).max(0) as u16, paint.max(0) as u16) + }; + // Thermometer encoding: bits 0-3 for my paint, bits 4-7 for opponent paint + for t in 0..my_paint { + bits |= 1 << t; + } + for t in 0..opp_paint { + bits |= 1 << (4 + t); + } + + if tile.is_wall() { + bits |= 1 << 8; + } + if board.powerups[coord] { + bits |= 1 << 9; + } + + match tile.beacon_owner() { + Some(Player::White) if is_white => bits |= 1 << 10, + Some(Player::Black) if !is_white => bits |= 1 << 10, + Some(Player::White) => bits |= 1 << 11, + Some(Player::Black) => bits |= 1 << 11, + None => {} + } + + let hill_id = board.hill_id[coord]; + if hill_id != u16::MAX { + match hill_meta[hill_id as usize].owner { + None => bits |= 1 << 12, + Some(Player::White) if is_white => bits |= 1 << 13, + Some(Player::Black) if !is_white => bits |= 1 << 13, + _ => bits |= 1 << 14, + } + } + + bits +} + +const WALL_BITS: u16 = 1 << 8; + +/// Write 25 u16 bitmasks for the manhattan-distance-3 diamond around `center`. +/// Out-of-bounds tiles default to the wall bit. +/// +/// Enumeration order: dx from -3..=3, dy from -3..=3, keeping |dx|+|dy| <= 3. +fn encode_local_window( + board: &Board, + is_white: bool, + hill_meta: &[HillData], + center: Coordinate, + out: &mut [u16], +) { + let cx = center.x as i16; + let cy = center.y as i16; + let mut i = 0; + for dx in -3i16..=3 { + for dy in -3i16..=3 { + if dx.abs() + dy.abs() > 3 { + continue; + } + let nx = cx + dx; + let ny = cy + dy; + out[i] = if (0..32).contains(&nx) && (0..32).contains(&ny) { + encode_tile( + board, + is_white, + hill_meta, + Coordinate::new(nx as u8, ny as u8), + ) + } else { + WALL_BITS + }; + i += 1; + } + } + debug_assert_eq!(i, 25); +} + +/// Find the minimum distance from `player_coord` to any tile in a hill +/// that is NOT controlled by the player identified by `is_white`. +fn nearest_non_controlled_hill_dist( + board: &Board, + player_coord: Coordinate, + is_white: bool, + hill_meta: &[HillData], +) -> u16 { + let my_owner = if is_white { + Some(Player::White) + } else { + Some(Player::Black) + }; + + let mut best = u16::MAX; + + for (i, hd) in hill_meta.iter().enumerate() { + // Skip hills we already control + if hd.owner == my_owner { + continue; + } + for &tile_coord in &board.hills[i] { + let d = board.dist[(player_coord, tile_coord)]; + if d < best { + best = d; + } + } + } + + best +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_words(fen: &str) -> Vec { + let board = Board::from_fen(fen).unwrap(); + let mut out = vec![0; OBS_WORDS]; + encode_into_slice(&board, &mut out); + out + } + + #[test] + fn obs_words_matches_expected_layout() { + assert_eq!(OBS_WORDS, 1024 + 50 + 4 + 20); + assert_eq!(OBS_WORDS, 1098); + } + + #[test] + fn tile_thermometer_paint_encoding() { + // White turn, white paint of strength 3 at (1,0) via 'B' = 2 strength + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:1B1/3/3"); + // White paint strength 2 at (1,0) -> bits 0,1 set (my paint >=1, >=2) + let tile_word = words[1 * 32 + 0]; + assert_eq!(tile_word & 0b1111, 0b0011); // bits 0,1 = my paint thermometer + assert_eq!((tile_word >> 4) & 0b1111, 0); // no opponent paint + } + + #[test] + fn tile_wall_and_powerup_bits() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:0,1|ps:-|bd:3/3/3"); + // (0,1) has powerup + let pu_word = words[0 * 32 + 1]; + assert_ne!(pu_word & (1 << 9), 0); + + // out-of-bounds cell should be wall + let oob_word = words[31 * 32 + 31]; + assert_ne!(oob_word & (1 << 8), 0); + } + + #[test] + fn global_features_in_correct_positions() { + let words = + encode_words("ap2|3x3|tc:0|cm:2|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + let g = &words[OFFSET_GLOBALS..]; + assert_eq!(g[0], 99); // my_stamina (white turn) + assert_eq!(g[1], 88); // opp_stamina + assert_eq!(g[16], 2); // consecutive_moves + assert_eq!(g[17], 0); // turn_count + assert_eq!(g[18], 3); // rows + assert_eq!(g[19], 3); // cols + } + + #[test] + fn local_window_has_25_tiles() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + // Local windows are at OFFSET_LOCALS..OFFSET_LOCALS+50 + let my_local = &words[OFFSET_LOCALS..OFFSET_LOCALS + 25]; + let opp_local = &words[OFFSET_LOCALS + 25..OFFSET_LOCALS + 50]; + // Both should have 25 entries, out-of-bounds tiles = WALL_BITS + assert_eq!(my_local.len(), 25); + assert_eq!(opp_local.len(), 25); + } + + #[test] + fn player_positions_encoded() { + let words = + encode_words("ap2|3x3|tc:0|cm:0|ep:0|w:0,0,99|b:2,2,88|h:n@1,1|pu:-|ps:-|bd:3/3/3"); + // White turn: my=(0,0), opp=(2,2) + assert_eq!(words[OFFSET_POSITIONS], 0); // my_x + assert_eq!(words[OFFSET_POSITIONS + 1], 0); // my_y + assert_eq!(words[OFFSET_POSITIONS + 2], 2); // opp_x + assert_eq!(words[OFFSET_POSITIONS + 3], 2); // opp_y + } + + #[test] + fn black_turn_flips_perspective() { + let words = + encode_words("ap2|3x3|tc:1|cm:0|ep:0|w:0,0,99|b:2,2,88|h:b@1,1|pu:-|ps:-|bd:1B1/3/3"); + // Black turn: my=(2,2), opp=(0,0) + assert_eq!(words[OFFSET_POSITIONS], 2); // my_x (black) + assert_eq!(words[OFFSET_POSITIONS + 1], 2); // my_y + assert_eq!(words[OFFSET_POSITIONS + 2], 0); // opp_x (white) + assert_eq!(words[OFFSET_POSITIONS + 3], 0); // opp_y + + // White paint at (1,0) should be opponent paint from black's perspective + let tile_word = words[1 * 32 + 0]; + assert_eq!(tile_word & 0b1111, 0); // no my paint (black has no paint here) + assert_ne!((tile_word >> 4) & 0b1111, 0); // opponent paint (white's paint) + + // Hill owned by black should be "my hill" from black's perspective + let hill_word = words[1 * 32 + 1]; + assert_ne!(hill_word & (1 << 13), 0); // hill current player + } +} diff --git a/training/src/queue.rs b/training/src/queue.rs new file mode 100644 index 0000000..1b2e685 --- /dev/null +++ b/training/src/queue.rs @@ -0,0 +1,600 @@ +//! Lock-free GPU job queue for batching inference requests. +//! +//! Uses atomic fetch_add for slot assignment and batch completion tracking. +//! Queue storage is sized from worker count at construction time. +//! +//! Observations are stored in a single contiguous array with shape +//! `(total_slots, ...obs_shape)`. +//! This enables zero-copy batch slicing for GPU dispatch. + +use std::cell::UnsafeCell; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +use event_listener::{Event, Listener}; +use ndarray::{Array, ArrayView, ArrayViewMut, Axis, Slice}; + +use crate::BatchDim; + +/// Number of jobs per batch. In production this would be 128. +/// Using a smaller value for tests to avoid deadlock with few workers. +#[cfg(test)] +pub const BATCH_SIZE: usize = 16; +#[cfg(not(test))] +pub const BATCH_SIZE: usize = 128; + +const SLOT_MULTIPLIER: usize = 32; + +/// Compute queue shape for a given worker count. +/// +/// Returns `(num_batches, total_slots)` where `total_slots` is rounded up to a +/// whole number of `BATCH_SIZE` lanes and is at least `2 * num_workers`. +pub fn queue_shape_for_workers(num_workers: usize) -> (usize, usize) { + assert!(num_workers > 0, "num_workers must be > 0"); + let min_slots = num_workers.saturating_mul(SLOT_MULTIPLIER).max(BATCH_SIZE); + let num_batches = min_slots.div_ceil(BATCH_SIZE); + let total_slots = num_batches * BATCH_SIZE; + (num_batches, total_slots) +} + +/// A lock-free queue for batching GPU inference jobs. +/// +/// Workers submit observations via callback and receive tickets. When a batch fills, +/// the completing worker triggers GPU dispatch with a zero-copy view of the batch. +pub struct GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + /// Monotonically increasing counter for slot assignment. + write_ticket: AtomicU64, + /// Number of batch slots in the ring. + num_batches: usize, + + /// Total number of observation/output slots. + total_slots: usize, + + /// Observation storage: shape is `(total_slots, ...obs_shape)`. + /// Single contiguous allocation for zero-copy batch slicing. + observations: UnsafeCell>, + + state: Arc>, + + /// Callback invoked when a batch is ready. + /// Receives the batch slot index, a view of batch observations, and a completion handle. + dispatch: Box, BatchCompletion) + Send + Sync>, +} + +struct QueueState +where + O: Copy + Default + Send + Sync, +{ + /// Count of completed writes per batch slot. + /// When this reaches BATCH_SIZE, the batch is ready for GPU dispatch. + batch_writes: Box<[AtomicU64]>, + + /// Ticket number at which each batch was completed (end of batch). + /// Workers check this to know if their result is ready. + batch_complete: Box<[AtomicU64]>, + + /// Output buffer. Size = `total_slots`. + outputs: Box<[UnsafeCell]>, + + /// Number of dispatched batches whose completion callback has not run yet. + inflight_batches: AtomicU64, + + /// Event for parking threads when waiting for GPU completion. + completion_event: Event, +} + +// SAFETY: Access to queue state is synchronized via ticket ownership and atomics. +unsafe impl Send for QueueState where O: Copy + Default + Send + Sync {} +unsafe impl Sync for QueueState where O: Copy + Default + Send + Sync {} + +/// Completion handle for a dispatched batch. +/// +/// The dispatch backend must call `complete` exactly once, either synchronously +/// or asynchronously (e.g. from a CUDA stream callback). +pub struct BatchCompletion +where + O: Copy + Default + Send + Sync, +{ + state: Arc>, + batch_idx: usize, + batch_start: usize, + batch_end_ticket: u64, +} + +// SAFETY: BatchCompletion only contains an Arc and plain integers. +unsafe impl Send for BatchCompletion where O: Copy + Default + Send + Sync {} + +impl BatchCompletion +where + O: Copy + Default + Send + Sync, +{ + #[inline] + pub fn complete(self, outputs: &[O]) { + debug_assert_eq!(outputs.len(), BATCH_SIZE); + for (i, output) in outputs.iter().copied().enumerate() { + unsafe { + *self.state.outputs[self.batch_start + i].get() = output; + } + } + + self.state.batch_complete[self.batch_idx].store(self.batch_end_ticket, Ordering::Release); + self.state.batch_writes[self.batch_idx].store(0, Ordering::Relaxed); + let previous = self.state.inflight_batches.fetch_sub(1, Ordering::AcqRel); + debug_assert!(previous > 0, "inflight batch counter underflowed"); + self.state.completion_event.notify(usize::MAX); + } +} + +// SAFETY: The queue is designed for concurrent access: +// - write_ticket ensures each slot is claimed by exactly one writer +// - batch_writes/batch_complete use atomic operations +// - observation slots are only written by their ticket owner, read after batch_complete +// - dispatch is Send + Sync +unsafe impl Send for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ +} +unsafe impl Sync for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ +} + +impl GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + fn compute_queue_shape(num_workers: usize) -> (usize, usize) { + queue_shape_for_workers(num_workers) + } + + /// Creates a new job queue with the given observation shape, worker count, + /// and dispatch callback. + /// + /// Queue storage is provisioned to at least `2 * num_workers` slots, + /// rounded up to a whole number of batches. + /// + /// The callback is invoked when a batch of BATCH_SIZE jobs is ready. + /// It receives a view of the batch observations (shape: BATCH_SIZE x obs_shape) + /// and should fill the outputs. + pub fn new(obs_shape: D, num_workers: usize, dispatch: F) -> Self + where + F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, + { + let (num_batches, total_slots) = Self::compute_queue_shape(num_workers); + + // Build the batched shape: `(total_slots, ...obs_shape)` + let full_shape = D::with_batch(total_slots, obs_shape); + let observations = Array::default(full_shape); + + let outputs: Box<[UnsafeCell]> = (0..total_slots) + .map(|_| UnsafeCell::new(O::default())) + .collect(); + + let batch_writes = (0..num_batches).map(|_| AtomicU64::new(0)).collect(); + let batch_complete = (0..num_batches).map(|_| AtomicU64::new(0)).collect(); + let state = Arc::new(QueueState { + batch_writes, + batch_complete, + outputs, + inflight_batches: AtomicU64::new(0), + completion_event: Event::new(), + }); + + Self { + write_ticket: AtomicU64::new(0), + num_batches, + total_slots, + observations: UnsafeCell::new(observations), + state, + dispatch: Box::new(dispatch), + } + } + + #[inline] + pub fn num_batches(&self) -> usize { + self.num_batches + } + + #[inline] + pub fn total_slots(&self) -> usize { + self.total_slots + } + + /// Submit a job by writing an observation via callback. + /// + /// The callback receives a mutable view into the queue's contiguous storage + /// for zero-copy observation writing. + /// + /// If this submission completes a batch, the current thread will + /// synchronously dispatch the batch (blocking until complete). + pub fn submit(&self, write_obs: F) -> u64 + where + F: FnOnce(ArrayViewMut), + { + // Claim a slot + let ticket = self.write_ticket.fetch_add(1, Ordering::Relaxed); + let slot_idx = (ticket as usize) % self.total_slots; + let batch_idx = ((ticket as usize) / BATCH_SIZE) % self.num_batches; + + // Get mutable view of our slot and let caller write the observation + // SAFETY: We own this slot exclusively until we increment batch_writes + // index_axis_mut on Array returns ArrayViewMut + // because D::BatchedDim::Smaller == D (guaranteed by BatchDim trait) + let slot_view = unsafe { (*self.observations.get()).index_axis_mut(Axis(0), slot_idx) }; + write_obs(slot_view); + + // AcqRel: Release our write, Acquire if we trigger dispatch to see others' writes + let writes_in_batch = self.state.batch_writes[batch_idx].fetch_add(1, Ordering::AcqRel) + 1; + + // If we completed the batch, dispatch it + if writes_in_batch == BATCH_SIZE as u64 { + self.dispatch_batch(batch_idx, ticket); + } + + ticket + } + + /// Dispatch a completed batch to the GPU. + fn dispatch_batch(&self, batch_idx: usize, trigger_ticket: u64) { + let batch_start = batch_idx * BATCH_SIZE; + + // Zero-copy slice of the batch observations + // SAFETY: All writes to this batch are complete (batch_writes == BATCH_SIZE) + let obs_array = unsafe { &*self.observations.get() }; + let batch_view = + obs_array.slice_axis(Axis(0), Slice::from(batch_start..batch_start + BATCH_SIZE)); + debug_assert!( + batch_view.is_standard_layout(), + "batch_view should be contiguous for efficient GPU transfer" + ); + + // Calculate the batch end ticket (first ticket of next batch) + let batch_number = trigger_ticket / BATCH_SIZE as u64; + let batch_end_ticket = (batch_number + 1) * BATCH_SIZE as u64; + + let completion = BatchCompletion { + state: self.state.clone(), + batch_idx, + batch_start, + batch_end_ticket, + }; + + self.state.inflight_batches.fetch_add(1, Ordering::AcqRel); + (self.dispatch)(batch_idx, batch_view, completion); + } + + fn wait_until_idle(&self) { + loop { + if self.state.inflight_batches.load(Ordering::Acquire) == 0 { + return; + } + let listener = self.state.completion_event.listen(); + if self.state.inflight_batches.load(Ordering::Acquire) == 0 { + return; + } + listener.wait(); + } + } + + /// Flush the open partial batch, if any, and wait for all in-flight GPU work. + /// + /// # Safety + /// + /// The caller must guarantee exclusive access to queue mutation while this + /// runs: no thread may call `submit`, `eval`, or any other quiesce function. + /// The intended callsite is after all session worker threads have left the + /// executor polling loop at a pause boundary. + pub unsafe fn quiesce_exclusive(&self) -> bool { + let next_ticket = self.write_ticket.load(Ordering::Acquire); + let remainder = (next_ticket % BATCH_SIZE as u64) as usize; + let mut flushed = false; + + if remainder != 0 { + let batch_number = next_ticket / BATCH_SIZE as u64; + let batch_idx = (batch_number as usize) % self.num_batches; + let writes = self.state.batch_writes[batch_idx].load(Ordering::Acquire) as usize; + + debug_assert_eq!( + writes, remainder, + "partial batch writes should match the open batch remainder" + ); + + if writes != 0 { + let batch_end_ticket = (batch_number + 1) * BATCH_SIZE as u64; + self.write_ticket.store(batch_end_ticket, Ordering::Release); + + self.dispatch_batch(batch_idx, batch_end_ticket - 1); + flushed = true; + } + } + + self.wait_until_idle(); + flushed + } + + /// Poll for a result. Returns Some(&O) if ready, None if still pending. + pub fn poll(&self, ticket: u64) -> Option<&O> { + let batch_idx = ((ticket as usize) / BATCH_SIZE) % self.num_batches; + let batch_end_ticket = ((ticket / BATCH_SIZE as u64) + 1) * BATCH_SIZE as u64; + + // Check if this batch is complete + if self.state.batch_complete[batch_idx].load(Ordering::Acquire) < batch_end_ticket { + return None; + } + + // Batch is complete, return reference to output + let slot_idx = (ticket as usize) % self.total_slots; + // SAFETY: batch_complete >= batch_end_ticket means output is written and won't change + Some(unsafe { &*self.state.outputs[slot_idx].get() }) + } + + /// Get a listener for the completion event. + /// Use this before polling to avoid missing notifications. + pub fn listen(&self) -> event_listener::EventListener { + self.state.completion_event.listen() + } + + /// Wake all waiters (used for GPU completion or external cancellation). + pub fn notify_all(&self) { + self.state.completion_event.notify(usize::MAX); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Ix0; + use std::sync::Arc; + + #[test] + fn test_single_batch_completion() { + // Use Ix0 (scalar) for simple tests + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + // Simple transform: output = input * 2 + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + // Submit BATCH_SIZE jobs + let tickets: Vec = (0..BATCH_SIZE as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // All should be complete now (last submit triggered dispatch) + for (i, &ticket) in tickets.iter().enumerate() { + let result = queue.poll(ticket); + assert!(result.is_some(), "ticket {} should be ready", ticket); + assert_eq!(*result.unwrap(), (i as u64) * 2); + } + } + + #[test] + fn test_partial_batch_not_ready() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + // Submit less than a full batch + let tickets: Vec = (0..BATCH_SIZE as u64 - 1) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // None should be ready + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_none(), + "partial batch should not be ready" + ); + } + + // Complete the batch + queue.submit(|mut out| out[()] = BATCH_SIZE as u64 - 1); + + // Now all should be ready + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_some(), + "batch should be ready after completion" + ); + } + } + + #[test] + fn test_quiesce_exclusive_completes_written_tickets() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input + 100; + } + completion.complete(&outputs); + }, + )); + + let tickets: Vec = (0..8u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + for &ticket in &tickets { + assert!( + queue.poll(ticket).is_none(), + "partial batch should not be ready" + ); + } + + unsafe { + assert!(queue.quiesce_exclusive()); + } + + for (i, &ticket) in tickets.iter().enumerate() { + let result = queue.poll(ticket).expect("flushed ticket should be ready"); + assert_eq!(*result, i as u64 + 100); + } + } + + #[test] + fn test_multiple_batches() { + let num_jobs = BATCH_SIZE * 3; + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + num_jobs, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input + 1000; + } + completion.complete(&outputs); + }, + )); + + // Submit 3 full batches + let all_tickets: Vec = (0..num_jobs as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // All should be ready + for (i, &ticket) in all_tickets.iter().enumerate() { + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!(*result, (i as u64) + 1000); + } + } + + #[test] + fn test_batch_slot_reuse() { + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + BATCH_SIZE, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = *input; + } + completion.complete(&outputs); + }, + )); + + let total_slots = queue.total_slots(); + + // Submit exactly enough jobs to fill every slot once. + let tickets_round1: Vec = (0..total_slots as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // Read all results from round 1 + for (i, &ticket) in tickets_round1.iter().enumerate() { + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!( + *result, i as u64, + "round 1 ticket {} has wrong value", + ticket + ); + } + + // Now submit another round (reusing slots) + let tickets_round2: Vec = (total_slots as u64..(total_slots * 2) as u64) + .map(|i| queue.submit(|mut out| out[()] = i)) + .collect(); + + // Read all results from round 2 + for (i, &ticket) in tickets_round2.iter().enumerate() { + let expected = (total_slots + i) as u64; + let result = queue.poll(ticket).expect("should be ready"); + assert_eq!( + *result, expected, + "round 2 ticket {} has wrong value", + ticket + ); + } + } + + #[test] + fn test_concurrent_submissions() { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::thread; + + let batch_count = Arc::new(AtomicUsize::new(0)); + let num_threads = 4; + let jobs_per_thread = BATCH_SIZE * 2; // Each thread submits 2 batches worth + + let queue: Arc> = Arc::new(GpuJobQueue::new( + Ix0(), + num_threads * jobs_per_thread, + |_batch_idx, inputs, completion| { + let mut outputs = vec![0u64; BATCH_SIZE]; + for (i, input) in inputs.iter().enumerate() { + outputs[i] = input * 2; + } + completion.complete(&outputs); + }, + )); + + let handles: Vec<_> = (0..num_threads) + .map(|thread_id| { + let queue = queue.clone(); + let batch_count = batch_count.clone(); + thread::spawn(move || { + let base = (thread_id * jobs_per_thread) as u64; + let mut results = Vec::new(); + + for i in 0..jobs_per_thread as u64 { + let val = base + i; + let ticket = queue.submit(|mut out| out[()] = val); + results.push((ticket, val)); + } + + // Wait for results + for (ticket, expected_input) in results { + loop { + if let Some(&result) = queue.poll(ticket) { + assert_eq!(result, expected_input * 2); + batch_count.fetch_add(1, Ordering::Relaxed); + break; + } + // Busy wait (in real code we'd use the event listener) + std::hint::spin_loop(); + } + } + }) + }) + .collect(); + + for handle in handles { + handle.join().expect("thread panicked"); + } + + assert_eq!( + batch_count.load(Ordering::Relaxed), + num_threads * jobs_per_thread + ); + } +} diff --git a/training/src/replay_buffer.rs b/training/src/replay_buffer.rs new file mode 100644 index 0000000..9637f1c --- /dev/null +++ b/training/src/replay_buffer.rs @@ -0,0 +1,280 @@ +//! Lock-free ring buffer for contiguous observation replay storage. + +use std::cell::UnsafeCell; +use std::sync::atomic::{AtomicU64, Ordering}; + +use ndarray::{Array, ArrayViewMut, Axis}; +use rand::seq::index::sample; +use rand::Rng; + +use crate::BatchDim; + +/// Batched sample output from [`ReplayBuffer::sample`]. +pub struct SampleBatch +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub observations: Array, + pub values: Vec, +} + +/// Lock-free ring buffer for storing observations and values. +/// +/// Observations are kept in single contiguous arrays: +/// - observations: `(capacity, ...obs_shape)` +/// - values: `(capacity,)` +pub struct ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + observations: UnsafeCell>, + values: UnsafeCell>, + capacity: usize, + obs_shape: D, + obs_elems_per_sample: usize, + head: AtomicU64, + writers: AtomicU64, +} + +// SAFETY: Each writer reserves unique slots through an atomic ticket and writes +// only to its owned slots until drop. Readers require no active writers. +unsafe impl Sync for ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} +unsafe impl Send for ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ +} + +/// RAII guard for writing to reserved slots. +pub struct ReserveGuard<'a, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + buffer: &'a ReplayBuffer, + start: u64, + len: usize, + written: usize, +} + +impl<'a, A, D> ReserveGuard<'a, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + /// Push one sample by writing observation directly into the reserved slot. + #[inline] + pub fn push_with_observation(&mut self, value: f32, write_observation: F) + where + F: FnOnce(ArrayViewMut), + { + assert!(self.written < self.len, "wrote more samples than reserved"); + let idx = (self.start + self.written as u64) as usize % self.buffer.capacity; + + unsafe { + let slot_view = (*self.buffer.observations.get()).index_axis_mut(Axis(0), idx); + write_observation(slot_view); + + (&mut *self.buffer.values.get())[idx] = value; + } + + self.written += 1; + } + + /// Push one sample from flattened observation data. + pub fn push(&mut self, observation: &[A], value: f32) { + assert_eq!( + observation.len(), + self.buffer.obs_elems_per_sample, + "observation length must match env observation size" + ); + + self.push_with_observation(value.clamp(-1., 1.), |mut out| { + for (dst, src) in out.iter_mut().zip(observation.iter()) { + *dst = src.clone(); + } + }); + } +} + +impl Drop for ReserveGuard<'_, A, D> +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + fn drop(&mut self) { + self.buffer.writers.fetch_sub(1, Ordering::Release); + } +} + +impl ReplayBuffer +where + A: Clone + Default + Send + Sync, + D: BatchDim, +{ + pub fn new(capacity: usize, obs_shape: D) -> Self { + assert!(capacity > 0, "capacity must be > 0"); + let obs_elems_per_sample = obs_shape.clone().size(); + + Self { + observations: UnsafeCell::new(Array::default(D::with_batch( + capacity, + obs_shape.clone(), + ))), + values: UnsafeCell::new(vec![0.0; capacity]), + capacity, + obs_shape, + obs_elems_per_sample, + head: AtomicU64::new(0), + writers: AtomicU64::new(0), + } + } + + pub fn reserve(&self, n: usize) -> ReserveGuard<'_, A, D> { + assert!( + n <= self.capacity, + "cannot reserve more samples than buffer capacity" + ); + + self.writers.fetch_add(1, Ordering::Acquire); + let start = self.head.fetch_add(n as u64, Ordering::AcqRel); + ReserveGuard { + buffer: self, + start, + len: n, + written: 0, + } + } + + #[inline] + fn valid_range(&self) -> (u64, u64) { + let head = self.head.load(Ordering::Acquire); + let tail = head.saturating_sub(self.capacity as u64); + (tail, head) + } + + #[inline] + pub fn len(&self) -> usize { + let (tail, head) = self.valid_range(); + (head - tail) as usize + } + + #[inline] + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn capacity(&self) -> usize { + self.capacity + } + + /// Sample `n` items uniformly. Panics if writers are active. + pub fn sample(&self, n: usize, rng: &mut impl Rng) -> SampleBatch { + assert_eq!( + self.writers.load(Ordering::Acquire), + 0, + "cannot sample while writers are active" + ); + + let (tail, head) = self.valid_range(); + let count = (head - tail) as usize; + if count == 0 || n == 0 { + return SampleBatch { + observations: Array::default(D::with_batch(0, self.obs_shape.clone())), + values: Vec::new(), + }; + } + + let sample_count = n.min(count); + let indices = sample(rng, count, sample_count); + + let observations = unsafe { &*self.observations.get() }; + let observation_slice = observations + .as_slice_memory_order() + .expect("observation storage must be contiguous"); + + let values = unsafe { &*self.values.get() }; + + let mut obs_data = Vec::with_capacity(sample_count * self.obs_elems_per_sample); + let mut value_data = Vec::with_capacity(sample_count); + + for offset in indices.iter() { + let idx = (tail + offset as u64) as usize % self.capacity; + + let obs_start = idx * self.obs_elems_per_sample; + obs_data.extend_from_slice( + &observation_slice[obs_start..obs_start + self.obs_elems_per_sample], + ); + + value_data.push(values[idx]); + } + + let observations = Array::from_shape_vec( + D::with_batch(sample_count, self.obs_shape.clone()), + obs_data, + ) + .expect("sample observation shape mismatch"); + + SampleBatch { + observations, + values: value_data, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use ndarray::Ix1; + use rand::SeedableRng; + use rand_chacha::ChaCha8Rng; + + #[test] + fn test_push_and_sample() { + let buffer = ReplayBuffer::::new(10, Ix1(4)); + + { + let mut guard = buffer.reserve(2); + guard.push(&[1, 2, 3, 4], 0.5); + guard.push(&[5, 6, 7, 8], -0.5); + } + + assert_eq!(buffer.len(), 2); + + let mut rng = ChaCha8Rng::seed_from_u64(7); + let batch = buffer.sample(2, &mut rng); + assert_eq!(batch.observations.shape(), &[2, 4]); + assert_eq!(batch.values.len(), 2); + } + + #[test] + fn test_wraparound_len() { + let buffer = ReplayBuffer::::new(3, Ix1(2)); + + for i in 0..10 { + let mut guard = buffer.reserve(1); + guard.push(&[i as i8, (i + 1) as i8], i as f32); + } + + assert_eq!(buffer.len(), 3); + } + + #[test] + #[should_panic(expected = "cannot sample while writers are active")] + fn test_sample_during_write_panics() { + let buffer = ReplayBuffer::::new(4, Ix1(2)); + let _guard = buffer.reserve(1); + + let mut rng = ChaCha8Rng::seed_from_u64(42); + let _ = buffer.sample(1, &mut rng); + } +} diff --git a/training/src/training.rs b/training/src/training.rs new file mode 100644 index 0000000..09c0262 --- /dev/null +++ b/training/src/training.rs @@ -0,0 +1,562 @@ +//! Training infrastructure - thread spawning and coordination. +//! +//! Provides `SelfPlaySession`: a persistent session with pause/resume semantics +//! that preserves in-progress game state across boundaries. + +use std::mem; +use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread; + +use ndarray::{ArrayView, Ix1, Ix2}; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; + +use crate::eval::GpuEvaluator; +use crate::executor::Executor; +use crate::queue::{BatchCompletion, GpuJobQueue}; +use crate::replay_buffer::ReplayBuffer; +use crate::worker::{worker_loop_forever, SelfPlayMetrics, WorkerConfig}; +use crate::BatchDim; + +/// Shared control state for the persistent self-play session. +struct SessionControl { + /// Whether workers should be actively polling. + running: AtomicBool, + /// Whether the session is being torn down. + shutdown: AtomicBool, + /// Absolute target: workers pause when `samples_collected >= target`. + target_samples: AtomicUsize, + /// Total samples collected (monotonic across the session lifetime). + samples_collected: Arc, + /// Total terminal leaf samples injected into replay via terminal mixing. + terminal_mix_samples_added: Arc, + /// Total games completed. + games_completed: Arc, + /// Total actions selected across all workers. + action_steps: Arc, + /// Total final actions selected across all workers. + final_actions: Arc, + /// Total non-final actions selected across all workers. + nonfinal_actions: Arc, + /// Total turns completed via selected final actions. + completed_turns: Arc, + /// Sum of `board.turn_count` observed before each selected action. + action_turn_count_total: Arc, + /// Largest `board.turn_count` seen in any in-progress game. + max_turn_count_seen: Arc, + /// Total selected actions consumed by completed games. + completed_game_actions_total: Arc, + /// Largest number of selected actions observed in a completed game. + max_actions_in_completed_game: Arc, + /// Total completed turn count accumulated across finished games. + completed_game_turn_count_total: Arc, + /// Largest completed turn count observed in a finished game. + max_turn_count_in_completed_game: Arc, + /// Completed game turn counts waiting to be drained by the trainer. + completed_game_turn_counts: Arc>>, + /// Total time spent building fresh search trees. + tree_build_nanos: Arc, + /// Total time spent running descent iterations. + descent_nanos: Arc, + /// Total time spent traversing trees to collect replay samples. + sample_collect_nanos: Arc, + /// Total time spent encoding and pushing replay samples. + replay_push_nanos: Arc, + /// Total CPU expansion prep time inside Descent. + descent_expand_cpu_nanos: Arc, + /// Total time spent applying actions during Descent. + descent_apply_action_nanos: Arc, + /// Total time spent submitting evals to the GPU queue. + descent_eval_submit_nanos: Arc, + /// Total time spent awaiting queued evals. + descent_eval_await_nanos: Arc, + /// Total time spent backing up/searching best children. + descent_backup_nanos: Arc, + /// Number of threads currently inside the executor polling loop. + active_pollers: AtomicUsize, + /// Condvar + mutex for coordinating start/pause/quiesce/shutdown. + condvar: Condvar, + condvar_mutex: Mutex<()>, +} + +impl SessionControl { + fn new() -> Self { + Self { + running: AtomicBool::new(false), + shutdown: AtomicBool::new(false), + target_samples: AtomicUsize::new(0), + samples_collected: Arc::new(AtomicUsize::new(0)), + terminal_mix_samples_added: Arc::new(AtomicU64::new(0)), + games_completed: Arc::new(AtomicUsize::new(0)), + action_steps: Arc::new(AtomicU64::new(0)), + final_actions: Arc::new(AtomicU64::new(0)), + nonfinal_actions: Arc::new(AtomicU64::new(0)), + completed_turns: Arc::new(AtomicU64::new(0)), + action_turn_count_total: Arc::new(AtomicU64::new(0)), + max_turn_count_seen: Arc::new(AtomicUsize::new(0)), + completed_game_actions_total: Arc::new(AtomicU64::new(0)), + max_actions_in_completed_game: Arc::new(AtomicU64::new(0)), + completed_game_turn_count_total: Arc::new(AtomicU64::new(0)), + max_turn_count_in_completed_game: Arc::new(AtomicUsize::new(0)), + completed_game_turn_counts: Arc::new(Mutex::new(Vec::new())), + tree_build_nanos: Arc::new(AtomicU64::new(0)), + descent_nanos: Arc::new(AtomicU64::new(0)), + sample_collect_nanos: Arc::new(AtomicU64::new(0)), + replay_push_nanos: Arc::new(AtomicU64::new(0)), + descent_expand_cpu_nanos: Arc::new(AtomicU64::new(0)), + descent_apply_action_nanos: Arc::new(AtomicU64::new(0)), + descent_eval_submit_nanos: Arc::new(AtomicU64::new(0)), + descent_eval_await_nanos: Arc::new(AtomicU64::new(0)), + descent_backup_nanos: Arc::new(AtomicU64::new(0)), + active_pollers: AtomicUsize::new(0), + condvar: Condvar::new(), + condvar_mutex: Mutex::new(()), + } + } + + fn wake_all(&self) { + self.condvar.notify_all(); + } + + fn should_pause(&self) -> bool { + if self.shutdown.load(Ordering::Acquire) { + return true; + } + if !self.running.load(Ordering::Acquire) { + return true; + } + self.samples_collected.load(Ordering::Acquire) + >= self.target_samples.load(Ordering::Acquire) + } +} + +/// Configuration for creating a persistent self-play session. +#[derive(Clone)] +pub struct SessionConfig { + /// Number of OS threads to spawn. + pub num_threads: usize, + /// Number of workers per thread. + pub workers_per_thread: usize, + /// Worker configuration. + pub worker: WorkerConfig, + /// Random seed for reproducibility. + pub seed: u64, + /// Stack size for each self-play OS thread. + pub thread_stack_size_bytes: usize, +} + +impl Default for SessionConfig { + fn default() -> Self { + Self { + num_threads: 32, + workers_per_thread: 8, + worker: WorkerConfig::default(), + seed: 42, + thread_stack_size_bytes: 128 * 1024 * 1024, + } + } +} + +/// Trait-object wrapper so we can call `notify_all()` on the queue. +trait QueueNotify: Send + Sync { + unsafe fn quiesce_exclusive(&self) -> bool; + fn notify_all(&self); +} + +impl QueueNotify for GpuJobQueue +where + A: Clone + Default + Send + Sync, + D: BatchDim, + O: Copy + Default + Send + Sync, +{ + unsafe fn quiesce_exclusive(&self) -> bool { + GpuJobQueue::quiesce_exclusive(self) + } + + fn notify_all(&self) { + GpuJobQueue::notify_all(self); + } +} + +/// A persistent self-play session that owns worker threads and preserves +/// in-progress game state across pause/resume boundaries. +pub struct SelfPlaySession { + control: Arc, + queue_notify: Arc, + threads: Option>>, +} + +impl SelfPlaySession { + /// Create a new persistent session. + /// + /// Uses concrete AlphaPaint types: packed u16 observations with Ix1 shape, f32 values. + pub fn new( + config: SessionConfig, + replay_buffer: Arc>, + dispatch: F, + ) -> Self + where + F: Fn(usize, ArrayView, BatchCompletion) + Send + Sync + 'static, + { + let total_workers = config + .num_threads + .checked_mul(config.workers_per_thread) + .expect("num_threads * workers_per_thread overflowed usize"); + + let obs_shape = Ix1(crate::cudagraph::OBS_WORDS); + let queue = Arc::new(GpuJobQueue::new(obs_shape, total_workers, dispatch)); + + let control = Arc::new(SessionControl::new()); + + let mut threads = Vec::with_capacity(config.num_threads); + for thread_id in 0..config.num_threads { + let queue = queue.clone(); + let control = control.clone(); + let config = config.clone(); + let replay_buffer = replay_buffer.clone(); + + let thread_name = format!("selfplay-{thread_id}"); + let stack_size = config.thread_stack_size_bytes; + let handle = thread::Builder::new() + .name(thread_name) + .stack_size(stack_size) + .spawn(move || { + session_thread_main(thread_id, queue, config, control, &replay_buffer); + }) + .expect("failed to spawn self-play thread"); + threads.push(handle); + } + + Self { + control, + queue_notify: queue, + threads: Some(threads), + } + } + + /// Start self-play with no sample limit. + pub fn start(&self) { + self.control + .target_samples + .store(usize::MAX, Ordering::Release); + self.control.running.store(true, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + } + + /// Block until at least `target_samples` absolute samples have been + /// collected, then pause and quiesce all workers. + pub fn wait_for(&self, target_samples: usize) -> usize { + self.control + .target_samples + .store(target_samples, Ordering::Release); + self.control.running.store(true, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + + { + let mut guard = self + .control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while self.control.samples_collected.load(Ordering::Acquire) < target_samples + && !self.control.shutdown.load(Ordering::Acquire) + { + guard = self + .control + .condvar + .wait(guard) + .expect("condvar wait failed"); + } + } + + self.control.running.store(false, Ordering::Release); + self.queue_notify.notify_all(); + self.control.wake_all(); + + { + let mut guard = self + .control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while self.control.active_pollers.load(Ordering::Acquire) > 0 + && !self.control.shutdown.load(Ordering::Acquire) + { + guard = self + .control + .condvar + .wait(guard) + .expect("condvar wait failed"); + } + } + + // SAFETY: `wait_for` has already stopped the session and waited until + // every session thread has left the executor loop (`active_pollers == 0`), + // so no thread can submit new queue writes or race another quiesce call. + unsafe { + self.queue_notify.quiesce_exclusive(); + } + + self.control.samples_collected.load(Ordering::Acquire) + } + + /// Return the current absolute sample count. + pub fn samples(&self) -> usize { + self.control.samples_collected.load(Ordering::Acquire) + } + + /// Return the current absolute game count. + pub fn games(&self) -> usize { + self.control.games_completed.load(Ordering::Acquire) + } + + /// Return the total number of terminal leaf samples injected into replay. + pub fn terminal_mix_samples_added(&self) -> u64 { + self.control + .terminal_mix_samples_added + .load(Ordering::Acquire) + } + + /// Return the total number of selected actions. + pub fn action_steps(&self) -> u64 { + self.control.action_steps.load(Ordering::Acquire) + } + + /// Return the total number of selected final actions. + pub fn final_actions(&self) -> u64 { + self.control.final_actions.load(Ordering::Acquire) + } + + /// Return the total number of selected non-final actions. + pub fn nonfinal_actions(&self) -> u64 { + self.control.nonfinal_actions.load(Ordering::Acquire) + } + + /// Return the total number of completed turns. + pub fn completed_turns(&self) -> u64 { + self.control.completed_turns.load(Ordering::Acquire) + } + + /// Return the sum of `board.turn_count` observed before each action. + pub fn action_turn_count_total(&self) -> u64 { + self.control.action_turn_count_total.load(Ordering::Acquire) + } + + /// Return the largest `board.turn_count` seen so far. + pub fn max_turn_count_seen(&self) -> usize { + self.control.max_turn_count_seen.load(Ordering::Acquire) + } + + /// Return the total number of selected actions in completed games. + pub fn completed_game_actions_total(&self) -> u64 { + self.control + .completed_game_actions_total + .load(Ordering::Acquire) + } + + /// Return the largest number of selected actions seen in a completed game. + pub fn max_actions_in_completed_game(&self) -> u64 { + self.control + .max_actions_in_completed_game + .load(Ordering::Acquire) + } + + /// Return the total completed turn count accumulated across finished games. + pub fn completed_game_turn_count_total(&self) -> u64 { + self.control + .completed_game_turn_count_total + .load(Ordering::Acquire) + } + + /// Return the largest completed turn count seen in a finished game. + pub fn max_turn_count_in_completed_game(&self) -> usize { + self.control + .max_turn_count_in_completed_game + .load(Ordering::Acquire) + } + + /// Drain completed game turn counts collected since the last call. + pub fn take_completed_game_turn_counts(&self) -> Vec { + let mut guard = self + .control + .completed_game_turn_counts + .lock() + .expect("completed game turn counts mutex poisoned"); + mem::take(&mut *guard) + } + + /// Return the total time spent building fresh search trees. + pub fn tree_build_nanos(&self) -> u64 { + self.control.tree_build_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent running descent iterations. + pub fn descent_nanos(&self) -> u64 { + self.control.descent_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent traversing trees to collect replay samples. + pub fn sample_collect_nanos(&self) -> u64 { + self.control.sample_collect_nanos.load(Ordering::Acquire) + } + + /// Return the total time spent encoding and pushing replay samples. + pub fn replay_push_nanos(&self) -> u64 { + self.control.replay_push_nanos.load(Ordering::Acquire) + } + + /// Return total CPU expansion prep time inside Descent. + pub fn descent_expand_cpu_nanos(&self) -> u64 { + self.control + .descent_expand_cpu_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent applying actions during Descent. + pub fn descent_apply_action_nanos(&self) -> u64 { + self.control + .descent_apply_action_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent submitting evals to the GPU queue. + pub fn descent_eval_submit_nanos(&self) -> u64 { + self.control + .descent_eval_submit_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent awaiting queued evals. + pub fn descent_eval_await_nanos(&self) -> u64 { + self.control + .descent_eval_await_nanos + .load(Ordering::Acquire) + } + + /// Return total time spent backing up/searching best children. + pub fn descent_backup_nanos(&self) -> u64 { + self.control.descent_backup_nanos.load(Ordering::Acquire) + } + + /// Shut down the session. Idempotent. + pub fn shutdown(&mut self) { + if let Some(threads) = self.threads.take() { + self.control.shutdown.store(true, Ordering::Release); + self.control.running.store(false, Ordering::Release); + self.control.wake_all(); + self.queue_notify.notify_all(); + + for handle in threads { + let _ = handle.join(); + } + } + } +} + +impl Drop for SelfPlaySession { + fn drop(&mut self) { + self.shutdown(); + } +} + +/// Main loop for a single thread in a persistent session. +fn session_thread_main( + thread_id: usize, + queue: Arc>, + config: SessionConfig, + control: Arc, + replay_buffer: &ReplayBuffer, +) { + let base_seed = config.seed.wrapping_add(thread_id as u64 * 1000); + let evaluator = GpuEvaluator::new(&*queue); + + let metrics = SelfPlayMetrics { + samples_collected: control.samples_collected.clone(), + terminal_mix_samples_added: control.terminal_mix_samples_added.clone(), + games_completed: control.games_completed.clone(), + action_steps: control.action_steps.clone(), + final_actions: control.final_actions.clone(), + nonfinal_actions: control.nonfinal_actions.clone(), + completed_turns: control.completed_turns.clone(), + action_turn_count_total: control.action_turn_count_total.clone(), + max_turn_count_seen: control.max_turn_count_seen.clone(), + completed_game_actions_total: control.completed_game_actions_total.clone(), + max_actions_in_completed_game: control.max_actions_in_completed_game.clone(), + completed_game_turn_count_total: control.completed_game_turn_count_total.clone(), + max_turn_count_in_completed_game: control.max_turn_count_in_completed_game.clone(), + completed_game_turn_counts: control.completed_game_turn_counts.clone(), + tree_build_nanos: control.tree_build_nanos.clone(), + descent_nanos: control.descent_nanos.clone(), + sample_collect_nanos: control.sample_collect_nanos.clone(), + replay_push_nanos: control.replay_push_nanos.clone(), + descent_expand_cpu_nanos: control.descent_expand_cpu_nanos.clone(), + descent_apply_action_nanos: control.descent_apply_action_nanos.clone(), + descent_eval_submit_nanos: control.descent_eval_submit_nanos.clone(), + descent_eval_await_nanos: control.descent_eval_await_nanos.clone(), + descent_backup_nanos: control.descent_backup_nanos.clone(), + }; + + let mut futures: Vec<_> = (0..config.workers_per_thread) + .map(|i| { + let metrics = metrics.clone(); + let mut rng = ChaCha8Rng::seed_from_u64(base_seed + i as u64); + let evaluator_ref = &evaluator; + let worker_config = &config.worker; + + let fut = async move { + worker_loop_forever( + evaluator_ref, + worker_config, + &mut rng, + metrics, + replay_buffer, + ) + .await; + }; + Box::pin(fut) + }) + .collect(); + + let executor = Executor::new(|| queue.listen()); + + loop { + { + let mut guard = control + .condvar_mutex + .lock() + .expect("condvar mutex poisoned"); + while !control.running.load(Ordering::Acquire) + && !control.shutdown.load(Ordering::Acquire) + { + guard = control.condvar.wait(guard).expect("condvar wait failed"); + } + } + + if control.shutdown.load(Ordering::Acquire) { + return; + } + + control.active_pollers.fetch_add(1, Ordering::AcqRel); + + let control_ref = &control; + let queue_ref = &queue; + executor.run(&mut futures, &mut || { + let should_pause = control_ref.should_pause(); + if should_pause { + queue_ref.notify_all(); + } + if control_ref.samples_collected.load(Ordering::Acquire) + >= control_ref.target_samples.load(Ordering::Acquire) + { + control_ref.wake_all(); + } + should_pause + }); + + control.active_pollers.fetch_sub(1, Ordering::AcqRel); + control.wake_all(); + } +} diff --git a/training/src/worker.rs b/training/src/worker.rs new file mode 100644 index 0000000..ffb16fb --- /dev/null +++ b/training/src/worker.rs @@ -0,0 +1,367 @@ +//! Worker loop for Athénan-style self-play with tree learning. +//! +//! Each worker runs Descent search, collects training samples from +//! internal tree nodes, and pushes them to the replay buffer. + +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use ndarray::Ix1; +use rand::rngs::SmallRng; +use rand::{Rng, RngExt, SeedableRng}; + +use alpha_paint::board::Board; +use alpha_paint::TRAINING_START_FENS; + +use crate::descent::{GameSearchTree, SearchTimingMetrics, TreeLearningSample}; +use crate::eval::{CountingEvaluator, Evaluator}; +use crate::observation; +use crate::replay_buffer::ReplayBuffer; + +const RANDOM_TERMINAL_MIX_PROB: f32 = 0.2; + +fn update_max_usize(max_value: &AtomicUsize, candidate: usize) { + let mut current = max_value.load(Ordering::Acquire); + while candidate > current { + match max_value.compare_exchange_weak( + current, + candidate, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => current = observed, + } + } +} + +fn update_max_u64(max_value: &AtomicU64, candidate: u64) { + let mut current = max_value.load(Ordering::Acquire); + while candidate > current { + match max_value.compare_exchange_weak( + current, + candidate, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => break, + Err(observed) => current = observed, + } + } +} + +fn elapsed_nanos(started_at: Instant) -> u64 { + started_at.elapsed().as_nanos().min(u64::MAX as u128) as u64 +} + +#[derive(Clone)] +pub struct SelfPlayMetrics { + pub samples_collected: Arc, + pub terminal_mix_samples_added: Arc, + pub games_completed: Arc, + pub action_steps: Arc, + pub final_actions: Arc, + pub nonfinal_actions: Arc, + pub completed_turns: Arc, + pub action_turn_count_total: Arc, + pub max_turn_count_seen: Arc, + pub completed_game_actions_total: Arc, + pub max_actions_in_completed_game: Arc, + pub completed_game_turn_count_total: Arc, + pub max_turn_count_in_completed_game: Arc, + pub completed_game_turn_counts: Arc>>, + pub tree_build_nanos: Arc, + pub descent_nanos: Arc, + pub sample_collect_nanos: Arc, + pub replay_push_nanos: Arc, + pub descent_expand_cpu_nanos: Arc, + pub descent_apply_action_nanos: Arc, + pub descent_eval_submit_nanos: Arc, + pub descent_eval_await_nanos: Arc, + pub descent_backup_nanos: Arc, +} + +/// Configuration for the worker. +#[derive(Clone)] +pub struct WorkerConfig { + /// Maximum GPU eval submissions spent on descent for a single move. + pub max_gpu_evals_per_move: u64, +} + +impl Default for WorkerConfig { + fn default() -> Self { + Self { + max_gpu_evals_per_move: 4 * 1024, + } + } +} + +fn push_samples_to_replay( + replay_buffer: &ReplayBuffer, + metrics: &SelfPlayMetrics, + samples: Vec, +) { + let num_samples = samples.len(); + if num_samples == 0 { + return; + } + + let terminal_mix_count = samples.iter().filter(|sample| sample.terminal_mix).count(); + + let started_at = Instant::now(); + + let mut guard = replay_buffer.reserve(num_samples); + for sample in samples { + guard.push_with_observation(sample.value, |mut out| { + observation::encode_into_slice( + &sample.board, + out.as_slice_mut() + .expect("replay observation slot must be contiguous"), + ); + }); + } + metrics + .samples_collected + .fetch_add(num_samples, Ordering::AcqRel); + metrics + .terminal_mix_samples_added + .fetch_add(terminal_mix_count as u64, Ordering::AcqRel); + metrics + .replay_push_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); +} + +fn collect_tree_learning_samples_timed( + tree: &GameSearchTree<'_, E>, + metrics: &SelfPlayMetrics, + rng: &mut R, +) -> Vec { + let started_at = Instant::now(); + let samples = + tree.collect_tree_learning_samples_with_terminal_mix(rng, RANDOM_TERMINAL_MIX_PROB); + metrics + .sample_collect_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); + samples +} + +fn step_tree_and_collect_dropped_samples_timed( + tree: &mut GameSearchTree<'_, E>, + new_board: &Board, + action_id: usize, + metrics: &SelfPlayMetrics, + rng: &mut R, +) -> Vec { + let started_at = Instant::now(); + let samples = tree.step_tree_and_collect_dropped_samples_with_terminal_mix( + new_board, + action_id, + rng, + RANDOM_TERMINAL_MIX_PROB, + ); + metrics + .sample_collect_nanos + .fetch_add(elapsed_nanos(started_at), Ordering::AcqRel); + samples +} + +fn finish_game_with_current_tree( + tree: &GameSearchTree<'_, E>, + replay_buffer: &ReplayBuffer, + metrics: &SelfPlayMetrics, + rng: &mut impl Rng, + action_steps_in_game: u64, + completed_turn_count_in_game: usize, +) { + push_samples_to_replay( + replay_buffer, + metrics, + collect_tree_learning_samples_timed(tree, metrics, rng), + ); + metrics + .completed_game_actions_total + .fetch_add(action_steps_in_game, Ordering::AcqRel); + update_max_u64(&metrics.max_actions_in_completed_game, action_steps_in_game); + metrics + .completed_game_turn_count_total + .fetch_add(completed_turn_count_in_game as u64, Ordering::AcqRel); + update_max_usize( + &metrics.max_turn_count_in_completed_game, + completed_turn_count_in_game, + ); + metrics + .completed_game_turn_counts + .lock() + .expect("completed game turn count mutex poisoned") + .push(completed_turn_count_in_game); +} + +/// Run a single self-play game with tree learning. +/// +/// At each move: +/// 1. Run Descent search +/// 2. Select action via ordinal distribution +/// 3. Collect the dropped portion of the tree with its latest values +/// 4. Apply action, reuse tree +/// +async fn play_game( + evaluator: &E, + config: &WorkerConfig, + rng: &mut R, + replay_buffer: &ReplayBuffer, + metrics: &SelfPlayMetrics, +) { + let board = + Board::from_fen(TRAINING_START_FENS[rng.random_range(0..TRAINING_START_FENS.len())]) + .expect("training start FEN must parse"); + + // Check if terminal before starting + // (board.get_valid_actions().len() == 0 would indicate terminal) + let actions = board.get_valid_actions(); + if actions.len() == 0 { + return; + } + + let counting_evaluator = CountingEvaluator::new(evaluator); + let search_timing = SearchTimingMetrics { + expand_cpu_nanos: metrics.descent_expand_cpu_nanos.clone(), + apply_action_nanos: metrics.descent_apply_action_nanos.clone(), + eval_submit_nanos: metrics.descent_eval_submit_nanos.clone(), + eval_await_nanos: metrics.descent_eval_await_nanos.clone(), + backup_nanos: metrics.descent_backup_nanos.clone(), + }; + let tree_rng = SmallRng::from_rng(rng); + let build_started_at = Instant::now(); + let mut tree = + GameSearchTree::new(&board, &counting_evaluator, tree_rng, search_timing.clone()).await; + metrics + .tree_build_nanos + .fetch_add(elapsed_nanos(build_started_at), Ordering::AcqRel); + let mut action_steps_in_game = 0u64; + + loop { + let current_turn_count = tree.root_state.turn_count; + metrics + .action_turn_count_total + .fetch_add(current_turn_count as u64, Ordering::AcqRel); + update_max_usize(&metrics.max_turn_count_seen, current_turn_count); + + // Run descent search + let descent_started_at = Instant::now(); + tree.run_descent_to_eval_limit(config.max_gpu_evals_per_move) + .await; + metrics + .descent_nanos + .fetch_add(elapsed_nanos(descent_started_at), Ordering::AcqRel); + + // Select action via ordinal distribution + if tree.root_node.children.is_empty() { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + rng, + action_steps_in_game, + tree.root_state.turn_count, + ); + break; + } + + let action_id = tree.ordinal_select(); + let action = tree.root_node.children[action_id].action; + action_steps_in_game += 1; + metrics.action_steps.fetch_add(1, Ordering::AcqRel); + if action.is_final() { + metrics.final_actions.fetch_add(1, Ordering::AcqRel); + metrics.completed_turns.fetch_add(1, Ordering::AcqRel); + } else { + metrics.nonfinal_actions.fetch_add(1, Ordering::AcqRel); + } + + // Apply action to get new board state + let mut new_board = tree.root_state.clone(); + let (outcome, _) = new_board.apply_action(action); + update_max_usize(&metrics.max_turn_count_seen, new_board.turn_count); + + // Check if game is over + match outcome { + alpha_paint::board::ApplyActionOutcome::Terminal { .. } => { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + rng, + action_steps_in_game, + new_board.turn_count, + ); + break; + } + alpha_paint::board::ApplyActionOutcome::Killshot { .. } => { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + rng, + action_steps_in_game, + tree.root_state.turn_count + 1, + ); + break; + } + alpha_paint::board::ApplyActionOutcome::PlayInstead { .. } => { + finish_game_with_current_tree( + &tree, + replay_buffer, + metrics, + rng, + action_steps_in_game, + tree.root_state.turn_count + 1, + ); + break; + } + alpha_paint::board::ApplyActionOutcome::Ongoing => { + if tree.root_node.children[action_id].node.is_some() { + let samples = step_tree_and_collect_dropped_samples_timed( + &mut tree, &new_board, action_id, metrics, rng, + ); + push_samples_to_replay(replay_buffer, metrics, samples); + } else { + push_samples_to_replay( + replay_buffer, + metrics, + collect_tree_learning_samples_timed(&tree, metrics, rng), + ); + let next_rng = SmallRng::from_rng(rng); + let build_started_at = Instant::now(); + tree = GameSearchTree::new( + &new_board, + &counting_evaluator, + next_rng, + search_timing.clone(), + ) + .await; + metrics + .tree_build_nanos + .fetch_add(elapsed_nanos(build_started_at), Ordering::AcqRel); + } + } + } + } +} + +/// Run a worker loop that plays games forever. +/// +/// Stopping is handled externally by the executor's cancel/pause mechanism. +/// Training samples (obs, value) are pushed to the shared replay buffer. +pub async fn worker_loop_forever( + evaluator: &E, + config: &WorkerConfig, + rng: &mut R, + metrics: SelfPlayMetrics, + replay_buffer: &ReplayBuffer, +) { + loop { + play_game(evaluator, config, rng, replay_buffer, &metrics).await; + metrics.games_completed.fetch_add(1, Ordering::AcqRel); + } +} diff --git a/uv.lock b/uv.lock index 9ff6113..749aad0 100644 --- a/uv.lock +++ b/uv.lock @@ -11,6 +11,9 @@ dependencies = [ { name = "numpy" }, { name = "psutil" }, { name = "py-cpuinfo" }, + { name = "torch" }, + { name = "triton", marker = "sys_platform == 'linux'" }, + { name = "wandb" }, ] [package.dev-dependencies] @@ -26,6 +29,9 @@ requires-dist = [ { name = "numpy", specifier = "==2.1.3" }, { name = "psutil", specifier = "==5.9.0" }, { name = "py-cpuinfo" }, + { name = "torch", specifier = "==2.10.0" }, + { name = "triton", marker = "sys_platform == 'linux'", specifier = ">=3.6.0" }, + { name = "wandb", specifier = ">=0.25.1" }, ] [package.metadata.requires-dev] @@ -34,6 +40,61 @@ dev = [ { name = "pytest-xdist", specifier = ">=3" }, ] +[[package]] +name = "annotated-types" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ee/67/531ea369ba64dcff5ec9c3402f9f51bf748cec26dde048a2f973a4eea7f5/annotated_types-0.7.0.tar.gz", hash = "sha256:aff07c09a53a08bc8cfccb9c85b05f1aa9a2a6f23728d790723543408344ce89", size = 16081, upload-time = "2024-05-20T21:33:25.928Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, +] + +[[package]] +name = "certifi" +version = "2026.2.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/2d/7bf41579a8986e348fa033a31cdd0e4121114f6bce2457e8876010b092dd/certifi-2026.2.25.tar.gz", hash = "sha256:e887ab5cee78ea814d3472169153c2d12cd43b14bd03329a39a9c6e2e80bfba7", size = 155029, upload-time = "2026-02-25T02:54:17.342Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/60/e3bec1881450851b087e301bedc3daa9377a4d45f1c26aa90b0b235e38aa/charset_normalizer-3.4.6.tar.gz", hash = "sha256:1ae6b62897110aa7c79ea2f5dd38d1abca6db663687c0b1ad9aed6f6bae3d9d6", size = 143363, upload-time = "2026-03-15T18:53:25.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e5/62/c0815c992c9545347aeea7859b50dc9044d147e2e7278329c6e02ac9a616/charset_normalizer-3.4.6-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:2ef7fedc7a6ecbe99969cd09632516738a97eeb8bd7258bf8a0f23114c057dab", size = 295154, upload-time = "2026-03-15T18:50:50.88Z" }, + { url = "https://files.pythonhosted.org/packages/a8/37/bdca6613c2e3c58c7421891d80cc3efa1d32e882f7c4a7ee6039c3fc951a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a4ea868bc28109052790eb2b52a9ab33f3aa7adc02f96673526ff47419490e21", size = 199191, upload-time = "2026-03-15T18:50:52.658Z" }, + { url = "https://files.pythonhosted.org/packages/6c/92/9934d1bbd69f7f398b38c5dae1cbf9cc672e7c34a4adf7b17c0a9c17d15d/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:836ab36280f21fc1a03c99cd05c6b7af70d2697e374c7af0b61ed271401a72a2", size = 218674, upload-time = "2026-03-15T18:50:54.102Z" }, + { url = "https://files.pythonhosted.org/packages/af/90/25f6ab406659286be929fd89ab0e78e38aa183fc374e03aa3c12d730af8a/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f1ce721c8a7dfec21fcbdfe04e8f68174183cf4e8188e0645e92aa23985c57ff", size = 215259, upload-time = "2026-03-15T18:50:55.616Z" }, + { url = "https://files.pythonhosted.org/packages/4e/ef/79a463eb0fff7f96afa04c1d4c51f8fc85426f918db467854bfb6a569ce3/charset_normalizer-3.4.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e28d62a8fc7a1fa411c43bd65e346f3bce9716dc51b897fbe930c5987b402d5", size = 207276, upload-time = "2026-03-15T18:50:57.054Z" }, + { url = "https://files.pythonhosted.org/packages/f7/72/d0426afec4b71dc159fa6b4e68f868cd5a3ecd918fec5813a15d292a7d10/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:530d548084c4a9f7a16ed4a294d459b4f229db50df689bfe92027452452943a0", size = 195161, upload-time = "2026-03-15T18:50:58.686Z" }, + { url = "https://files.pythonhosted.org/packages/bf/18/c82b06a68bfcb6ce55e508225d210c7e6a4ea122bfc0748892f3dc4e8e11/charset_normalizer-3.4.6-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:30f445ae60aad5e1f8bdbb3108e39f6fbc09f4ea16c815c66578878325f8f15a", size = 203452, upload-time = "2026-03-15T18:51:00.196Z" }, + { url = "https://files.pythonhosted.org/packages/44/d6/0c25979b92f8adafdbb946160348d8d44aa60ce99afdc27df524379875cb/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:ac2393c73378fea4e52aa56285a3d64be50f1a12395afef9cce47772f60334c2", size = 202272, upload-time = "2026-03-15T18:51:01.703Z" }, + { url = "https://files.pythonhosted.org/packages/2e/3d/7fea3e8fe84136bebbac715dd1221cc25c173c57a699c030ab9b8900cbb7/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:90ca27cd8da8118b18a52d5f547859cc1f8354a00cd1e8e5120df3e30d6279e5", size = 195622, upload-time = "2026-03-15T18:51:03.526Z" }, + { url = "https://files.pythonhosted.org/packages/57/8a/d6f7fd5cb96c58ef2f681424fbca01264461336d2a7fc875e4446b1f1346/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:8e5a94886bedca0f9b78fecd6afb6629142fd2605aa70a125d49f4edc6037ee6", size = 220056, upload-time = "2026-03-15T18:51:05.269Z" }, + { url = "https://files.pythonhosted.org/packages/16/50/478cdda782c8c9c3fb5da3cc72dd7f331f031e7f1363a893cdd6ca0f8de0/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:695f5c2823691a25f17bc5d5ffe79fa90972cc34b002ac6c843bb8a1720e950d", size = 203751, upload-time = "2026-03-15T18:51:06.858Z" }, + { url = "https://files.pythonhosted.org/packages/75/fc/cc2fcac943939c8e4d8791abfa139f685e5150cae9f94b60f12520feaa9b/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:231d4da14bcd9301310faf492051bee27df11f2bc7549bc0bb41fef11b82daa2", size = 216563, upload-time = "2026-03-15T18:51:08.564Z" }, + { url = "https://files.pythonhosted.org/packages/a8/b7/a4add1d9a5f68f3d037261aecca83abdb0ab15960a3591d340e829b37298/charset_normalizer-3.4.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a056d1ad2633548ca18ffa2f85c202cfb48b68615129143915b8dc72a806a923", size = 209265, upload-time = "2026-03-15T18:51:10.312Z" }, + { url = "https://files.pythonhosted.org/packages/6c/18/c094561b5d64a24277707698e54b7f67bd17a4f857bbfbb1072bba07c8bf/charset_normalizer-3.4.6-cp312-cp312-win32.whl", hash = "sha256:c2274ca724536f173122f36c98ce188fd24ce3dad886ec2b7af859518ce008a4", size = 144229, upload-time = "2026-03-15T18:51:11.694Z" }, + { url = "https://files.pythonhosted.org/packages/ab/20/0567efb3a8fd481b8f34f739ebddc098ed062a59fed41a8d193a61939e8f/charset_normalizer-3.4.6-cp312-cp312-win_amd64.whl", hash = "sha256:c8ae56368f8cc97c7e40a7ee18e1cedaf8e780cd8bc5ed5ac8b81f238614facb", size = 154277, upload-time = "2026-03-15T18:51:13.004Z" }, + { url = "https://files.pythonhosted.org/packages/15/57/28d79b44b51933119e21f65479d0864a8d5893e494cf5daab15df0247c17/charset_normalizer-3.4.6-cp312-cp312-win_arm64.whl", hash = "sha256:899d28f422116b08be5118ef350c292b36fc15ec2daeb9ea987c89281c7bb5c4", size = 142817, upload-time = "2026-03-15T18:51:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/2a/68/687187c7e26cb24ccbd88e5069f5ef00eba804d36dde11d99aad0838ab45/charset_normalizer-3.4.6-py3-none-any.whl", hash = "sha256:947cf925bc916d90adba35a64c82aace04fa39b46b52d4630ece166655905a69", size = 61455, upload-time = "2026-03-15T18:53:23.833Z" }, +] + +[[package]] +name = "click" +version = "8.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3d/fa/656b739db8587d7b5dfa22e22ed02566950fbfbcdc20311993483657a5c0/click-8.3.1.tar.gz", hash = "sha256:12ff4785d337a1bb490bb7e9c2b1ee5da3112e94a8622f26a6c77f5d2fc6842a", size = 295065, upload-time = "2025-11-15T20:45:42.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/98/78/01c019cdb5d6498122777c1a43056ebb3ebfeef2076d9d026bfe15583b2b/click-8.3.1-py3-none-any.whl", hash = "sha256:981153a64e25f12d547d3426c367a4857371575ee7ad18df2a6183ab0545b2a6", size = 108274, upload-time = "2025-11-15T20:45:41.139Z" }, +] + [[package]] name = "colorama" version = "0.4.6" @@ -43,6 +104,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "cuda-bindings" +version = "12.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, +] + +[[package]] +name = "cuda-pathfinder" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/66/0c02bd330e7d976f83fa68583d6198d76f23581bcbb5c0e98a6148f326e5/cuda_pathfinder-1.5.0-py3-none-any.whl", hash = "sha256:498f90a9e9de36044a7924742aecce11c50c49f735f1bc53e05aa46de9ea4110", size = 49739, upload-time = "2026-03-24T21:14:30.869Z" }, +] + [[package]] name = "cython" version = "3.0.11" @@ -69,6 +149,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/84/02fc1827e8cdded4aa65baef11296a9bbe595c474f0d6d758af082d849fd/execnet-2.1.2-py3-none-any.whl", hash = "sha256:67fba928dd5a544b783f6056f449e5e3931a5c378b128bc18501f7ea79e296ec", size = 40708, upload-time = "2025-11-12T09:56:36.333Z" }, ] +[[package]] +name = "filelock" +version = "3.25.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, +] + +[[package]] +name = "gitdb" +version = "4.0.12" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "smmap" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/72/94/63b0fc47eb32792c7ba1fe1b694daec9a63620db1e313033d18140c2320a/gitdb-4.0.12.tar.gz", hash = "sha256:5ef71f855d191a3326fcfbc0d5da835f26b13fbcba60c32c21091c349ffdb571", size = 394684, upload-time = "2025-01-02T07:20:46.413Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/61/5c78b91c3143ed5c14207f463aecfc8f9dbb5092fb2869baf37c273b2705/gitdb-4.0.12-py3-none-any.whl", hash = "sha256:67073e15955400952c6565cc3e707c554a4eea2e428946f7a4c162fab9bd9bcf", size = 62794, upload-time = "2025-01-02T07:20:43.624Z" }, +] + +[[package]] +name = "gitpython" +version = "3.1.46" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "gitdb" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/b5/59d16470a1f0dfe8c793f9ef56fd3826093fc52b3bd96d6b9d6c26c7e27b/gitpython-3.1.46.tar.gz", hash = "sha256:400124c7d0ef4ea03f7310ac2fbf7151e09ff97f2a3288d64a440c584a29c37f", size = 215371, upload-time = "2026-01-01T15:37:32.073Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/09/e21df6aef1e1ffc0c816f0522ddc3f6dcded766c3261813131c78a704470/gitpython-3.1.46-py3-none-any.whl", hash = "sha256:79812ed143d9d25b6d176a10bb511de0f9c67b1fa641d82097b0ab90398a2058", size = 208620, upload-time = "2026-01-01T15:37:30.574Z" }, +] + +[[package]] +name = "idna" +version = "3.11" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6f/6d/0703ccc57f3a7233505399edb88de3cbd678da106337b9fcde432b65ed60/idna-3.11.tar.gz", hash = "sha256:795dafcc9c04ed0c1fb032c2aa73654d8e8c5023a7df64a53f39190ada629902", size = 194582, upload-time = "2025-10-12T14:55:20.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -78,6 +209,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + [[package]] name = "llvmlite" version = "0.44.0" @@ -91,6 +234,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e2/3b/a9a17366af80127bd09decbe2a54d8974b6d8b274b39bf47fbaedeec6307/llvmlite-0.44.0-cp312-cp312-win_amd64.whl", hash = "sha256:eae7e2d4ca8f88f89d315b48c6b741dcb925d6a1042da694aa16ab3dd4cbd3a1", size = 30332380, upload-time = "2025-01-20T11:14:02.442Z" }, ] +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + [[package]] name = "numba" version = "0.61.0" @@ -126,6 +306,140 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a6/84/fa11dad3404b7634aaab50733581ce11e5350383311ea7a7010f464c0170/numpy-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:0d30c543f02e84e92c4b1f415b7c6b5326cbe45ee7882b6b77db7195fb971e3a", size = 12566858, upload-time = "2024-11-02T17:40:08.851Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12" }, + { name = "nvidia-cusparse-cu12" }, + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.27.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + [[package]] name = "packaging" version = "26.0" @@ -135,6 +449,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, ] +[[package]] +name = "platformdirs" +version = "4.9.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/19/56/8d4c30c8a1d07013911a8fdbd8f89440ef9f08d07a1b50ab8ca8be5a20f9/platformdirs-4.9.4.tar.gz", hash = "sha256:1ec356301b7dc906d83f371c8f487070e99d3ccf9e501686456394622a01a934", size = 28737, upload-time = "2026-03-05T18:34:13.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/d7/97f7e3a6abb67d8080dd406fd4df842c2be0efaf712d1c899c32a075027c/platformdirs-4.9.4-py3-none-any.whl", hash = "sha256:68a9a4619a666ea6439f2ff250c12a853cd1cbd5158d258bd824a7df6be2f868", size = 21216, upload-time = "2026-03-05T18:34:12.172Z" }, +] + [[package]] name = "pluggy" version = "1.6.0" @@ -144,6 +467,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "protobuf" +version = "6.33.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/66/70/e908e9c5e52ef7c3a6c7902c9dfbb34c7e29c25d2f81ade3856445fd5c94/protobuf-6.33.6.tar.gz", hash = "sha256:a6768d25248312c297558af96a9f9c929e8c4cee0659cb07e780731095f38135", size = 444531, upload-time = "2026-03-18T19:05:00.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fc/9f/2f509339e89cfa6f6a4c4ff50438db9ca488dec341f7e454adad60150b00/protobuf-6.33.6-cp310-abi3-win32.whl", hash = "sha256:7d29d9b65f8afef196f8334e80d6bc1d5d4adedb449971fefd3723824e6e77d3", size = 425739, upload-time = "2026-03-18T19:04:48.373Z" }, + { url = "https://files.pythonhosted.org/packages/76/5d/683efcd4798e0030c1bab27374fd13a89f7c2515fb1f3123efdfaa5eab57/protobuf-6.33.6-cp310-abi3-win_amd64.whl", hash = "sha256:0cd27b587afca21b7cfa59a74dcbd48a50f0a6400cfb59391340ad729d91d326", size = 437089, upload-time = "2026-03-18T19:04:50.381Z" }, + { url = "https://files.pythonhosted.org/packages/5c/01/a3c3ed5cd186f39e7880f8303cc51385a198a81469d53d0fdecf1f64d929/protobuf-6.33.6-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:9720e6961b251bde64edfdab7d500725a2af5280f3f4c87e57c0208376aa8c3a", size = 427737, upload-time = "2026-03-18T19:04:51.866Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/b3c01fdec7d2f627b3a6884243ba328c1217ed2d978def5c12dc50d328a3/protobuf-6.33.6-cp39-abi3-manylinux2014_aarch64.whl", hash = "sha256:e2afbae9b8e1825e3529f88d514754e094278bb95eadc0e199751cdd9a2e82a2", size = 324610, upload-time = "2026-03-18T19:04:53.096Z" }, + { url = "https://files.pythonhosted.org/packages/9b/ca/25afc144934014700c52e05103c2421997482d561f3101ff352e1292fb81/protobuf-6.33.6-cp39-abi3-manylinux2014_s390x.whl", hash = "sha256:c96c37eec15086b79762ed265d59ab204dabc53056e3443e702d2681f4b39ce3", size = 339381, upload-time = "2026-03-18T19:04:54.616Z" }, + { url = "https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl", hash = "sha256:e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593", size = 323436, upload-time = "2026-03-18T19:04:55.768Z" }, + { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, +] + [[package]] name = "psutil" version = "5.9.0" @@ -159,6 +497,50 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e0/a9/023730ba63db1e494a271cb018dcd361bd2c917ba7004c3e49d5daf795a2/py_cpuinfo-9.0.0-py3-none-any.whl", hash = "sha256:859625bc251f64e21f077d099d4162689c762b5d6a4c3c97553d56241c9674d5", size = 22335, upload-time = "2022-10-25T20:38:27.636Z" }, ] +[[package]] +name = "pydantic" +version = "2.12.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types" }, + { name = "pydantic-core" }, + { name = "typing-extensions" }, + { name = "typing-inspection" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.41.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/5d/5f6c63eebb5afee93bcaae4ce9a898f3373ca23df3ccaef086d0233a35a7/pydantic_core-2.41.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:f41a7489d32336dbf2199c8c0a215390a751c5b014c2c1c5366e817202e9cdf7", size = 2110990, upload-time = "2025-11-04T13:39:58.079Z" }, + { url = "https://files.pythonhosted.org/packages/aa/32/9c2e8ccb57c01111e0fd091f236c7b371c1bccea0fa85247ac55b1e2b6b6/pydantic_core-2.41.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:070259a8818988b9a84a449a2a7337c7f430a22acc0859c6b110aa7212a6d9c0", size = 1896003, upload-time = "2025-11-04T13:39:59.956Z" }, + { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, + { url = "https://files.pythonhosted.org/packages/38/de/8c36b5198a29bdaade07b5985e80a233a5ac27137846f3bc2d3b40a47360/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed2e99c456e3fadd05c991f8f437ef902e00eedf34320ba2b0842bd1c3ca3a75", size = 2052578, upload-time = "2025-11-04T13:40:04.401Z" }, + { url = "https://files.pythonhosted.org/packages/00/b5/0e8e4b5b081eac6cb3dbb7e60a65907549a1ce035a724368c330112adfdd/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:65840751b72fbfd82c3c640cff9284545342a4f1eb1586ad0636955b261b0b05", size = 2208504, upload-time = "2025-11-04T13:40:06.072Z" }, + { url = "https://files.pythonhosted.org/packages/77/56/87a61aad59c7c5b9dc8caad5a41a5545cba3810c3e828708b3d7404f6cef/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e536c98a7626a98feb2d3eaf75944ef6f3dbee447e1f841eae16f2f0a72d8ddc", size = 2335816, upload-time = "2025-11-04T13:40:07.835Z" }, + { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, + { url = "https://files.pythonhosted.org/packages/d3/43/ebef01f69baa07a482844faaa0a591bad1ef129253ffd0cdaa9d8a7f72d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d38548150c39b74aeeb0ce8ee1d8e82696f4a4e16ddc6de7b1d8823f7de4b9b5", size = 2171698, upload-time = "2025-11-04T13:40:12.004Z" }, + { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, + { url = "https://files.pythonhosted.org/packages/49/7d/4c00df99cb12070b6bccdef4a195255e6020a550d572768d92cc54dba91a/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:482c982f814460eabe1d3bb0adfdc583387bd4691ef00b90575ca0d2b6fe2294", size = 2329591, upload-time = "2025-11-04T13:40:15.672Z" }, + { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, + { url = "https://files.pythonhosted.org/packages/49/3b/774f2b5cd4192d5ab75870ce4381fd89cf218af999515baf07e7206753f0/pydantic_core-2.41.5-cp312-cp312-win32.whl", hash = "sha256:b74557b16e390ec12dca509bce9264c3bbd128f8a2c376eaa68003d7f327276d", size = 1985908, upload-time = "2025-11-04T13:40:19.309Z" }, + { url = "https://files.pythonhosted.org/packages/86/45/00173a033c801cacf67c190fef088789394feaf88a98a7035b0e40d53dc9/pydantic_core-2.41.5-cp312-cp312-win_amd64.whl", hash = "sha256:1962293292865bca8e54702b08a4f26da73adc83dd1fcf26fbc875b35d81c815", size = 2020145, upload-time = "2025-11-04T13:40:21.548Z" }, + { url = "https://files.pythonhosted.org/packages/f9/22/91fbc821fa6d261b376a3f73809f907cec5ca6025642c463d3488aad22fb/pydantic_core-2.41.5-cp312-cp312-win_arm64.whl", hash = "sha256:1746d4a3d9a794cacae06a5eaaccb4b8643a131d45fbc9af23e353dc0a5ba5c3", size = 1976179, upload-time = "2025-11-04T13:40:23.393Z" }, + { url = "https://files.pythonhosted.org/packages/09/32/59b0c7e63e277fa7911c2fc70ccfb45ce4b98991e7ef37110663437005af/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:7da7087d756b19037bc2c06edc6c170eeef3c3bafcb8f532ff17d64dc427adfd", size = 2110495, upload-time = "2025-11-04T13:42:49.689Z" }, + { url = "https://files.pythonhosted.org/packages/aa/81/05e400037eaf55ad400bcd318c05bb345b57e708887f07ddb2d20e3f0e98/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:aabf5777b5c8ca26f7824cb4a120a740c9588ed58df9b2d196ce92fba42ff8dc", size = 1915388, upload-time = "2025-11-04T13:42:52.215Z" }, + { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, + { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, +] + [[package]] name = "pygments" version = "2.19.2" @@ -196,3 +578,186 @@ sdist = { url = "https://files.pythonhosted.org/packages/78/b4/439b179d1ff526791 wheels = [ { url = "https://files.pythonhosted.org/packages/ca/31/d4e37e9e550c2b92a9cbc2e4d0b7420a27224968580b5a447f420847c975/pytest_xdist-3.8.0-py3-none-any.whl", hash = "sha256:202ca578cfeb7370784a8c33d6d05bc6e13b4f25b5053c30a152269fd10f0b88", size = 46396, upload-time = "2025-07-01T13:30:56.632Z" }, ] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/33/422b98d2195232ca1826284a76852ad5a86fe23e31b009c9886b2d0fb8b2/pyyaml-6.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:7f047e29dcae44602496db43be01ad42fc6f1cc0d8cd6c83d342306c32270196", size = 182063, upload-time = "2025-09-25T21:32:11.445Z" }, + { url = "https://files.pythonhosted.org/packages/89/a0/6cf41a19a1f2f3feab0e9c0b74134aa2ce6849093d5517a0c550fe37a648/pyyaml-6.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:fc09d0aa354569bc501d4e787133afc08552722d3ab34836a80547331bb5d4a0", size = 173973, upload-time = "2025-09-25T21:32:12.492Z" }, + { url = "https://files.pythonhosted.org/packages/ed/23/7a778b6bd0b9a8039df8b1b1d80e2e2ad78aa04171592c8a5c43a56a6af4/pyyaml-6.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9149cad251584d5fb4981be1ecde53a1ca46c891a79788c0df828d2f166bda28", size = 775116, upload-time = "2025-09-25T21:32:13.652Z" }, + { url = "https://files.pythonhosted.org/packages/65/30/d7353c338e12baef4ecc1b09e877c1970bd3382789c159b4f89d6a70dc09/pyyaml-6.0.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:5fdec68f91a0c6739b380c83b951e2c72ac0197ace422360e6d5a959d8d97b2c", size = 844011, upload-time = "2025-09-25T21:32:15.21Z" }, + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/05/c0/b3be26a015601b822b97d9149ff8cb5ead58c66f981e04fedf4e762f4bd4/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8dc52c23056b9ddd46818a57b78404882310fb473d63f17b07d5c40421e47f8e", size = 761089, upload-time = "2025-09-25T21:32:17.56Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, + { url = "https://files.pythonhosted.org/packages/74/93/7baea19427dcfbe1e5a372d81473250b379f04b1bd3c4c5ff825e2327202/pyyaml-6.0.3-cp312-cp312-win32.whl", hash = "sha256:96b533f0e99f6579b3d4d4995707cf36df9100d67e0c8303a0c55b27b5f99bc5", size = 137658, upload-time = "2025-09-25T21:32:20.209Z" }, + { url = "https://files.pythonhosted.org/packages/86/bf/899e81e4cce32febab4fb42bb97dcdf66bc135272882d1987881a4b519e9/pyyaml-6.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:5fcd34e47f6e0b794d17de1b4ff496c00986e1c83f7ab2fb8fcfe9616ff7477b", size = 154003, upload-time = "2025-09-25T21:32:21.167Z" }, + { url = "https://files.pythonhosted.org/packages/1a/08/67bd04656199bbb51dbed1439b7f27601dfb576fb864099c7ef0c3e55531/pyyaml-6.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:64386e5e707d03a7e172c0701abfb7e10f0fb753ee1d773128192742712a98fd", size = 140344, upload-time = "2025-09-25T21:32:22.617Z" }, +] + +[[package]] +name = "requests" +version = "2.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "charset-normalizer" }, + { name = "idna" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/34/64/8860370b167a9721e8956ae116825caff829224fbca0ca6e7bf8ddef8430/requests-2.33.0.tar.gz", hash = "sha256:c7ebc5e8b0f21837386ad0e1c8fe8b829fa5f544d8df3b2253bff14ef29d7652", size = 134232, upload-time = "2026-03-25T15:10:41.586Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/5d/c814546c2333ceea4ba42262d8c4d55763003e767fa169adc693bd524478/requests-2.33.0-py3-none-any.whl", hash = "sha256:3324635456fa185245e24865e810cecec7b4caf933d7eb133dcde67d48cee69b", size = 65017, upload-time = "2026-03-25T15:10:40.382Z" }, +] + +[[package]] +name = "sentry-sdk" +version = "2.56.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/df/5008954f5466085966468612a7d1638487596ee6d2fd7fb51783a85351bf/sentry_sdk-2.56.0.tar.gz", hash = "sha256:fdab72030b69625665b2eeb9738bdde748ad254e8073085a0ce95382678e8168", size = 426820, upload-time = "2026-03-24T09:56:36.575Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cd/1a/b3a3e9f6520493fed7997af4d2de7965d71549c62f994a8fd15f2ecd519e/sentry_sdk-2.56.0-py2.py3-none-any.whl", hash = "sha256:5afafb744ceb91d22f4cc650c6bd048ac6af5f7412dcc6c59305a2e36f4dbc02", size = 451568, upload-time = "2026-03-24T09:56:34.807Z" }, +] + +[[package]] +name = "setuptools" +version = "81.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/0d/1c/73e719955c59b8e424d015ab450f51c0af856ae46ea2da83eba51cc88de1/setuptools-81.0.0.tar.gz", hash = "sha256:487b53915f52501f0a79ccfd0c02c165ffe06631443a886740b91af4b7a5845a", size = 1198299, upload-time = "2026-02-06T21:10:39.601Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl", hash = "sha256:fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", size = 1062021, upload-time = "2026-02-06T21:10:37.175Z" }, +] + +[[package]] +name = "smmap" +version = "5.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1f/ea/49c993d6dfdd7338c9b1000a0f36817ed7ec84577ae2e52f890d1a4ff909/smmap-5.0.3.tar.gz", hash = "sha256:4d9debb8b99007ae47165abc08670bd74cb74b5227dda7f643eccc4e9eb5642c", size = 22506, upload-time = "2026-03-09T03:43:26.1Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "torch" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools" }, + { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, + { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, + { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, + { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, + { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, +] + +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/17/5d/08201db32823bdf77a0e2b9039540080b2e5c23a20706ddba942924ebcd6/triton-3.6.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4", size = 176128243, upload-time = "2026-01-20T16:16:07.857Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, +] + +[[package]] +name = "urllib3" +version = "2.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c7/24/5f1b3bdffd70275f6661c76461e25f024d5a38a46f04aaca912426a2b1d3/urllib3-2.6.3.tar.gz", hash = "sha256:1b62b6884944a57dbe321509ab94fd4d3b307075e0c2eae991ac71ee15ad38ed", size = 435556, upload-time = "2026-01-07T16:24:43.925Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, +] + +[[package]] +name = "wandb" +version = "0.25.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click" }, + { name = "gitpython" }, + { name = "packaging" }, + { name = "platformdirs" }, + { name = "protobuf" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "requests" }, + { name = "sentry-sdk" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/60/bb/eb579bf9abac70934a014a9d4e45346aab307994f3021d201bebe5fa25ec/wandb-0.25.1.tar.gz", hash = "sha256:b2a95cd777ecbe7499599a43158834983448a0048329bc7210ef46ca18d21994", size = 43983308, upload-time = "2026-03-10T23:51:44.227Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/d8/873553b6818499d1b1de314067d528b892897baf0dc81fedc0e845abc2dd/wandb-0.25.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:9bb0679a3e2dcd96db9d9b6d3e17d046241d8d122974b24facb85cc93309a8c9", size = 23615900, upload-time = "2026-03-10T23:51:06.278Z" }, + { url = "https://files.pythonhosted.org/packages/71/ea/b131f319aaa5d0bf7572b6bfcff3dd89e1cf92b17eee443bbab71d12d74c/wandb-0.25.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:0fb13ed18914027523e7b4fc20380c520e0d10da0ee452f924a13f84509fbe12", size = 25576144, upload-time = "2026-03-10T23:51:11.527Z" }, + { url = "https://files.pythonhosted.org/packages/70/5f/81508581f0bb77b0495665c1c78e77606a48e66e855ca71ba7c8ae29efa4/wandb-0.25.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:cc4521eb5223429ddab5e8eee9b42fdf4caabdf0bc4e0e809042720e5fbef0ed", size = 23070425, upload-time = "2026-03-10T23:51:15.71Z" }, + { url = "https://files.pythonhosted.org/packages/f2/c7/445155ef010e2e35d190797d7c36ff441e062a5b566a6da4778e22233395/wandb-0.25.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:e73b4c55b947edae349232d5845204d30fac88e18eb4ad1d4b96bf7cf898405a", size = 25628142, upload-time = "2026-03-10T23:51:19.326Z" }, + { url = "https://files.pythonhosted.org/packages/d5/63/f5c55ee00cf481ef1ccd3c385a0585ad52e7840d08419d4f82ddbeeea959/wandb-0.25.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:22b84065aa398e1624d2e5ad79e08bc4d2af41a6db61697b03b3aaba332977c6", size = 23123172, upload-time = "2026-03-10T23:51:23.418Z" }, + { url = "https://files.pythonhosted.org/packages/3e/d9/19eb7974c0e9253bcbaee655222c0f0e1a52e63e9479ee711b4208f8ac31/wandb-0.25.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:005c4c6b5126ef8f4b4110e5372d950918b00637d6dc4b615ad17445f9739478", size = 25714479, upload-time = "2026-03-10T23:51:27.421Z" }, + { url = "https://files.pythonhosted.org/packages/11/19/466c1d03323a4a0ed7d4036a59b18d6b6f67cb5032e444205927e226b18d/wandb-0.25.1-py3-none-win32.whl", hash = "sha256:8f2d04f16b88d65bfba9d79fb945f6c64e2686215469a841936e0972be8ec6a5", size = 24967338, upload-time = "2026-03-10T23:51:31.833Z" }, + { url = "https://files.pythonhosted.org/packages/89/22/680d34c1587f3a979c701b66d71aa7c42b4ef2fdf0774f67034e618e834e/wandb-0.25.1-py3-none-win_amd64.whl", hash = "sha256:62db5166de14456156d7a85953a58733a631228e6d4248a753605f75f75fb845", size = 24967343, upload-time = "2026-03-10T23:51:36.026Z" }, + { url = "https://files.pythonhosted.org/packages/c4/e8/76836b75d401ff5912aaf513176e64557ceaec4c4946bfd38a698ff84d48/wandb-0.25.1-py3-none-win_arm64.whl", hash = "sha256:cc7c34b70cf4b7be4d395541e82e325fd9d2be978d62c9ec01f1a7141523b6bb", size = 22080774, upload-time = "2026-03-10T23:51:40.196Z" }, +]